test(fx): add NPU tests for symbolic_shapes.lru_cache and PropagateUnbackedSymInts-pytorch-AtomGit

已合并

test(fx): add NPU tests for symbolic_shapes.lru_cache and PropagateUnbackedSymInts #38674

Nokstella创建于 6月16日从wangshiqi-2026/pytorch： test-fx-symbolic-shapes-unbacked-symints-v2.12.0 合入到 Ascend/pytorch： v2.12.0

test(fx): add NPU tests for symbolic_shapes.lru_cache and PropagateUnbackedSymInts #38674

已合并

Nokstella创建于 6月16日

从wangshiqi-2026/pytorch test-fx-symbolic-shapes-unbacked-symints-v2.12.0合入到Ascend/pytorch v2.12.0

共 2 个文件变更+130-3

Mtest/fx/test_fx_propagate_unbacked_symints.py+92-2

	@@ -1,16 +1,39 @@
		1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd
		2	+# All rights reserved.
		3	+#
		4	+# Licensed under the BSD 3-Clause License (the "License");
		5	+# you may not use this file except in compliance with the License.
		6	+# You may obtain a copy of the License at
		7	+#
		8	+# https://opensource.org/licenses/BSD-3-Clause
		9	+#
		10	+# Unless required by applicable law or agreed to in writing, software
		11	+# distributed under the License is distributed on an "AS IS" BASIS,
		12	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
		13	+# See the License for the specific language governing permissions and
		14	+# limitations under the License.
		15	+
1	"""	16	"""
2	Add validation cases for torch.fx.experimental.symbolic_shapes.PropagateUnbackedSymInts on NPU.	17	Add validation cases for torch.fx.experimental.symbolic_shapes.PropagateUnbackedSymInts on NPU.
3		18
4	1. PyTorch community lacks sufficient and direct API validations for PropagateUnbackedSymInts,	19	1. PyTorch community lacks sufficient and direct API validations for PropagateUnbackedSymInts,
5	so this file is added.	20	so this file is added.
6	2. This file validates PropagateUnbackedSymInts.run, PropagateUnbackedSymInts.run_node,	21	2. This file validates PropagateUnbackedSymInts.run, PropagateUnbackedSymInts.run_node,
7	- PropagateUnbackedSymInts.placeholder, PropagateUnbackedSymInts.output, and rebind_unbacked.	22	+ PropagateUnbackedSymInts.placeholder, PropagateUnbackedSymInts.output,
		23	+ PropagateUnbackedSymInts.boxed_run, PropagateUnbackedSymInts.call_function,
		24	+ PropagateUnbackedSymInts.call_method, and rebind_unbacked.
8	"""	25	"""
9		26
10	import torch	27	import torch
11	from torch.testing._internal.common_utils import TestCase, run_tests	28	from torch.testing._internal.common_utils import TestCase, run_tests
12	from torch._dynamo.utils import detect_fake_mode	29	from torch._dynamo.utils import detect_fake_mode
13	-from torch.fx.experimental.symbolic_shapes import PropagateUnbackedSymInts, rebind_unbacked	30	+from torch.fx import Interpreter, symbolic_trace
		31	+from torch.fx.experimental.symbolic_shapes import PropagateUnbackedSymInts
		32	+from torch_npu.utils._dynamo import _dynamo_register_interface_for_device
		33	+
		34	+# Ensure NPU Dynamo device interface is registered before torch.export(strict=True).
		35	+# has_triton() may query "npu" before lazy inductor/init registration runs.
		36	+_dynamo_register_interface_for_device()
14		37
15	device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"	38	device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
16	torch.zeros(3, 4).to(device_type)	39	torch.zeros(3, 4).to(device_type)
	@@ -127,6 +150,73 @@ class TestPropagateUnbackedSymInts(TestCase):
