[torchgen] Fix multiple backends with custom namespace (#82133)
Summary:
Some quantized operators needs `QuantizedCPU` backend, due to an issue in namespace checking, currently if we have two backends as well as a custom namespaces in native function, codegen will hit assertion error. This PR fixes this issue
The root cause is that codegen right now asserts that a native function should only have one namespace. The current behavior is that If a native function is not found in a `BackendIndex`, we will use default namespace for that backend, for fallback kernels. However that default namespace may not be listed in the yaml file and it should not be counted when checking if we have two different namespaces for that backend. In our error case, we have 2 `BackendIndex`, one for `QuantizedCPU` and one for `CPU`. The native function doesn't have a kernel in `QuantizedCPU` but we still use a default namespace (`at::native`) for it. Since we have a custom namespace for dispatch key `CPU`, we ran into the assertion error.
This PR changes the assertion criteria. We only error out if a namespace has two or more kernels and they have two or more different namespaces.
Test Plan: rely on newly added unit test
Differential Revision: D38101345
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82133
Approved by: https://github.com/iseeyuan
diff --git a/tools/BUCK.bzl b/tools/BUCK.bzl
index 959a73d..2ca59ee 100644
--- a/tools/BUCK.bzl
+++ b/tools/BUCK.bzl
@@ -261,3 +261,16 @@
":gen_operators_yaml_lib",
],
)
+
+ python_test(
+ name = "test_codegen",
+ srcs = [
+ "test/test_codegen.py",
+ ],
+ contacts = contacts,
+ visibility = ["PUBLIC"],
+ deps = [
+ torchgen_deps,
+ ":autograd",
+ ],
+ )
