Shard Operators.cpp (pytorch#62185)

Summary:
Pull Request resolved: pytorch#62185

This file can take 5 minutes on its own to compile, and is the single limiting
factor for compile time of `libtorch_cpu` on a 32-core threadripper. Instead,
sharding into 5 files that take around 1 minute each cuts a full minute off the
overall build time.

This also factors out the `.findSchemaOrThrow(...).typed` step so the code can
be shared between `call` and `redispatch`.

Test Plan: Imported from OSS

Reviewed By: bdhirsh

Differential Revision: D29962049

Pulled By: albanD

fbshipit-source-id: be5df05fbea09ada0d825855f1618c25a11abbd8

BUILD.bazel

@@ -147,7 +147,11 @@ genrule(
         "aten/src/ATen/Functions.cpp",
         "aten/src/ATen/RedispatchFunctions.h",
         "aten/src/ATen/Operators.h",
-        "aten/src/ATen/Operators.cpp",
+        "aten/src/ATen/Operators_0.cpp",
+        "aten/src/ATen/Operators_1.cpp",
+        "aten/src/ATen/Operators_2.cpp",
+        "aten/src/ATen/Operators_3.cpp",
+        "aten/src/ATen/Operators_4.cpp",
         "aten/src/ATen/NativeFunctions.h",
         "aten/src/ATen/MetaFunctions.h",
         "aten/src/ATen/MetaFunctions_inl.h",

aten/src/ATen/templates/Operators.cpp

@@ -2,6 +2,8 @@
 #include <ATen/Tensor.h>
 #include <ATen/core/dispatch/Dispatcher.h>
 
+// NOTE See [Sharded File] comment in VariableType
+
 namespace at { namespace _ops {
 
 ${definitions}

tools/code_analyzer/op_deps_pass.cpp

@@ -770,6 +770,12 @@ class OpDependency : public ModulePass {
 
   static void extractStringValue(
       Value* V, const std::function<void(const std::string&)>& CB) {
+    if (isa<UndefValue>(V)) {
+      // UndefValue inherits from ConstantValue, but don't contain any data
+      // See: https://llvm.org/docs/LangRef.html#undefined-values
+      return;
+    }
+
     if (auto array = dyn_cast<ConstantDataArray>(V)) {
       // Normal case for c-style string literal and "std::basic_string".
       if (array->isCString()) {

tools/codegen/gen.py

@@ -215,7 +215,7 @@ class ComputeOperators:
     ]
 
     @method_with_native_function
-    def __call__(self, f: NativeFunction) -> Optional[str]:
+    def __call__(self, f: NativeFunction) -> str:
         sig = DispatcherSignature.from_schema(f.func)
         name = f.func.name.unambiguous_name()
         call_method_name = 'call'
@@ -260,7 +260,15 @@ def __call__(self, f: NativeFunction) -> Optional[str]:
             defns = f"""
 STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, name, "aten::{str(f.func.name)}")
 STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, overload_name, "{f.func.name.overload_name}")
-STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, schema_str, {cpp_string(str(f.func))})"""
+STATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, schema_str, {cpp_string(str(f.func))})
+
+// aten::{f.func}
+static C10_NOINLINE c10::TypedOperatorHandle<{name}::schema> create_{name}_typed_handle() {{
+  return c10::Dispatcher::singleton()
+      .findSchemaOrThrow("aten::{f.func.name.name}", "{f.func.name.overload_name}")
+      .typed<{name}::schema>();
+}}
+"""
 
             for is_redispatching_fn in [False, True]:
                 if is_redispatching_fn:
@@ -275,9 +283,7 @@ def __call__(self, f: NativeFunction) -> Optional[str]:
                 defns += f"""
 // aten::{f.func}
 {sig.defn(name=method_name, is_redispatching_fn=is_redispatching_fn)} {{
-    static auto op = c10::Dispatcher::singleton()
-        .findSchemaOrThrow("aten::{f.func.name.name}", "{f.func.name.overload_name}")
-        .typed<{sig.type()}>();
+    static auto op = create_{name}_typed_handle();
     return op.{dispatcher_call}({dispatcher_exprs_str});
 }}
 """
@@ -1175,10 +1181,18 @@ def make_file_manager(install_dir: str) -> FileManager:
         'schema_registrations': list(mapMaybe(RegisterSchema(schema_selector), native_functions)),
     })
 
-    cpu_fm.write('Operators.cpp', lambda: {
-        'definitions': list(mapMaybe(ComputeOperators(
-            Target.DEFINITION), native_functions)),
-    })
+    def key_func(fn: NativeFunction) -> str:
+        return fn.func.name.unambiguous_name()
+
+    cpu_fm.write_sharded(
+        'Operators.cpp',
+        native_functions,
+        key_fn=key_func,
+        env_callable=lambda fn: {
+            'definitions': [ComputeOperators(Target.DEFINITION)(fn)]},
+        num_shards=5,
+        sharded_keys={'definitions'}
+    )
     cpu_fm.write('Operators.h', lambda: {
         'declarations': list(mapMaybe(ComputeOperators(
             Target.DECLARATION), native_functions)),