已合并
[inductor] add pass ut #43217
lvjiangdong创建于 26 天前
[inductor] add pass ut #43217
已合并
lvjiangdong创建于 26 天前
4 个文件变更+810-0
Atest/_inductor/test_fold_add_pass.py+204-0
@@ -0,0 +1,204 @@
1+import torch
2+import torch.fx as fx
3+from torch.fx.passes.shape_prop import ShapeProp
4+from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests
5+from testutils import TestUtils
6+import torch_npu
7+import torch_npu._inductor
8+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_four_op_pass
9+ 
10+ 
11+class FoldAddModel(torch.nn.Module):
12+ def forward(self, first_element):
13+ add = torch.ops.aten.add.Tensor(first_element, 0)
14+ add_output = torch.ops.aten.relu.default(add)
15+ return add_output
16+ 
17+ 
18+class TestFoldAddPass(TestUtils):
19+ def op_calc(self, first_element):
20+ add = torch.add(first_element, 0)
21+ add_output = torch.relu(add)
22+ return add_output
23+ 
24+ 
25+ @parametrize('shape', [(1, 2, 3)])
26+ @parametrize('dtype', ['float32'])
27+ def test_compile_cases(self, shape, dtype):
28+ first_element = self._generate_tensor(shape, dtype)
29+ std_result = self.op_calc(first_element)
30+ with torch.no_grad():
31+ compiled_op_calc = torch.compile(self.op_calc, backend="inductor")
32+ inductor_result = compiled_op_calc(first_element)
33+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
34+ 
35+ 
36+ @parametrize('shape', [(1, 2, 3)])
37+ @parametrize('dtype', ['float32'])
38+ def test_ut_cases(self, shape, dtype):
39+ first_element = self._generate_tensor(shape, dtype)
40+ model = FoldAddModel()
41+ graph_module = fx.symbolic_trace(model)
42+ ShapeProp(graph_module).propagate(first_element)
43+ 
44+ # 应用优化 Pass
45+ fold_four_op_pass(graph_module.graph)
46+ graph_module.recompile()
47+ 
48+ # 验证输出是否一致
49+ std_result = model(first_element)
50+ inductor_result = graph_module(first_element)
51+ 
52+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
53+ 
54+ 
55+ @parametrize('shape', [(4, 8)])
56+ @parametrize('dtype', ['float32', 'float16'])
57+ def test_add_zero_left_fold(self, shape, dtype):
58+ """ 0 + x 应该折叠成 x """
59+ class M(torch.nn.Module):
60+ def forward(self, x):
61+ y = torch.ops.aten.add.Tensor(0, x) # 左边是 0 (scalar)
62+ return torch.relu(y + 1.0)
63+ 
64+ t = self._generate_tensor(shape, dtype)
65+ m = M()
66+ gm = fx.symbolic_trace(m)
67+ ShapeProp(gm).propagate(t)
68+ 
69+ fold_four_op_pass(gm.graph)
70+ gm.recompile()
71+ 
72+ # 检查 pass 是否删除了 add 节点
73+ self.assertFalse(
74+ any(n.target in (torch.add, torch.ops.aten.add.Tensor, torch.ops.aten.add.Scalar)
75+ for n in gm.graph.nodes),
76+ "add + 0 should be folded away"
77+ )
atomgit-bot
atomgit-botatomgit-bot26 天前

🟠 High Priority

测试 test_add_zero_left_fold 在第 73-77 行断言图中不应存在任何 add 节点:self.assertFalse(any(n.target in (torch.add, torch.ops.aten.add.Tensor, torch.ops.aten.add.Scalar) for n in gm.graph.nodes))

但模型定义(第 59-62 行)中:y = torch.ops.aten.add.Tensor(0, x) 折叠后 y→x,而 torch.relu(y + 1.0) 中的 y + 1.0 仍会生成一个新的 add 节点(add.Tensor(x, 1.0)add.Scalar(x, 1.0),因为 1.0 ≠ 0,不会被 fold_four_op_pass 折叠)。因此图中必然残留一个 add 节点,assertFalse 将失败。

触发条件fold_four_op_pass 正确折叠 0 + x 后,x + 1.0 的 add 节点仍然存在。

修复方向:将断言改为验证折叠后的图中仅剩余 x + 1.0 这一个 add 节点(即 len(add_nodes) == 1),或修改模型使 y + 1.0 不产生 add 节点(例如改为 torch.mul(y, 2.0))。

建议:将断言从 assertFalse(零 add 节点)改为 assertEqual(len(add_nodes), 1),只验证 0 + x 被折叠,而 x + 1.0 正常保留。

