Add segment alignment option to export.py, default constant_segment t…

…o True (#1965)

Summary:
Ios uses 0x4000 alignment

Pull Request resolved: #1965

Reviewed By: larryliu0820

Differential Revision: D53733852

Pulled By: lucylq

fbshipit-source-id: 2eade519c4c7c553d7db00385492179a8f6495ed

build/test_ios_ci.sh

@@ -41,7 +41,7 @@ say "Installing MPS Backend Requirements"
 
 say "Exporting Models"
 
-python3 -m examples.portable.scripts.export --model_name="$MODEL_NAME"
+python3 -m examples.portable.scripts.export --model_name="$MODEL_NAME" --segment_alignment=0x4000
 python3 -m examples.apple.coreml.scripts.export_and_delegate --model_name="$MODEL_NAME"
 python3 -m examples.apple.mps.scripts.mps_example --model_name="$MODEL_NAME"
 python3 -m examples.xnnpack.aot_compiler --model_name="$MODEL_NAME" --delegate

examples/portable/scripts/export.py

@@ -28,6 +28,13 @@ def main() -> None:
         required=True,
         help=f"provide a model name. Valid ones: {list(MODEL_NAME_TO_MODEL.keys())}",
     )
+
+    parser.add_argument(
+        "-a",
+        "--segment_alignment",
+        required=False,
+        help="specify segment alignment in hex. Default is 0x1000. Use 0x4000 for iOS",
+    )
     parser.add_argument("-o", "--output_dir", default=".", help="output directory")
 
     args = parser.parse_args()
@@ -42,6 +49,9 @@ def main() -> None:
         *MODEL_NAME_TO_MODEL[args.model_name]
     )
 
+    backend_config = ExecutorchBackendConfig(extract_constant_segment=True)
+    if args.segment_alignment is not None:
+        backend_config.segment_alignment = int(args.segment_alignment, 16)
     if (
         dynamic_shapes is not None
     ):  # capture_pre_autograd_graph does not work with dynamic shapes
@@ -53,11 +63,14 @@ def main() -> None:
                 _check_ir_validity=False,
             ),
         )
-        prog = edge_manager.to_executorch(
-            config=ExecutorchBackendConfig(extract_constant_segment=False)
-        )
+        prog = edge_manager.to_executorch(config=backend_config)
     else:
-        prog = export_to_exec_prog(model, example_inputs, dynamic_shapes=dynamic_shapes)
+        prog = export_to_exec_prog(
+            model,
+            example_inputs,
+            dynamic_shapes=dynamic_shapes,
+            backend_config=backend_config,
+        )
     save_pte_program(prog.buffer, args.model_name, args.output_dir)
 
 

exir/capture/_config.py

@@ -59,7 +59,7 @@ class ExecutorchBackendConfig:
     # 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 = False
+    extract_constant_segment: bool = True
 
     # When extracting segments, the starting offset of each segment will be
     # aligned to this value (in bytes). When using mmap() to load segments, this