127	shape_prop_gm.propagate(*fake_inputs)	150	shape_prop_gm.propagate(*fake_inputs)
128	self.assertEqual(len(fake_mode.shape_env.pending_fresh_unbacked_symbols), 0)	151	self.assertEqual(len(fake_mode.shape_env.pending_fresh_unbacked_symbols), 0)
129		152
		153	+ def test_propagate_unbacked_symints_boxed_run(self):
		154	+ """Test PropagateUnbackedSymInts.boxed_run with NPU tensor."""
		155	+
		156	+ class M(torch.nn.Module):
		157	+ def forward(self, x: torch.Tensor):
		158	+ return torch.nonzero(x)
		159	+
		160	+ inp = (torch.tensor([1, 0, 1, 0]).to(device_type),)
		161	+ gm = torch.export.export(M(), inp, strict=True).module()
		162	+ fake_inputs = [
		163	+ node.meta.get("val") for node in gm.graph.nodes if node.op == "placeholder"
		164	+ ]
		165	+ fake_mode = detect_fake_mode(fake_inputs)
		166	+ with fake_mode:
		167	+ interpreter = PropagateUnbackedSymInts(gm)
		168	+ args_list = list(fake_inputs)
		169	+ result = interpreter.boxed_run(args_list)
		170	+ self.assertIsNotNone(result)
		171	+
		172	+ def test_propagate_unbacked_symints_call_function(self):
		173	+ """Test PropagateUnbackedSymInts.call_function with NPU tensor."""
		174	+ self.assertIs(
		175	+ PropagateUnbackedSymInts.call_function,
		176	+ Interpreter.call_function,
		177	+ )
		178	+
		179	+ class M(torch.nn.Module):
		180	+ def forward(self, x: torch.Tensor):
		181	+ return torch.add(x, x)
		182	+
		183	+ gm = symbolic_trace(M())
		184	+ placeholder = next(node for node in gm.graph.nodes if node.op == "placeholder")
		185	+ call_function = next(
		186	+ node for node in gm.graph.nodes if node.op == "call_function"
		187	+ )
		188	+
		189	+ interpreter = PropagateUnbackedSymInts(gm)
		190	+ interpreter.env[placeholder] = torch.ones(2, 3).to(device_type)
		191	+
		192	+ args, kwargs = interpreter.fetch_args_kwargs_from_env(call_function)
		193	+ result = interpreter.call_function(call_function.target, args, kwargs)
		194	+ self.assertEqual(tuple(result.shape), (2, 3))
		195	+ self.assertEqual(result.device.type, device_type)
		196	+
		197	+ def test_propagate_unbacked_symints_call_method(self):
		198	+ """Test PropagateUnbackedSymInts.call_method with NPU tensor."""
		199	+ self.assertIs(
		200	+ PropagateUnbackedSymInts.call_method,
		201	+ Interpreter.call_method,
		202	+ )
		203	+
		204	+ class M(torch.nn.Module):
		205	+ def forward(self, x: torch.Tensor):
		206	+ return x.relu()
		207	+
		208	+ gm = symbolic_trace(M())
		209	+ placeholder = next(node for node in gm.graph.nodes if node.op == "placeholder")
		210	+ call_method = next(node for node in gm.graph.nodes if node.op == "call_method")
		211	+
		212	+ interpreter = PropagateUnbackedSymInts(gm)
		213	+ interpreter.env[placeholder] = torch.randn(2, 3).to(device_type)
		214	+
		215	+ args, kwargs = interpreter.fetch_args_kwargs_from_env(call_method)
		216	+ result = interpreter.call_method(call_method.target, args, kwargs)
		217	+ self.assertEqual(tuple(result.shape), (2, 3))
		218	+ self.assertEqual(result.device.type, device_type)
		219	+
130		220
131	if __name__ == "__main__":	221	if __name__ == "__main__":
132	run_tests()	222	run_tests()

