ATen DerivedType is dead, long live ATen RegisterDispatchKey (#47011)

ezyang · facebook-github-bot · commit 809660ffa49e · 2020-11-12T09:53:48.000-08:00
Summary: Pull Request resolved: #47011 smessmer has complained about how it is difficult to find generated code. Well hopefully this diffs helps a bit with that. There are three components to this refactor: - Rename TypeDerived (CPUType) to RegisterDispatchKey (RegisterCPU). The 'Type' nomenclature is vestigial and I think Register says what these files do a lot more clearly. I also got rid of the CPUType namespace; everything just goes in anonymous namespace now, less moving parts this way. - Give Math and DefaultBackend their own files (RegisterMath and RegisterDefaultBackend) - Restructure code generation so that schema definition is done completely separately from RegisterDispatchKey I decided to name the files RegisterCPU rather than the old convention BackendSelectRegister, because it seems better to me if these files clump together in an alphabetical listing rather than being spread out everywhere. There are a few manual registration files which should probably get similar renaming. I also did a little garden cleaning about how we identify if a dispatch key is a cuda key or a generic key (previously called KEYWORD_ALL_BACKENDS but I like my naming better). Signed-off-by: Edward Z. Yang <ezyang@fb.com> Differential Revision: D24600806 Test Plan: Imported from OSS Reviewed By: smessmer Pulled By: ezyang fbshipit-source-id: c1b510dd7515bd95e3ad25b8edf961b2fb30a25a
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -125,15 +125,17 @@ genrule(
     ] + glob(["aten/src/ATen/templates/**"]),
     outs = [
         "aten/src/ATen/Declarations.yaml",
-        "aten/src/ATen/BackendSelectRegister.cpp",
-        "aten/src/ATen/CPUType.cpp",
+        "aten/src/ATen/RegisterBackendSelect.cpp",
+        "aten/src/ATen/RegisterCPU.cpp",
+        "aten/src/ATen/RegisterMkldnnCPU.cpp",
+        "aten/src/ATen/RegisterQuantizedCPU.cpp",
+        "aten/src/ATen/RegisterSparseCPU.cpp",
+        "aten/src/ATen/RegisterMath.cpp",
+        "aten/src/ATen/RegisterDefaultBackend.cpp",
+        "aten/src/ATen/RegisterSchema.cpp",
         "aten/src/ATen/Functions.h",
         "aten/src/ATen/Functions.cpp",
         "aten/src/ATen/NativeFunctions.h",
-        "aten/src/ATen/MkldnnCPUType.cpp",
-        "aten/src/ATen/QuantizedCPUType.cpp",
-        "aten/src/ATen/SparseCPUType.cpp",
-        "aten/src/ATen/TypeDefault.cpp",
         "aten/src/ATen/core/TensorBody.h",
         "aten/src/ATen/core/TensorMethods.cpp",
         "aten/src/ATen/core/ATenOpList.cpp",
diff --git a/aten/src/ATen/templates/RegisterBackendSelect.cpp b/aten/src/ATen/templates/RegisterBackendSelect.cpp
diff --git a/aten/src/ATen/templates/RegisterDispatchKey.cpp b/aten/src/ATen/templates/RegisterDispatchKey.cpp
@@ -33,21 +33,14 @@
 
 namespace at {
 
-/* example
-Tensor * ${Type}::add(Tensor & a, Tensor & b) {
-  std::cout << "add Tensor with backend ${Backend}\n";
-  return &a;
-}
-*/
-
-namespace ${Type} {
+namespace {
 
-${type_derived_method_definitions}
+${dispatch_definitions}
 
-}  // namespace ${Type}
-
-TORCH_LIBRARY_IMPL(aten, ${Backend}, m) {
-  ${function_registrations}
+TORCH_LIBRARY_IMPL(aten, ${DispatchKey}, m) {
+  ${dispatch_registrations}
 }
 
+} // anonymous namespace
+
 } // namespace at
diff --git a/aten/src/ATen/templates/RegisterSchema.cpp b/aten/src/ATen/templates/RegisterSchema.cpp
@@ -17,14 +17,8 @@
 #include <torch/library.h>
 
 namespace at {
-namespace TypeDefault {
-
-${type_method_definitions}
-
-}  // namespace TypeDefault
-
 TORCH_LIBRARY(aten, m) {
-  ${function_registrations};
+  ${schema_registrations};
 
   // String Ops
   // Implementations located in torch/csrc/jit/runtime/register_prim_ops.cpp
@@ -63,12 +57,4 @@ TORCH_LIBRARY(aten, m) {
   // Implementations located in torch/csrc/jit/runtime/register_distributed_ops.cpp
   m.def("get_gradients(int context_id) -> Dict(Tensor, Tensor)");
 }
-
-TORCH_LIBRARY_IMPL(aten, Math, m) {
-  ${math_function_registrations};
-}
-
-TORCH_LIBRARY_IMPL(aten, DefaultBackend, m) {
-  ${default_backend_function_registrations};
-}
 }  // namespace at
diff --git a/c10/core/DispatchKey.h b/c10/core/DispatchKey.h
@@ -51,8 +51,8 @@ enum class DispatchKey : uint8_t {
 
   // Here are backends which you think of as traditionally specifying
   // how to implement operations on some device.
-  CPU, // registered at build/aten/src/ATen/CPUType.cpp
-  CUDA, // registered at build/aten/src/ATen/CUDAType.cpp
+  CPU, // registered at build/aten/src/ATen/RegisterCPU.cpp
+  CUDA, // registered at build/aten/src/ATen/RegisterCUDA.cpp
   HIP, // NB: I think this is not actually used, due to Note [Masquerading as
        // CUDA]
   FPGA, // Xilinx support lives out of tree at https://gitlab.com/pytorch-complex/vitis_kernels
@@ -73,8 +73,8 @@ enum class DispatchKey : uint8_t {
 
   // Here are backends which specify more specialized operators
   // based on the dtype of the tensor.
-  QuantizedCPU, // registered at build/aten/src/ATen/QuantizedCPUType.cpp
-  QuantizedCUDA, // registered at build/aten/src/ATen/QuantizedCUDAType.cpp
+  QuantizedCPU, // registered at build/aten/src/ATen/RegisterQuantizedCPU.cpp
+  QuantizedCUDA, // registered at build/aten/src/ATen/RegisterQuantizedCUDA.cpp
   ComplexCPU, // lives out of tree at
               // https://gitlab.com/pytorch-complex/pytorch-cpu-strided-complex
   ComplexCUDA, // and
@@ -97,10 +97,10 @@ enum class DispatchKey : uint8_t {
   // based on the layout of the tensor.  Note that the sparse backends
   // are one case where ordering matters: sparse multi-dispatches with
   // the corresponding dense tensors, and must be handled before them.
-  MkldnnCPU, // registered at build/aten/src/ATen/MkldnnCPUType.cpp
+  MkldnnCPU, // registered at build/aten/src/ATen/RegisterMkldnnCPU.cpp
   // NB: not to be confused with MKLDNN, which is Caffe2 only
-  SparseCPU, // registered at build/aten/src/ATen/SparseCPUType.cpp
-  SparseCUDA, // registered at build/aten/src/ATen/SparseCUDAType.cpp
+  SparseCPU, // registered at build/aten/src/ATen/RegisterSparseCPU.cpp
+  SparseCUDA, // registered at build/aten/src/ATen/RegisterSparseCUDA.cpp
   SparseHIP, // TODO: I think this is not actually used, due to Note
              // [Masquerading as CUDA]
 
@@ -276,8 +276,8 @@ enum class DispatchKey : uint8_t {
 
   // See Note [Alias Dispatch Key : Autograd]
   Autograd,
-  Math,
-  DefaultBackend,
+  Math,  // registered at build/aten/src/ATen/RegisterMath.cpp
+  DefaultBackend,  // registered at build/aten/src/ATen/RegisterDefaultBackend.cpp
 
   // Define an alias key to represent end of alias dispatch keys.
   // If you add new alias keys after Autograd, please also update it here.
diff --git a/tools/codegen/gen.py b/tools/codegen/gen.py
