Skip to content

Commit 082e69f

Browse files
lucylqlucylq
committed
Use PTEFile class in serialize_pte_binary
Pull Request resolved: #15801 Take in a PTE file for serialization instead of program, mutable_segments, named_data segments separately. ghstack-source-id: 323125958 @exported-using-ghexport Differential Revision: [D86908241](https://our.internmc.facebook.com/intern/diff/D86908241/)
1 parent 8bb5c50 commit 082e69f

File tree

6 files changed

+48
-40
lines changed

6 files changed

+48
-40
lines changed

‎devtools/bundled_program/test/test_bundle_data.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
from executorch.devtools.bundled_program.util.test_util import (
1919
get_common_executorch_program,
2020
)
21-
from executorch.exir._serialize import _serialize_pte_binary
21+
from executorch.exir._serialize import _PTEFile, _serialize_pte_binary
2222

2323

2424
class TestBundle(unittest.TestCase):
@@ -72,7 +72,11 @@ def test_bundled_program(self) -> None:
7272

7373
self.assertEqual(
7474
bundled_program.serialize_to_schema().program,
75-
bytes(_serialize_pte_binary(executorch_program.executorch_program)),
75+
bytes(
76+
_serialize_pte_binary(
77+
pte_file=_PTEFile(program=executorch_program.executorch_program)
78+
)
79+
),
7680
)
7781

7882
def test_bundled_program_from_pte(self) -> None:

‎exir/_serialize/_program.py‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -419,19 +419,17 @@ def _extract_named_data(
419419

420420

421421
def serialize_pte_binary(
422-
program: Program,
422+
pte_file: PTEFile,
423423
*,
424-
mutable_data: Optional[List[Buffer]] = None,
425424
extract_delegate_segments: bool = False,
426425
segment_alignment: int = 128,
427426
constant_tensor_alignment: Optional[int] = None,
428427
delegate_alignment: Optional[int] = None,
429-
named_data: Optional[NamedDataStoreOutput] = None,
430428
) -> Cord:
431429
"""Returns the runtime binary representation of the given Program.
432430
433431
Args:
434-
program: The Program to serialize.
432+
pte_file: PTEFile class containing the program and segments.
435433
extract_delegate_segments: Whether to move delegate data blobs from the
436434
Program into separate segments, rather than encoding those blobs
437435
in the flatbuffer data. When true, will also:
@@ -446,8 +444,6 @@ def serialize_pte_binary(
446444
delegate_alignment: If provided, the minimum alignment of delegate data
447445
in the program. Must be a power of 2. If not provided, uses the
448446
value in the schema file.
449-
named_data: If provided, named blobs to be stored in segments
450-
after the PTE file.
451447
Returns:
452448
The serialized form of the Program, ready for execution by the runtime.
453449
"""
@@ -458,7 +454,7 @@ def serialize_pte_binary(
458454
# Don't modify the original program.
459455
# TODO(T144120904): Could avoid yet more huge copies with a more shallow
460456
# copy, reusing the actual data blobs.
461-
program = copy.deepcopy(program)
457+
program = copy.deepcopy(pte_file.program)
462458

463459
# Store extracted segment data, with any buffer-specific alignment.
464460
# This may be constant data, delegate data or named data.
@@ -482,9 +478,9 @@ def serialize_pte_binary(
482478
# Add to the aggregate segments cord.
483479
segments.append(AlignedData(constant_segment_data))
484480

485-
if mutable_data is not None:
481+
if pte_file.mutable_data is not None:
486482
mutable_segment_data, mutable_segment_offsets = _extract_constant_segment(
487-
mutable_data,
483+
pte_file.mutable_data,
488484
tensor_alignment=None, # data is copied at Method load so no need to align.
489485
)
490486
if len(mutable_segment_data) > 0:
@@ -499,8 +495,10 @@ def serialize_pte_binary(
499495

500496
if extract_delegate_segments:
501497
_extract_delegate_segments(program, segments)
502-
if named_data is not None:
503-
_extract_named_data(program, segments, named_data.buffers, named_data.pte_data)
498+
if pte_file.named_data is not None:
499+
_extract_named_data(
500+
program, segments, pte_file.named_data.buffers, pte_file.named_data.pte_data
501+
)
504502

505503
# Append all segments into a single Cord, adding any necessary padding to ensure that
506504
# each segment begins at the required alignment.

‎exir/_serialize/_serialize.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,10 @@
88

99
from typing import Dict, Optional, Set, Tuple
1010

11-
from executorch.exir._serialize import _serialize_pte_binary
12-
1311
from executorch.exir._serialize._cord import Cord
1412
from executorch.exir._serialize._named_data_store import NamedDataStoreOutput
13+
14+
from executorch.exir._serialize._program import PTEFile, serialize_pte_binary
1515
from executorch.exir._serialize.data_serializer import (
1616
DataEntry,
1717
DataPayload,
@@ -46,14 +46,16 @@ def serialize_for_executorch(
4646
pte_data=named_data_store.pte_data,
4747
external_data={},
4848
)
49-
pte: Cord = _serialize_pte_binary(
50-
program=emitter_output.program,
51-
mutable_data=emitter_output.mutable_data,
49+
pte: Cord = serialize_pte_binary(
50+
pte_file=PTEFile(
51+
program=emitter_output.program,
52+
mutable_data=emitter_output.mutable_data,
53+
named_data=pte_named_data,
54+
),
5255
extract_delegate_segments=config.extract_delegate_segments,
5356
segment_alignment=config.segment_alignment,
5457
constant_tensor_alignment=config.constant_tensor_alignment,
5558
delegate_alignment=config.delegate_alignment,
56-
named_data=pte_named_data,
5759
)
5860

5961
# Serialize PTD files.

‎exir/_serialize/test/test_program.py‎

Lines changed: 17 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
_json_to_program,
2525
_program_to_json,
2626
deserialize_pte_binary,
27+
PTEFile,
2728
serialize_pte_binary,
2829
)
2930
from executorch.exir._serialize.data_serializer import DataEntry
@@ -173,7 +174,7 @@ def constant_segment_with_tensor_alignment(
173174
# Extract blobs into constant segment during serialization.
174175
pte_data = bytes(
175176
serialize_pte_binary(
176-
program,
177+
PTEFile(program=program),
177178
segment_alignment=SEGMENT_ALIGNMENT,
178179
constant_tensor_alignment=constant_tensor_alignment,
179180
)
@@ -446,7 +447,7 @@ def test_round_trip_no_header_no_segments(self) -> None:
446447
deserializing.
447448
"""
448449
program = get_test_program()
449-
pte_data = bytes(serialize_pte_binary(program))
450+
pte_data = bytes(serialize_pte_binary(pte_file=PTEFile(program)))
450451
self.assertGreater(len(pte_data), 16)
451452

452453
# File magic should be present at the expected offset.
@@ -471,7 +472,7 @@ def test_round_trip_large_buffer_sizes(self) -> None:
471472
"""
472473
program = get_test_program()
473474
program.execution_plan[0].non_const_buffer_sizes = [0, 2**48]
474-
flatbuffer_from_py = bytes(serialize_pte_binary(program))
475+
flatbuffer_from_py = bytes(serialize_pte_binary(pte_file=PTEFile(program)))
475476
self.assert_programs_equal(
476477
program, deserialize_pte_binary(flatbuffer_from_py).program
477478
)
@@ -483,7 +484,11 @@ def test_round_trip_no_segments_and_no_header(self) -> None:
483484
the same after serializing and deserializing.
484485
"""
485486
program = get_test_program()
486-
pte_data = bytes(serialize_pte_binary(program, extract_delegate_segments=True))
487+
pte_data = bytes(
488+
serialize_pte_binary(
489+
pte_file=PTEFile(program), extract_delegate_segments=True
490+
)
491+
)
487492
self.assertGreater(len(pte_data), 16)
488493

489494
# File magic should be present at the expected offset.
@@ -533,7 +538,7 @@ def test_round_trip_with_segments(self) -> None:
533538
# Extract the blobs into segments during serialization.
534539
pte_data = bytes(
535540
serialize_pte_binary(
536-
program,
541+
PTEFile(program=program),
537542
extract_delegate_segments=True,
538543
segment_alignment=SEGMENT_ALIGNMENT,
539544
)
@@ -647,7 +652,7 @@ def test_no_constants(self) -> None:
647652

648653
pte_data = bytes(
649654
serialize_pte_binary(
650-
program,
655+
PTEFile(program=program),
651656
extract_delegate_segments=True,
652657
segment_alignment=SEGMENT_ALIGNMENT,
653658
constant_tensor_alignment=CONSTANT_TENSOR_ALIGNMENT,
@@ -679,7 +684,7 @@ def test_unused_inline_delegate_blobs_with_segments(self) -> None:
679684
# Extract the blobs into segments should succeeed.
680685
pte_data = bytes(
681686
serialize_pte_binary(
682-
program,
687+
PTEFile(program=program),
683688
extract_delegate_segments=True,
684689
segment_alignment=SEGMENT_ALIGNMENT,
685690
)
@@ -694,7 +699,7 @@ def test_unused_inline_delegate_blobs_with_segments(self) -> None:
694699
# Should cause serialization to fail.
695700
with self.assertRaises(ValueError):
696701
serialize_pte_binary(
697-
program,
702+
PTEFile(program=program),
698703
extract_delegate_segments=True,
699704
segment_alignment=SEGMENT_ALIGNMENT,
700705
)
@@ -715,7 +720,7 @@ def test_constant_segment_tensor_alignment_non_power_of_2_fails(self) -> None:
715720
# Expect failure as tensor alignment 14 is not a power of 2.
716721
with self.assertRaises(ValueError):
717722
serialize_pte_binary(
718-
program,
723+
PTEFile(program=program),
719724
segment_alignment=SEGMENT_ALIGNMENT,
720725
constant_tensor_alignment=constant_tensor_alignment,
721726
)
@@ -750,11 +755,10 @@ def test_constant_delegate_and_named_data_segments(self) -> None:
750755
# Extract the blobs into segments during serialization.
751756
pte_data = bytes(
752757
serialize_pte_binary(
753-
program,
758+
PTEFile(program=program, named_data=named_data),
754759
extract_delegate_segments=True,
755760
segment_alignment=SEGMENT_ALIGNMENT,
756761
constant_tensor_alignment=CONSTANT_TENSOR_ALIGNMENT,
757-
named_data=named_data,
758762
)
759763
)
760764

@@ -961,11 +965,10 @@ def test_named_data_segments(self) -> None:
961965
# Serialize the program with named data segments.
962966
pte_data = bytes(
963967
serialize_pte_binary(
964-
program,
968+
PTEFile(program=program, named_data=named_data),
965969
extract_delegate_segments=True,
966970
segment_alignment=SEGMENT_ALIGNMENT,
967971
constant_tensor_alignment=CONSTANT_TENSOR_ALIGNMENT,
968-
named_data=named_data,
969972
)
970973
)
971974

@@ -1046,11 +1049,10 @@ def test_named_data_segments(self) -> None:
10461049

10471050
# Test re-serialize
10481051
pte_data2 = serialize_pte_binary(
1049-
deserialized.program,
1052+
PTEFile(program=deserialized.program, named_data=deserialized.named_data),
10501053
extract_delegate_segments=True,
10511054
segment_alignment=SEGMENT_ALIGNMENT,
10521055
constant_tensor_alignment=CONSTANT_TENSOR_ALIGNMENT,
1053-
named_data=deserialized.named_data,
10541056
)
10551057
# pte_data2 is not going to be the same as pte_data due to alignment;
10561058
# directly test the deserialized one.

‎exir/backend/test/test_compatibility.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
import torch
1010
from executorch.exir import to_edge
11-
from executorch.exir._serialize import _serialize_pte_binary
11+
from executorch.exir._serialize import _PTEFile, _serialize_pte_binary
1212
from executorch.exir.backend.backend_api import to_backend
1313
from executorch.exir.backend.canonical_partitioners.all_node_partitioner import (
1414
AllNodePartitioner,
@@ -58,7 +58,7 @@ def forward(self, x):
5858
# Generate the .pte file with the wrong version.
5959
buff = bytes(
6060
_serialize_pte_binary(
61-
program=prog,
61+
pte_file=_PTEFile(program=prog),
6262
)
6363
)
6464

@@ -105,7 +105,7 @@ def forward(self, x):
105105
# Generate the .pte file with the wrong version.
106106
buff = bytes(
107107
_serialize_pte_binary(
108-
program=prog,
108+
pte_file=_PTEFile(program=prog),
109109
)
110110
)
111111

‎exir/lowered_backend_module.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313

1414
import torch
1515
import torch.utils._pytree as pytree
16-
from executorch.exir._serialize import _serialize_pte_binary
16+
from executorch.exir._serialize import _PTEFile, _serialize_pte_binary
1717
from executorch.exir._serialize._named_data_store import NamedDataStoreOutput
1818
from executorch.exir.backend.compile_spec_schema import CompileSpec
1919
from executorch.exir.delegate import executorch_call_delegate, get_lowered_module_name
@@ -164,12 +164,14 @@ def buffer(
164164
# TODO(T181463742): avoid calling bytes(..) which incurs large copies.
165165
out = bytes(
166166
_serialize_pte_binary(
167-
program=self.program(memory_planning=memory_planning),
167+
pte_file=_PTEFile(
168+
program=self.program(memory_planning=memory_planning),
169+
named_data=self.named_data_store_output,
170+
),
168171
extract_delegate_segments=extract_delegate_segments,
169172
segment_alignment=segment_alignment,
170173
constant_tensor_alignment=constant_tensor_alignment,
171174
delegate_alignment=delegate_alignment,
172-
named_data=self.named_data_store_output,
173175
)
174176
)
175177
return out

0 commit comments

Comments
 (0)