Skip to content

Commit 5a491dc

Browse files
lucylqlucylq
committed
Refactor export_delegated_program
Currently, we generate every combination of inputs for each module with the export_delegate_program script: - extract_segments=True, False - delegate_alignment=None,1024 Planning to add another flag, 'external_constants', which will move constants into a separate file to test program-data separation for delegated programs. This test only requires pte, ptd, with default settings. So refactoring the export script to only generate based on the args, and update genrule to generate what the test requires. Differential Revision: [D73278562](https://our.internmc.facebook.com/intern/diff/D73278562/) ghstack-source-id: 279015190 Pull Request resolved: #10303
1 parent 605bfa6 commit 5a491dc

File tree

3 files changed

+46
-40
lines changed

3 files changed

+46
-40
lines changed

‎runtime/executor/test/targets.bzl‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -233,9 +233,9 @@ def define_common_targets(is_fbcode = False):
233233
# Uses an fbcode target path because the authoring/export tools
234234
# intentionally don't work in xplat (since they're host-only
235235
# tools).
236-
"ET_MODULE_ADD_MUL_NOSEGMENTS_DA1024_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul-nosegments-da1024.pte])",
237-
"ET_MODULE_ADD_MUL_NOSEGMENTS_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul-nosegments.pte])",
238-
"ET_MODULE_ADD_MUL_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul.pte])",
236+
"ET_MODULE_ADD_MUL_NOSEGMENTS_DA1024_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul-nosegments-da1024.pte])",
237+
"ET_MODULE_ADD_MUL_NOSEGMENTS_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul-nosegments.pte])",
238+
"ET_MODULE_ADD_MUL_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul.pte])",
239239
},
240240
)
241241

‎test/models/export_delegated_program.py‎

Lines changed: 28 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from executorch.exir.backend.test.backend_with_compiler_demo import (
2020
BackendWithCompilerDemo,
2121
)
22+
from executorch.exir.program import ExecutorchProgramManager
2223
from torch import nn
2324
from torch.export import export
2425

@@ -111,10 +112,10 @@ def export_module_to_program(
111112
*,
112113
backend_id: str,
113114
extract_delegate_segments: bool,
114-
constant_tensor_alignemnt: Optional[int] = None,
115+
constant_tensor_alignment: Optional[int] = None,
115116
delegate_alignment: Optional[int] = None,
116117
method: str = "forward",
117-
) -> bytes:
118+
) -> ExecutorchProgramManager:
118119
eager_module = module_class().eval()
119120
inputs = ()
120121
if hasattr(eager_module, "get_random_inputs"):
@@ -135,7 +136,7 @@ def forward(self, *args, **kwargs):
135136
edge_config = EdgeCompileConfig(_check_ir_validity=False)
136137
et_config = exir.ExecutorchBackendConfig(
137138
extract_delegate_segments=extract_delegate_segments,
138-
constant_tensor_alignment=constant_tensor_alignemnt,
139+
constant_tensor_alignment=constant_tensor_alignment,
139140
delegate_alignment=delegate_alignment,
140141
)
141142

@@ -170,7 +171,7 @@ def forward(self, *args, **kwargs):
170171
export(composite_module, args=inputs, strict=True)
171172
).to_executorch(config=et_config)
172173

173-
return executorch_program.buffer
174+
return executorch_program
174175

175176

176177
def main() -> None:
@@ -199,6 +200,14 @@ def main() -> None:
199200
help="ID of the backend to use for delegation; "
200201
+ f"one of {known_backend_ids}",
201202
)
203+
parser.add_argument(
204+
"--inline_delegate_segments",
205+
action="store_true",
206+
help="Store delegate data inside the flatbuffer.",
207+
)
208+
parser.add_argument(
209+
"--delegate_alignment", type=int, default=None, help="Delegate alignment."
210+
)
202211
parser.add_argument(
203212
"--outdir",
204213
type=str,
@@ -219,25 +228,22 @@ def main() -> None:
219228

220229
# Export and write to the output files.
221230
os.makedirs(args.outdir, exist_ok=True)
231+
suffix = ""
222232
for module_name, module_class in module_names_to_classes.items():
223-
for extract_delegate_segments in (True, False):
224-
suffix = "" if extract_delegate_segments else "-nosegments"
225-
# Create files with the default alignment, and a large alignment.
226-
# This alignment should be so large that it's extremely unlikely for
227-
# the data to accidentally be aligned to it in the default case.
228-
for delegate_alignment in (None, 1024):
229-
suffix += f"-da{delegate_alignment}" if delegate_alignment else ""
230-
outfile = os.path.join(args.outdir, f"{module_name}{suffix}.pte")
231-
with open(outfile, "wb") as fp:
232-
fp.write(
233-
export_module_to_program(
234-
module_class,
235-
backend_id=args.backend_id,
236-
extract_delegate_segments=extract_delegate_segments,
237-
delegate_alignment=delegate_alignment,
238-
)
239-
)
240-
print(f"Exported {module_name} and wrote program data to {outfile}")
233+
if args.inline_delegate_segments:
234+
suffix += f"-nosegments"
235+
if args.delegate_alignment is not None:
236+
suffix += f"-da{args.delegate_alignment}"
237+
outfile = os.path.join(args.outdir, f"{module_name}{suffix}.pte")
238+
executorch_program = export_module_to_program(
239+
module_class,
240+
backend_id=args.backend_id,
241+
extract_delegate_segments=not args.inline_delegate_segments,
242+
delegate_alignment=args.delegate_alignment,
243+
)
244+
with open(outfile, "wb") as fp:
245+
fp.write(executorch_program.buffer)
246+
print(f"Exported {module_name} and wrote program data to {outfile}")
241247

242248

243249
if __name__ == "__main__":

‎test/models/targets.bzl‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ def define_common_targets():
150150
visibility = [], # Private
151151
)
152152

153-
# Class names of nn.Modules for :exported_delegated_programs to export.
153+
# Class names of nn.Modules available in export_delegated_program.py.
154154
DELEGATED_MODULES_TO_EXPORT = [
155155
"ModuleAddMul",
156156
"ModuleAddLarge",
@@ -160,23 +160,23 @@ def define_common_targets():
160160
# Name of the backend to use when exporting delegated programs.
161161
BACKEND_ID = "StubBackend"
162162

163-
# Generates Executorch .pte program files for various modules at build time.
163+
# Generates Executorch .pte program files for the AddMul module at build time.
164164
# To use one, depend on a target like
165-
# ":exported_delegated_programs[ModuleAdd.pte]" or
166-
# ":exported_delegated_programs[ModuleAdd-nosegments.pte]" (which does not
165+
# ":exported_delegated_add_mul[ModuleAdd.pte]" or
166+
# ":exported_delegated_add_mul[ModuleAdd-nosegments.pte]" (which does not
167167
# extract the delegate data blobs into segments).
168168
runtime.genrule(
169-
name = "exported_delegated_programs",
170-
cmd = "$(exe :export_delegated_program)" +
171-
" --modules " + ",".join(DELEGATED_MODULES_TO_EXPORT) +
172-
" --backend_id " + BACKEND_ID +
173-
" --outdir $OUT",
169+
name = "exported_delegated_add_mul",
170+
cmd = "$(exe :export_delegated_program) --modules ModuleAddMul --backend_id " + BACKEND_ID + " --outdir $OUT" +
171+
" && $(exe :export_delegated_program) --modules ModuleAddMul --backend_id " + BACKEND_ID + " --inline_delegate_segments --outdir $OUT" +
172+
# Create files with a large alignment as well as the default.
173+
# This alignment should be so large that it's extremely unlikely for
174+
# the data to accidentally be aligned to it in the default case.
175+
" && $(exe :export_delegated_program) --modules ModuleAddMul --backend_id " + BACKEND_ID + " --inline_delegate_segments --delegate_alignment 1024 --outdir $OUT",
174176
outs = {
175-
fname + seg_suffix + da_suffix + ".pte": [fname + seg_suffix + da_suffix + ".pte"]
176-
for fname in DELEGATED_MODULES_TO_EXPORT
177-
for seg_suffix in ["", "-nosegments"]
178-
# "da" = delegate alignment
179-
for da_suffix in ["", "-da1024"]
177+
"ModuleAddMul.pte": ["ModuleAddMul.pte"],
178+
"ModuleAddMul-nosegments.pte": ["ModuleAddMul-nosegments.pte"],
179+
"ModuleAddMul-nosegments-da1024.pte": ["ModuleAddMul-nosegments-da1024.pte"],
180180
},
181181
default_outs = ["."],
182182
visibility = [
@@ -188,7 +188,7 @@ def define_common_targets():
188188
runtime.genrule(
189189
name = "exported_xnnp_delegated_programs",
190190
cmd = "$(exe :export_delegated_program)" +
191-
" --modules " + ",".join(DELEGATED_MODULES_TO_EXPORT) +
191+
" --modules ModuleAddLarge,ModuleSubLarge" +
192192
" --backend_id " + "XnnpackBackend" +
193193
" --outdir $OUT",
194194
outs = {

0 commit comments

Comments
 (0)