Mtest/fx/test_symbolic_shapes_api.py+38-1

	@@ -1,3 +1,18 @@
		1	+# Copyright (c) 2026 Huawei Technologies Co., Ltd
		2	+# All rights reserved.
		3	+#
		4	+# Licensed under the BSD 3-Clause License (the "License");
		5	+# you may not use this file except in compliance with the License.
		6	+# You may obtain a copy of the License at
		7	+#
		8	+# https://opensource.org/licenses/BSD-3-Clause
		9	+#
		10	+# Unless required by applicable law or agreed to in writing, software
		11	+# distributed under the License is distributed on an "AS IS" BASIS,
		12	+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
		13	+# See the License for the specific language governing permissions and
		14	+# limitations under the License.
		15	+
1	# Owner(s): ["module: fx"]	16	# Owner(s): ["module: fx"]
2	"""	17	"""
3	Add validation cases for torch.fx symbolic_shapes related APIs on NPU:	18	Add validation cases for torch.fx symbolic_shapes related APIs on NPU:
	@@ -19,6 +34,7 @@ Add validation cases for torch.fx symbolic_shapes related APIs on NPU:
19	- torch.fx.experimental.symbolic_shapes.StatefulSymbolicContext	34	- torch.fx.experimental.symbolic_shapes.StatefulSymbolicContext
20	- torch.fx.experimental.symbolic_shapes.StatelessSymbolicContext	35	- torch.fx.experimental.symbolic_shapes.StatelessSymbolicContext
21	- symbolic_shapes._lru_cache	36	- symbolic_shapes._lru_cache
		37	+ - symbolic_shapes.lru_cache
22	- symbolic_shapes.CallMethodKey	38	- symbolic_shapes.CallMethodKey
23	- symbolic_shapes.CallMethodKey.get	39	- symbolic_shapes.CallMethodKey.get
24	- symbolic_shapes.canonicalize_bool_expr	40	- symbolic_shapes.canonicalize_bool_expr
	@@ -35,7 +51,6 @@ import inspect
35	import sympy	51	import sympy
36	import torch	52	import torch
37		53
38	-import torch_npu
39	from torch._dynamo.source import ConstantSource	54	from torch._dynamo.source import ConstantSource
40	from torch.export import Dim	55	from torch.export import Dim
41	from torch.fx.experimental import symbolic_shapes	56	from torch.fx.experimental import symbolic_shapes
	@@ -45,6 +60,7 @@ from torch.utils._sympy.value_ranges import ValueRanges
45		60
46		61
47	device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"	62	device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
		63	+torch.zeros(3, 4).to(device_type)
48		64
49		65
50	class TestSymbolicShapesAPI(TestCase):	66	class TestSymbolicShapesAPI(TestCase):
	@@ -354,6 +370,27 @@ class TestSymbolicShapesTargetApiNPU(TestCase):
354	self.assertEqual(limited.cache_info().maxsize, 1)	370	self.assertEqual(limited.cache_info().maxsize, 1)
355	self.assertEqual(limited_calls["count"], 3)	371	self.assertEqual(limited_calls["count"], 3)
356		372
		373	+ def test_public_lru_cache(self):
		374	+ calls = {"count": 0}
		375	+
		376	+ @symbolic_shapes.lru_cache(128)
		377	+ def cached(value):
		378	+ calls["count"] += 1
		379	+ return value + 1
		380	+
		381	+ self.assertEqual(cached(3), 4)
		382	+ self.assertEqual(cached(3), 4)
		383	+ self.assertEqual(calls["count"], 1)
		384	+
		385	+ cache_info = cached.cumulative_cache_info()
		386	+ self.assertEqual(cache_info.hits, 1)
		387	+ self.assertEqual(cache_info.misses, 1)
		388	+
		389	+ cached.cache_clear()
		390	+ self.assertEqual(cached.cache_info().currsize, 0)
		391	+ self.assertEqual(cached(3), 4)
		392	+ self.assertEqual(calls["count"], 2)
		393	+
357	def test_call_method_key_get_on_npu_tensor(self):	394	def test_call_method_key_get_on_npu_tensor(self):
358	tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4).to(device_type)	395	tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4).to(device_type)
359	size_key = symbolic_shapes.CallMethodKey("size")	396	size_key = symbolic_shapes.CallMethodKey("size")