[Pytorch Edge] Make RegisterBackendSelect Selective (#59096)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/59096
RegisterBackendSelect is bringing in ~100 extra ops to the runtime. This messes with the compatibility api, and also adds a nontrivial amount of size.
Test Plan: Model Unittests/CI
Reviewed By: iseeyuan
Differential Revision: D28588100
fbshipit-source-id: ffd0b5b9cbe20f27dbf3be418a6c1f80c7396fdb
diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py
index cb99a57..24f6f95 100644
--- a/tools/codegen/gen.py
+++ b/tools/codegen/gen.py
@@ -187,7 +187,9 @@
def __call__(self, f: NativeFunction) -> Optional[str]:
if not self.selector.is_native_function_selected(f):
return None
- return f'm.def({cpp_string(str(f.func))});\n'
+ schema_str = cpp_string(str(f.func))
+ schema_str = '"' + "aten::" + schema_str[1:]
+ return f'm.def({schema_str});\n'
def _num_leading_spaces(line: str) -> int:
@@ -454,6 +456,10 @@
Literal[Target.REGISTRATION]
]
+ # Selector object to determine which operators to generate
+ # registration code for.
+ selector: SelectiveBuilder
+
@method_with_native_function
def __call__(self, f: NativeFunction) -> Optional[str]:
if str(f.func.name.name).endswith('_like') or str(f.func.name.name).startswith('new_'):
@@ -465,6 +471,9 @@
if not any(isinstance(a.argument, TensorOptionsArguments) for a in native_sig.arguments()):
return None
+ if not self.selector.is_native_function_selected(f):
+ return None
+
native_tensor_args = [
a for a in native_sig.arguments()
if isinstance(a.argument, Argument) and a.argument.type.is_tensor_like()
@@ -1057,9 +1066,9 @@
# BackendSelect is generated specially
cpu_fm.write('RegisterBackendSelect.cpp', lambda: {
'backend_select_method_definitions':
- list(mapMaybe(ComputeBackendSelect(Target.DEFINITION), native_functions)),
+ list(mapMaybe(ComputeBackendSelect(Target.DEFINITION, selector), native_functions)),
'backend_select_function_registrations':
- list(mapMaybe(ComputeBackendSelect(Target.REGISTRATION), native_functions)),
+ list(mapMaybe(ComputeBackendSelect(Target.REGISTRATION, selector), native_functions)),
})
cpu_fm.write('MetaFunctions.h', lambda: {