likedislike
78+ 
79+ self.assertEqual(m(t), gm(t))
80+ 
81+ 
82+ @parametrize('shape', [(5, 5)])
83+ def test_add_scalar_zero_right_fold(self, shape):
84+ """ x + 0.0 (scalar) 应该折叠 """
85+ class M(torch.nn.Module):
86+ def forward(self, x):
87+ y = torch.add(x, 0.0) # 使用 python torch.add → 应转为 add.Scalar
88+ return y * 2.0
89+ 
90+ t = torch.randn(shape)
91+ m = M()
92+ gm = fx.symbolic_trace(m)
93+ ShapeProp(gm).propagate(t)
94+ 
95+ fold_four_op_pass(gm.graph)
96+ gm.recompile()
97+ 
98+ self.assertFalse(
99+ any(n.target in (torch.add, torch.ops.aten.add.Tensor, torch.ops.aten.add.Scalar)
100+ for n in gm.graph.nodes)
101+ )
102+ torch.testing.assert_close(m(t), gm(t), atol=1e-5, rtol=1e-5)
103+ 
104+ 
105+ @parametrize('shape', [(2, 3, 4)])
106+ def test_add_int_zero_right_fold(self, shape):
107+ """ x + 0 (int literal) """
108+ class M(torch.nn.Module):
109+ def forward(self, x):
110+ return torch.ops.aten.add.Tensor(x, 0) + x
111+ 
112+ t = self._generate_tensor(shape, 'float32')
113+ m = M()
114+ gm = fx.symbolic_trace(m)
115+ ShapeProp(gm).propagate(t)
116+ 
117+ fold_four_op_pass(gm.graph)
118+ gm.recompile()
119+ 
120+ self.assertEqual(m(t), gm(t))
121+ 
122+ 
123+ @parametrize('shape', [(10,)])
124+ def test_add_zero_not_fold_nonzero(self, shape):
125+ """ x + 1.5 不应该折叠 """
126+ class M(torch.nn.Module):
127+ def forward(self, x):
128+ y = torch.add(x, 1.5)
129+ return torch.sigmoid(y)
130+ 
131+ t = self._generate_tensor(shape, 'float32')
132+ m = M()
133+ gm = fx.symbolic_trace(m)
134+ ShapeProp(gm).propagate(t)
135+ 
136+ before_nodes = len([n for n in gm.graph.nodes if n.op == 'call_function'])
137+ 
138+ fold_four_op_pass(gm.graph)
139+ gm.recompile()
140+ 
141+ after_nodes = len([n for n in gm.graph.nodes if n.op == 'call_function'])
142+ self.assertEqual(before_nodes, after_nodes, "不应折叠非零常量")
143+ 
144+ 
145+ @parametrize('shape', [(3, 4)])
146+ def test_add_zero_chain_fold(self, shape):
147+ """ 多个连续的 +0 折叠 """
148+ class M(torch.nn.Module):
149+ def forward(self, x):
150+ a = torch.ops.aten.add.Tensor(x, 0)
151+ b = torch.ops.aten.add.Tensor(a, 0.0)
152+ c = torch.ops.aten.add.Tensor(b, torch.zeros_like(x))
153+ return torch.ops.aten.relu.default(c)
154+ 
155+ t = self._generate_tensor(shape, 'float32')
156+ m = M()
157+ gm = fx.symbolic_trace(m)
158+ ShapeProp(gm).propagate(t)
159+ fold_four_op_pass(gm.graph)
160+ gm.recompile()
161+ 
162+ add_nodes = [n for n in gm.graph.nodes if 'add' in str(n.target)]
163+ self.assertEqual(len(add_nodes), 0, "所有 +0 都应该被折叠掉")
164+ 
165+ 
166+ def test_no_add_op_no_change(self):
167+ """ 图中没有 add 节点,pass 应安全通过 """
168+ class M(torch.nn.Module):
169+ def forward(self, x):
170+ return torch.mul(x, 2.0) + torch.pow(x, 2)
171+ 
172+ t = torch.randn(8)
173+ gm = fx.symbolic_trace(M())
174+ ShapeProp(gm).propagate(t)
175+ 
176+ orig_graph_str = str(gm.graph)
177+ fold_four_op_pass(gm.graph)
178+ self.assertEqual(orig_graph_str, str(gm.graph), "无 add 节点不应修改图")
179+ 
180+ 
181+ @parametrize('shape', [(1, 128)])
182+ def test_add_zero_with_downstream_users(self, shape):
183+ """ +0 后有多个下游用户,都应正确替换 """
184+ class M(torch.nn.Module):
185+ def forward(self, x):
186+ y = torch.add(x, 0)
187+ return (y * 2, y + 3, torch.relu(y))
188+ 
189+ t = self._generate_tensor(shape, 'float32')
190+ m = M()
191+ gm = fx.symbolic_trace(m)
192+ ShapeProp(gm).propagate(t)
193+ 
194+ fold_four_op_pass(gm.graph)
195+ gm.recompile()
196+ 
197+ self.assertEqual(m(t), gm(t))
198+ 
199+ 
200+instantiate_parametrized_tests(TestFoldAddPass)
201+ 
202+ 
203+if __name__ == "__main__":
204+ run_tests()
Atest/_inductor/test_fold_cat_pass.py+204-0
@@ -0,0 +1,204 @@
1+import torch
2+import torch.fx as fx
3+from torch.fx.passes.shape_prop import ShapeProp
4+from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests
5+from testutils import TestUtils
6+import torch_npu
7+import torch_npu._inductor
8+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_cat
9+ 
10+ 
11+class FoldCatModel(torch.nn.Module):
12+ def forward(self, t1, t2, t3, t4, t5):
13+ cat1 = torch.ops.aten.cat.default([t1, t2], 1)
14+ cat2 = torch.ops.aten.cat.default([cat1, t3], 1)
15+ cat3 = torch.ops.aten.cat.default([cat2, t4], 1)
16+ cat4 = torch.ops.aten.cat.default([cat3, t5], 1)
17+ return cat4
18+ 
19+ 
20+class TestFoldCatPass(TestUtils):
21+ def op_calc(self, t1, t2, t3, t4, t5):
22+ cat1 = torch.cat([t1, t2], dim=1)
23+ cat2 = torch.cat([cat1, t3], dim=1)
24+ cat3 = torch.cat([cat2, t4], dim=1)
25+ cat4 = torch.cat([cat3, t5], dim=1)
26+ return cat4
27+ 
28+ 
29+ @parametrize('shape', [(2, 4)])
30+ @parametrize('dtype', ['float32'])
31+ def test_compile_cases(self, shape, dtype):
32+ t1 = self._generate_tensor(shape, dtype)
33+ t2 = self._generate_tensor(shape, dtype)
34+ t3 = self._generate_tensor(shape, dtype)
35+ t4 = self._generate_tensor(shape, dtype)
36+ t5 = self._generate_tensor(shape, dtype)
37+ std_result = self.op_calc(t1, t2, t3, t4, t5)
38+ with torch.no_grad():
39+ compiled_op_calc = torch.compile(self.op_calc, backend="inductor")
40+ inductor_result = compiled_op_calc(t1, t2, t3, t4, t5)
41+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
42+ 
43+ 
44+ @parametrize('shape', [(2, 4)])
45+ @parametrize('dtype', ['float32'])
46+ def test_ut_cases(self, shape, dtype):
47+ t1 = self._generate_tensor(shape, dtype)
48+ t2 = self._generate_tensor(shape, dtype)
49+ t3 = self._generate_tensor(shape, dtype)
50+ t4 = self._generate_tensor(shape, dtype)
51+ t5 = self._generate_tensor(shape, dtype)
52+ model = FoldCatModel()
53+ graph_module = fx.symbolic_trace(model)
54+ ShapeProp(graph_module).propagate(t1, t2, t3, t4, t5)
55+ 
56+ # 应用优化 Pass
57+ fold_cat(graph_module.graph)
58+ graph_module.recompile()
59+ 
60+ # 验证输出是否一致
61+ std_result = model(t1, t2, t3, t4, t5)
62+ inductor_result = graph_module(t1, t2, t3, t4, t5)
63+ 
64+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
65+ 
66+ 
67+ @parametrize('shape', [(4, 8, 16)])
68+ @parametrize('dtype', ['float32'])
69+ def test_deep_nested_same_axis_fold(self, shape, dtype):
70+ """多层嵌套同轴 cat,应完全折叠为单层 cat"""
71+ class M(torch.nn.Module):
72+ def forward(self, a, b, c, d, e):
73+ c1 = torch.ops.aten.cat.default([a, b], dim=1)
74+ c2 = torch.ops.aten.cat.default([c1, c], dim=1)
75+ c3 = torch.ops.aten.cat.default([c2, d], dim=1)
76+ c4 = torch.ops.aten.cat.default([c3, e], dim=1)
77+ return c4
78+ 
79+ tensors = [self._generate_tensor(shape, dtype) for _ in range(5)]
80+ model = M()
81+ gm = fx.symbolic_trace(model)
82+ ShapeProp(gm).propagate(*tensors)
83+ gm.graph.print_tabular()
84+ fold_cat(gm.graph)
85+ gm.graph.print_tabular()
86+ gm.recompile()
87+ 
88+ cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default]
89+ self.assertEqual(len(cat_nodes), 1, "应折叠为单层 cat")
90+ self.assertEqual(len(cat_nodes[0].args[0]), 5, "输入应包含所有 5 个原始 tensor")
91+ 
92+ self.assertEqual(model(*tensors), gm(*tensors))
93+ 
94+ 
95+ @parametrize('shape', [(3, 5)])
96+ @parametrize('dtype', ['float32'])
97+ def test_different_axis_no_fold(self, shape, dtype):
98+ """不同轴的 cat 不应折叠"""
99+ class M(torch.nn.Module):
100+ def forward(self, a, b, c):
101+ c1 = torch.ops.aten.cat.default([a, b], 0) # dim=0
102+ c2 = torch.ops.aten.cat.default([c1, c], 1) # dim=1 → 不同
103+ return c2
104+ 
105+ t1, t2 = [self._generate_tensor(shape, dtype) for _ in range(2)]
106+ shape1 = (shape[0] * 2, shape[1])
107+ t3 = self._generate_tensor(shape1, dtype)
108+ model = M()
109+ gm = fx.symbolic_trace(model)
110+ ShapeProp(gm).propagate(t1, t2, t3)
111+ 
112+ fold_cat(gm.graph)
113+ 
114+ cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default]
115+ self.assertEqual(len(cat_nodes), 2, "不同轴不应折叠")
116+ 
117+ 
118+ @parametrize('shape', [(2, 6)])
119+ @parametrize('dtype', ['float32'])
120+ def test_input_cat_has_multiple_users_no_fold(self, shape, dtype):
121+ """被多个下游使用的 cat 输入不应折叠"""
122+ class M(torch.nn.Module):
123+ def forward(self, a, b, c):
124+ inner = torch.ops.aten.cat.default([a, b], 1)
125+ out1 = inner + 1.0
126+ out2 = inner * 2.0
127+ out3 = torch.ops.aten.cat.default([inner, c], 1)
128+ return out1, out2, out3
129+ 
130+ t1, t2, t3 = [self._generate_tensor(shape, dtype) for _ in range(3)]
131+ model = M()
132+ gm = fx.symbolic_trace(model)
133+ ShapeProp(gm).propagate(t1, t2, t3)
134+ 
135+ fold_cat(gm.graph)
136+ 
137+ cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default]
138+ self.assertEqual(len(cat_nodes), 2, "有多个用户,不应折叠 inner cat")
139+ 
140+ 
141+ @parametrize('shape', [(4, 3, 5)])
142+ @parametrize('dtype', ['float32'])
143+ def test_last_dim_negative_and_positive_axis(self, shape, dtype):
144+ """最后一维使用 -1 和正数轴应视为相同"""
145+ class M(torch.nn.Module):
146+ def forward(self, a, b, c):
147+ inner = torch.ops.aten.cat.default([a, b], -1) # -1
148+ outer = torch.ops.aten.cat.default([inner, c], 2) # 2
149+ return outer
150+ 
151+ t1, t2, t3 = [self._generate_tensor(shape, dtype) for _ in range(3)]
152+ model = M()
153+ gm = fx.symbolic_trace(model)
154+ ShapeProp(gm).propagate(t1, t2, t3)
155+ fold_cat(gm.graph)
156+ cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default]
157+ self.assertEqual(len(cat_nodes), 1, "最后一维 -1 和 2 应折叠")
158+ self.assertEqual(cat_nodes[0].args[1], -1)
159+ 
160+ 
161+ @parametrize('shape', [(2, 7)])
162+ @parametrize('dtype', ['float32'])
163+ def test_mixed_cat_and_non_cat_inputs(self, shape, dtype):
164+ """cat 输入中混有非 cat 节点"""
165+ class M(torch.nn.Module):
166+ def forward(self, a, b, c):
167+ inner = torch.ops.aten.cat.default([a, b], 0)
168+ outer = torch.ops.aten.cat.default([inner, c * 1.5, torch.ops.aten.relu.default(c)], 0)
169+ return outer
170+ 
171+ t1, t2, t3 = [self._generate_tensor(shape, dtype) for _ in range(3)]
172+ model = M()
173+ gm = fx.symbolic_trace(model)
174+ ShapeProp(gm).propagate(t1, t2, t3)
175+ 
176+ fold_cat(gm.graph)
177+ 
178+ cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default]
179+ self.assertEqual(len(cat_nodes), 1, "混有非 cat 输入,应保留一层 cat")
180+ self.assertEqual(len(cat_nodes[0].args[0]), 4, "应包含 inner + c*1.5 + relu(c)")
181+ 
182+ 
183+ def test_single_input_cat_no_fold(self):
184+ """只有一个输入的 cat 不应折叠"""
185+ class M(torch.nn.Module):
186+ def forward(self, x):
187+ return torch.ops.aten.cat.default([x], dim=0)
188+ 
189+ t = self._generate_tensor((4, 5), 'float32')
190+ model = M()
191+ gm = fx.symbolic_trace(model)
192+ ShapeProp(gm).propagate(t)
193+ 
194+ fold_cat(gm.graph)
195+ 
196+ cat_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.cat.default]
197+ self.assertEqual(len(cat_nodes), 1, "单输入 cat 不应被移除")
198+ 
199+ 
200+instantiate_parametrized_tests(TestFoldCatPass)
201+ 
202+ 
203+if __name__ == "__main__":
204+ run_tests()
Atest/_inductor/test_fold_clone_pass.py+201-0
@@ -0,0 +1,201 @@
1+import torch
2+import torch.fx as fx
3+from torch.fx.passes.shape_prop import ShapeProp
4+from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests
5+from testutils import TestUtils
6+import torch_npu
7+import torch_npu._inductor
8+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_clone
9+ 
10+ 
11+class FoldCloneModel(torch.nn.Module):
12+ def forward(self, t1):
13+ clone_1 = torch.ops.aten.clone.default(t1)
14+ relu_1 = torch.ops.aten.relu.default(clone_1)
15+ return relu_1
16+ 
17+ 
18+class TestFoldClonePass(TestUtils):
19+ def op_calc(self, t1):
20+ clone_1 = torch.clone(t1)
21+ output = torch.relu(clone_1)
22+ return output
23+ 
24+ 
25+ @parametrize('shape', [(1, 2, 3)])
26+ @parametrize('dtype', ['float32'])
27+ def test_compile_cases(self, shape, dtype):
28+ t1 = self._generate_tensor(shape, dtype)
29+ std_result = self.op_calc(t1)
30+ with torch.no_grad():
31+ compiled_op_calc = torch.compile(self.op_calc, backend="inductor")
32+ inductor_result = compiled_op_calc(t1)
33+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
34+ 
35+ 
36+ @parametrize('shape', [(1, 2, 3)])
37+ @parametrize('dtype', ['float32'])
38+ def test_ut_cases(self, shape, dtype):
39+ t1 = self._generate_tensor(shape, dtype)
40+ model = FoldCloneModel()
41+ graph_module = fx.symbolic_trace(model)
42+ ShapeProp(graph_module).propagate(t1)
43+ 
44+ # 应用优化 Pass
45+ fold_clone(graph_module.graph)
46+ graph_module.recompile()
47+ 
48+ # 验证输出是否一致
49+ std_result = model(t1)
50+ inductor_result = graph_module(t1)
51+ 
52+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
53+ 
54+ 
55+ @parametrize('shape', [(4, 8)])
56+ @parametrize('dtype', ['float32'])
57+ def test_clone_same_memory_format_should_fold(self, shape, dtype):
58+ """memory_format 相同 → 应折叠删除 clone"""
59+ class M(torch.nn.Module):
60+ def forward(self, x):
61+ cloned = torch.ops.aten.clone.default(x) # 默认 memory_format 相同
62+ return cloned + 1.0
63+ 
64+ x = self._generate_tensor(shape, dtype)
65+ model = M()
66+ gm = fx.symbolic_trace(model)
67+ ShapeProp(gm).propagate(x)
68+ fold_clone(gm.graph)
69+ gm.recompile()
70+ 
71+ clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default]
72+ self.assertEqual(len(clone_nodes), 0, "相同 memory_format 的 clone 应被删除")
73+ 
74+ self.assertEqual(model(x), gm(x))
75+ 
76+ 
77+ @parametrize('shape', [(5, 5)])
78+ @parametrize('dtype', ['float32'])
79+ def test_clone_different_memory_format_no_fold(self, shape, dtype):
80+ """memory_format 不同 → 不折叠"""
81+ class M(torch.nn.Module):
82+ def forward(self, x):
83+ # 显式指定不同 memory_format(假设原 x 是 contiguous)
84+ cloned = torch.ops.aten.clone.default(x, memory_format=torch.preserve_format)
85+ # 如果想强制不同,可以用 channels_last,但需确保 dtype 支持
86+ return cloned
87+ 
88+ x = self._generate_tensor(shape, dtype)
89+ model = M()
90+ gm = fx.symbolic_trace(model)
91+ ShapeProp(gm).propagate(x)
92+ 
93+ fold_clone(gm.graph)
94+ 
95+ clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default]
96+ self.assertEqual(len(clone_nodes), 1, "memory_format 不同不应折叠")
atomgit-bot
atomgit-botatomgit-bot26 天前

