1# Owner(s): ["module: codegen"] 2 3from __future__ import annotations 4 5import os 6import tempfile 7import unittest 8 9import expecttest 10 11from torchgen.gen import _GLOBAL_PARSE_NATIVE_YAML_CACHE # noqa: F401 12from torchgen.gen_backend_stubs import run 13 14 15path = os.path.dirname(os.path.realpath(__file__)) 16gen_backend_stubs_path = os.path.join(path, "../torchgen/gen_backend_stubs.py") 17 18 19# gen_backend_stubs.py is an integration point that is called directly by external backends. 20# The tests here are to confirm that badly formed inputs result in reasonable error messages. 21class TestGenBackendStubs(expecttest.TestCase): 22 def setUp(self) -> None: 23 global _GLOBAL_PARSE_NATIVE_YAML_CACHE 24 _GLOBAL_PARSE_NATIVE_YAML_CACHE.clear() 25 26 def assert_success_from_gen_backend_stubs(self, yaml_str: str) -> None: 27 with tempfile.NamedTemporaryFile(mode="w") as fp: 28 fp.write(yaml_str) 29 fp.flush() 30 run(fp.name, "", True) 31 32 def get_errors_from_gen_backend_stubs( 33 self, yaml_str: str, *, kernels_str: str | None = None 34 ) -> str: 35 with tempfile.NamedTemporaryFile(mode="w") as fp: 36 fp.write(yaml_str) 37 fp.flush() 38 try: 39 if kernels_str is None: 40 run(fp.name, "", True) 41 else: 42 with tempfile.NamedTemporaryFile(mode="w") as kernel_file: 43 kernel_file.write(kernels_str) 44 kernel_file.flush() 45 run(fp.name, "", True, impl_path=kernel_file.name) 46 except AssertionError as e: 47 # Scrub out the temp file name from any error messages to simplify assertions. 48 return str(e).replace(fp.name, "") 49 self.fail( 50 "Expected gen_backend_stubs to raise an AssertionError, but it did not." 51 ) 52 53 def test_valid_single_op(self) -> None: 54 yaml_str = """\ 55backend: XLA 56cpp_namespace: torch_xla 57supported: 58- abs""" 59 self.assert_success_from_gen_backend_stubs(yaml_str) 60 61 def test_valid_multiple_ops(self) -> None: 62 yaml_str = """\ 63backend: XLA 64cpp_namespace: torch_xla 65supported: 66- add.Tensor 67- abs""" 68 self.assert_success_from_gen_backend_stubs(yaml_str) 69 70 def test_valid_zero_ops(self) -> None: 71 yaml_str = """\ 72backend: XLA 73cpp_namespace: torch_xla 74supported:""" 75 self.assert_success_from_gen_backend_stubs(yaml_str) 76 77 def test_valid_zero_ops_doesnt_require_backend_dispatch_key(self) -> None: 78 yaml_str = """\ 79backend: BAD_XLA 80cpp_namespace: torch_xla 81supported:""" 82 # External codegen on a yaml file with no operators is effectively a no-op, 83 # so there's no reason to parse the backend 84 self.assert_success_from_gen_backend_stubs(yaml_str) 85 86 def test_valid_with_autograd_ops(self) -> None: 87 yaml_str = """\ 88backend: XLA 89cpp_namespace: torch_xla 90supported: 91- abs 92autograd: 93- add.Tensor""" 94 # External codegen on a yaml file with no operators is effectively a no-op, 95 # so there's no reason to parse the backend 96 self.assert_success_from_gen_backend_stubs(yaml_str) 97 98 def test_missing_backend(self) -> None: 99 yaml_str = """\ 100cpp_namespace: torch_xla 101supported: 102- abs""" 103 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 104 self.assertExpectedInline( 105 output_error, '''You must provide a value for "backend"''' 106 ) 107 108 def test_empty_backend(self) -> None: 109 yaml_str = """\ 110backend: 111cpp_namespace: torch_xla 112supported: 113- abs""" 114 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 115 self.assertExpectedInline( 116 output_error, '''You must provide a value for "backend"''' 117 ) 118 119 def test_backend_invalid_dispatch_key(self) -> None: 120 yaml_str = """\ 121backend: NOT_XLA 122cpp_namespace: torch_xla 123supported: 124- abs""" 125 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 126 self.assertExpectedInline( 127 output_error, 128 """\ 129unknown dispatch key NOT_XLA 130 The provided value for "backend" must be a valid DispatchKey, but got NOT_XLA.""", 131 ) # noqa: B950 132 133 def test_missing_cpp_namespace(self) -> None: 134 yaml_str = """\ 135backend: XLA 136supported: 137- abs""" 138 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 139 self.assertExpectedInline( 140 output_error, '''You must provide a value for "cpp_namespace"''' 141 ) 142 143 def test_whitespace_cpp_namespace(self) -> None: 144 yaml_str = """\ 145backend: XLA 146cpp_namespace:\t 147supported: 148- abs""" 149 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 150 self.assertExpectedInline( 151 output_error, '''You must provide a value for "cpp_namespace"''' 152 ) 153 154 # supported is a single item (it should be a list) 155 def test_nonlist_supported(self) -> None: 156 yaml_str = """\ 157backend: XLA 158cpp_namespace: torch_xla 159supported: abs""" 160 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 161 self.assertExpectedInline( 162 output_error, 163 """expected "supported" to be a list, but got: abs (of type <class 'str'>)""", 164 ) 165 166 # supported contains an op that isn't in native_functions.yaml 167 def test_supported_invalid_op(self) -> None: 168 yaml_str = """\ 169backend: XLA 170cpp_namespace: torch_xla 171supported: 172- abs_BAD""" 173 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 174 self.assertExpectedInline( 175 output_error, """Found an invalid operator name: abs_BAD""" 176 ) 177 178 # The backend is valid, but doesn't have a valid autograd key. They can't override autograd kernels in that case. 179 # Only using Vulkan here because it has a valid backend key but not an autograd key- if this changes we can update the test. 180 def test_backend_has_no_autograd_key_but_provides_entries(self) -> None: 181 yaml_str = """\ 182backend: Vulkan 183cpp_namespace: torch_vulkan 184supported: 185- add 186autograd: 187- sub""" 188 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 189 self.assertExpectedInline( 190 output_error, """Found an invalid operator name: add""" 191 ) # noqa: B950 192 193 # in an operator group, currently all operators must either be registered to the backend or autograd kernel. 194 # Here, functional and out mismatch 195 def test_backend_autograd_kernel_mismatch_out_functional(self) -> None: 196 yaml_str = """\ 197backend: XLA 198cpp_namespace: torch_xla 199supported: 200- add.Tensor 201autograd: 202- add.out""" 203 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 204 self.assertExpectedInline( 205 output_error, 206 """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 207 ) 208 209 # in an operator group, currently all operators must either be registered to the backend or autograd kernel. 210 # Here, functional and inplace mismatch 211 def test_backend_autograd_kernel_mismatch_functional_inplace(self) -> None: 212 yaml_str = """\ 213backend: XLA 214cpp_namespace: torch_xla 215supported: 216- add.Tensor 217autograd: 218- add_.Tensor""" 219 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 220 self.assertExpectedInline( 221 output_error, 222 """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 223 ) 224 225 # Currently, the same operator can't be listed under both 'supported' and 'autograd', which would 226 # involve registering the same kernel to both the XLA and AutogradXLA keys. 227 # If we need that functionality in the future, we'll need to augment the codegen. 228 def test_op_appears_in_supported_and_autograd_lists(self) -> None: 229 yaml_str = """\ 230backend: XLA 231cpp_namespace: torch_xla 232supported: 233- add.Tensor 234autograd: 235- add.Tensor""" 236 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 237 self.assertExpectedInline( 238 output_error, 239 """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 240 ) 241 242 # unrecognized extra yaml key 243 def test_unrecognized_key(self) -> None: 244 yaml_str = """\ 245backend: XLA 246cpp_namespace: torch_xla 247supported: 248- abs 249invalid_key: invalid_val""" 250 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 251 self.assertExpectedInline( 252 output_error, 253 """ contains unexpected keys: invalid_key. Only the following keys are supported: backend, class_name, cpp_namespace, extra_headers, supported, autograd, full_codegen, non_native, ir_gen, symint""", # noqa: B950 254 ) 255 256 # if use_out_as_primary is provided, it must be a bool 257 def test_use_out_as_primary_non_bool(self) -> None: 258 yaml_str = """\ 259backend: XLA 260cpp_namespace: torch_xla 261use_out_as_primary: frue 262supported: 263- abs""" 264 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 265 self.assertExpectedInline( 266 output_error, 267 """You must provide either True or False for use_out_as_primary. Provided: frue""", 268 ) # noqa: B950 269 270 # if device_guard is provided, it must be a bool 271 def test_device_guard_non_bool(self) -> None: 272 yaml_str = """\ 273backend: XLA 274cpp_namespace: torch_xla 275device_guard: frue 276supported: 277- abs""" 278 output_error = self.get_errors_from_gen_backend_stubs(yaml_str) 279 self.assertExpectedInline( 280 output_error, 281 """You must provide either True or False for device_guard. Provided: frue""", 282 ) # noqa: B950 283 284 def test_incorrect_kernel_name(self) -> None: 285 yaml_str = """\ 286backend: XLA 287cpp_namespace: torch_xla 288supported: 289- abs 290autograd: 291- add.Tensor""" 292 # Codegen will expect two kernel names (and try to parse them with regex): 293 # XLANativeFunctions::abs(...) 294 # XLANativeFunctions::add(...) 295 kernels_str = """\ 296at::Tensor& XLANativeFunctions::absWRONG(at::Tensor& self) {} 297at::Tensor& XLANativeFunctions::add(at::Tensor& self) {}""" 298 output_error = self.get_errors_from_gen_backend_stubs( 299 yaml_str, kernels_str=kernels_str 300 ) 301 self.assertExpectedInline( 302 output_error, 303 """\ 304 305XLANativeFunctions is missing a kernel definition for abs. We found 0 kernel(s) with that name, 306but expected 1 kernel(s). The expected function schemas for the missing operator are: 307at::Tensor abs(const at::Tensor & self) 308 309""", 310 ) 311 312 313if __name__ == "__main__": 314 unittest.main() 315