已合并
test: 新增torch.fx.experimental.symbolic_shapes NPU 适配验证与统一运行脚本 #35552
yuhongming-2026创建于 5月13日
test: 新增torch.fx.experimental.symbolic_shapes NPU 适配验证与统一运行脚本 #35552
已合并
共 1 个文件变更+81-0
| @@ -0,0 +1,81 @@ | |||
| 1 | +# Owner(s): ["module: fx"] | ||
| 2 | +"""NPU compatibility tests for torch.fx.experimental.symbolic_shapes APIs. | ||
| 3 | + | ||
| 4 | +This test verifies: | ||
| 5 | +- torch.fx.experimental.symbolic_shapes.compute_unbacked_bindings | ||
| 6 | +- torch.fx.experimental.symbolic_shapes.constrain_range | ||
| 7 | +- torch.fx.experimental.symbolic_shapes.constrain_unify | ||
| 8 | +- torch.fx.experimental.symbolic_shapes.ConvertIntKey | ||
| 9 | +- torch.fx.experimental.symbolic_shapes.ConvertIntKey.get | ||
| 10 | +""" | ||
| 11 | + | ||
| 12 | +import torch | ||
| 13 | +import torch_npu | ||
| 14 | +import torch.fx.experimental.symbolic_shapes as symbolic_shapes | ||
| 15 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +class TestFXExperimentalSymbolic(TestCase): | ||
| 19 | + """Verify symbolic_shapes helpers work on NPU and remain importable.""" | ||
| 20 | + | ||
| 21 | + def test_compute_unbacked_bindings(self): | ||
| 22 | + # Call compute_unbacked_bindings with a plain NPU tensor input. | ||
| 23 | + x = torch.randn(2, 3, device="npu") | ||
| 24 | + result = symbolic_shapes.compute_unbacked_bindings(None, x) | ||
| 25 | + self.assertTrue(result is None or isinstance(result, dict)) | ||
| 26 | + | ||
| 27 | + def test_compute_unbacked_bindings_npu_context(self): | ||
| 28 | + # Same API inside an explicit torch.npu.device context. | ||
| 29 | + with torch.npu.device(0): | ||
| 30 | + x = torch.randn(4, 5, device="npu") | ||
| 31 | + result = symbolic_shapes.compute_unbacked_bindings(None, x) | ||
| 32 | + self.assertTrue(result is None or isinstance(result, dict)) | ||
| 33 | + | ||
| 34 | + def test_compute_unbacked_bindings_to_npu(self): | ||
| 35 | + # Call compute_unbacked_bindings after transferring the input tensor to NPU. | ||
| 36 | + x = torch.randn(2, 3).to("npu") | ||
| 37 | + result = symbolic_shapes.compute_unbacked_bindings(None, x) | ||
| 38 | + self.assertTrue(result is None or isinstance(result, dict)) | ||
| 39 | + | ||
| 40 | + def test_constrain_range(self): | ||
| 41 | + # constrain_range accepts in-range plain ints without a ShapeEnv. | ||
| 42 | + symbolic_shapes.constrain_range(5, min=2, max=10) | ||
| 43 | + with self.assertRaises(ValueError): | ||
| 44 | + symbolic_shapes.constrain_range(1, min=2, max=10) | ||
| 45 | + | ||
| 46 | + def test_constrain_range_npu_context(self): | ||
| 47 | + with torch.npu.device(0): | ||
| 48 | + symbolic_shapes.constrain_range(5, min=2, max=10) | ||
| 49 | + | ||
| 50 | + def test_constrain_unify(self): | ||
| 51 | + symbolic_shapes.constrain_unify(5, 5) | ||
| 52 | + with self.assertRaises(AssertionError): | ||
| 53 | + symbolic_shapes.constrain_unify(5, 6) | ||
| 54 | + | ||
| 55 | + def test_constrain_unify_npu_context(self): | ||
| 56 | + with torch.npu.device(0): | ||
| 57 | + symbolic_shapes.constrain_unify(7, 7) | ||
| 58 | + | ||
| 59 | + def test_convert_int_key_singleton(self): | ||
| 60 | + # ConvertIntKey maps bool conditions to integer constants 1/0. | ||
| 61 | + cik = symbolic_shapes.ConvertIntKey() | ||
| 62 | + self.assertEqual(cik.get(True), 1) | ||
| 63 | + self.assertEqual(cik.get(False), 0) | ||
| 64 | + | ||
| 65 | + def test_convert_int_key_get_npu_derived(self): | ||
| 66 | + # ConvertIntKey.get accepts bools materialized from NPU tensor compares. | ||
| 67 | + cik = symbolic_shapes.ConvertIntKey() | ||
| 68 | + x = torch.tensor([1, 2, 3], device="npu") | ||
| 69 | + self.assertEqual(cik.get(bool((x[0] == x[0]).item())), 1) | ||
| 70 | + self.assertEqual(cik.get(bool((x[0] != x[1]).item())), 1) | ||
| 71 | + self.assertEqual(cik.get(bool((x[1] <= x[0]).item())), 0) | ||
| 72 | + | ||
| 73 | + def test_convert_int_key_npu_context(self): | ||
| 74 | + with torch.npu.device(0): | ||
| 75 | + cik = symbolic_shapes.ConvertIntKey() | ||
| 76 | + self.assertEqual(cik.get(True), 1) | ||
| 77 | + self.assertEqual(cik.get(False), 0) | ||
| 78 | + | ||
| 79 | + | ||
| 80 | +if __name__ == "__main__": | ||
| 81 | + run_tests() | ||