test(fx): add NPU coverage for PropagateUnbackedSymInts and rebind

已合并

test(fx): add NPU coverage for PropagateUnbackedSymInts and rebind_unbacked #37265

lihaofei-2026创建于 5月31日从已删除 :symbolic-shapes-propagate-unbacked-symints-v2.9.0 合入到 Ascend/pytorch： v2.9.0

test(fx): add NPU coverage for PropagateUnbackedSymInts and rebind_unbacked #37265

已合并

lihaofei-2026创建于 5月31日

从已删除 :symbolic-shapes-propagate-unbacked-symints-v2.9.0合入到Ascend/pytorch v2.9.0

共 1 个文件变更+131-0

Atest/fx/test_fx_propagate_unbacked_symints.py+131-0

@@ -0,0 +1,131 @@
	1	+"""
	2	+Add validation cases for torch.fx.experimental.symbolic_shapes.PropagateUnbackedSymInts on NPU.
	3	+
	4	+1. PyTorch community lacks sufficient and direct API validations for PropagateUnbackedSymInts,
	5	+ so this file is added.
	6	+2. This file validates PropagateUnbackedSymInts.run, PropagateUnbackedSymInts.run_node,
	7	+ PropagateUnbackedSymInts.placeholder, PropagateUnbackedSymInts.output, and rebind_unbacked.
	8	+"""
	9	+
	10	+import torch
	11	+from torch.testing._internal.common_utils import TestCase, run_tests
	12	+from torch._dynamo.utils import detect_fake_mode
	13	+from torch.fx.experimental.symbolic_shapes import PropagateUnbackedSymInts, rebind_unbacked
	14	+
	15	+device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
	16	+torch.zeros(3, 4).to(device_type)
	17	+
	18	+class TestPropagateUnbackedSymInts(TestCase):
	19	+
	20	+ def test_propagate_unbacked_symints_run(self):
	21	+ """Test PropagateUnbackedSymInts.run with NPU tensor."""
	22	+
	23	+ class M(torch.nn.Module):
	24	+ def forward(self, x: torch.Tensor):
	25	+ return torch.nonzero(x)
	26	+
	27	+ inp = (torch.tensor([1, 0, 1, 0]).to(device_type),)
	28	+ gm = torch.export.export(M(), inp, strict=True).module()
	29	+ fake_inputs = [
	30	+ node.meta.get("val") for node in gm.graph.nodes if node.op == "placeholder"
	31	+ ]
	32	+ fake_mode = detect_fake_mode(fake_inputs)
	33	+ with fake_mode:
	34	+ result = PropagateUnbackedSymInts(gm).run(*fake_inputs)
	35	+ self.assertIsNotNone(result)
	36	+
	37	+ def test_propagate_unbacked_symints_run_node(self):
	38	+ """Test PropagateUnbackedSymInts.run_node with NPU tensor."""
	39	+
	40	+ class RunNodeCapturingInterpreter(PropagateUnbackedSymInts):
	41	+ def __init__(self, args, *kwargs):
	42	+ super().__init__(args, *kwargs)
	43	+ self.captured_results = {}
	44	+
	45	+ def run_node(self, n):
	46	+ result = super().run_node(n)
	47	+ self.captured_results[n] = result
	48	+ return result
	49	+
	50	+ class M(torch.nn.Module):
	51	+ def forward(self, x: torch.Tensor):
	52	+ return torch.nonzero(x)
	53	+
	54	+ inp = (torch.tensor([1, 0, 1, 0]).to(device_type),)
	55	+ gm = torch.export.export(M(), inp, strict=True).module()
	56	+ fake_inputs = [
	57	+ node.meta.get("val") for node in gm.graph.nodes if node.op == "placeholder"
	58	+ ]
	59	+ fake_mode = detect_fake_mode(fake_inputs)
	60	+ with fake_mode:
	61	+ interpreter = RunNodeCapturingInterpreter(gm)
	62	+ result = interpreter.run(*fake_inputs)
	63	+ self.assertIsNotNone(result)
	64	+ for node in gm.graph.nodes:
	65	+ if node.op == "call_function":
	66	+ self.assertIn(node, interpreter.captured_results)
	67	+
	68	+ def test_propagate_unbacked_symints_placeholder(self):
	69	+ """Test PropagateUnbackedSymInts.placeholder with NPU tensor."""
	70	+
	71	+ class M(torch.nn.Module):
	72	+ def forward(self, x: torch.Tensor):
	73	+ return torch.nonzero(x)
	74	+
	75	+ inp = (torch.tensor([1, 0, 1, 0]).to(device_type),)
	76	+ gm = torch.export.export(M(), inp, strict=True).module()
	77	+ fake_inputs = [
	78	+ node.meta.get("val") for node in gm.graph.nodes if node.op == "placeholder"
	79	+ ]
	80	+ fake_mode = detect_fake_mode(fake_inputs)
	81	+ with fake_mode:
	82	+ interpreter = PropagateUnbackedSymInts(gm)
	83	+ interpreter.args_iter = iter(fake_inputs)
	84	+ for node in gm.graph.nodes:
	85	+ if node.op == "placeholder":
	86	+ result = interpreter.placeholder(node.target, node.args, node.kwargs)
	87	+ self.assertIsNotNone(result)
	88	+
	89	+ def test_propagate_unbacked_symints_output(self):
	90	+ """Test PropagateUnbackedSymInts.output with NPU tensor."""
	91	+
	92	+ class M(torch.nn.Module):
	93	+ def forward(self, x: torch.Tensor):
	94	+ return torch.nonzero(x)
	95	+
	96	+ inp = (torch.tensor([1, 0, 1, 0]).to(device_type),)
	97	+ gm = torch.export.export(M(), inp, strict=True).module()
	98	+ fake_inputs = [
	99	+ node.meta.get("val") for node in gm.graph.nodes if node.op == "placeholder"
	100	+ ]
	101	+ fake_mode = detect_fake_mode(fake_inputs)
	102	+ with fake_mode:
	103	+ interpreter = PropagateUnbackedSymInts(gm)
	104	+ interpreter.run(*fake_inputs)
	105	+ for node in gm.graph.nodes:
	106	+ if node.op == "output":
	107	+ result = interpreter.output(node.target, node.args, node.kwargs)
	108	+ self.assertIsNotNone(result)
	109	+
	110	+ def test_rebind_unbacked(self):
	111	+ """Test rebind_unbacked with NPU tensor."""
	112	+
	113	+ class M(torch.nn.Module):
	114	+ def forward(self, x: torch.Tensor):
	115	+ return torch.nonzero(x)
	116	+
	117	+ inp = (torch.tensor([1, 0, 1, 0]).to(device_type),)
	118	+ gm = torch.export.export(M(), inp, strict=True).module()
	119	+ fake_inputs = [
	120	+ node.meta.get("val") for node in gm.graph.nodes if node.op == "placeholder"
	121	+ ]
	122	+ fake_mode = detect_fake_mode(fake_inputs)
	123	+ shape_prop_gm = torch.fx.passes.shape_prop.ShapeProp(
	124	+ gm=gm, fake_mode=fake_mode
	125	+ )
	126	+ shape_prop_gm.propagate(*fake_inputs)
	127	+ self.assertEqual(len(fake_mode.shape_env.pending_fresh_unbacked_symbols), 0)
	128	+
	129	+
	130	+if __name__ == "__main__":
	131	+ run_tests()