Star846
1323
代码
代码
Issues1501
Pull Requests2286
流水线
项目讨论
Wiki
分析
Star846
1323
已合并
[inductor]sync PR:38117 #39786
rain-666创建于 7月1日
[inductor]sync PR:38117 #39786
已合并
rain-666创建于 7月1日
从rain-666/pytorchv2.10.0_simd_codegen_bug合入到Ascend/pytorchv2.10.0
共 1 个文件变更+277-0
Atest/_inductor/test_reduction_multi_output.py+277-0
@@ -0,0 +1,277 @@
1+"""
2+Regression tests for NPU Inductor Triton Kernel DSL generation bugs.
3+ 
4+Covers the following fixed issues:
5+ Bug 1: NameError('tmp15 is not defined') — reduction result variable used in
6+ non-reduction store placed inside the loop instead of outside.
7+ Bug 2: NameError('ps1 is not defined') — CSE variable definition order error
8+ in dynamic shape mode.
9+ Bug 3: NameError('x1'/'x3' is not defined') — upstream outside_loop_vars
10+ pollution causing misplacement of stores.
11+ Bug 4: IndentationError / NameError regression — IndentedBuffer _indent
12+ inconsistency and reduction_result_vars cleared across nodes.
13+ 
14+Key trigger condition: A fused kernel containing a persistent reduction whose
15+result is referenced by a subsequent (epilogue) node's store operation.
16+This exercises the reduction_result_vars + _deferred_reduction_stores path.
17+"""
18+ 
19+import torch
20+from torch.testing._internal.common_utils import (
21+ run_tests,
22+ parametrize,
23+ instantiate_parametrized_tests,
24+)
25+from testutils import TestUtils
26+import torch_npu
27+ 
28+ 
29+class TestReductionMultiOutput(TestUtils):
30+ """
31+ Tests for reduction kernels where the reduction result variable is stored
32+ to multiple output buffers (one via store_reduction, one or more via regular
33+ store referencing the same reduction result).
34+ 
35+ This directly exercises the fix for:
36+ - reduction_result_vars independent set (avoids outside_loop_vars pollution)
37+ - _deferred_reduction_stores plain list (avoids IndentedBuffer _indent issues)
38+ - deferred output only in last_tiling branch (correct position in static/dynamic)
39+ - reduction_result_vars NOT cleared at end of codegen_body (cross-node survival)
40+ """
41+ 
42+ # ------------------------------------------------------------------ #
43+ # Test Group 1: Reduction result stored to two outputs (static mode) #
44+ # ------------------------------------------------------------------ #
45+ 
46+ def _reduction_dual_store(self, x, dim):
47+ """Sum reduction, result stored to two separate outputs."""
48+ r = x.sum(dim)
49+ # The second return value references the same reduction result 'r',
50+ # triggering the deferred store path in codegen.
51+ return r, r
52+ 
53+ @parametrize("shape", [(128, 38), (64, 128, 16), (32, 8, 128, 4)])
54+ @parametrize("dim", [(-1,), (1,)])
55+ @parametrize("dtype", ["float32"])
56+ def test_reduction_dual_store_static(self, shape, dim, dtype):
57+ """
58+ Bug #1 / #3 / #4 regression: reduction result 'tmp15' used in a second
59+ store (out_ptr1) must appear AFTER the reduction definition and with
60+ correct indentation (outside the reduction loop).
61+ Uses small reduction dimension (e.g., 38) to force loops_r2 > 1 when
62+ R2BLOCK_SUB=4, which was the exact trigger for op49_op321.
63+ """
64+ input_tensor = self._generate_tensor(shape, dtype)
65+ 
66+ ref_out0, ref_out1 = self._reduction_dual_store(input_tensor, *dim)
67+ 
68+ compiled_fn = torch.compile(
69+ self._reduction_dual_store, backend="inductor", dynamic=False
70+ )
71+ ind_out0, ind_out1 = compiled_fn(input_tensor, *dim)
72+ 
73+ self.assertEqual(ref_out0, ind_out0, atol=1e-1, rtol=1e-1)
74+ self.assertEqual(ref_out1, ind_out1, atol=1e-1, rtol=1e-1)
75+ 
76+ # ------------------------------------------------------------------ #
77+ # Test Group 2: Same pattern with dynamic=True #
78+ # ------------------------------------------------------------------ #
79+ 
80+ @parametrize("shape", [(128, 38), (64, 128, 16), (32, 8, 128, 4)])
81+ @parametrize("dim", [(-1,), (1,)])
82+ @parametrize("dtype", ["float32"])
83+ def test_reduction_dual_store_dynamic(self, shape, dim, dtype):
84+ """
85+ Bug #2 / #4 regression: dynamic shape mode must also correctly place
86+ the deferred store outside the loop with proper indentation.
87+ Also exercises the ps* variable ordering fix in wrapper.py.
88+ """
89+ input_tensor = self._generate_tensor(shape, dtype)
90+ 
91+ ref_out0, ref_out1 = self._reduction_dual_store(input_tensor, *dim)
92+ 
93+ compiled_fn = torch.compile(
94+ self._reduction_dual_store, backend="inductor", dynamic=True
95+ )
96+ ind_out0, ind_out1 = compiled_fn(input_tensor, *dim)
97+ 
98+ self.assertEqual(ref_out0, ind_out0, atol=1e-1, rtol=1e-1)
99+ self.assertEqual(ref_out1, ind_out1, atol=1e-1, rtol=1e-1)
100+ 
101+ # ------------------------------------------------------------------ #
102+ # Test Group 3: Reduction + epilogue (unsqueeze/broadcast) #
103+ # ------------------------------------------------------------------ #
104+ 
105+ def _reduction_with_epilogue(self, x, dim):
106+ """Reduction followed by an epilogue op that uses the result."""
107+ r = x.sum(dim)
108+ # unsqueeze creates a new node whose store references 'r'
109+ v = r.unsqueeze(dim[0] if isinstance(dim, tuple) else dim)
110+ return r, v
111+ 
112+ @parametrize("shape", [(128, 38), (256, 64, 10), (16, 16, 64, 8)])
113+ @parametrize("dim", [(-1,), (1,)])
114+ @parametrize("dtype", ["float32"])
115+ def test_reduction_epilogue_static(self, shape, dim, dtype):
116+ """
117+ Regression for the full op49_op321 pattern:
118+ - reduction node generates tmp15 via tl.sum()
119+ - epilogue (pointwise) node stores tmp15 to out_ptr1
120+ - Both stores must be outside the reduction loop with matching indent.
121+ """
122+ input_tensor = self._generate_tensor(shape, dtype)
123+ 
124+ ref_r, ref_v = self._reduction_with_epilogue(input_tensor, *dim)
125+ 
126+ compiled_fn = torch.compile(
127+ self._reduction_with_epilogue, backend="inductor", dynamic=False
128+ )
129+ ind_r, ind_v = compiled_fn(input_tensor, *dim)
130+ 
131+ self.assertEqual(ref_r, ind_r, atol=1e-1, rtol=1e-1)
132+ self.assertEqual(ref_v, ind_v, atol=1e-1, rtol=1e-1)
133+ 
134+ @parametrize("shape", [(128, 38), (256, 64, 10), (16, 16, 64, 8)])
135+ @parametrize("dim", [(-1,), (1,)])
136+ @parametrize("dtype", ["float32"])
137+ def test_reduction_epilogue_dynamic(self, shape, dim, dtype):
138+ """Same as above but in dynamic mode."""
139+ input_tensor = self._generate_tensor(shape, dtype)
140+ 
141+ ref_r, ref_v = self._reduction_with_epilogue(input_tensor, *dim)
142+ 
143+ compiled_fn = torch.compile(
144+ self._reduction_with_epilogue, backend="inductor", dynamic=True
145+ )
146+ ind_r, ind_v = compiled_fn(input_tensor, *dim)
147+ 
148+ self.assertEqual(ref_r, ind_r, atol=1e-1, rtol=1e-1)
149+ self.assertEqual(ref_v, ind_v, atol=1e-1, rtol=1e-1)
150+ 
151+ # ------------------------------------------------------------------ #
152+ # Test Group 4: Various reduction types #
153+ # ------------------------------------------------------------------ #
154+ 
155+ def _reduction_mean_dual(self, x, dim):
156+ """Mean reduction with dual output."""
157+ m = x.mean(dim)
158+ return m, m
159+ 
160+ @parametrize("shape", [(128, 50), (64, 100, 12)])
161+ @parametrize("dim", [(-1,), (1,)])
162+ @parametrize("dtype", ["float32"])
163+ def test_mean_dual_store_static(self, shape, dim, dtype):
164+ """Mean reduction (non-sum) with dual store, static mode."""
165+ input_tensor = self._generate_tensor(shape, dtype)
166+ 
167+ ref0, ref1 = self._reduction_mean_dual(input_tensor, *dim)
168+ 
169+ compiled_fn = torch.compile(
170+ self._reduction_mean_dual, backend="inductor", dynamic=False
171+ )
172+ ind0, ind1 = compiled_fn(input_tensor, *dim)
173+ 
174+ self.assertEqual(ref0, ind0, atol=1e-1, rtol=1e-1)
175+ self.assertEqual(ref1, ind1, atol=1e-1, rtol=1e-1)
176+ 
177+ def _reduction_var_mean_dual(self, x, dim):
178+ """Var_mean reduction with dual output on both results."""
179+ var, mean = torch.var_mean(x, dim)
180+ return var, mean, var, mean
181+ 
182+ @parametrize("shape", [(64, 48), (32, 32, 20)])
183+ @parametrize("dim", [(-1,)])
184+ @parametrize("dtype", ["float32"])
185+ def test_var_mean_multi_store_static(self, shape, dim, dtype):
186+ """
187+ Var_mean produces TWO reduction results (var and mean), each stored
188+ twice. Exercises multiple entries in reduction_result_vars.
189+ """
190+ input_tensor = self._generate_tensor(shape, dtype)
191+ 
192+ ref = self._reduction_var_mean_dual(input_tensor, *dim)
193+ 
194+ compiled_fn = torch.compile(
195+ self._reduction_var_mean_dual, backend="inductor", dynamic=False
196+ )
197+ ind = compiled_fn(input_tensor, *dim)
198+ 
199+ for i in range(len(ref)):
200+ self.assertEqual(ref[i], ind[i], atol=1e-1, rtol=1e-1)
201+ 
202+ # ------------------------------------------------------------------ #
203+ # Test Group 5: Edge cases — shapes that exercise different #
204+ # codegen_range branches #
205+ # ------------------------------------------------------------------ #
206+ 
207+ def _sum_3d_last_dim(self, x):
208+ """3D sum over last dim, result used in two stores."""
209+ s = x.sum(2)
210+ return s, s
211+ 
212+ @parametrize("shape", [(8, 128, 38), (4, 64, 100), (2, 32, 200)])
213+ @parametrize("dtype", ["float32"])
214+ def test_3d_reduction_last_dim_static(self, shape, dtype):
215+ """
216+ 3D tensor with reduction on last dimension. Exercises tiling_axis +
217+ reduction axis combination where the reduction axis is NOT the last
218+ tiling axis, potentially hitting a different codegen_range branch.
219+ """
220+ input_tensor = self._generate_tensor(shape, dtype)
221+ 
222+ ref0, ref1 = self._sum_3d_last_dim(input_tensor)
223+ 
224+ compiled_fn = torch.compile(
225+ self._sum_3d_last_dim, backend="inductor", dynamic=False
226+ )
227+ ind0, ind1 = compiled_fn(input_tensor)
228+ 
229+ self.assertEqual(ref0, ind0, atol=1e-1, rtol=1e-1)
230+ self.assertEqual(ref1, ind1, atol=1e-1, rtol=1e-1)
231+ 
232+ @parametrize("shape", [(8, 128, 38), (4, 64, 100), (2, 32, 200)])
233+ @parametrize("dtype", ["float32"])
234+ def test_3d_reduction_last_dim_dynamic(self, shape, dtype):
235+ """Same 3D pattern but in dynamic mode."""
236+ input_tensor = self._generate_tensor(shape, dtype)
237+ 
238+ ref0, ref1 = self._sum_3d_last_dim(input_tensor)
239+ 
240+ compiled_fn = torch.compile(
241+ self._sum_3d_last_dim, backend="inductor", dynamic=True
242+ )
243+ ind0, ind1 = compiled_fn(input_tensor)
244+ 
245+ self.assertEqual(ref0, ind0, atol=1e-1, rtol=1e-1)
246+ self.assertEqual(ref1, ind1, atol=1e-1, rtol=1e-1)
247+ 
248+ def _sum_4d_mid_dim(self, x):
249+ """4D sum over a middle dimension, dual store."""
250+ s = x.sum(1)
251+ return s, s
252+ 
253+ @parametrize("shape", [(16, 38, 64, 8), (8, 50, 32, 4)])
254+ @parametrize("dtype", ["float32"])
255+ def test_4d_reduction_mid_dim_static(self, shape, dtype):
256+ """
257+ 4D tensor reducing a middle dimension. This exercises the axis
258+ reordering logic in codegen_body() where non-tiling axes are moved
259+ to the front of sorted_axis.
260+ """
261+ input_tensor = self._generate_tensor(shape, dtype)
262+ 
263+ ref0, ref1 = self._sum_4d_mid_dim(input_tensor)
264+ 
265+ compiled_fn = torch.compile(
266+ self._sum_4d_mid_dim, backend="inductor", dynamic=False
267+ )
268+ ind0, ind1 = compiled_fn(input_tensor)
269+ 
270+ self.assertEqual(ref0, ind0, atol=1e-1, rtol=1e-1)
271+ self.assertEqual(ref1, ind1, atol=1e-1, rtol=1e-1)
272+ 
273+ 
274+instantiate_parametrized_tests(TestReductionMultiOutput)
275+ 
276+if __name__ == "__main__":
277+ run_tests()