Skip to content

Commit 269aaf2

Browse files
lucylqlucylq
committed
Refactor export_delegated_program
Pull Request resolved: #10303 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. ghstack-source-id: 279308036 Differential Revision: [D73278562](https://our.internmc.facebook.com/intern/diff/D73278562/)
1 parent 06f912d commit 269aaf2

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
@@ -234,9 +234,9 @@ def define_common_targets(is_fbcode = False):
234234
# Uses an fbcode target path because the authoring/export tools
235235
# intentionally don't work in xplat (since they're host-only
236236
# tools).
237-
"ET_MODULE_ADD_MUL_NOSEGMENTS_DA1024_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul-nosegments-da1024.pte])",
238-
"ET_MODULE_ADD_MUL_NOSEGMENTS_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul-nosegments.pte])",
239-
"ET_MODULE_ADD_MUL_PATH": "$(location fbcode//executorch/test/models:exported_delegated_programs[ModuleAddMul.pte])",
237+
"ET_MODULE_ADD_MUL_NOSEGMENTS_DA1024_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul-nosegments-da1024.pte])",
238+
"ET_MODULE_ADD_MUL_NOSEGMENTS_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul-nosegments.pte])",
239+
"ET_MODULE_ADD_MUL_PATH": "$(location fbcode//executorch/test/models:exported_delegated_add_mul[ModuleAddMul.pte])",
240240
},
241241
)
242242

‎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 += "-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
@@ -151,7 +151,7 @@ def define_common_targets():
151151
visibility = [], # Private
152152
)
153153

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

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

0 commit comments

Comments
 (0)