allow external backend codegen to toggle whether to generate out= and inplace kernels (#68530)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/68530
Test Plan: Imported from OSS
Reviewed By: albanD
Differential Revision: D32498569
Pulled By: bdhirsh
fbshipit-source-id: ebd932d042b988e19c71aa04a21677db9bdc9f04
diff --git a/test/run_test.py b/test/run_test.py
index 1707220..171383c 100644
--- a/test/run_test.py
+++ b/test/run_test.py
@@ -97,7 +97,6 @@
'test_bundled_images',
'test_cpp_extensions_aot',
'test_determination',
- 'test_gen_backend_stubs',
'test_jit_fuser',
'test_jit_simple',
'test_jit_string',
diff --git a/tools/codegen/gen_backend_stubs.py b/tools/codegen/gen_backend_stubs.py
index a860ff6..96eed45 100644
--- a/tools/codegen/gen_backend_stubs.py
+++ b/tools/codegen/gen_backend_stubs.py
@@ -43,6 +43,11 @@
cpp_namespace = yaml_values.pop('cpp_namespace', None)
assert cpp_namespace is not None, 'You must provide a value for "cpp_namespace"'
+ # Mostly just defaulting to false to stick with LazyTensor convention.
+ use_out_as_primary = yaml_values.pop('use_out_as_primary', False)
+ assert isinstance(use_out_as_primary, bool), \
+ f'You must provide either True or False for use_out_as_primary. Provided: {use_out_as_primary}'
+
supported = yaml_values.pop('supported', [])
if supported is None:
supported = [] # Allow an empty list of supported ops
@@ -59,7 +64,7 @@
f'{backend_yaml_path} contains unexpected keys: {", ".join(yaml_values.keys())}. \
Only the following keys are supported: {", ".join(valid_keys)}'
- def create_backend_index(backend_ops: List[str], dispatch_key: DispatchKey) -> BackendIndex:
+ def create_backend_index(backend_ops: List[str], dispatch_key: DispatchKey, *, use_out_as_primary: bool) -> BackendIndex:
metadata: Dict[OperatorName, BackendMetadata] = {}
for op in backend_ops:
op_name = OperatorName.parse(op)
@@ -69,11 +74,9 @@
# TODO: allow structured external backends later.
m = BackendMetadata(kernel=kernel_name, structured=False)
metadata[op_name] = m
- # TODO: currently hardcoding the fact that XLA implements out/inplace in terms of functional ops,
- # this should eventually be toggleable per-backend.
return BackendIndex(
dispatch_key=dispatch_key,
- use_out_as_primary=False,
+ use_out_as_primary=use_out_as_primary,
external=True,
index=metadata)
@@ -82,7 +85,7 @@
with context(lambda: f'The provided value for "backend" must be a valid DispatchKey, but got {backend}.'):
backend_key = DispatchKey.parse(backend)
- backend_idx = create_backend_index(supported, backend_key)
+ backend_idx = create_backend_index(supported, backend_key, use_out_as_primary=use_out_as_primary)
assert backend_key not in backend_indices
backend_indices[backend_key] = backend_idx
@@ -92,7 +95,7 @@
the behavior of autograd for some operators on your backend. However "Autograd{backend}" is not a valid DispatchKey.'):
autograd_key = DispatchKey.parse(f'Autograd{backend}')
- autograd_idx = create_backend_index(supported_autograd, autograd_key)
+ autograd_idx = create_backend_index(supported_autograd, autograd_key, use_out_as_primary=use_out_as_primary)
assert autograd_key not in backend_indices
backend_indices[autograd_key] = autograd_idx
@@ -262,7 +265,7 @@
)),
})
-def run(source_yaml: str, output_dir: str, dry_run: bool, impl_path: Optional[str]) -> None:
+def run(source_yaml: str, output_dir: str, dry_run: bool, impl_path: Optional[str] = None) -> None:
# Assumes that this file lives at PYTORCH_ROOT/tools/codegen/gen_backend_stubs.py
pytorch_root = pathlib.Path(__file__).parent.parent.parent.absolute()
@@ -286,7 +289,10 @@
selector = SelectiveBuilder.get_nop_selector()
- assert backend_key is not None
+ if backend_key is None:
+ # This could be useful if a backend wants to quickly set up a noop yaml file but doesn't have any kernels ready yet.
+ return
+
class_name = backend_indices[backend_key].native_function_class_name()
if impl_path is not None:
diff --git a/test/test_gen_backend_stubs.py b/tools/test/test_gen_backend_stubs.py
similarity index 81%
rename from test/test_gen_backend_stubs.py
rename to tools/test/test_gen_backend_stubs.py
index f84a1c1..5b829be 100644
--- a/test/test_gen_backend_stubs.py
+++ b/tools/test/test_gen_backend_stubs.py
@@ -2,35 +2,42 @@
import os
import tempfile
+import unittest
+import expecttest
-from torch.testing._internal.common_utils import TestCase, run_tests
-import tools.codegen.gen_backend_stubs
+from tools.codegen.gen_backend_stubs import run
+from tools.codegen.gen import _GLOBAL_PARSE_NATIVE_YAML_CACHE # noqa: F401
path = os.path.dirname(os.path.realpath(__file__))
gen_backend_stubs_path = os.path.join(path, '../tools/codegen/gen_backend_stubs.py')
# gen_backend_stubs.py is an integration point that is called directly by external backends.
# The tests here are to confirm that badly formed inputs result in reasonable error messages.
-class TestGenBackendStubs(TestCase):
+class TestGenBackendStubs(expecttest.TestCase):
- def assert_success_from_gen_backend_stubs(self, yaml_str: str) -> str:
+ def setUp(self) -> None:
+ global _GLOBAL_PARSE_NATIVE_YAML_CACHE
+ _GLOBAL_PARSE_NATIVE_YAML_CACHE.clear()
+
+
+ def assert_success_from_gen_backend_stubs(self, yaml_str: str) -> None:
with tempfile.NamedTemporaryFile(mode='w') as fp:
fp.write(yaml_str)
fp.flush()
- tools.codegen.gen_backend_stubs.run(fp.name, '', True)
+ run(fp.name, '', True)
def get_errors_from_gen_backend_stubs(self, yaml_str: str) -> str:
with tempfile.NamedTemporaryFile(mode='w') as fp:
fp.write(yaml_str)
fp.flush()
try:
- tools.codegen.gen_backend_stubs.run(fp.name, '', True)
+ run(fp.name, '', True)
except AssertionError as e:
# Scrub out the temp file name from any error messages to simplify assertions.
return str(e).replace(fp.name, '')
self.fail('Expected gen_backend_stubs to raise an AssertionError, but it did not.')
- def test_valid_single_op(self):
+ def test_valid_single_op(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -38,7 +45,7 @@
- abs'''
self.assert_success_from_gen_backend_stubs(yaml_str)
- def test_valid_multiple_ops(self):
+ def test_valid_multiple_ops(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -47,14 +54,14 @@
- abs'''
self.assert_success_from_gen_backend_stubs(yaml_str)
- def test_valid_zero_ops(self):
+ def test_valid_zero_ops(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
supported:'''
self.assert_success_from_gen_backend_stubs(yaml_str)
- def test_valid_zero_ops_doesnt_require_backend_dispatch_key(self):
+ def test_valid_zero_ops_doesnt_require_backend_dispatch_key(self) -> None:
yaml_str = '''\
backend: BAD_XLA
cpp_namespace: torch_xla
@@ -63,7 +70,7 @@
# so there's no reason to parse the backend
self.assert_success_from_gen_backend_stubs(yaml_str)
- def test_valid_with_autograd_ops(self):
+ def test_valid_with_autograd_ops(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -75,7 +82,7 @@
# so there's no reason to parse the backend
self.assert_success_from_gen_backend_stubs(yaml_str)
- def test_missing_backend(self):
+ def test_missing_backend(self) -> None:
yaml_str = '''\
cpp_namespace: torch_xla
supported:
@@ -83,7 +90,7 @@
output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
self.assertExpectedInline(output_error, '''You must provide a value for "backend"''')
- def test_empty_backend(self):
+ def test_empty_backend(self) -> None:
yaml_str = '''\
backend:
cpp_namespace: torch_xla
@@ -92,7 +99,7 @@
output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
self.assertExpectedInline(output_error, '''You must provide a value for "backend"''')
- def test_backend_invalid_dispatch_key(self):
+ def test_backend_invalid_dispatch_key(self) -> None:
yaml_str = '''\
backend: NOT_XLA
cpp_namespace: torch_xla
@@ -103,7 +110,7 @@
unknown dispatch key NOT_XLA
The provided value for "backend" must be a valid DispatchKey, but got NOT_XLA.''') # noqa: B950
- def test_missing_cpp_namespace(self):
+ def test_missing_cpp_namespace(self) -> None:
yaml_str = '''\
backend: XLA
supported:
@@ -111,7 +118,7 @@
output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
self.assertExpectedInline(output_error, '''You must provide a value for "cpp_namespace"''')
- def test_whitespace_cpp_namespace(self):
+ def test_whitespace_cpp_namespace(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace:\t
@@ -121,7 +128,7 @@
self.assertExpectedInline(output_error, '''You must provide a value for "cpp_namespace"''')
# supported is a single item (it should be a list)
- def test_nonlist_supported(self):
+ def test_nonlist_supported(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -130,7 +137,7 @@
self.assertExpectedInline(output_error, '''expected "supported" to be a list, but got: abs (of type <class 'str'>)''')
# supported contains an op that isn't in native_functions.yaml
- def test_supported_invalid_op(self):
+ def test_supported_invalid_op(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -141,7 +148,7 @@
# The backend is valid, but doesn't have a valid autograd key. They can't override autograd kernels in that case.
# Only using Vulkan here because it has a valid backend key but not an autograd key- if this changes we can update the test.
- def test_backend_has_no_autograd_key_but_provides_entries(self):
+ def test_backend_has_no_autograd_key_but_provides_entries(self) -> None:
yaml_str = '''\
backend: Vulkan
cpp_namespace: torch_vulkan
@@ -154,7 +161,7 @@
# in an operator group, currently all operators must either be registered to the backend or autograd kernel.
# Here, functional and out mismatch
- def test_backend_autograd_kernel_mismatch_out_functional(self):
+ def test_backend_autograd_kernel_mismatch_out_functional(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -167,7 +174,7 @@
# in an operator group, currently all operators must either be registered to the backend or autograd kernel.
# Here, functional and inplace mismatch
- def test_backend_autograd_kernel_mismatch_functional_inplace(self):
+ def test_backend_autograd_kernel_mismatch_functional_inplace(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -181,7 +188,7 @@
# Currently, the same operator can't be listed under both 'supported' and 'autograd', which would
# involve registering the same kernel to both the XLA and AutogradXLA keys.
# If we need that functionality in the future, we'll need to augment the codegen.
- def test_op_appears_in_supported_and_autograd_lists(self):
+ def test_op_appears_in_supported_and_autograd_lists(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -193,7 +200,7 @@
self.assertExpectedInline(output_error, '''Currently, all variants of an op must either be registered to a backend key, or to a backend's autograd key. They cannot be mix and matched. If this is something you need, feel free to create an issue! add is listed under "supported", but add is listed under "autograd".''') # noqa: B950
# unrecognized extra yaml key
- def test_unrecognized_key(self):
+ def test_unrecognized_key(self) -> None:
yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
@@ -201,8 +208,19 @@
- abs
invalid_key: invalid_val'''
output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
- self.assertExpectedInline(output_error, ''' contains unexpected keys: invalid_key. Only the following keys are supported: backend, cpp_namespace, extra_headers, supported, autograd''') # noqa: B950
+ self.assertExpectedInline(output_error, ''' contains unexpected keys: invalid_key. Only the following keys are supported: backend, cpp_namespace, extra_headers, supported, autograd, full_codegen''') # noqa: B950
+
+ # if use_out_as_primary is provided, it must be a bool
+ def test_use_out_as_primary_non_bool(self) -> None:
+ yaml_str = '''\
+backend: XLA
+cpp_namespace: torch_xla
+use_out_as_primary: frue
+supported:
+- abs'''
+ output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
+ self.assertExpectedInline(output_error, '''You must provide either True or False for use_out_as_primary. Provided: frue''') # noqa: B950
if __name__ == '__main__':
- run_tests()
+ unittest.main()
