Add operator selection ability to gen_unboxing (#74271)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74271

Add operator selection ability to gen_unboxing:
- Adds ability to pass in selected_operators.yaml to gen_aten_unboxing_files target
- Creates a selective builder using the selected_operators.yaml in gen_unboxing.py and uses the selective builder to select native functions for codegen

Test Plan: Used selected_operators.yaml used in oacr_service target and ensured the codegen only generated code for included ops

Reviewed By: larryliu0820

Differential Revision: D34809772

fbshipit-source-id: c8d58e9693730b3d5b55f77ea057ee14a5ee1337
(cherry picked from commit 96f8494c21239ba4da7e61df38ddcf47c0f4db9f)

tools/jit/gen_unboxing.py

@@ -10,6 +10,7 @@ from tools.codegen.api.unboxing import convert_arguments
 from tools.codegen.context import method_with_native_function
 from tools.codegen.gen import parse_native_yaml, cpp_string
 from tools.codegen.model import NativeFunction, NativeFunctionsGroup, Variant
+from tools.codegen.selective_build.selector import SelectiveBuilder
 from tools.codegen.utils import Target, FileManager, mapMaybe, make_file_manager
 from typing import Union, Sequence
 from typing_extensions import Literal
@@ -19,9 +20,12 @@ from typing_extensions import Literal
 @dataclass(frozen=True)
 class ComputeUnboxingFunctions:
     target: Union[Literal[Target.DECLARATION], Literal[Target.DEFINITION]]
+    selector: SelectiveBuilder
 
     @method_with_native_function
     def __call__(self, f: NativeFunction) -> str:
+        if not self.selector.is_native_function_selected(f):
+            return ""
 
         if self.target is Target.DECLARATION:
             # Note [The ATen Codegen Unboxing API]
@@ -78,8 +82,12 @@ TORCH_API void {f.func.name.unambiguous_name()}(Stack & stack) {{
 # Generates RegisterCodegenUnboxedKernels.cpp.
 @dataclass(frozen=True)
 class ComputeCodegenUnboxedKernels:
+    selector: SelectiveBuilder
+
     @method_with_native_function
     def __call__(self, f: NativeFunction) -> str:
+        if not self.selector.is_native_function_selected(f):
+            return ""
         # We unconditionally generate function wrappers,
         sig_group = CppSignatureGroup.from_native_function(
             f, method=False
@@ -131,6 +139,7 @@ def gen_unboxing(
         *,
         native_functions: Sequence[NativeFunction],
         cpu_fm: FileManager,
+        selector: SelectiveBuilder,
 ) -> None:
     def key_func(fn: Union[NativeFunction, NativeFunctionsGroup]) -> str:
         return fn.root_name
@@ -140,7 +149,7 @@ def gen_unboxing(
         native_functions,
         key_fn=key_func,
         env_callable=lambda fn: {
-            "definitions": [ComputeUnboxingFunctions(Target.DEFINITION)(fn)]
+            "definitions": [ComputeUnboxingFunctions(Target.DEFINITION, selector)(fn)]
         },
         num_shards=5,
         sharded_keys={"definitions"},
@@ -149,7 +158,7 @@ def gen_unboxing(
         "UnboxingFunctions.h",
         lambda: {
             "declarations": list(
-                mapMaybe(ComputeUnboxingFunctions(Target.DECLARATION), native_functions)
+                mapMaybe(ComputeUnboxingFunctions(Target.DECLARATION, selector), native_functions)
             ),
         },
     )
@@ -157,7 +166,7 @@ def gen_unboxing(
         "RegisterCodegenUnboxedKernels.cpp",
         native_functions,
         key_fn=key_func,
-        env_callable=lambda fn: {"unboxed_ops": [ComputeCodegenUnboxedKernels()(fn)]},
+        env_callable=lambda fn: {"unboxed_ops": [ComputeCodegenUnboxedKernels(selector)(fn)]},
         num_shards=10,
         sharded_keys={"unboxed_ops"},
     )
@@ -181,9 +190,21 @@ def main() -> None:
     parser.add_argument(
         '--dry-run', action='store_true',
         help='run without writing any files (still updates outputs)')
+    parser.add_argument(
+        '--op_selection_yaml_path',
+        help='Provide a path to the operator selection (for custom build) YAML '
+             'that contains the information about the set of selected operators '
+             'and their categories (training, ...). Each operator is either a '
+             'full operator name with overload or just a bare operator name. '
+             'The operator names also contain the namespace prefix (e.g. aten::)')
 
     options = parser.parse_args()
 
+    if options.op_selection_yaml_path is not None:
+        selector = SelectiveBuilder.from_yaml_path(options.op_selection_yaml_path)
+    else:
+        selector = SelectiveBuilder.get_nop_selector()
+
     native_yaml_path = os.path.join(options.source_path, "native/native_functions.yaml")
     parsed_yaml = parse_native_yaml(native_yaml_path)
     native_functions, backend_indices = (
@@ -192,7 +213,7 @@ def main() -> None:
     )
 
     cpu_fm = make_file_manager(options=options)
-    gen_unboxing(native_functions=native_functions, cpu_fm=cpu_fm)
+    gen_unboxing(native_functions=native_functions, cpu_fm=cpu_fm, selector=selector)
 
     if options.output_dependencies:
         depfile_path = pathlib.Path(options.output_dependencies).resolve()