From 1711bbcadc1d90555a86c251ea82445689c1f7fd Mon Sep 17 00:00:00 2001
From: Lucy Qiu <lfq@meta.com>
Date: Tue, 13 Feb 2024 17:04:23 -0800
Subject: [PATCH] Add segment alignment option to export.py, default
 constant_segment to True (#1965)

Summary:
Ios uses 0x4000 alignment

Pull Request resolved: https://github.com/pytorch/executorch/pull/1965

Reviewed By: larryliu0820

Differential Revision: D53733852

Pulled By: lucylq

fbshipit-source-id: 2eade519c4c7c553d7db00385492179a8f6495ed
---
 build/test_ios_ci.sh                |  2 +-
 examples/portable/scripts/export.py | 21 +++++++++++++++++----
 exir/capture/_config.py             |  2 +-
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/build/test_ios_ci.sh b/build/test_ios_ci.sh
index ad81767d6e..fc893d4828 100755
--- a/build/test_ios_ci.sh
+++ b/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
diff --git a/examples/portable/scripts/export.py b/examples/portable/scripts/export.py
index 2f67846599..728f51fc20 100644
--- a/examples/portable/scripts/export.py
+++ b/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)
 
 
diff --git a/exir/capture/_config.py b/exir/capture/_config.py
index c0a33c2424..1a0ffd3b5f 100644
--- a/exir/capture/_config.py
+++ b/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