[pytorch] register c10 ops for static dispatch to unblock c10 boxing

Summary:
PR #32521 broke static dispatch because some ops are no longer
registered in register_aten_ops_*.cpp - it expects the c10 registers in
TypeDefault.cpp / CPUType.cpp / etc to register these ops. However, all
c10 registers are inside `#ifndef USE_STATIC_DISPATCH` section.

To measure the OSS mobile build size impact of this PR:
```
 # default build: SELECTED_OP_LIST=MobileNetV2.yaml scripts/build_pytorch_android.sh armeabi-v7a
 # mobilenetv2 custom build: scripts/build_pytorch_android.sh armeabi-v7a
```

- Before this PR, Android AAR size for arm-v7:
* default build: 5.5M;
* mobilenetv2 custom build: 3.2M;

- After this PR:
* default build: 6.4M;
* mobilenetv2 custom build: 3.3M;

It regressed default build size by ~1M because more root ops are
registered by c10 registers, e.g. backward ops which are filtered out by
gen_jit_dispatch.py for inference-only mobile build.

mobilenetv2 custom build size regressed by ~100k presumably because
the op whitelist is not yet applied to things like BackendSelectRegister.

Differential Revision: D20266240

Test Plan: Imported from OSS

Pulled By: ljk53

fbshipit-source-id: 97a9a06779f8c62fe3ff5cce089aa7fa9dee3c4a

aten/src/ATen/templates/BackendSelectRegister.cpp

@@ -9,8 +9,6 @@
 #include <ATen/core/op_registration/op_registration.h>
 #include <c10/core/TensorOptions.h>
 
-
-#ifndef USE_STATIC_DISPATCH
 namespace at {
 
 namespace {
@@ -23,4 +21,3 @@ static auto registry = torch::RegisterOperators()
 
 } // namespace
 } // at
-#endif

aten/src/ATen/templates/PerOpRegistration.cpp

@@ -7,11 +7,9 @@ $extra_headers
 
 namespace at {
 
-#ifndef USE_STATIC_DISPATCH
 namespace {
 auto registerer = torch::import()
   ${function_registrations};
 }
-#endif
 
 }  // namespace at

aten/src/ATen/templates/SparseTypeDerived.cpp

@@ -42,11 +42,9 @@ ${type_derived_method_definitions}
 
 }  // namespace ${Type}
 
-#ifndef USE_STATIC_DISPATCH
 namespace {
 static auto registerer = torch::import()
   ${function_registrations};
 }
-#endif
 
 }

aten/src/ATen/templates/TypeDefault.cpp

@@ -29,11 +29,9 @@ ${type_method_definitions}
 
 }  // namespace TypeDefault
 
-#ifndef USE_STATIC_DISPATCH
 namespace {
 auto registerer = torch::import()
   ${function_registrations};
 }
-#endif
 
 }  // namespace at

aten/src/ATen/templates/TypeDerived.cpp

@@ -53,11 +53,9 @@ ${type_derived_method_definitions}
 
 }  // namespace ${Type}
 
-#ifndef USE_STATIC_DISPATCH
 namespace {
 auto registerer = torch::import()
   ${function_registrations};
 }
-#endif
 
 }

cmake/Codegen.cmake

@@ -152,21 +152,21 @@ if (INTERN_BUILD_ATEN_OPS)
   endif()
 
   set(CUSTOM_BUILD_FLAG)
-  if (SELECTED_OP_LIST AND NOT USE_STATIC_DISPATCH)
-    if (NOT OP_DEPENDENCY)
+  if (SELECTED_OP_LIST)
+    if (NOT USE_STATIC_DISPATCH AND NOT OP_DEPENDENCY)
       message(FATAL_ERROR "Must provide op dependency graph .yaml file for custom build with dynamic dispatch!")
     endif()
     EXECUTE_PROCESS(
       COMMAND
-      "${PYTHON_EXECUTABLE}" ${CMAKE_CURRENT_LIST_DIR}/../tools/code_analyzer/gen_transitive_deps.py
+      "${PYTHON_EXECUTABLE}" ${CMAKE_CURRENT_LIST_DIR}/../tools/code_analyzer/gen_op_registration_whitelist.py
       --op-dependency "${OP_DEPENDENCY}"
       --root-ops "${SELECTED_OP_LIST}"
-      OUTPUT_VARIABLE SELECTED_OP_CLOSURE
+      OUTPUT_VARIABLE OP_REGISTRATION_WHITELIST
     )
-    separate_arguments(SELECTED_OP_CLOSURE)
-    message(STATUS "Custom build with selected op closure: ${SELECTED_OP_CLOSURE}")
+    separate_arguments(OP_REGISTRATION_WHITELIST)
+    message(STATUS "Custom build with op registration whitelist: ${OP_REGISTRATION_WHITELIST}")
     set(CUSTOM_BUILD_FLAG
-      --op_registration_whitelist ${SELECTED_OP_CLOSURE})
+      --op_registration_whitelist ${OP_REGISTRATION_WHITELIST})
   endif()
 
   SET(GEN_COMMAND

tools/code_analyzer/gen_op_registration_whitelist.py

@@ -1,7 +1,11 @@
 """
-This util takes the op dependency graph of ATen and the list of root ops, and
-outputs all transitive dependencies of the root ops. It is invoked from cmake
-for custom mobile build.
+This util is invoked from cmake to produce the op registration whitelist param
+for `ATen/gen.py` for custom mobile build.
+For custom build with dynamic dispatch, it takes the op dependency graph of ATen
+and the list of root ops, and outputs all transitive dependencies of the root
+ops as the whitelist.
+For custom build with static dispatch, the op dependency graph will be omitted,
+and it will directly output root ops as the whitelist.
 """
 
 import argparse
@@ -53,12 +57,14 @@ if __name__ == "__main__":
         description='Util to produce transitive dependencies for custom build')
     parser.add_argument(
         '--op-dependency',
-        help='input yaml file of op dependency graph')
+        help='input yaml file of op dependency graph '
+             '- can be omitted for custom build with static dispatch')
     parser.add_argument(
         '--root-ops',
+        required=True,
         help='input yaml file of root (directly used) operators')
     args = parser.parse_args()
 
-    deps = load_op_dep_graph(args.op_dependency)
+    deps = load_op_dep_graph(args.op_dependency) if args.op_dependency else {}
     root_ops = load_root_ops(args.root_ops)
     print(gen_transitive_closure(deps, root_ops))