Add sparse, quantized and nested tensor meta support to codegen (pyto…

…rch#81793) Signed-off-by: Edward Z. Yang <[email protected]> Pull Request resolved: pytorch#81793 Approved by: https://github.com/cpuhrsch, https://github.com/bdhirsh
juhofuriosa · Jul 21, 2022 · 6f0c253 · 6f0c253
1 parent a33c5ac
commit 6f0c253
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 0 deletions.
diff --git a/.bazelrc b/.bazelrc
@@ -107,6 +107,9 @@ build --per_file_copt='//:aten/src/ATen/RegisterNestedTensorCPU\.cpp$'@-Wno-erro
 build --per_file_copt='//:aten/src/ATen/RegisterQuantizedCPU\.cpp$'@-Wno-error=unused-function
 build --per_file_copt='//:aten/src/ATen/RegisterSparseCPU\.cpp$'@-Wno-error=unused-function
 build --per_file_copt='//:aten/src/ATen/RegisterSparseCsrCPU\.cpp$'@-Wno-error=unused-function
+build --per_file_copt='//:aten/src/ATen/RegisterNestedTensorMeta\.cpp$'@-Wno-error=unused-function
+build --per_file_copt='//:aten/src/ATen/RegisterSparseMeta\.cpp$'@-Wno-error=unused-function
+build --per_file_copt='//:aten/src/ATen/RegisterQuantizedMeta\.cpp$'@-Wno-error=unused-function
 build --per_file_copt='//:aten/src/ATen/RegisterZeroTensor\.cpp$'@-Wno-error=unused-function
 build --per_file_copt='//:torch/csrc/lazy/generated/RegisterAutogradLazy\.cpp$'@-Wno-error=unused-function
 build --per_file_copt='//:torch/csrc/lazy/generated/RegisterLazy\.cpp$'@-Wno-error=unused-function
diff --git a/BUILD.bazel b/BUILD.bazel
@@ -50,6 +50,9 @@ generated_cpu_cpp = [
     "aten/src/ATen/RegisterCompositeExplicitAutograd.cpp",
     "aten/src/ATen/RegisterCompositeExplicitAutogradNonFunctional.cpp",
     "aten/src/ATen/RegisterMeta.cpp",
+    "aten/src/ATen/RegisterSparseMeta.cpp",
+    "aten/src/ATen/RegisterQuantizedMeta.cpp",
+    "aten/src/ATen/RegisterNestedTensorMeta.cpp",
     "aten/src/ATen/RegisterSchema.cpp",
     "aten/src/ATen/CPUFunctions.h",
     "aten/src/ATen/CPUFunctions_inl.h",

diff --git a/build.bzl b/build.bzl
@@ -193,6 +193,9 @@ GENERATED_CPP = [
     "RegisterCompositeImplicitAutograd.cpp",
     "RegisterZeroTensor.cpp",
     "RegisterMeta.cpp",
+    "RegisterQuantizedMeta.cpp",
+    "RegisterNestedTensorMeta.cpp",
+    "RegisterSparseMeta.cpp",
     "RegisterCompositeExplicitAutograd.cpp",
     "RegisterCompositeExplicitAutogradNonFunctional.cpp",
     "CompositeViewCopyKernels.cpp",

diff --git a/torchgen/model.py b/torchgen/model.py
@@ -218,6 +218,9 @@ def codegen_per_backend_entries() -> str:
     # Meta is a magic key: it is automatically generated for structured
     # kernels
     DispatchKey.Meta,
+    DispatchKey.SparseMeta,
+    DispatchKey.QuantizedMeta,
+    DispatchKey.NestedTensorMeta,
     DispatchKey.ZeroTensor,
 ]