🟡 Medium Priority

变更行:第 79-96 行(test_clone_different_memory_format_no_fold)。

证据链

  1. 测试模型的 forward 直接把 cloned 作为返回值(第 86 行 return cloned),即 clone 节点是 output 节点的直接输入。
  2. fold_clone 在构建候选列表时(ascend_graph_pass.py 第 370 行),要求 get_node_unique_id(node) not in output_storages。该 clone 在 output_storages 中,因此永远不会成为候选节点
  3. 实际的 memory_format 比较(ascend_graph_pass.py 第 377-378 行)对此 clone 节点从未被执行
  4. 测试断言 len(clone_nodes) == 1 通过,但原因不是"memory_format 不同",而是"clone 在输出路径上被保护"。

失败模式:如果后续有人误修改了 fold_clone 的 memory_format 比较逻辑(例如错误地把 preserve_format 视为与 contiguous_format 等同),该测试不会捕获到回归——因为 clone 从未进入 memory_format 检查分支。这造成了对 memory_format 守卫逻辑的虚假覆盖率

建议:将模型改为 clone 不直接作为输出,让 clone 进入 fold_clone 的候选列表,真正触发 memory_format 比较路径。例如:return cloned + 1.0 代替 return cloned

likedislike
97+ 
98+ 
99+ @parametrize('shape', [(3, 7)])
100+ def test_clone_in_output_path_no_fold(self, shape):
101+ """clone 是 output 的直接/间接输入 → 不折叠"""
102+ class M(torch.nn.Module):
103+ def forward(self, x):
104+ cloned = torch.ops.aten.clone.default(x)
105+ return cloned # 直接输出 clone
106+ 
107+ x = self._generate_tensor(shape, 'float32')
108+ model = M()
109+ gm = fx.symbolic_trace(model)
110+ ShapeProp(gm).propagate(x)
111+ fold_clone(gm.graph)
112+ clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default]
113+ self.assertEqual(len(clone_nodes), 1, "clone 在 output 路径上不应折叠")
114+ 
115+ 
116+ @parametrize('shape', [(2, 16)])
117+ def test_multi_clone_some_fold_some_not(self, shape):
118+ """多个 clone,只有部分可折叠"""
119+ class M(torch.nn.Module):
120+ def forward(self, x, y):
121+ c1 = torch.ops.aten.clone.default(x)
122+ c2 = torch.ops.aten.clone.default(y, memory_format=torch.channels_last)
123+ c3 = torch.ops.aten.clone.default(c1)
124+ return c1 + c2 + c3
125+ 
126+ x = self._generate_tensor(shape, 'float32')
127+ y = self._generate_tensor(shape, 'float32')
128+ model = M()
129+ gm = fx.symbolic_trace(model)
130+ ShapeProp(gm).propagate(x, y)
131+ fold_clone(gm.graph)
132+ clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default]
133+ # 预期:c1 和 c3 折叠,c2 保留(或根据你的 memory_format 判断)
134+ self.assertLessEqual(len(clone_nodes), 1)
135+ 
136+ 
137+ @parametrize('shape', [(1, 64)])
138+ def test_clone_with_downstream_users(self, shape):
139+ """clone 被折叠后,下游多个用户应正确指向原输入"""
140+ class M(torch.nn.Module):
141+ def forward(self, x):
142+ cloned = torch.ops.aten.clone.default(x)
143+ a = cloned * 2
144+ b = cloned + 3
145+ return a, b
146+ 
147+ x = self._generate_tensor(shape, 'float32')
148+ model = M()
149+ gm = fx.symbolic_trace(model)
150+ ShapeProp(gm).propagate(x)
151+ 
152+ fold_clone(gm.graph)
153+ gm.recompile()
154+ 
155+ clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default]
156+ self.assertEqual(len(clone_nodes), 0)
157+ 
158+ self.assertEqual(model(x), gm(x))
159+ 
160+ 
161+ def test_no_clone_node_no_change(self):
162+ """图中无 clone 节点,pass 不应修改图"""
163+ class M(torch.nn.Module):
164+ def forward(self, x):
165+ return x + x * 2.0
166+ 
167+ x = torch.randn(8, 8)
168+ gm = fx.symbolic_trace(M())
169+ ShapeProp(gm).propagate(x)
170+ 
171+ orig_graph = str(gm.graph)
172+ fold_clone(gm.graph)
173+ self.assertEqual(orig_graph, str(gm.graph))
174+ 
175+ 
176+ @parametrize('shape', [(4, 4)])
177+ def test_clone_without_tensor_meta_no_fold(self, shape):
178+ """输入没有 tensor_meta → 不折叠(安全检查)"""
179+ class M(torch.nn.Module):
180+ def forward(self, x):
181+ cloned = torch.ops.aten.clone.default(x)
182+ return cloned
183+ 
184+ x = self._generate_tensor(shape, 'float32')
185+ gm = fx.symbolic_trace(M())
186+ # 故意移除 tensor_meta
187+ for node in gm.graph.nodes:
188+ if "tensor_meta" in node.meta:
189+ del node.meta["tensor_meta"]
190+ 
191+ fold_clone(gm.graph)
192+ 
193+ clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default]
194+ self.assertEqual(len(clone_nodes), 1, "缺少 tensor_meta 不应折叠")
atomgit-bot
atomgit-botatomgit-bot26 天前

