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