test_gen_backend_stubs.py

import os
import tempfile

from torch.testing._internal.common_utils import TestCase, run_tests
import tools.codegen.gen_backend_stubs

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):

    def assert_success_from_gen_backend_stubs(self, yaml_str: str) -> str:
        with tempfile.NamedTemporaryFile(mode='w') as fp:
            fp.write(yaml_str)
            fp.flush()
            tools.codegen.gen_backend_stubs.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)
            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):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
supported:
- abs'''
        self.assert_success_from_gen_backend_stubs(yaml_str)

    def test_valid_multiple_ops(self):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
supported:
- add.Tensor
- abs'''
        self.assert_success_from_gen_backend_stubs(yaml_str)

    def test_valid_zero_ops(self):
        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):
        yaml_str = '''\
backend: BAD_XLA
cpp_namespace: torch_xla
supported:'''
        # External codegen on a yaml file with no operators is effectively a no-op,
        # 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):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
supported:
- abs
autograd:
- add.Tensor'''
        # External codegen on a yaml file with no operators is effectively a no-op,
        # so there's no reason to parse the backend
        self.assert_success_from_gen_backend_stubs(yaml_str)

    def test_missing_backend(self):
        yaml_str = '''\
cpp_namespace: torch_xla
supported:
- abs'''
        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):
        yaml_str = '''\
backend:
cpp_namespace: torch_xla
supported:
- abs'''
        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):
        yaml_str = '''\
backend: NOT_XLA
cpp_namespace: torch_xla
supported:
- abs'''
        output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
        self.assertExpectedInline(output_error, '''\
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):
        yaml_str = '''\
backend: XLA
supported:
- abs'''
        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):
        yaml_str = '''\
backend: XLA
cpp_namespace:\t
supported:
- abs'''
        output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
        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):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
supported: abs'''
        output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
        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):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
supported:
- abs_BAD'''
        output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
        self.assertExpectedInline(output_error, '''Found an invalid operator name: abs_BAD''')

    # The backend is valid, but doesn't have a valid autograd key. They can't override autograd kernels in that case.
    # Only using MSNPU 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):
        yaml_str = '''\
backend: MSNPU
cpp_namespace: torch_msnpu
supported:
- add
autograd:
- sub'''
        output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
        self.assertExpectedInline(output_error, '''Found an invalid operator name: add''')  # noqa: B950

    # 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):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_msnpu
supported:
- add.Tensor
autograd:
- add.out'''
        output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
        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_out is listed under "autograd".''')  # noqa: B950

    # 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):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_msnpu
supported:
- add.Tensor
autograd:
- add_.Tensor'''
        output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
        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

    # 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):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_msnpu
supported:
- add.Tensor
autograd:
- add.Tensor'''
        output_error = self.get_errors_from_gen_backend_stubs(yaml_str)
        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):
        yaml_str = '''\
backend: XLA
cpp_namespace: torch_xla
supported:
- 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


if __name__ == '__main__':
    run_tests()