🟡 Medium Priority

变更行:第 176-194 行(test_clone_without_tensor_meta_no_fold)。

证据链

  1. 测试模型的 forward 直接把 cloned 作为返回值(第 182 行 return cloned),即 clone 节点是 output 节点的直接输入。
  2. fold_clone 在构建候选列表时(ascend_graph_pass.py 第 370 行),要求 get_node_unique_id(node) not in output_storages。该 clone 在 output_storages 中,因此永远不会成为候选节点
  3. 对候选节点的 tensor_meta 检查(ascend_graph_pass.py 第 374 行)对此 clone 节点从未被执行
  4. 测试在第 187-189 行删除 tensor_meta,然后在第 191 行调用 fold_clone,断言 len(clone_nodes) == 1。断言通过,但原因不是"缺少 tensor_meta 导致跳过折叠",而是"clone 在输出路径上被保护"——与第 187-189 行的 tensor_meta 删除操作完全无关。

失败模式:如果后续有人误修改了 fold_clone 的 tensor_meta 守卫(例如错误地删除了第 374 行的 continue),该测试不会捕获到回归——因为 clone 从未进入 tensor_meta 检查分支。这造成了对 tensor_meta 守卫逻辑的虚假覆盖率。此外,第 187-189 行的 "故意移除 tensor_meta" 注释和实际代码暗示作者意图测试 tensor_meta 守卫,但测试结构无法达到此目的。

