已合并
test: 补齐torch.fx.experimental.symbolic_shapes NPU 适配验证与统一运行脚本 #36866
cuiyunhao-2026创建于 5月27日
test: 补齐torch.fx.experimental.symbolic_shapes NPU 适配验证与统一运行脚本 #36866
已合并
共 1 个文件变更+84-1
| @@ -9,8 +9,8 @@ import unittest | |||
| 9 | 9 | ||
| 10 | import sympy | 10 | import sympy |
| 11 | import torch | 11 | import torch |
| 12 | -import torch_npu | ||
| 13 | from torch import nn | 12 | from torch import nn |
| 13 | +from torch._guards import ShapeGuard, SLoc | ||
| 14 | from torch._subclasses.fake_tensor import FakeTensorMode | 14 | from torch._subclasses.fake_tensor import FakeTensorMode |
| 15 | from torch.fx import Graph, Interpreter, symbolic_trace | 15 | from torch.fx import Graph, Interpreter, symbolic_trace |
| 16 | from torch.fx.experimental.symbolic_shapes import ( | 16 | from torch.fx.experimental.symbolic_shapes import ( |
| @@ -305,6 +305,89 @@ class TestShapeEnvNPU(TestCase): | |||
| 305 | a, b = sympy.symbols("a b") | 305 | a, b = sympy.symbols("a b") |
| 306 | self.assertEqual(env.simplify((a + b) - b), a) | 306 | self.assertEqual(env.simplify((a + b) - b), a) |
| 307 | 307 | ||
| 308 | + @unittest.skipUnless(torch.npu.is_available(), "requires npu") | ||
| 309 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_format_guards(self): | ||
| 310 | + """Verify format_guards returns formatted guard strings.""" | ||
| 311 | + se = ShapeEnv() | ||
| 312 | + self.assertEqual(se.format_guards(), "") | ||
| 313 | + | ||
| 314 | + s0 = sympy.Symbol('s0', integer=True, positive=True) | ||
| 315 | + s1 = sympy.Symbol('s1', integer=True, positive=True) | ||
| 316 | + sloc = SLoc("framework_loc", "user_loc") | ||
| 317 | + se.guards.append(ShapeGuard( | ||
| 318 | + sympy.Lt(s0, s1, evaluate=False), sloc, False)) | ||
| 319 | + se.guards.append(ShapeGuard( | ||
| 320 | + sympy.Ge(s0, sympy.Integer(1), evaluate=False), sloc, False)) | ||
| 321 | + | ||
| 322 | + result = se.format_guards() | ||
| 323 | + self.assertIn("s0 < s1", result) | ||
| 324 | + self.assertIn("s0 >= 1", result) | ||
| 325 | + | ||
| 326 | + result_verbose = se.format_guards(verbose=True) | ||
| 327 | + self.assertIn("s0 < s1", result_verbose) | ||
| 328 | + self.assertIn("user_loc", result_verbose) | ||
| 329 | + | ||
| 330 | + @unittest.skipUnless(torch.npu.is_available(), "requires npu") | ||
| 331 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_freeze(self): | ||
| 332 | + """Verify freeze toggles the frozen state of ShapeEnv.""" | ||
| 333 | + se = ShapeEnv() | ||
| 334 | + self.assertFalse(se.frozen) | ||
| 335 | + se.freeze() | ||
| 336 | + self.assertTrue(se.frozen) | ||
| 337 | + | ||
| 338 | + @unittest.skipUnless(torch.npu.is_available(), "requires npu") | ||
| 339 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_freeze_runtime_asserts(self): | ||
| 340 | + """Verify freeze_runtime_asserts toggles runtime_asserts_frozen.""" | ||
| 341 | + se = ShapeEnv() | ||
| 342 | + self.assertFalse(se.runtime_asserts_frozen) | ||
| 343 | + se.freeze_runtime_asserts() | ||
| 344 | + self.assertTrue(se.runtime_asserts_frozen) | ||
| 345 | + | ||
| 346 | + @unittest.skipUnless(torch.npu.is_available(), "requires npu") | ||
| 347 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_get_axioms(self): | ||
| 348 | + """Verify get_axioms returns axioms tuple, optionally filtered by symbols.""" | ||
| 349 | + se = ShapeEnv() | ||
| 350 | + self.assertIsInstance(se.get_axioms(), tuple) | ||
| 351 | + | ||
| 352 | + s0 = sympy.Symbol('s0', integer=True, positive=True) | ||
| 353 | + s1 = sympy.Symbol('s1', integer=True, positive=True) | ||
| 354 | + sloc = SLoc("framework_loc", "user_loc") | ||
| 355 | + se.guards.append(ShapeGuard( | ||
| 356 | + sympy.Lt(s0, s1, evaluate=False), sloc, False)) | ||
| 357 | + | ||
| 358 | + axioms = se.get_axioms(symbols=(s0,)) | ||
| 359 | + self.assertIn(sympy.Lt(s0, s1, evaluate=False), axioms) | ||
| 360 | + | ||
| 361 | + @unittest.skipUnless(torch.npu.is_available(), "requires npu") | ||
| 362 | + def test_torch_fx_experimental_symbolic_shapes_ShapeEnv_get_implications(self): | ||
| 363 | + """Verify get_implications returns implications for Eq/Lt/Ne/Le expressions.""" | ||
| 364 | + se = ShapeEnv() | ||
| 365 | + s0 = sympy.Symbol('s0', integer=True, positive=True) | ||
| 366 | + s1 = sympy.Symbol('s1', integer=True, positive=True) | ||
| 367 | + | ||
| 368 | + # Eq implies equality | ||
| 369 | + impl_eq = dict(se.get_implications( | ||
| 370 | + sympy.Eq(s0, s1, evaluate=False))) | ||
| 371 | + self.assertIn(sympy.Eq(s0, s1, evaluate=False), impl_eq) | ||
| 372 | + | ||
| 373 | + # Lt implies Le and Ne | ||
| 374 | + impl_lt = dict(se.get_implications( | ||
| 375 | + sympy.Lt(s0, s1, evaluate=False))) | ||
| 376 | + self.assertIn(sympy.Lt(s0, s1, evaluate=False), impl_lt) | ||
| 377 | + self.assertIn(sympy.Le(s0, s1, evaluate=False), impl_lt) | ||
| 378 | + self.assertIn(sympy.Ne(s0, s1, evaluate=False), impl_lt) | ||
| 379 | + | ||
| 380 | + # Ne | ||
| 381 | + impl_ne = dict(se.get_implications( | ||
| 382 | + sympy.Ne(s0, s1, evaluate=False))) | ||
| 383 | + self.assertIn(sympy.Ne(s0, s1, evaluate=False), impl_ne) | ||
| 384 | + | ||
| 385 | + # Le implies Lt(a, b+1) | ||
| 386 | + impl_le = dict(se.get_implications( | ||
| 387 | + sympy.Le(s0, s1, evaluate=False))) | ||
| 388 | + self.assertIn(sympy.Le(s0, s1, evaluate=False), impl_le) | ||
| 389 | + self.assertIn(sympy.Lt(s0, s1 + 1, evaluate=False), impl_le) | ||
| 390 | + | ||
| 308 | 391 | ||
| 309 | if __name__ == "__main__": | 392 | if __name__ == "__main__": |
| 310 | run_tests() | 393 | run_tests() |