Remove extract_constant_segment from config (#5096)

Summary:
Pull Request resolved: #5096

All usages of extract_constant_segment should be removed.

Update program and method tests, and update test artifacts to only generate .pte with constant segment.

Differential Revision: D61996249

Author: lucylq

exir/_serialize/_program.py

@@ -347,7 +347,6 @@ def serialize_pte_binary(
     *,
     mutable_data: Optional[List[Buffer]] = None,
     extract_delegate_segments: bool = False,
-    extract_constant_segment: bool = True,
     segment_alignment: int = 128,
     constant_tensor_alignment: Optional[int] = None,
     delegate_alignment: Optional[int] = None,
@@ -363,8 +362,6 @@ def serialize_pte_binary(
               and the starting segment offset.
             - Update the Program.segments field with the offsets and lengths
               of each segment.
-        extract_constant_segment: Whether to move the constant data from the Program
-            into a separate segment.
         segment_alignment: Alignment in bytes. The starting offset of each
             segment will be aligned to this value in the output data.
         constant_tensor_alignment: The minimum alignment of tensor
@@ -387,19 +384,18 @@ def serialize_pte_binary(
     # Store extracted segment data; this may be constant data or delegate data.
     segments: List[Cord] = []
 
-    if extract_constant_segment:
-        constant_segment_data, constant_segment_offsets = _extract_constant_segment(
-            program.constant_buffer, tensor_alignment=constant_tensor_alignment
+    constant_segment_data, constant_segment_offsets = _extract_constant_segment(
+        program.constant_buffer, tensor_alignment=constant_tensor_alignment
+    )
+    if len(constant_segment_data) > 0:
+        # Update program.constant_segment with constant subsegment offset information.
+        program.constant_segment = SubsegmentOffsets(
+            segment_index=len(segments), offsets=constant_segment_offsets
         )
-        if len(constant_segment_data) > 0:
-            # Update program.constant_segment with constant subsegment offset information.
-            program.constant_segment = SubsegmentOffsets(
-                segment_index=len(segments), offsets=constant_segment_offsets
-            )
-            # Clear the constant buffer, as constant data will be stored in segments.
-            program.constant_buffer = []
-            # Add to the aggregate segments cord.
-            segments.append(constant_segment_data)
+        # Clear the constant buffer, as constant data will be stored in segments.
+        program.constant_buffer = []
+        # Add to the aggregate segments cord.
+        segments.append(constant_segment_data)
 
     if mutable_data is not None:
         mutable_segment_data, mutable_segment_offsets = _extract_constant_segment(

exir/capture/_config.py

@@ -65,12 +65,6 @@ class ExecutorchBackendConfig:
     # This makes it possible to free those blobs at runtime.
     extract_delegate_segments: bool = True
 
-    # Whether to extract constants from the Program into separate segments,
-    # rather than encoding those constants in the flatbuffer data.
-    # This reduces the memory overhead of creating the .pte file for models with
-    # large constant data.
-    extract_constant_segment: bool = True
-
     # When extracting segments, the starting offset of each segment will be
     # aligned to this value (in bytes). Must be a power of two.
     segment_alignment: int = 128

exir/program/_program.py

@@ -439,7 +439,6 @@ def to_executorch(
             new_prog,
             emit_stacktrace=config.emit_stacktrace,
             extract_delegate_segments=config.extract_delegate_segments,
-            extract_constant_segment=config.extract_constant_segment,
             segment_alignment=config.segment_alignment,
             constant_tensor_alignment=config.constant_tensor_alignment,
             delegate_alignment=config.delegate_alignment,
@@ -468,7 +467,6 @@ def __init__(
         exir_exported_program: ExirExportedProgram,
         emit_stacktrace: bool,
         extract_delegate_segments: bool,
-        extract_constant_segment: bool,
         segment_alignment: int,
         constant_tensor_alignment: Optional[int] = None,
         delegate_alignment: Optional[int] = None,
@@ -483,7 +481,6 @@ def __init__(
         self._emitter_output: Optional[EmitterOutput] = None
         self._emit_stacktrace: bool = emit_stacktrace
         self._extract_delegate_segments: bool = extract_delegate_segments
-        self._extract_constant_segment: bool = extract_constant_segment
         self._segment_alignment: int = segment_alignment
         self._constant_tensor_alignment: Optional[int] = constant_tensor_alignment
         self._delegate_alignment: Optional[int] = delegate_alignment
@@ -493,7 +490,6 @@ def _get_pte_data(self) -> Cord:
             self._pte_data = _serialize_pte_binary(
                 program=self.program,
                 extract_delegate_segments=self._extract_delegate_segments,
-                extract_constant_segment=self._extract_constant_segment,
                 segment_alignment=self._segment_alignment,
                 constant_tensor_alignment=self._constant_tensor_alignment,
                 delegate_alignment=self._delegate_alignment,
@@ -1351,7 +1347,6 @@ def __init__(
             program=self._emitter_output.program,
             mutable_data=self._emitter_output.mutable_data,
             extract_delegate_segments=backend_config.extract_delegate_segments,
-            extract_constant_segment=backend_config.extract_constant_segment,
             segment_alignment=backend_config.segment_alignment,
             constant_tensor_alignment=backend_config.constant_tensor_alignment,
             delegate_alignment=backend_config.delegate_alignment,

runtime/executor/test/method_test.cpp

@@ -59,11 +59,9 @@ class MethodTest : public ::testing::Test {
     load_program(std::getenv("ET_MODULE_INDEX_PATH"), "index");
     load_program(
         std::getenv("ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH"), "cat");
+    load_program(std::getenv("ET_MODULE_LINEAR_PATH"), "linear");
     load_program(
-        std::getenv("ET_MODULE_LINEAR_CONSTANT_SEGMENT_PATH"),
-        "linear_constant_segment");
-    load_program(
-        std::getenv("ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH"),
+        std::getenv("DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH"),
         "linear_constant_buffer");
   }
 
@@ -274,7 +272,7 @@ TEST_F(MethodTest, ConstantSegmentTest) {
   // Execute model with constants stored in segment.
   ManagedMemoryManager mmm(kDefaultNonConstMemBytes, kDefaultRuntimeMemBytes);
   Result<Method> method =
-      programs_["linear_constant_segment"]->load_method("forward", &mmm.get());
+      programs_["linear"]->load_method("forward", &mmm.get());
   ASSERT_EQ(method.error(), Error::Ok);
 
   // Can execute the method.

runtime/executor/test/program_test.cpp

@@ -382,8 +382,7 @@ TEST_F(ProgramTest, DEPRECATEDLoad) {
 TEST_F(ProgramTest, LoadConstantSegment) {
   // Load the serialized ModuleLinear data, with constants in the segment and no
   // constants in the flatbuffer.
-  const char* linear_path =
-      std::getenv("ET_MODULE_LINEAR_CONSTANT_SEGMENT_PATH");
+  const char* linear_path = std::getenv("ET_MODULE_LINEAR_PATH");
   Result<FileDataLoader> linear_loader = FileDataLoader::from(linear_path);
   ASSERT_EQ(linear_loader.error(), Error::Ok);
 
@@ -424,11 +423,11 @@ TEST_F(ProgramTest, LoadConstantSegment) {
   EXPECT_GE(flatbuffer_program->constant_segment()->offsets()->size(), 1);
 }
 
-TEST_F(ProgramTest, LoadConstantSegmentWithNoConstantSegment) {
+TEST_F(ProgramTest, LoadConstantSegmentWhenConstantBufferExists) {
   // Load the serialized ModuleLinear data, with constants in the flatbuffer and
   // no constants in the segment.
   const char* linear_path =
-      std::getenv("ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH");
+      std::getenv("DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH");
   Result<FileDataLoader> linear_loader = FileDataLoader::from(linear_path);
   ASSERT_EQ(linear_loader.error(), Error::Ok);
 

runtime/executor/test/targets.bzl

@@ -97,15 +97,16 @@ def define_common_targets(is_fbcode = False):
     # file in fbcode. See https://fburl.com/9esapdmd
     if not runtime.is_oss and is_fbcode:
         modules_env = {
+            # Deprecated model that still works with ExecuTorch runtime.
+            "DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH": "$(location fbcode//executorch/test/models/deprecated:ModuleLinear-no-constant-segment.pte)",
             # The tests use this var to find the program file to load. This uses
             # an fbcode target path because the authoring/export tools
             # intentionally don't work in xplat (since they're host-only tools).
             "ET_MODULE_ADD_HALF_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleAddHalf.pte])",
             "ET_MODULE_ADD_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleAdd.pte])",
             "ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleDynamicCatUnallocatedIO.pte])",
             "ET_MODULE_INDEX_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleIndex.pte])",
-            "ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleLinear-no-constant-segment.pte])",
-            "ET_MODULE_LINEAR_CONSTANT_SEGMENT_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleLinear.pte])",
+            "ET_MODULE_LINEAR_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleLinear.pte])",
             "ET_MODULE_MULTI_ENTRY_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleMultipleEntry.pte])",
             "ET_MODULE_SIMPLE_TRAIN_PATH": "$(location fbcode//executorch/test/models:exported_programs[ModuleSimpleTrain.pte])",
         }

schema/program.fbs

@@ -429,6 +429,7 @@ table Program {
   // Each constant is assigned an index into the table which are each individually aligned.
   // 0 index is reserved to be pointed to by non-constant Tensors.
   // If this field is non-empty, constant_segment.offsets must be empty.
+  // DEPRECATED: After D61996249 on 2024-09-05, no new PTE files will use this field.
   constant_buffer:[Buffer];
 
   // List of delegate data. Pointed to by BackendDelegateDataReference.

test/end2end/exported_module.py

@@ -67,7 +67,6 @@ def export(
         ignore_to_out_var_failure: bool = False,
         dynamic_memory_planning_mode: DynamicMemoryPlanningMode = DynamicMemoryPlanningMode.UPPER_BOUND,
         capture_config=None,
-        extract_constant_segment: bool = True,
         skip_type_promotion: bool = False,
         export_joint_graph: bool = False,
     ) -> "ExportedModule":
@@ -206,7 +205,6 @@ def __init__(self, method):
                 dynamic_memory_planning_mode=dynamic_memory_planning_mode,
                 memory_planning_pass=memory_planning_pass,
                 to_out_var_pass=ToOutVarPass(ignore_to_out_var_failure),
-                extract_constant_segment=extract_constant_segment,
             )
         )
 

test/models/deprecated/ModuleLinear-no-constant-segment.pte

@@ -0,0 +1,11 @@
+## Deprecated Models
+
+This readme documents deprecated models that remain compatible with versions of the ExecuTorch runtime.
+
+ModuleLinear-no-constant-segment.pte
+- This file contains constants stored in the constant_buffer, which was deprecated in D61996249 on 09-04-2024. Now, constants are stored in a separate segment.
+- Works with ExecuTorch runtime ET12.
+- Generated using commit hash e7c60c17d2d3, with command:
+    ```
+    buck2 build fbcode//executorch/test/models:exported_programs[ModuleLinear-no-constant-segment.pte] --show-output
+    ```

test/models/deprecated/README.md

@@ -0,0 +1,12 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+runtime.export_file(
+    name = "ModuleLinear-no-constant-segment.pte",
+    src = "ModuleLinear-no-constant-segment.pte",
+    visibility = [
+        "//executorch/runtime/executor/test/...",
+        "//executorch/test/...",
+    ],
+)

test/models/deprecated/TARGETS

@@ -190,7 +190,6 @@ def export_joint():
 
 def export_module_to_program(
     module_class: Type[nn.Module],
-    extract_constant_segment: bool,
     skip_type_promotion: bool,
 ):
     """Exports the module and returns the serialized program data."""
@@ -211,7 +210,6 @@ def export_module_to_program(
     module = ExportedModule.export(
         module_class,
         methods,
-        extract_constant_segment=extract_constant_segment,
         skip_type_promotion=skip_type_promotion,
         export_joint_graph=export_joint,
         **export_kwargs,
@@ -259,18 +257,15 @@ def main() -> None:
             # Skip type promotion to keep the model in fp16.
             # Type promotion will convert to fp32.
             skip_type_promotion = True
-        for extract_constant_segment in (True, False):
-            suffix = "" if extract_constant_segment else "-no-constant-segment"
-            outfile = os.path.join(args.outdir, f"{module_name}{suffix}.pte")
-            with open(outfile, "wb") as fp:
-                fp.write(
-                    export_module_to_program(
-                        module_class,
-                        extract_constant_segment=extract_constant_segment,
-                        skip_type_promotion=skip_type_promotion,
-                    )
+        outfile = os.path.join(args.outdir, f"{module_name}.pte")
+        with open(outfile, "wb") as fp:
+            fp.write(
+                export_module_to_program(
+                    module_class,
+                    skip_type_promotion=skip_type_promotion,
                 )
-            print(f"Exported {module_name} and wrote program data to {outfile}")
+            )
+        print(f"Exported {module_name} and wrote program data to {outfile}")
 
 
 if __name__ == "__main__":

test/models/export_program.py

@@ -56,23 +56,23 @@ export_test_model() {
   python3 -m test.models.export_program --modules "ModuleAdd,ModuleAddHalf,ModuleDynamicCatUnallocatedIO,ModuleIndex,ModuleLinear,ModuleMultipleEntry,ModuleSimpleTrain" --outdir "cmake-out" 2> /dev/null
   python3 -m test.models.export_delegated_program --modules "ModuleAddMul" --backend_id "StubBackend" --outdir "cmake-out" || true
 
+  DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH="./tests/models/deprecated/ModuleLinear-no-constant-segment.pte"
   ET_MODULE_ADD_HALF_PATH="$(realpath cmake-out/ModuleAddHalf.pte)"
   ET_MODULE_ADD_PATH="$(realpath cmake-out/ModuleAdd.pte)"
   ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH="$(realpath cmake-out/ModuleDynamicCatUnallocatedIO.pte)"
   ET_MODULE_INDEX_PATH="$(realpath cmake-out/ModuleIndex.pte)"
-  ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH="$(realpath cmake-out/ModuleLinear-no-constant-segment.pte)"
-  ET_MODULE_LINEAR_CONSTANT_SEGMENT_PATH="$(realpath cmake-out/ModuleLinear.pte)"
+  ET_MODULE_LINEAR_PATH="$(realpath cmake-out/ModuleLinear.pte)"
   ET_MODULE_MULTI_ENTRY_PATH="$(realpath cmake-out/ModuleMultipleEntry.pte)"
   ET_MODULE_ADD_MUL_NOSEGMENTS_DA1024_PATH="$(realpath cmake-out/ModuleAddMul-nosegments-da1024.pte)"
   ET_MODULE_ADD_MUL_NOSEGMENTS_PATH="$(realpath cmake-out/ModuleAddMul-nosegments.pte)"
   ET_MODULE_ADD_MUL_PATH="$(realpath cmake-out/ModuleAddMul.pte)"
   ET_MODULE_SIMPLE_TRAIN_PATH="$(realpath cmake-out/ModuleSimpleTrain.pte)"
+  export DEPRECATED_ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH
   export ET_MODULE_ADD_HALF_PATH
   export ET_MODULE_ADD_PATH
   export ET_MODULE_DYNAMIC_CAT_UNALLOCATED_IO_PATH
   export ET_MODULE_INDEX_PATH
-  export ET_MODULE_LINEAR_CONSTANT_BUFFER_PATH
-  export ET_MODULE_LINEAR_CONSTANT_SEGMENT_PATH
+  export ET_MODULE_LINEAR_PATH
   export ET_MODULE_MULTI_ENTRY_PATH
   export ET_MODULE_ADD_MUL_NOSEGMENTS_DA1024_PATH
   export ET_MODULE_ADD_MUL_NOSEGMENTS_PATH