建议:将模型改为 clone 不直接作为输出(如 return cloned + 1.0),并先用 ShapeProp 添加 tensor_meta 再删除,确保 clone 进入候选列表后真正触发 tensor_meta 缺失的守卫逻辑。

改动建议
194
+ @parametrize('shape', [(4, 4)])
195
+ def test_clone_without_tensor_meta_no_fold(self, shape):
196
+ """输入没有 tensor_meta → 不折叠(安全检查)"""
197
+ class M(torch.nn.Module):
198
+ def forward(self, x):
199
+ cloned = torch.ops.aten.clone.default(x)
200
+ return cloned + 1.0 # clone 不作为直接输出,使其进入候选列表
201
+
202
+ x = self._generate_tensor(shape, 'float32')
203
+ gm = fx.symbolic_trace(M())
204
+ ShapeProp(gm).propagate(x) # 先添加 tensor_meta
205
+ # 故意移除 tensor_meta
206
+ for node in gm.graph.nodes:
207
+ if "tensor_meta" in node.meta:
208
+ del node.meta["tensor_meta"]
209
+
210
+ fold_clone(gm.graph)
211
+
212
+ clone_nodes = [n for n in gm.graph.nodes if n.target == torch.ops.aten.clone.default]
194
213
  self.assertEqual(len(clone_nodes), 1, "缺少 tensor_meta 不应折叠")
