Skip to content

Commit 85a8e88

Browse files
pytorchbotpytorchbot
authored
Refactor export_delegated_program (#10334)
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/)
1 parent 369792e commit 85a8e88

File tree

3 files changed

+48
-40
lines changed

3 files changed

+48
-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: 30 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7+
# pyre-unsafe
8+
79
import argparse
810
import inspect
911
import os
@@ -19,6 +21,7 @@
1921
from executorch.exir.backend.test.backend_with_compiler_demo import (
2022
BackendWithCompilerDemo,
2123
)
24+
from executorch.exir.program import ExecutorchProgramManager
2225
from torch import nn
2326
from torch.export import export
2427

@@ -111,10 +114,10 @@ def export_module_to_program(
111114
*,
112115
backend_id: str,
113116
extract_delegate_segments: bool,
114-
constant_tensor_alignemnt: Optional[int] = None,
117+
constant_tensor_alignment: Optional[int] = None,
115118
delegate_alignment: Optional[int] = None,
116119
method: str = "forward",
117-
) -> bytes:
120+
) -> ExecutorchProgramManager:
118121
eager_module = module_class().eval()
119122
inputs = ()
120123
if hasattr(eager_module, "get_random_inputs"):
@@ -135,7 +138,7 @@ def forward(self, *args, **kwargs):
135138
edge_config = EdgeCompileConfig(_check_ir_validity=False)
136139
et_config = exir.ExecutorchBackendConfig(
137140
extract_delegate_segments=extract_delegate_segments,
138-
constant_tensor_alignment=constant_tensor_alignemnt,
141+
constant_tensor_alignment=constant_tensor_alignment,
139142
delegate_alignment=delegate_alignment,
140143
)
141144

@@ -170,7 +173,7 @@ def forward(self, *args, **kwargs):
170173
export(composite_module, args=inputs, strict=True)
171174
).to_executorch(config=et_config)
172175

173-
return executorch_program.buffer
176+
return executorch_program
174177

175178

176179
def main() -> None:
@@ -199,6 +202,14 @@ def main() -> None:
199202
help="ID of the backend to use for delegation; "
200203
+ f"one of {known_backend_ids}",
201204
)
205+
parser.add_argument(
206+
"--inline_delegate_segments",
207+
action="store_true",
208+
help="Store delegate data inside the flatbuffer.",
209+
)
210+
parser.add_argument(
211+
"--delegate_alignment", type=int, default=None, help="Delegate alignment."
212+
)
202213
parser.add_argument(
203214
"--outdir",
204215
type=str,
@@ -219,25 +230,22 @@ def main() -> None:
219230

220231
# Export and write to the output files.
221232
os.makedirs(args.outdir, exist_ok=True)
233+
suffix = ""
222234
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}")
235+
if args.inline_delegate_segments:
236+
suffix += "-nosegments"
237+
if args.delegate_alignment is not None:
238+
suffix += f"-da{args.delegate_alignment}"
239+
outfile = os.path.join(args.outdir, f"{module_name}{suffix}.pte")
240+
executorch_program = export_module_to_program(
241+
module_class,
242+
backend_id=args.backend_id,
243+
extract_delegate_segments=not args.inline_delegate_segments,
244+
delegate_alignment=args.delegate_alignment,
245+
)
246+
with open(outfile, "wb") as fp:
247+
fp.write(executorch_program.buffer)
248+
print(f"Exported {module_name} and wrote program data to {outfile}")
241249

242250

243251
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)