pytorch · dzhulgakov · Sep 25, 2019 · Sep 25, 2019 · Sep 26, 2019
diff --git a/aten/src/ATen/core/OpsAlreadyMovedToC10.cpp b/aten/src/ATen/core/OpsAlreadyMovedToC10.cpp
diff --git a/aten/src/ATen/core/TensorMethods.h b/aten/src/ATen/core/TensorMethods.h
diff --git a/aten/src/ATen/function_wrapper.py b/aten/src/ATen/function_wrapper.py
@@ -113,16 +113,10 @@ def TypedDict(name, attrs, total=True):  # type: ignore
 """)
 
 DEFAULT_UNBOXEDONLY_FUNCTION_REGISTRATION = CodeTemplate("""\
-.op(torch::RegisterOperators::options()
-  .schema("${schema_string}")
-  .impl_unboxedOnlyCatchAllKernel<${return_type} (${formals_types}), &TypeDefault::${api_name}>()
-  .aliasAnalysis(c10::AliasAnalysisKind::FROM_SCHEMA))
+.registerOp<${return_type} (${formals_types})>(TensorTypeId::UndefinedTensorId, "${schema_string}", &TypeDefault::${api_name})
 """)
 BACKEND_UNBOXEDONLY_FUNCTION_REGISTRATION = CodeTemplate("""\
-.op(torch::RegisterOperators::options()
-  .schema("${schema_string}")
-  .impl_unboxedOnlyKernel<${return_type} (${formals_types}), &${Type}::${api_name}>(TensorTypeId::${Backend}TensorId)
-  .aliasAnalysis(c10::AliasAnalysisKind::FROM_SCHEMA))
+.registerOp<${return_type} (${formals_types})>(TensorTypeId::${Backend}TensorId, "${schema_string}", &${Type}::${api_name})
 """)
 DEFAULT_FUNCTION_REGISTRATION = CodeTemplate("""\
 .op(torch::RegisterOperators::options()

diff --git a/aten/src/ATen/native/README.md b/aten/src/ATen/native/README.md
@@ -303,6 +303,9 @@ set of reviewers.
 
 ### `use_c10_dispatcher`
 
+Note: this functionality is temporary disabled because of code size bloat,
+annotations are ignored.
+
 ```
 use_c10_dispatcher: 'no'
 use_c10_dispatcher: 'unboxed_only'

diff --git a/aten/src/ATen/native_parse.py b/aten/src/ATen/native_parse.py
@@ -415,6 +415,10 @@ def run(paths):
                 declaration['supports_named_tensor'] = func.get('supports_named_tensor', False)
                 declaration['use_c10_dispatcher'] = func.get('use_c10_dispatcher', 'no')
                 assert declaration['use_c10_dispatcher'] in ['no', 'unboxed_only', 'full']
+                # TODO: remove this override once code size issues are fixed (see
+                # https://github.com/pytorch/pytorch/pull/26821 for more context)
+                declaration['use_c10_dispatcher'] = 'no'
+
                 declaration['category_override'] = func.get('category_override', '')
                 declaration['arguments'] = func.get('arguments', arguments)
                 declaration['type_method_definition_dispatch'] = func.get('dispatch', declaration['name'])

diff --git a/aten/src/ATen/templates/SparseTypeDerived.cpp b/aten/src/ATen/templates/SparseTypeDerived.cpp
@@ -17,7 +17,6 @@
 #include <c10/util/Half.h>
 #include <c10/core/UndefinedTensorImpl.h>
 #include <c10/util/Optional.h>
-#include <ATen/core/op_registration/op_registration.h>
 #include <ATen/core/EnableNamedTensor.h>
 
 #include <cstddef>
@@ -33,7 +32,7 @@ namespace at {
 ${type_derived_method_definitions}
 
 #ifndef USE_STATIC_DISPATCH
-static auto registerer = torch::RegisterOperators()
+static auto& registerer = globalATenDispatch()
   ${function_registrations};
 #endif
 }
diff --git a/aten/src/ATen/templates/TypeDefault.cpp b/aten/src/ATen/templates/TypeDefault.cpp
@@ -14,15 +14,14 @@
 #include <ATen/DeviceGuard.h>
 #include <ATen/SparseTensorUtils.h>
 #include <ATen/core/ATenDispatch.h>
-#include <ATen/core/op_registration/op_registration.h>
 #include <ATen/core/EnableNamedTensor.h>
 
 namespace at {
 
 ${type_method_definitions}
 
 #ifndef USE_STATIC_DISPATCH
-static auto registerer = torch::RegisterOperators()
+static auto& registerer = globalATenDispatch()
   ${function_registrations};
 #endif
 }
diff --git a/aten/src/ATen/templates/TypeDerived.cpp b/aten/src/ATen/templates/TypeDerived.cpp
@@ -29,7 +29,6 @@
 #include <utility>
 
 #include <ATen/Config.h>
-#include <ATen/core/op_registration/op_registration.h>
 $extra_cuda_headers
 $legacy_th_headers
 
@@ -45,7 +44,7 @@ Tensor * ${Type}::add(Tensor & a, Tensor & b) {
 ${type_derived_method_definitions}
 
 #ifndef USE_STATIC_DISPATCH
-static auto registerer = torch::RegisterOperators()
+static auto& registerer = globalATenDispatch()
   ${function_registrations};
 #endif
 }
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
@@ -154,10 +154,7 @@
 """)
 
 UNBOXEDONLY_WRAPPER_REGISTRATION = CodeTemplate("""\
-.op(torch::RegisterOperators::options()
-  .schema("${schema_string}")
-  .impl_unboxedOnlyKernel<${return_type} (${formal_types}), &VariableType::${api_name}>(TensorTypeId::VariableTensorId)
-  .aliasAnalysis(c10::AliasAnalysisKind::FROM_SCHEMA))
+.registerOp<${return_type} (${formal_types})>(TensorTypeId::VariableTensorId, "${schema_string}", &VariableType::${api_name})
 """)
 
 WRAPPER_REGISTRATION = CodeTemplate("""\

diff --git a/tools/autograd/templates/VariableType.cpp b/tools/autograd/templates/VariableType.cpp
@@ -1,7 +1,6 @@
 #include "torch/csrc/autograd/VariableTypeUtils.h"
 
 #include <ATen/TypeDefault.h>
-#include <ATen/core/op_registration/op_registration.h>
 
 // ${generated_comment}
 
@@ -32,11 +31,6 @@ namespace torch { namespace autograd {
 
 ${type_derived_method_definitions}
 
-namespace {
-
-static auto registerer = torch::RegisterOperators()
+static auto& registerer = globalATenDispatch()
   ${wrapper_registrations};
-
-}
-
 }} // namespace torch::autograd