应用建议
likedislike
195+ 
196+ 
197+instantiate_parametrized_tests(TestFoldClonePass)
198+ 
199+ 
200+if __name__ == "__main__":
201+ run_tests()
Atest/_inductor/test_fold_div_pass.py+201-0
@@ -0,0 +1,201 @@
1+import torch
2+import torch.fx as fx
3+from torch.fx.passes.shape_prop import ShapeProp
4+from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests
5+from testutils import TestUtils
6+import torch_npu
7+import torch_npu._inductor
8+from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_four_op_pass
9+ 
10+ 
11+class FoldDivModel(torch.nn.Module):
12+ def forward(self, t1, t2, t3):
13+ div_1 = torch.ops.aten.div(t1, 1)
14+ div_output = torch.relu(div_1)
15+ return div_output
16+ 
17+ 
18+class TestFoldDivPass(TestUtils):
19+ def op_calc(self, t1):
20+ div_1 = torch.ops.aten.div(t1, 1)
21+ div_output = torch.relu(div_1)
22+ return div_output
23+ 
24+ 
25+ @parametrize('shape', [(1, 2, 3)])
26+ @parametrize('dtype', ['float32'])
27+ def test_compile_cases(self, shape, dtype):
28+ t1 = self._generate_tensor(shape, dtype)
29+ std_result = self.op_calc(t1)
30+ with torch.no_grad():
31+ compiled_op_calc = torch.compile(self.op_calc, backend="inductor")
32+ inductor_result = compiled_op_calc(t1)
33+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
34+ 
35+ 
36+ @parametrize('shape', [(2, 4, 8)])
37+ @parametrize('dtype', ['float32'])
38+ def test_ut_cases(self, shape, dtype):
39+ t1 = self._generate_tensor(shape, dtype)
40+ t2 = self._generate_tensor(shape, dtype)
41+ t3 = self._generate_tensor(shape, dtype)
42+ model = FoldDivModel()
43+ graph_module = fx.symbolic_trace(model)
44+ ShapeProp(graph_module).propagate(t1, t2, t3)
45+ 
46+ # 应用优化 Pass
47+ fold_four_op_pass(graph_module.graph)
48+ graph_module.recompile()
49+ 
50+ # 验证输出是否一致
51+ std_result = model(t1, t2, t3)
52+ inductor_result = graph_module(t1, t2, t3)
53+ 
54+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
55+ 
56+ 
57+ @parametrize('shape', [(4, 16)])
58+ @parametrize('dtype', ['float32', 'float16'])
59+ def test_div_by_one_right_fold(self, shape, dtype):
60+ """ x / 1.0 → should be x """
61+ class M(torch.nn.Module):
62+ def forward(self, x):
63+ y = torch.div(x, 1.0)
64+ return torch.relu(y * 2.0)
65+ 
66+ t = self._generate_tensor(shape, dtype)
67+ m = M()
68+ gm = fx.symbolic_trace(m)
69+ ShapeProp(gm).propagate(t)
70+ 
71+ fold_four_op_pass(gm.graph)
72+ gm.recompile()
73+ 
74+ self.assertFalse(
75+ any(n.target in (torch.ops.aten.div.Tensor, torch.ops.aten.div.Scalar)
76+ for n in gm.graph.nodes),
77+ "x / 1.0 should be folded to x"
78+ )
79+ torch.testing.assert_close(m(t), gm(t), atol=1e-5, rtol=1e-5)
80+ 
81+ 
82+ @parametrize('shape', [(3, 8)])
83+ def test_div_by_one_int_fold(self, shape):
84+ """ x / 1 (int) should be folded """
85+ class M(torch.nn.Module):
86+ def forward(self, x):
87+ y = torch.div(x, 1)
88+ return y + x.mean()
89+ 
90+ t = self._generate_tensor(shape, 'float32')
91+ m = M()
92+ gm = fx.symbolic_trace(m)
93+ ShapeProp(gm).propagate(t)
94+ 
95+ fold_four_op_pass(gm.graph)
96+ gm.recompile()
97+ 
98+ self.assertFalse(any('div' in str(n.target) for n in gm.graph.nodes if n.op == 'call_function'))
99+ self.assertEqual(m(t), gm(t), "x / 1 (int) should be folded")
100+ 
101+ 
102+ @parametrize('shape', [(5, 5)])
103+ def test_rdiv_one_left_fold(self, shape):
104+ """ 1.0 / x should be folded """
105+ class M(torch.nn.Module):
106+ def forward(self, x):
107+ y = torch.div(1.0, x) # div.Scalar(1.0, x) 或 rdiv
108+ return torch.sigmoid(y)
109+ 
110+ t = self._generate_tensor(shape, 'float32') + 0.1 # 避免除零
111+ m = M()
112+ gm = fx.symbolic_trace(m)
113+ ShapeProp(gm).propagate(t)
114+ 
115+ fold_four_op_pass(gm.graph)
116+ gm.recompile()
117+ 
118+ torch.testing.assert_close(m(t), gm(t), atol=1e-5, rtol=1e-5)
atomgit-bot
atomgit-botatomgit-bot26 天前

🟡 Medium Priority

测试名和 docstring 都声称 "1.0 / x should be folded",但测试体中仅通过 torch.testing.assert_close(m(t), gm(t)) 验证正确性,未对图中是否仍有 div 节点做任何断言。

查看 fold_four_op_pass 实现:div 分支仅处理右操作数为 1 的情况(try_match(inp0, inp1, is_one_like, "right")),不处理左操作数为 1 的 1.0 / x 场景(无 rdiv_ops 分支)。因此该 pass 实际上不会折叠 torch.div(1.0, x)

后果:测试会通过(pass 未破坏正确性),但并未验证其声称的折叠行为,造成虚假的测试覆盖信心。

修复方向:要么在 pass 中补充 rdiv 折叠逻辑并加入对应的节点断言;要么将测试改名/改 docstring 为验证"不对 1.0/x 做错误折叠"。

建议:要么为 fold_four_op_pass 补充 rdiv 折叠能力并在测试中加入节点计数断言;要么将测试名和 docstring 改为表示验证"1.0/x 不被错误折叠"。

likedislike
119+ 
120+ 
121+ @parametrize('shape', [(2, 32)])
122+ def test_div_one_chain_fold(self, shape):
123+ """ all / 1.0 should be folded """
124+ class M(torch.nn.Module):
125+ def forward(self, x):
126+ a = torch.div(x, 1.0)
127+ b = torch.div(a, 1)
128+ c = torch.div(b, torch.ones_like(x))
129+ return torch.relu(c + x)
130+ 
131+ t = self._generate_tensor(shape, 'float32')
132+ m = M()
133+ gm = fx.symbolic_trace(m)
134+ ShapeProp(gm).propagate(t)
135+ 
136+ fold_four_op_pass(gm.graph)
137+ gm.recompile()
138+ 
139+ div_nodes = [n for n in gm.graph.nodes if 'div' in str(n.target)]
140+ self.assertEqual(len(div_nodes), 0, "all / 1.0 should be folded")
141+ 
142+ 
143+ @parametrize('shape', [(6, 6)])
144+ def test_div_non_one_no_fold(self, shape):
145+ """ should not fold since div not 1 """
146+ class M(torch.nn.Module):
147+ def forward(self, x):
148+ y = torch.div(x, 2.0)
149+ return torch.tanh(y)
150+ 
151+ t = self._generate_tensor(shape, 'float32')
152+ m = M()
153+ gm = fx.symbolic_trace(m)
154+ ShapeProp(gm).propagate(t)
155+ 
156+ before = len([n for n in gm.graph.nodes if n.op == 'call_function'])
157+ fold_four_op_pass(gm.graph)
158+ after = len([n for n in gm.graph.nodes if n.op == 'call_function'])
159+ 
160+ self.assertEqual(before, after, "should not fold since div not 1")
161+ 
162+ 
163+ @parametrize('shape', [(1, 128)])
164+ def test_div_one_multi_users(self, shape):
165+ """ there are multiple usrs after div /1 all of them should be replaced """
166+ class M(torch.nn.Module):
167+ def forward(self, x):
168+ y = torch.div(x, 1.0)
169+ return y * 3.0, torch.relu(y), y + x
170+ 
171+ t = self._generate_tensor(shape, 'float32')
172+ m = M()
173+ gm = fx.symbolic_trace(m)
174+ ShapeProp(gm).propagate(t)
175+ 
176+ fold_four_op_pass(gm.graph)
177+ gm.recompile()
178+ 
179+ self.assertEqual(m(t), gm(t), "there are multiple usrs after div /1 all of them should be replaced")
180+ 
181+ 
182+ def test_no_div_op_no_change(self):
183+ """ no div node, pass should not modify graph """
184+ class M(torch.nn.Module):
185+ def forward(self, x):
186+ return torch.mul(x, x) + torch.add(x, 1.0)
187+ 
188+ t = torch.randn(4, 4)
189+ gm = fx.symbolic_trace(M())
190+ ShapeProp(gm).propagate(t)
191+ 
192+ orig_str = str(gm.graph)
193+ fold_four_op_pass(gm.graph)
194+ self.assertEqual(orig_str, str(gm.graph), "no div node, pass should not modify graph")
195+ 
196+ 
197+instantiate_parametrized_tests(TestFoldDivPass)
198+ 
199+ 
200+if __name__ == "__main__":
201+ run_tests()