已合并
[test] ShapeEnv APIs NPU兼容性验证与资料更新 #36139
Yhw050920创建于 5月19日
[test] ShapeEnv APIs NPU兼容性验证与资料更新 #36139
已合并
共 1 个文件变更+100-0
| @@ -1,10 +1,20 @@ | |||
| 1 | +""" | ||
| 2 | +Add validation cases for torch.fx.experimental.symbolic_shapes.ShapeEnv APIs on NPU: | ||
| 3 | + | ||
| 4 | +PyTorch community lacks sufficient and direct API validations for some APIs, so this file is added. | ||
| 5 | +This file validates ShapeEnv.produce_guards_expression, ShapeEnv.produce_guards_verbose, ShapeEnv.replace, ShapeEnv.set_unbacked_var_to_val, ShapeEnv.simplify (extendable). | ||
| 6 | +""" | ||
| 7 | + | ||
| 1 | import unittest | 8 | import unittest |
 | |||
| 2 | 9 | ||
| 10 | +import sympy | ||
| 3 | import torch | 11 | import torch |
| 4 | import torch_npu | 12 | import torch_npu |
| 13 | +from torch._subclasses.fake_tensor import FakeTensorMode | ||
| 5 | from torch.fx import Graph | 14 | from torch.fx import Graph |
| 6 | from torch.fx.experimental.symbolic_shapes import ( | 15 | from torch.fx.experimental.symbolic_shapes import ( |
| 7 | ShapeEnv, | 16 | ShapeEnv, |
| 17 | + Source, | ||
| 8 | is_accessor_node, | 18 | is_accessor_node, |
| 9 | is_concrete_bool, | 19 | is_concrete_bool, |
| 10 | is_concrete_float, | 20 | is_concrete_float, |
| @@ -14,6 +24,37 @@ from torch.fx.experimental.symbolic_shapes import ( | |||
| 14 | from torch.testing._internal.common_utils import TestCase, run_tests | 24 | from torch.testing._internal.common_utils import TestCase, run_tests |
| 15 | 25 | ||
| 16 | 26 | ||
| 27 | +def _shape_env_has(name: str) -> bool: | ||
| 28 | + """Return True if ShapeEnv exports a callable method with the given name.""" | ||
| 29 | + return callable(getattr(ShapeEnv, name, None)) | ||
| 30 | + | ||
| 31 | + | ||
| 32 | +def _produce_guards_verbose_works() -> bool: | ||
| 33 | + """Probe produce_guards_verbose with bare Source(); skip when upstream rejects it.""" | ||
| 34 | + if not _shape_env_has("produce_guards_verbose"): | ||
| 35 | + return False | ||
| 36 | + try: | ||
| 37 | + env = ShapeEnv() | ||
| 38 | + fake_mode = FakeTensorMode(shape_env=env) | ||
| 39 | + placeholder = fake_mode.from_tensor(torch.randn(2, 2)) | ||
| 40 | + env.produce_guards_verbose([placeholder], [Source()]) | ||
| 41 | + return True | ||
| 42 | + except (NotImplementedError, TypeError): | ||
| 43 | + return False | ||
| 44 | + | ||
| 45 | + | ||
| 46 | +def _shape_env_set_unbacked_var_to_val_works() -> bool: | ||
| 47 | + """Probe set_unbacked_var_to_val(create_unbacked_symint(), val); skip on known breakage.""" | ||
| 48 | + if not _shape_env_has("set_unbacked_var_to_val"): | ||
| 49 | + return False | ||
| 50 | + try: | ||
| 51 | + env = ShapeEnv() | ||
| 52 | + env.set_unbacked_var_to_val(env.create_unbacked_symint(), 4) | ||
| 53 | + return True | ||
| 54 | + except TypeError: | ||
| 55 | + return False | ||
| 56 | + | ||
| 57 | + | ||
| 17 | class TestSymbolicShapes(TestCase): | 58 | class TestSymbolicShapes(TestCase): |
| 18 | 59 | ||
| 19 | def test_is_accessor_node_with_call_method(self): | 60 | def test_is_accessor_node_with_call_method(self): |
| @@ -62,5 +103,64 @@ class TestSymbolicShapes(TestCase): | |||
| 62 | self.assertTrue(is_symbolic(sym_bool)) | 103 | self.assertTrue(is_symbolic(sym_bool)) |
| 63 | 104 | ||
| 64 | 105 | ||
| 106 | +class TestShapeEnvNPU(TestCase): | ||
| 107 | + """Issue #1627: direct NPU unit tests for ShapeEnv guard / sympy methods (v2.7.1).""" | ||
| 108 | + | ||
| 109 | + def _shape_env_with_fake_placeholders(self, shape=(3, 4)): | ||
| 110 | + """Build a ShapeEnv with NPU FakeTensor placeholders for guard-generation tests.""" | ||
| 111 | + env = ShapeEnv() | ||
| 112 | + fake_mode = FakeTensorMode(shape_env=env) | ||
| 113 | + fake_tensor = fake_mode.from_tensor(torch.randn(*shape, device="npu")) | ||
| 114 | + return env, [fake_tensor] | ||
| 115 | + | ||
| 116 | + | ||
| 117 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_produce_guards_expression(self): | ||
| 118 | + """Verify produce_guards_expression returns a guard expression string on NPU.""" | ||
| 119 | + env, placeholders = self._shape_env_with_fake_placeholders() | ||
| 120 | + guards = env.produce_guards_expression(placeholders) | ||
| 121 | + self.assertIsInstance(guards, str) | ||
| 122 | + | ||
| 123 | + | ||
| 124 | + | ||
| 125 | + _produce_guards_verbose_works(), | ||
| 126 | + "torch.fx.experimental.symbolic_shapes.ShapeEnv.produce_guards_verbose " | ||
| 127 | + "cannot run with bare Source() on this PyTorch build", | ||
| 128 | + ) | ||
| 129 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_produce_guards_verbose(self): | ||
| 130 | + """Verify produce_guards_verbose with Source list on NPU FakeTensor placeholders.""" | ||
| 131 | + env, placeholders = self._shape_env_with_fake_placeholders() | ||
| 132 | + source = Source() | ||
| 133 | + sources = [source] * len(placeholders) | ||
| 134 | + guards = env.produce_guards_verbose(placeholders, sources) | ||
| 135 | + self.assertIsNotNone(guards) | ||
| 136 | + | ||
| 137 | + | ||
| 138 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_replace(self): | ||
| 139 | + """Verify replace is identity when no substitution rules are registered.""" | ||
| 140 | + env = ShapeEnv() | ||
| 141 | + a, b = sympy.symbols("a b") | ||
| 142 | + original_expr = a + b | ||
| 143 | + self.assertEqual(env.replace(original_expr), original_expr) | ||
| 144 | + | ||
| 145 | + | ||
| 146 | + | ||
| 147 | + _shape_env_set_unbacked_var_to_val_works(), | ||
| 148 | + "torch.fx.experimental.symbolic_shapes.ShapeEnv.set_unbacked_var_to_val " | ||
| 149 | + "is broken with create_unbacked_symint() on this PyTorch build", | ||
| 150 | + ) | ||
| 151 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_set_unbacked_var_to_val(self): | ||
| 152 | + """Verify set_unbacked_var_to_val binds a concrete value to an unbacked SymInt.""" | ||
| 153 | + env = ShapeEnv() | ||
| 154 | + unbacked_sym = env.create_unbacked_symint() | ||
| 155 | + env.set_unbacked_var_to_val(unbacked_sym, 4) | ||
| 156 | + | ||
| 157 | + | ||
| 158 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_simplify(self): | ||
| 159 | + """Verify simplify reduces sympy expressions inside the ShapeEnv context.""" | ||
| 160 | + env = ShapeEnv() | ||
| 161 | + a, b = sympy.symbols("a b") | ||
| 162 | + self.assertEqual(env.simplify((a + b) - b), a) | ||
| 163 | + | ||
| 164 | + | ||
| 65 | if __name__ == "__main__": | 165 | if __name__ == "__main__": |
| 66 | run_tests() | 166 | run_tests() |
加一下一下信息吧，把这个文件测试的 API 梳理一下
""" Add validation cases for torch.nn APIs on NPU: