Star846
1321
代码
代码
Issues1500
Pull Requests2287
流水线
项目讨论
Wiki
分析
Star846
1321
已合并
[inductor]sync graph optimization and multi stream ut #44482
zzll创建于 13 天前
[inductor]sync graph optimization and multi stream ut #44482
已合并
zzll创建于 13 天前
从zhoujinning/pytorchv210合入到Ascend/pytorchv2.10.0
共 2 个文件变更+1618-0
Atest/_inductor/test_ascend_graph_pass.py+864-0
@@ -0,0 +1,864 @@
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+ 
9+ 
10+class CastStandardModel(torch.nn.Module):
11+ def forward(self, arg1_1):
12+ cast_1 = torch.ops.npu._npu_dtype_cast.default(arg1_1, torch.int64)
13+ output = torch.ops.aten.relu.default(cast_1)
14+ return output
15+ 
16+ 
17+class CatSliceCatModel(torch.nn.Module):
18+ def forward(self, first_element):
19+ cat1 = torch.cat(
20+ [
21+ first_element[:, 0:1, :, :],
22+ first_element[:, 1:2, :, :],
23+ first_element[:, 2:3, :, :],
24+ ],
25+ dim=1,
26+ )
27+ cat2 = torch.cat(
28+ [
29+ cat1[:, 0:1, :, :], # 0~1
30+ cat1[:, 1:2, :, :], # 1~2
31+ cat1[:, 2:3, :, :], # 2~3
32+ ],
33+ dim=1,
34+ )
35+ return cat2
36+ 
37+ 
38+class FoldAddModel(torch.nn.Module):
39+ def forward(self, first_element):
40+ add = torch.ops.aten.add.Tensor(first_element, 0)
41+ add_output = torch.ops.aten.relu.default(add)
42+ return add_output
43+ 
44+ 
45+class FoldCatModel(torch.nn.Module):
46+ def forward(self, t1, t2, t3, t4, t5):
47+ cat1 = torch.ops.aten.cat.default([t1, t2], 1)
48+ cat2 = torch.ops.aten.cat.default([cat1, t3], 1)
49+ cat3 = torch.ops.aten.cat.default([cat2, t4], 1)
50+ cat4 = torch.ops.aten.cat.default([cat3, t5], 1)
51+ return cat4
52+ 
53+ 
54+class FoldCloneModel(torch.nn.Module):
55+ def forward(self, t1):
56+ clone_1 = torch.ops.aten.clone.default(t1)
57+ relu_1 = torch.ops.aten.relu.default(clone_1)
58+ return relu_1
59+ 
60+ 
61+class FoldDetachModel(torch.nn.Module):
62+ def forward(self, t1):
63+ detach_x = torch.ops.aten.detach.default(t1)
64+ output = torch.ops.aten.relu.default(detach_x)
65+ return output
66+ 
67+ 
68+class FoldDivModel(torch.nn.Module):
69+ def forward(self, t1, t2, t3):
70+ div_1 = torch.ops.aten.div(t1, 1)
71+ div_output = torch.relu(div_1)
72+ return div_output
73+ 
74+ 
75+class FoldExpandModel(torch.nn.Module):
76+ def forward(self, t1):
77+ add = torch.ops.aten.expand.default(t1, [256, 128, 1])
78+ add_output = torch.ops.aten.relu.default(add)
79+ return add_output
80+ 
81+ 
82+class FoldMulModel(torch.nn.Module):
83+ def forward(self, t1):
84+ add = torch.ops.aten.mul.Tensor(t1, 1)
85+ add_output = torch.ops.aten.relu.default(add)
86+ return add_output
87+ 
88+ 
89+class FoldReduceModel(torch.nn.Module):
90+ def forward(self, t1):
91+ sum_1 = torch.ops.aten.sum.dim_IntList(t1, [1, 3])
92+ return sum_1
93+ 
94+ 
95+class FoldMultiShapeUnchangeModel(torch.nn.Module):
96+ def forward(self, arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1):
97+ embedding = torch.ops.aten.embedding.default(arg0_1, arg1_1)
98+ view = torch.ops.aten.view.default(embedding, [-1, 1, 64])
99+ squeeze = torch.ops.aten.squeeze.dim(view, 1)
100+ embedding_1 = torch.ops.aten.embedding.default(arg2_1, arg3_1)
101+ view_1 = torch.ops.aten.view.default(embedding_1, [-1, 1, 32])
102+ squeeze_1 = torch.ops.aten.squeeze.dim(view_1, 1)
103+ embedding_2 = torch.ops.aten.embedding.default(arg4_1, arg5_1)
104+ view_2 = torch.ops.aten.view.default(embedding_2, [-1, 1, 16])
105+ squeeze_2 = torch.ops.aten.squeeze.dim(view_2, 1)
106+ permute_1 = torch.ops.aten.permute.default(arg6_1, [1, 0])
107+ permute_2 = torch.ops.aten.permute.default(permute_1, [1, 0])
108+ relu = torch.ops.aten.relu.default(arg7_1)
109+ addmm_1 = torch.ops.aten.addmm.default(arg8_1, relu, permute_2)
110+ relu_1 = torch.ops.aten.relu.default(addmm_1)
111+ return {"squeeze": squeeze, "squeeze_1": squeeze_1, "squeeze_2": squeeze_2, "permute_2": permute_2, "relu_1": relu_1}
112+ 
113+ 
114+class FoldSinkViewModel(torch.nn.Module):
115+ def forward(self, t1):
116+ view_1 = torch.ops.aten.view.default(t1, [1, -1])
117+ output = torch.ops.aten.relu.default(view_1)
118+ return output
119+ 
120+ 
121+class FoldSliceModel(torch.nn.Module):
122+ def forward(self, base, view, t1, t2, t3):
123+ end = 16
124+ slice_1 = torch.ops.aten.slice_scatter.default(base, view, 1, 0, end)
125+ result = view + slice_1
126+ slice_2 = torch.ops.aten.slice_scatter.default(t1, t2, 1, 0, 3)
127+ b = torch.ops.aten.slice.Tensor(t3, 1, 0, None)
128+ result_c = torch.ops.aten.add.Tensor(b, b)
129+ return {"result": result, "slice_2": slice_2, "result_c": result_c}
130+ 
131+ 
132+class FoldSqueezeModel(torch.nn.Module):
133+ def forward(self, t1, t2,):
134+ squeeze_1 = torch.ops.aten.squeeze.default(t1)
135+ squeeze_2 = torch.ops.aten.squeeze.default(squeeze_1)
136+ unsqueeze_1 = torch.ops.aten.unsqueeze.default(t2, 1)
137+ squeeze_3 = torch.ops.aten.squeeze.dim(unsqueeze_1, 1)
138+ return {"squeeze_2": squeeze_2, "squeeze_3": squeeze_3}
139+ 
140+ 
141+class FoldSubModel(torch.nn.Module):
142+ def forward(self, t1):
143+ sub = torch.ops.aten.sub.Tensor(t1, torch.ops.aten.zeros_like.default(t1))
144+ sub_output = torch.ops.aten.relu.default(sub)
145+ rsub = torch.ops.aten.rsub.Tensor(torch.ops.aten.zeros_like.default(t1), t1)
146+ rsub_output = torch.ops.aten.relu.default(rsub)
147+ return sub_output + rsub_output
148+ 
149+ 
150+class FoldToCopyModel(torch.nn.Module):
151+ def forward(self, t1):
152+ copy_1 = torch.ops.aten._to_copy.default(t1)
153+ result = torch.ops.aten.add.Tensor(t1, copy_1)
154+ return result
155+ 
156+ 
157+class FoldViewModel(torch.nn.Module):
158+ def forward(self, t1, t2):
159+ squeeze_1 = torch.ops.aten.squeeze.dim(t1, 2)
160+ unsqueeze_1 = torch.ops.aten.unsqueeze.default(squeeze_1, 0)
161+ view_1 = torch.ops.aten.view.default(unsqueeze_1, [1, -1])
162+ output = torch.ops.aten.view.default(t2, [128, 64])
163+ return {"view_1": view_1, "output": output}
164+ 
165+ 
166+class FoldWhereModel(torch.nn.Module):
167+ def forward(self, t1):
168+ mask = t1 > 0
169+ return torch.ops.aten.where.self(mask, t1, t1)
170+ 
171+ 
172+class FoldPadSliceModel(torch.nn.Module):
173+ def forward(self, t1):
174+ inputPad = torch._C._nn.pad(t1, [0, 0, 0, 50], "constant", 0.0)
175+ inputSlice = inputPad[:, :50]
176+ output = torch.relu(inputSlice)
177+ return output
178+ 
179+ 
180+class TestAscendGraphPass(TestUtils):
181+ def cast_standard_op_calc(self, first_element):
182+ cast_1 = torch.ops.npu._npu_dtype_cast.default(first_element, torch.int64)
183+ output = torch.ops.aten.relu.default(cast_1)
184+ return output
185+ 
186+ 
187+ @parametrize('shape', [(256, 5)])
188+ @parametrize('dtype', ['int64'])
189+ def test_cast_standard_compile_cases(self, shape, dtype):
190+ first_element = self._generate_tensor(shape, dtype)
191+ std_result = self.cast_standard_op_calc(first_element)
192+ with torch.no_grad():
193+ compiled_op_calc = torch.compile(self.cast_standard_op_calc, backend="inductor")
194+ inductor_result = compiled_op_calc(first_element)
195+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
196+ 
197+ 
198+ @parametrize('shape', [(256, 5)])
199+ @parametrize('dtype', ['int64'])
200+ def test_cast_standard_ut_cases(self, shape, dtype):
201+ first_element = self._generate_tensor(shape, dtype)
202+ model = CastStandardModel()
203+ graph_module = fx.symbolic_trace(model)
204+ ShapeProp(graph_module).propagate(first_element)
205+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_cast
206+ fold_cast(graph_module.graph)
207+ graph_module.recompile()
208+ std_result = model(first_element)
209+ inductor_result = graph_module(first_element)
210+ 
211+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
212+ 
213+ 
214+ def cat_slice_cat_op_calc(self, first_element):
215+ # 第一次 cat：在 dim=1 上拼接 3 个部分 → [1, 3, 64, 1024]
216+ cat1 = torch.cat(
217+ [
218+ first_element[:, 0:1, :, :], # 第0个channel
219+ first_element[:, 1:2, :, :], # 第1个channel
220+ first_element[:, 2:3, :, :], # 第2个channel
221+ ],
222+ dim=1,
223+ )
224+ # 第二次 cat：在同一个 dim=1 上，对 cat1 的连续切片再做一次 cat
225+ # → 完全等价于 cat1，Inductor 应该直接 erase 掉第二次 cat 和所有 getitem
226+ cat2 = torch.cat(
227+ [
228+ cat1[:, 0:1, :, :], # 0~1
229+ cat1[:, 1:2, :, :], # 1~2
230+ cat1[:, 2:3, :, :], # 2~3
231+ ],
232+ dim=1,
233+ )
234+ return cat2
235+ 
236+ 
237+ @parametrize('shape', [(1, 3, 64, 1024)])
238+ @parametrize('dtype', ['float32'])
239+ def test_cat_slice_cat_compile_cases(self, shape, dtype):
240+ first_element = self._generate_tensor(shape, dtype)
241+ 
242+ std_result = self.cat_slice_cat_op_calc(first_element)
243+ with torch.no_grad():
244+ compiled_op_calc = torch.compile(self.cat_slice_cat_op_calc, backend="inductor")
245+ inductor_result = compiled_op_calc(first_element)
246+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
247+ 
248+ 
249+ @parametrize('shape', [(1, 3, 64, 1024)])
250+ @parametrize('dtype', ['float32'])
251+ def test_cat_slice_cat_ut_cases(self, shape, dtype):
252+ first_element = self._generate_tensor(shape, dtype)
253+ model = CatSliceCatModel()
254+ graph_module = fx.symbolic_trace(model)
255+ ShapeProp(graph_module).propagate(first_element)
256+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import cat_slice_cat_fold_pass
257+ cat_slice_cat_fold_pass(graph_module.graph)
258+ graph_module.recompile()
259+ std_result = model(first_element)
260+ inductor_result = graph_module(first_element)
261+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
262+ 
263+ 
264+ def fold_add_op_calc(self, first_element):
265+ add = torch.add(first_element, 0)
266+ add_output = torch.relu(add)
267+ return add_output
268+ 
269+ 
270+ @parametrize('shape', [(1, 2, 3)])
271+ @parametrize('dtype', ['float32'])
272+ def test_fold_add_compile_cases(self, shape, dtype):
273+ first_element = self._generate_tensor(shape, dtype)
274+ 
275+ std_result = self.fold_add_op_calc(first_element)
276+ with torch.no_grad():
277+ compiled_op_calc = torch.compile(self.fold_add_op_calc, backend="inductor")
278+ inductor_result = compiled_op_calc(first_element)
279+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
280+ 
281+ 
282+ @parametrize('shape', [(1, 2, 3)])
283+ @parametrize('dtype', ['float32'])
284+ def test_fold_add_ut_cases(self, shape, dtype):
285+ first_element = self._generate_tensor(shape, dtype)
286+ model = FoldAddModel()
287+ graph_module = fx.symbolic_trace(model)
288+ ShapeProp(graph_module).propagate(first_element)
289+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_four_op_pass
290+ fold_four_op_pass(graph_module.graph)
291+ graph_module.recompile()
292+ std_result = model(first_element)
293+ inductor_result = graph_module(first_element)
294+ 
295+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
296+ 
297+ 
298+ def fold_cat_op_calc(self, t1, t2, t3, t4, t5):
299+ cat1 = torch.cat([t1, t2], dim=1)
300+ cat2 = torch.cat([cat1, t3], dim=1)
301+ cat3 = torch.cat([cat2, t4], dim=1)
302+ cat4 = torch.cat([cat3, t5], dim=1)
303+ return cat4
304+ 
305+ 
306+ @parametrize('shape', [(2, 4)])
307+ @parametrize('dtype', ['float32'])
308+ def test_fold_cat_compile_cases(self, shape, dtype):
309+ t1 = self._generate_tensor(shape, dtype)
310+ t2 = self._generate_tensor(shape, dtype)
311+ t3 = self._generate_tensor(shape, dtype)
312+ t4 = self._generate_tensor(shape, dtype)
313+ t5 = self._generate_tensor(shape, dtype)
314+ 
315+ std_result = self.fold_cat_op_calc(t1, t2, t3, t4, t5)
316+ with torch.no_grad():
317+ compiled_op_calc = torch.compile(self.fold_cat_op_calc, backend="inductor")
318+ inductor_result = compiled_op_calc(t1, t2, t3, t4, t5)
319+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
320+ 
321+ 
322+ @parametrize('shape', [(2, 4)])
323+ @parametrize('dtype', ['float32'])
324+ def test_fold_cat_ut_cases(self, shape, dtype):
325+ t1 = self._generate_tensor(shape, dtype)
326+ t2 = self._generate_tensor(shape, dtype)
327+ t3 = self._generate_tensor(shape, dtype)
328+ t4 = self._generate_tensor(shape, dtype)
329+ t5 = self._generate_tensor(shape, dtype)
330+ model = FoldCatModel()
331+ graph_module = fx.symbolic_trace(model)
332+ ShapeProp(graph_module).propagate(t1, t2, t3, t4, t5)
333+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_cat
334+ fold_cat(graph_module.graph)
335+ graph_module.recompile()
336+ std_result = model(t1, t2, t3, t4, t5)
337+ inductor_result = graph_module(t1, t2, t3, t4, t5)
338+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
339+ 
340+ 
341+ def fold_clone_op_calc(self, t1):
342+ clone_1 = torch.clone(t1)
343+ output = torch.relu(clone_1)
344+ return output
345+ 
346+ 
347+ @parametrize('shape', [(1, 2, 3)])
348+ @parametrize('dtype', ['float32'])
349+ def test_fold_clone_compile_cases(self, shape, dtype):
350+ t1 = self._generate_tensor(shape, dtype)
351+ 
352+ std_result = self.fold_clone_op_calc(t1)
353+ with torch.no_grad():
354+ compiled_op_calc = torch.compile(self.fold_clone_op_calc, backend="inductor")
355+ inductor_result = compiled_op_calc(t1)
356+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
357+ 
358+ 
359+ @parametrize('shape', [(1, 2, 3)])
360+ @parametrize('dtype', ['float32'])
361+ def test_fold_clone_ut_cases(self, shape, dtype):
362+ t1 = self._generate_tensor(shape, dtype)
363+ model = FoldCloneModel()
364+ graph_module = fx.symbolic_trace(model)
365+ ShapeProp(graph_module).propagate(t1)
366+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_clone
367+ fold_clone(graph_module.graph)
368+ graph_module.recompile()
369+ std_result = model(t1)
370+ inductor_result = graph_module(t1)
371+ 
372+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
373+ 
374+ 
375+ def fold_detach_op_calc(self, t1):
376+ detach_x = torch.ops.aten.detach.default(t1)
377+ output = torch.ops.aten.relu.default(detach_x)
378+ return output
379+ 
380+ 
381+ @parametrize('shape', [(3, 3)])
382+ @parametrize('dtype', ['float32'])
383+ def test_fold_detach_compile_cases(self, shape, dtype):
384+ t1 = self._generate_tensor(shape, dtype)
385+ std_result = self.fold_detach_op_calc(t1)
386+ with torch.no_grad():
387+ compiled_op_calc = torch.compile(self.fold_detach_op_calc, backend="inductor")
388+ inductor_result = compiled_op_calc(t1)
389+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
390+ 
391+ 
392+ @parametrize('shape', [(3, 3)])
393+ @parametrize('dtype', ['float32'])
394+ def test_fold_detach_ut_cases(self, shape, dtype):
395+ t1 = self._generate_tensor(shape, dtype)
396+ model = FoldDetachModel()
397+ graph_module = fx.symbolic_trace(model)
398+ ShapeProp(graph_module).propagate(t1)
399+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_detach
400+ fold_detach(graph_module.graph)
401+ graph_module.recompile()
402+ std_result = model(t1)
403+ inductor_result = graph_module(t1)
404+ 
405+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
406+ 
407+ 
408+ def fold_div_op_calc(self, t1):
409+ div_1 = torch.ops.aten.div(t1, 1)
410+ div_output = torch.relu(div_1)
411+ return div_output
412+ 
413+ 
414+ @parametrize('shape', [(1, 2, 3)])
415+ @parametrize('dtype', ['float32'])
416+ def test_fold_div_compile_cases(self, shape, dtype):
417+ t1 = self._generate_tensor(shape, dtype)
418+ std_result = self.fold_div_op_calc(t1)
419+ with torch.no_grad():
420+ compiled_op_calc = torch.compile(self.fold_div_op_calc, backend="inductor")
421+ inductor_result = compiled_op_calc(t1)
422+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
423+ 
424+ 
425+ @parametrize('shape', [(2, 4, 8)])
426+ @parametrize('dtype', ['float32'])
427+ def test_fold_div_ut_cases(self, shape, dtype):
428+ t1 = self._generate_tensor(shape, dtype)
429+ t2 = self._generate_tensor(shape, dtype)
430+ t3 = self._generate_tensor(shape, dtype)
431+ model = FoldDivModel()
432+ graph_module = fx.symbolic_trace(model)
433+ ShapeProp(graph_module).propagate(t1, t2, t3)
434+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_four_op_pass
435+ fold_four_op_pass(graph_module.graph)
436+ graph_module.recompile()
437+ std_result = model(t1, t2, t3)
438+ inductor_result = graph_module(t1, t2, t3)
439+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
440+ 
441+ 
442+ def fold_expand_op_calc(self, t1):
443+ expand = torch.ops.aten.expand.default(t1, [256, 128, 1])
444+ relu = torch.relu(expand)
445+ return relu
446+ 
447+ 
448+ @parametrize('shape', [(256, 128, 1)])
449+ @parametrize('dtype', ['float32'])
450+ def test_fold_expand_compile_cases(self, shape, dtype):
451+ t1 = self._generate_tensor(shape, dtype)
452+ 
453+ std_result = self.fold_expand_op_calc(t1)
454+ with torch.no_grad():
455+ compiled_op_calc = torch.compile(self.fold_expand_op_calc, backend="inductor")
456+ inductor_result = compiled_op_calc(t1)
457+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
458+ 
459+ 
460+ @parametrize('shape', [(256, 128, 1)])
461+ @parametrize('dtype', ['float32'])
462+ def test_fold_expand_ut_cases(self, shape, dtype):
463+ t1 = self._generate_tensor(shape, dtype)
464+ model = FoldExpandModel()
465+ graph_module = fx.symbolic_trace(model)
466+ ShapeProp(graph_module).propagate(t1)
467+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_expand
468+ fold_expand(graph_module.graph)
469+ graph_module.recompile()
470+ std_result = model(t1)
471+ inductor_result = graph_module(t1)
472+ 
473+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
474+ 
475+ 
476+ def fold_mul_op_calc(self, t1):
477+ mul_1 = torch.ops.aten.mul.Tensor(t1, 1)
478+ mul_output = torch.relu(mul_1)
479+ return mul_output
480+ 
481+ 
482+ @parametrize('shape', [(1, 2, 3)])
483+ @parametrize('dtype', ['float32'])
484+ def test_fold_mul_compile_cases(self, shape, dtype):
485+ t1 = self._generate_tensor(shape, dtype)
486+ std_result = self.fold_mul_op_calc(t1)
487+ with torch.no_grad():
488+ compiled_op_calc = torch.compile(self.fold_mul_op_calc, backend="inductor")
489+ inductor_result = compiled_op_calc(t1)
490+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
491+ 
492+ 
493+ @parametrize('shape', [(1, 2, 3)])
494+ @parametrize('dtype', ['float32'])
495+ def test_fold_mul_ut_cases(self, shape, dtype):
496+ t1 = self._generate_tensor(shape, dtype)
497+ model = FoldMulModel()
498+ graph_module = fx.symbolic_trace(model)
499+ ShapeProp(graph_module).propagate(t1)
500+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_four_op_pass
501+ fold_four_op_pass(graph_module.graph)
502+ graph_module.recompile()
503+ std_result = model(t1)
504+ inductor_result = graph_module(t1)
505+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
506+ 
507+ 
508+ def fold_reduce_op_calc(self, t1):
509+ output = torch.sum(t1, dim=(1, 3), keepdim=False)
510+ return output
511+ 
512+ 
513+ @parametrize('shape', [(128, 1, 64, 1)])
514+ @parametrize('dtype', ['float32'])
515+ def test_fold_reduce_compile_cases(self, shape, dtype):
516+ t1 = self._generate_tensor(shape, dtype)
517+ std_result = self.fold_reduce_op_calc(t1)
518+ with torch.no_grad():
519+ compiled_op_calc = torch.compile(self.fold_reduce_op_calc, backend="inductor")
520+ inductor_result = compiled_op_calc(t1)
521+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
522+ 
523+ 
524+ @parametrize('shape', [(128, 1, 64, 1)])
525+ @parametrize('dtype', ['float32'])
526+ def test_fold_reduce_ut_cases(self, shape, dtype):
527+ t1 = self._generate_tensor(shape, dtype)
528+ model = FoldReduceModel()
529+ graph_module = fx.symbolic_trace(model)
530+ ShapeProp(graph_module).propagate(t1)
531+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_reduce
532+ fold_reduce(graph_module.graph)
533+ graph_module.recompile()
534+ std_result = model(t1)
535+ inductor_result = graph_module(t1)
536+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
537+ 
538+ 
539+ def fold_redundant_op_calc(self, arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1):
540+ embedding = torch.ops.aten.embedding.default(arg0_1, arg1_1)
541+ view = torch.ops.aten.view.default(embedding, [-1, 1, 64])
542+ squeeze = torch.ops.aten.squeeze.dim(view, 1)
543+ embedding_1 = torch.ops.aten.embedding.default(arg2_1, arg3_1)
544+ view_1 = torch.ops.aten.view.default(embedding_1, [-1, 1, 32])
545+ squeeze_1 = torch.ops.aten.squeeze.dim(view_1, 1)
546+ embedding_2 = torch.ops.aten.embedding.default(arg4_1, arg5_1)
547+ view_2 = torch.ops.aten.view.default(embedding_2, [-1, 1, 16])
548+ squeeze_2 = torch.ops.aten.squeeze.dim(view_2, 1)
549+ permute_1 = torch.ops.aten.permute.default(arg6_1, [1, 0])
550+ permute_2 = torch.ops.aten.permute.default(permute_1, [1, 0])
551+ relu = torch.ops.aten.relu.default(arg7_1)
552+ addmm_1 = torch.ops.aten.addmm.default(arg8_1, relu, permute_2)
553+ relu_1 = torch.ops.aten.relu.default(addmm_1)
554+ return {"squeeze": squeeze, "squeeze_1": squeeze_1, "squeeze_2": squeeze_2, "permute_2": permute_2, "relu_1": relu_1}
555+ 
556+ 
557+ def test_fold_redundant_compile_cases(self):
558+ arg0_1 = torch.randn(289094, 64, dtype=torch.float32)
559+ arg1_1 = torch.randint(0, 289094, (128,), dtype=torch.int64)
560+ arg2_1 = torch.randn(98, 32, dtype=torch.float32)
561+ arg3_1 = torch.randint(0, 98, (128,), dtype=torch.int64)
562+ arg4_1 = torch.randn(14, 16, dtype=torch.float32)
563+ arg5_1 = torch.randint(0, 14, (128,), dtype=torch.int64)
564+ arg6_1 = torch.randn(6144, 6144, dtype=torch.float32)
565+ arg7_1 = torch.randn(128, 6144, dtype=torch.float32)
566+ arg8_1 = torch.randn(6144, dtype=torch.float32)
567+ std_result = self.fold_redundant_op_calc(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1)
568+ with torch.no_grad():
569+ compiled_op_calc = torch.compile(self.fold_redundant_op_calc, backend="inductor")
570+ inductor_result = compiled_op_calc(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1)
571+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
572+ 
573+ 
574+ def test_fold_redundant_ut_cases(self):
575+ arg0_1 = torch.randn(289094, 64, dtype=torch.float32)
576+ arg1_1 = torch.randint(0, 289094, (128,), dtype=torch.int64)
577+ arg2_1 = torch.randn(98, 32, dtype=torch.float32)
578+ arg3_1 = torch.randint(0, 98, (128,), dtype=torch.int64)
579+ arg4_1 = torch.randn(14, 16, dtype=torch.float32)
580+ arg5_1 = torch.randint(0, 14, (128,), dtype=torch.int64)
581+ arg6_1 = torch.randn(6144, 6144, dtype=torch.float32)
582+ arg7_1 = torch.randn(128, 6144, dtype=torch.float32)
583+ arg8_1 = torch.randn(6144, dtype=torch.float32)
584+ model = FoldMultiShapeUnchangeModel()
585+ graph_module = fx.symbolic_trace(model)
586+ ShapeProp(graph_module).propagate(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1)
587+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_redundant_ops
588+ fold_redundant_ops(graph_module.graph)
589+ graph_module.recompile()
590+ std_result = model(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1)
591+ inductor_result = graph_module(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1)
592+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
593+ 
594+ 
595+ def fold_sink_view_op_calc(self, t1):
596+ x = t1.view(1, -1) # reshape
597+ return torch.relu(x)
598+ 
599+ 
600+ @parametrize('shape', [(2, 3, 4)])
601+ @parametrize('dtype', ['float32'])
602+ def test_fold_sink_viewcompile_cases(self, shape, dtype):
603+ t1 = self._generate_tensor(shape, dtype)
604+ std_result = self.fold_sink_view_op_calc(t1)
605+ with torch.no_grad():
606+ compiled_op_calc = torch.compile(self.fold_sink_view_op_calc, backend="inductor")
607+ inductor_result = compiled_op_calc(t1)
608+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
609+ 
610+ 
611+ @parametrize('shape', [(2, 3, 4)])
612+ @parametrize('dtype', ['float32'])
613+ def test_fold_sink_viewut_cases(self, shape, dtype):
614+ t1 = self._generate_tensor(shape, dtype)
615+ model = FoldSinkViewModel()
616+ graph_module = fx.symbolic_trace(model)
617+ ShapeProp(graph_module).propagate(t1)
618+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_sink_view
619+ fold_sink_view(graph_module.graph)
620+ graph_module.recompile()
621+ std_result = model(t1)
622+ inductor_result = graph_module(t1)
623+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
624+ 
625+ 
626+ def fold_slice_op_calc(self, base, view, t1, t2, t3):
627+ end = view.shape[1]
628+ result = torch.slice_scatter(base, view, 1, 0, end)
629+ result = result + view
630+ data = t1.slice_scatter(t2, dim=1, start=0, end=t2.shape[1])
631+ b = t3[:, 0:]
632+ result_c = b + b
633+ return {"result": result, "data": data, "result_c": result_c}
634+ 
635+ 
636+ def test_fold_slice_compile_cases(self):
637+ base = torch.randn(8, 16, 32)
638+ view = torch.ones(8, 16, 32)
639+ t1 = torch.tensor([[1, 2, 3], [4, 5, 6]])
640+ t2 = torch.tensor([[9, 9, 9], [8, 8, 8]])
641+ t3 = torch.randn(4, 16, 32, 64)
642+ std_result = self.fold_slice_op_calc(base, view, t1, t2, t3)
643+ with torch.no_grad():
644+ compiled_op_calc = torch.compile(self.fold_slice_op_calc, backend="inductor")
645+ inductor_result = compiled_op_calc(base, view, t1, t2, t3)
646+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
647+ 
648+ 
649+ def test_fold_slice_ut_cases(self):
650+ base = torch.randn(8, 16, 32)
651+ view = torch.ones(8, 16, 32)
652+ t1 = torch.tensor([[1, 2, 3], [4, 5, 6]])
653+ t2 = torch.tensor([[9, 9, 9], [8, 8, 8]])
654+ t3 = torch.randn(4, 16, 32, 64)
655+ model = FoldSliceModel()
656+ graph_module = fx.symbolic_trace(model)
657+ ShapeProp(graph_module).propagate(base, view, t1, t2, t3)
658+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_slice
659+ fold_slice(graph_module.graph)
660+ graph_module.recompile()
661+ std_result = model(base, view, t1, t2, t3)
662+ inductor_result = graph_module(base, view, t1, t2, t3)
663+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
664+ 
665+ 
666+ def fold_squeeze_op_calc(self, t1, t2):
667+ squeeze_1 = torch.squeeze(t1)
668+ squeeze_2 = torch.squeeze(squeeze_1)
669+ unsqueeze_1 = torch.unsqueeze(t2, 1)
670+ squeeze_3 = torch.squeeze(unsqueeze_1, 1)
671+ return {"squeeze_2": squeeze_2, "squeeze_3": squeeze_3}
672+ 
673+ 
674+ def test_fold_squeeze_compile_cases(self):
675+ t1 = torch.randn(2, 4)
676+ t2 = torch.randn(2, 1, 1, 4)
677+ std_result = self.fold_squeeze_op_calc(t1, t2)
678+ with torch.no_grad():
679+ compiled_op_calc = torch.compile(self.fold_squeeze_op_calc, backend="inductor")
680+ inductor_result = compiled_op_calc(t1, t2)
681+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
682+ 
683+ 
684+ def test_fold_squeeze_ut_cases(self):
685+ t1 = torch.randn(2, 4)
686+ t2 = torch.randn(2, 1, 1, 4)
687+ model = FoldSqueezeModel()
688+ graph_module = fx.symbolic_trace(model)
689+ ShapeProp(graph_module).propagate(t1, t2)
690+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_squeeze
691+ fold_squeeze(graph_module.graph)
692+ graph_module.recompile()
693+ std_result = model(t1, t2)
694+ inductor_result = graph_module(t1, t2)
695+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
696+ 
697+ 
698+ def fold_sub_op_calc(self, t1):
699+ sub = torch.sub(t1, torch.ops.aten.zeros_like.default(t1))
700+ sub_output = torch.relu(sub)
701+ rsub = torch.rsub(torch.ops.aten.zeros_like.default(t1), t1)
702+ rsub_output = torch.relu(rsub)
703+ return sub_output + rsub_output
704+ 
705+ 
706+ @parametrize('shape', [(1, 2, 3)])
707+ @parametrize('dtype', ['float32'])
708+ def test_fold_sub_compile_cases(self, shape, dtype):
709+ t1 = self._generate_tensor(shape, dtype)
710+ std_result = self.fold_sub_op_calc(t1)
711+ with torch.no_grad():
712+ compiled_op_calc = torch.compile(self.fold_sub_op_calc, backend="inductor")
713+ inductor_result = compiled_op_calc(t1)
714+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
715+ 
716+ 
717+ @parametrize('shape', [(1, 2, 3)])
718+ @parametrize('dtype', ['float32'])
719+ def test_fold_sub_ut_cases(self, shape, dtype):
720+ t1 = self._generate_tensor(shape, dtype)
721+ model = FoldSubModel()
722+ graph_module = fx.symbolic_trace(model)
723+ ShapeProp(graph_module).propagate(t1)
724+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_four_op_pass
725+ fold_four_op_pass(graph_module.graph)
726+ graph_module.recompile()
727+ std_result = model(t1)
728+ inductor_result = graph_module(t1)
729+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
730+ 
731+ 
732+ def fold_to_copy_op_calc(self, t1):
733+ copy_1 = torch.ops.aten._to_copy.default(t1)
734+ result = copy_1 + t1
735+ return result
736+ 
737+ 
738+ @parametrize('shape', [(2, 4)])
739+ @parametrize('dtype', ['float32'])
740+ def test_fold_to_copy_compile_cases(self, shape, dtype):
741+ t1 = self._generate_tensor(shape, dtype)
742+ std_result = self.fold_to_copy_op_calc(t1)
743+ with torch.no_grad():
744+ compiled_op_calc = torch.compile(self.fold_to_copy_op_calc, backend="inductor")
745+ inductor_result = compiled_op_calc(t1)
746+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
747+ 
748+ 
749+ @parametrize('shape', [(2, 4)])
750+ @parametrize('dtype', ['float32'])
751+ def test_fold_to_copy_ut_cases(self, shape, dtype):
752+ t1 = self._generate_tensor(shape, dtype)
753+ model = FoldToCopyModel()
754+ graph_module = fx.symbolic_trace(model)
755+ ShapeProp(graph_module).propagate(t1)
756+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_to_copy
757+ fold_to_copy(graph_module.graph)
758+ graph_module.recompile()
759+ std_result = model(t1)
760+ inductor_result = graph_module(t1)
761+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
762+ 
763+ 
764+ def fold_view_op_calc(self, t1, t2):
765+ squeeze_1 = t1.squeeze(2)
766+ unsqueeze_1 = squeeze_1.unsqueeze(0)
767+ view_1 = unsqueeze_1.view(1, -1)
768+ output = t2.view(128, 64)
769+ return {"view_1": view_1, "output": output}
770+ 
771+ 
772+ def test_fold_view_compile_cases(self):
773+ t1 = torch.randn(1, 3, 1, 5)
774+ t2 = torch.randn(128, 64)
775+ std_result = self.fold_view_op_calc(t1, t2)
776+ with torch.no_grad():
777+ compiled_op_calc = torch.compile(self.fold_view_op_calc, backend="inductor")
778+ inductor_result = compiled_op_calc(t1, t2)
779+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
780+ 
781+ 
782+ def test_fold_view_ut_cases(self):
783+ t1 = torch.randn(1, 3, 1, 5)
784+ t2 = torch.randn(128, 64)
785+ model = FoldViewModel()
786+ graph_module = fx.symbolic_trace(model)
787+ ShapeProp(graph_module).propagate(t1, t2)
788+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import view_fold_pass
789+ view_fold_pass(graph_module.graph)
790+ graph_module.recompile()
791+ std_result = model(t1, t2)
792+ inductor_result = graph_module(t1, t2)
793+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
794+ 
795+ 
796+ def fold_where_op_calc(self, t1):
797+ mask = t1 > 0
798+ return torch.where(mask, t1, t1) # 两个分支完全相同
799+ 
800+ 
801+ @parametrize('shape', [(3, 4)])
802+ @parametrize('dtype', ['float32'])
803+ def test_fold_where_compile_cases(self, shape, dtype):
804+ t1 = self._generate_tensor(shape, dtype)
805+ std_result = self.fold_where_op_calc(t1)
806+ with torch.no_grad():
807+ compiled_op_calc = torch.compile(self.fold_where_op_calc, backend="inductor")
808+ inductor_result = compiled_op_calc(t1)
809+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
810+ 
811+ 
812+ @parametrize('shape', [(3, 4)])
813+ @parametrize('dtype', ['float32'])
814+ def test_fold_where_ut_cases(self, shape, dtype):
815+ t1 = self._generate_tensor(shape, dtype)
816+ model = FoldWhereModel()
817+ graph_module = fx.symbolic_trace(model)
818+ ShapeProp(graph_module).propagate(t1)
819+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import fold_where
820+ fold_where(graph_module.graph)
821+ graph_module.recompile()
822+ std_result = model(t1)
823+ inductor_result = graph_module(t1)
824+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
825+ 
826+ 
827+ def pad_slice_op_calc(self, t1):
828+ inputPad = torch._C._nn.pad(t1, [0, 0, 0, 50], "constant", 0.0)
829+ inputSlice = inputPad[:, :50]
830+ output = torch.relu(inputSlice)
831+ return output
832+ 
833+ 
834+ @parametrize('shape', [(128, 50, 128)])
835+ @parametrize('dtype', ['float32'])
836+ def test_pad_slice_compile_cases(self, shape, dtype):
837+ t1 = self._generate_tensor(shape, dtype)
838+ std_result = self.pad_slice_op_calc(t1)
839+ with torch.no_grad():
840+ compiled_op_calc = torch.compile(self.pad_slice_op_calc, backend="inductor")
841+ inductor_result = compiled_op_calc(t1)
842+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
843+ 
844+ 
845+ @parametrize('shape', [(128, 50, 128)])
846+ @parametrize('dtype', ['float32'])
847+ def test_pad_slice_ut_cases(self, shape, dtype):
848+ t1 = self._generate_tensor(shape, dtype)
849+ model = FoldPadSliceModel()
850+ graph_module = fx.symbolic_trace(model)
851+ ShapeProp(graph_module).propagate(t1)
852+ from torch_npu._inductor.fx_passes.ascend_custom_passes.ascend_graph_pass import pad_slice_fold
853+ pad_slice_fold(graph_module.graph)
854+ graph_module.recompile()
855+ std_result = model(t1)
856+ inductor_result = graph_module(t1)
857+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
858+ 
859+ 
860+instantiate_parametrized_tests(TestAscendGraphPass)
861+ 
862+ 
863+if __name__ == "__main__":
864+ run_tests()
Atest/_inductor/test_multi_stream.py+754-0
@@ -0,0 +1,754 @@
1+import os
2+import itertools
3+import torch
4+from torch.testing._internal.common_utils import run_tests, instantiate_parametrized_tests
5+from testutils import TestUtils
6+import torch_npu
7+import torch_npu._inductor
8+from unittest.mock import patch, Mock, MagicMock
9+from torch._inductor.virtualized import V
10+from torch._inductor.codegen.wrapper import WorkspaceArg
11+from torch_npu._inductor.codegen.catlass.catlass_kernel import CATLASSTemplateKernel
12+from torch_npu._inductor.codegen.catlass.catlass_scheduling import CATLASSScheduling
13+from torch_npu._inductor.codegen.scheduling import NPUTritonScheduling
14+from torch_npu._inductor.codegen.triton import NPUIndexTritonKernel
15+import torch._inductor.scheduler as sch
16+ 
17+ 
18+class FakeWrapper:
19+ def __init__(self):
20+ self.write_triton_header_once = MagicMock()
21+ self.generate_workspace_allocation = MagicMock()
22+ self.generate_workspace_deallocation = MagicMock()
23+ self.generate_kernel_call = MagicMock()
24+ 
25+ 
26+class FakeGraph:
27+ def __init__(self, cpp_wrapper=False):
28+ self.wrapper_code = FakeWrapper()
29+ self.cpp_wrapper = cpp_wrapper
30+ self.workspace_id = itertools.count()
31+ self.sizevars = FakeSizeVars()
32+ 
33+ def get_current_device_or_throw(self):
34+ return "npu:0"
35+ 
36+ def is_unspec_arg(self, arg):
37+ return False
38+ 
39+ 
40+class FakeArgs:
41+ def python_argdefs(self):
42+ return [], ["input_ptr"], [], []
43+ 
44+ def cpp_argdefs(self):
45+ return [], ["input_ptr"], []
46+ 
47+ 
48+class FakeNode:
49+ def __init__(self, workspace_size, name):
50+ self.workspace_size = workspace_size
51+ self._name = name
52+ self.last_usage = []
53+ 
54+ def get_workspace_size(self):
55+ return self.workspace_size
56+ 
57+ def get_name(self):
58+ return self._name
59+ 
60+ def get_estimated_runtime(self):
61+ return 1.0
62+ 
63+ def get_device(self):
64+ return None
65+ 
66+ def is_extern(self):
67+ return False
68+ 
69+ def is_template(self):
70+ return False
71+ 
72+ def is_foreach(self):
73+ return False
74+ 
75+ def get_buffer_names(self):
76+ return []
77+ 
78+ def get_operation_names(self):
79+ return []
80+ 
81+ def mark_run(self):
82+ pass
83+ 
84+ def get_nodes(self):
85+ return []
86+ 
87+ 
88+class FakeDebugPrinterManager:
89+ def __init__(self):
90+ self.set_printer_args = MagicMock()
91+ 
92+ def __enter__(self):
93+ return self
94+ 
95+ def __exit__(self, exc_type, exc_val, exc_tb):
96+ return False
97+ 
98+ 
99+class FakeWrapperCode:
100+ def __init__(self):
101+ self.debug_printer = FakeDebugPrinterManager()
102+ 
103+ 
104+class FakeSizeVars:
105+ def statically_known_leq(self, a, b):
106+ return False
107+ 
108+ def size_hints(self, x, fallback=8192):
109+ return x
110+ 
111+ def simplify(self, x):
112+ return x
113+ 
114+ 
115+class FakeScheduleGraph:
116+ def __init__(self):
117+ self.sizevars = FakeSizeVars()
118+ 
119+class FakeKernel:
120+ def __init__(self):
121+ self.removed_buffers = set()
122+ self.args = MagicMock()
123+ self.args.python_argdefs.return_value = (
124+ [],
125+ ["arg0", "arg1"],
126+ ["sig0", "sig1"],
127+ [],
128+ )
129+ self.get_layout_args = MagicMock(return_value=[1, 2])
130+ self.call_kernel = MagicMock()
131+ 
132+ def __enter__(self):
133+ return self
134+ 
135+ def __exit__(self, exc_type, exc_val, exc_tb):
136+ return False
137+ 
138+ 
139+class FakeTemplateBuffer:
140+ def make_kernel_render(self, ctb, epilogue_nodes=None):
141+ kernel = FakeKernel()
142+ 
143+ def render():
144+ return "fake_src_code"
145+ 
146+ return kernel, render
147+ 
148+ def emulate_store_fn(self):
149+ pass
150+ 
151+ 
152+class FakeCatlassSchedulingNode:
153+ def __init__(self):
154+ self.node = FakeTemplateBuffer()
155+ self.group = (None, (1, 1))
156+ 
157+ def mark_run(self):
158+ pass
159+ 
160+ 
161+class FakeCatlassScheduleGraph:
162+ def __init__(self):
163+ self.wrapper_code = FakeWrapperCode()
164+ self.removed_buffers = set()
165+ self.sizevars = FakeSizeVars()
166+ 
167+ 
168+class FakeNPUTritonGraph:
169+ def __init__(self):
170+ self.removed_buffers = set()
171+ self.inplaced_to_remove = set()
172+ 
173+ self.wrapper_code = MagicMock()
174+ self.wrapper_code.supports_intermediate_hooks = False
175+ 
176+ 
177+class FakeNPUTritonFeatures:
178+ def __init__(self, node_schedule):
179+ self.node_schedule = node_schedule
180+ self.numel = 1024
181+ self.reduction_numel = 1
182+ 
183+ def scheduler_nodes(self):
184+ return self.node_schedule
185+ 
186+ 
187+class FakeNPUTritonNode:
188+ def mark_run(self):
189+ pass
190+ 
191+ def get_name(self):
192+ return "fake_node"
193+ 
194+ @property
195+ def node(self):
196+ return None
197+ 
198+ 
199+class FakeNPUTritonKernel:
200+ def __init__(self):
201+ self.kernel_name = "fake_kernel"
202+ self.removed_buffers = set()
203+ self.inplaced_to_remove = set()
204+ 
205+ self.args = MagicMock()
206+ self.args.live_output_buffers.return_value = set()
207+ 
208+ def codegen_kernel(self):
209+ return "fake_src_code"
210+ 
211+ def codegen_nan_check(self):
212+ pass
213+ 
214+ def warn_mix_layout(self, kernel_name):
215+ pass
216+ 
217+ def __enter__(self):
218+ return self
219+ 
220+ def __exit__(self, exc_type, exc_val, exc_tb):
221+ return False
222+ 
223+ def call_kernel(self, name, origin_node=None):
224+ pass
225+ 
226+ 
227+class FakeSchedulerNode:
228+ def __init__(self, is_reduction=False):
229+ self._is_reduction = is_reduction
230+ self.group = ("group_key", (1024, 1))
231+ 
232+ def is_reduction(self):
233+ return self._is_reduction
234+ 
235+ 
236+class FakeFusedNode:
237+ def __init__(self, nodes):
238+ self._nodes = nodes
239+ 
240+ def get_nodes(self):
241+ return self._nodes
242+ 
243+ 
244+class FakeKernelFeatures:
245+ def __init__(self, node_schedule, numel, rnumel):
246+ self.node_schedule = node_schedule
247+ self.numel = numel
248+ self.rnumel = rnumel
249+ 
250+ 
251+class FakeTritonArgs:
252+ def python_argdefs(self):
253+ return (
254+ None,
255+ ["a", "b"], # call_args
256+ None,
257+ ["int", "int"] # arg_types
258+ )
259+ 
260+ @property
261+ def workspace_args(self):
262+ return ["ws1", "ws2"]
263+ 
264+ 
265+class TestMultiStreamPass(TestUtils):
266+ 
267+ def define_catlass_template_kernel(self):
268+ kernel = CATLASSTemplateKernel(kernel_name="test_catlass_kernel")
269+ kernel.args = Mock()
270+ kernel.args.python_argdefs.return_value = (
271+ None, ["input", "weight"], None, [torch.float16, torch.float16]
272+ )
273+ kernel.args.cpp_argdefs.return_value = (
274+ None, ["input", "weight"], ["half*", "half*"]
275+ )
276+ kernel.get_layout_args = Mock(return_value=["dim0", "dim1", "dim2"])
277+ return kernel
278+ 
279+ 
280+ def define_catlass_scheduling(self):
281+ fake_scheduler = MagicMock()
282+ scheduler = CATLASSScheduling(fake_scheduler)
283+ scheduler.is_catlass_template = MagicMock(return_value=True)
284+ scheduler.define_kernel = MagicMock(
285+ return_value="generated_kernel"
286+ )
287+ scheduler.free_buffers_in_scheduler = MagicMock()
288+ return scheduler
289+ 
290+ 
291+ def define_scheduling(self):
292+ fake_scheduler = MagicMock()
293+ scheduling = NPUTritonScheduling(fake_scheduler)
294+ scheduling.select_tiling = MagicMock(return_value="fake_tiling")
295+ scheduling.codegen_node_schedule_with_kernel = MagicMock()
296+ scheduling.make_ttir_for_check = MagicMock()
297+ scheduling.codegen_comment = MagicMock()
298+ scheduling.scheduler = MagicMock()
299+ scheduling.scheduler.free_buffers = MagicMock()
300+ scheduling.generate_node_schedule = MagicMock(
301+ return_value=["node_a", "node_b"]
302+ )
303+ scheduling.codegen_node_schedule = MagicMock(
304+ return_value="codegen_result"
305+ )
306+ return scheduling
307+ 
308+ def define_codegen_node_schedule(self):
309+ sched = MagicMock()
310+ kernel = MagicMock()
311+ kernel.kernel_name = "k0"
312+ kernel.call_kernel = MagicMock()
313+ kernel.removed_buffers = set()
314+ kernel.inplaced_to_remove = set()
315+ kernel.codegen_kernel.return_value = "src"
316+ sched.create_kernel_choices.return_value = [kernel]
317+ sched.select_tiling.return_value = "tiling"
318+ sched.define_kernel.return_value = ("k0", "src")
319+ sched.codegen_node_schedule_with_kernel = MagicMock()
320+ features = MagicMock()
321+ features.node_schedule = []
322+ features.numel = 10
323+ features.reduction_numel = 1
324+ features.scheduler_nodes.return_value = []
325+ nodes = []
326+ origin_node = MagicMock()
327+ return sched, kernel, features, nodes, origin_node
328+ 
329+ @patch("torch_npu._inductor.codegen.catlass.catlass_kernel.is_multi_stream", return_value=False)
330+ def test_catlass_workspace_single_stream(self, mock_multi_stream):
331+ """
332+ case CATLASSTemplateKernel call_kernel:
333+ 1.workspace > 0
334+ 2.single stream
335+ 3.call generate_workspace_allocation/deallocation
336+ 4.kernel_call origin_node=None
337+ """
338+ fake_graph = FakeGraph(cpp_wrapper=False)
339+ node = FakeNode(workspace_size=1024, name="test_01")
340+ with V.set_graph_handler(fake_graph):
341+ kernel = self.define_catlass_template_kernel()
342+ kernel.call_kernel(
343+ name="test_kernel",
344+ node=node,
345+ origin_node="origin",
346+ )
347+ wrapper = fake_graph.wrapper_code
348+ # workspace allocation
349+ wrapper.generate_workspace_allocation.assert_called_once()
350+ alloc_args = wrapper.generate_workspace_allocation.call_args[0]
351+ self.assertIsInstance(alloc_args[0], WorkspaceArg)
352+ # kernel call
353+ wrapper.generate_kernel_call.assert_called_once()
354+ _, kwargs = wrapper.generate_kernel_call.call_args
355+ self.assertEqual(kwargs.get("origin_node", None), None)
356+ self.assertEqual(kwargs["triton"], False)
357+ # workspace deallocation
358+ wrapper.generate_workspace_deallocation.assert_called_once()
359+ 
360+ 
361+ @patch("torch_npu._inductor.codegen.catlass.catlass_kernel.is_multi_stream", return_value=True)
362+ def test_catlass_workspace_multi_stream(self, mock_multi_stream):
363+ """
364+ case CATLASSTemplateKernel call_kernel:
365+ 1. workspace > 0
366+ 2. multi stream
367+ 3. allocation/deallocation with origin_node
368+ 4. kernel_call with origin_node
369+ """
370+ fake_graph = FakeGraph(cpp_wrapper=False)
371+ node = FakeNode(workspace_size=2048, name="test_02")
372+ with V.set_graph_handler(fake_graph):
373+ kernel = self.define_catlass_template_kernel()
374+ kernel.call_kernel(
375+ name="test_kernel",
376+ node=node,
377+ origin_node="origin_node_x",
378+ )
379+ wrapper = fake_graph.wrapper_code
380+ # allocation
381+ wrapper.generate_workspace_allocation.assert_called_once()
382+ alloc_args = wrapper.generate_workspace_allocation.call_args[0]
383+ self.assertEqual(alloc_args[1], "origin_node_x")
384+ # kernel call
385+ wrapper.generate_kernel_call.assert_called_once()
386+ call_args, call_kwargs = wrapper.generate_kernel_call.call_args
387+ self.assertEqual(call_args[0], "test_kernel")
388+ self.assertEqual(call_args[2], "origin_node_x")
389+ # deallocation
390+ wrapper.generate_workspace_deallocation.assert_called_once()
391+ dealloc_args = wrapper.generate_workspace_deallocation.call_args[0]
392+ self.assertEqual(dealloc_args[1], "origin_node_x")
393+ 
394+ 
395+ @patch("torch_npu._inductor.codegen.catlass.catlass_scheduling.is_multi_stream", return_value=True)
396+ def test_catlass_codegen_template_multi_stream(self, mock_multi_stream):
397+ """
398+ case CATLASSScheduling codegen_template:
399+ 1. multi stream
400+ 2. kernel.call_kernel(kernel_name, ctb, template_node)
401+ """
402+ fake_graph = FakeCatlassScheduleGraph()
403+ template_node = FakeCatlassSchedulingNode()
404+ kernel = FakeKernel()
405+ 
406+ def fake_make_kernel_render(ctb, epilogue_nodes=None):
407+ def render():
408+ return "fake_src"
409+ return kernel, render
410+ 
411+ template_node.node.make_kernel_render = fake_make_kernel_render
412+ with V.set_graph_handler(fake_graph):
413+ scheduler = self.define_catlass_scheduling()
414+ scheduler.codegen_template(
415+ template_node=template_node,
416+ epilogue_nodes=[],
417+ prologue_nodes=[],
418+ only_src_code=False,
419+ )
420+ kernel.call_kernel.assert_called_once_with(
421+ "generated_kernel",
422+ template_node.node,
423+ template_node,
424+ )
425+ 
426+ @patch("torch_npu._inductor.codegen.catlass.catlass_scheduling.is_multi_stream", return_value=False)
427+ def test_catlass_codegen_template_single_stream(self, mock_multi_stream):
428+ """
429+ case CATLASSScheduling codegen_template:
430+ 1. single stream
431+ 2. kernel.call_kernel(kernel_name, ctb, None)
432+ """
433+ fake_graph = FakeCatlassScheduleGraph()
434+ template_node = FakeCatlassSchedulingNode()
435+ kernel = FakeKernel()
436+ ctb = template_node.node
437+ 
438+ def fake_make_kernel_render(ctb, epilogue_nodes=None):
439+ def render():
440+ return "fake_src"
441+ return kernel, render
442+ 
443+ ctb.make_kernel_render = fake_make_kernel_render
444+ with V.set_graph_handler(fake_graph):
445+ scheduler = self.define_catlass_scheduling()
446+ 
447+ scheduler.codegen_template(
448+ template_node=template_node,
449+ epilogue_nodes=[],
450+ prologue_nodes=[],
451+ only_src_code=False,
452+ )
453+ 
454+ kernel.call_kernel.assert_called_once_with(
455+ "generated_kernel",
456+ ctb,
457+ None,
458+ )
459+ 
460+ 
461+ @patch("torch_npu._inductor.codegen.scheduling.is_multi_stream", return_value=True)
462+ @patch("torch_npu._inductor.codegen.scheduling.V")
463+ def test_codegen_node_schedule_multi_stream(self, mock_V, mock_multi_stream):
464+ """
465+ case:
466+ 1. multi stream
467+ 2. final_kernel.call_kernel(..., origin_node=origin_node)
468+ """
469+ mock_V.graph.removed_buffers = set()
470+ mock_V.graph.inplaced_to_remove = set()
471+ sched, kernel, features, nodes, origin_node = self.define_codegen_node_schedule()
472+ NPUTritonScheduling.codegen_node_schedule(
473+ sched,
474+ features,
475+ nodes,
476+ origin_node=origin_node
477+ )
478+ kernel.call_kernel.assert_called_once()
479+ _, kwargs = kernel.call_kernel.call_args
480+ assert kwargs["name"] == "k0"
481+ assert kwargs["origin_node"] is origin_node
482+ 
483+ 
484+ @patch("torch_npu._inductor.codegen.scheduling.is_multi_stream", return_value=False)
485+ @patch("torch_npu._inductor.codegen.scheduling.V")
486+ def test_codegen_node_schedule_single_stream(self, mock_V, mock_multi_stream):
487+ """
488+ case:
489+ 1. single stream
490+ 2. final_kernel.call_kernel(..., origin_node=None)
491+ """
492+ mock_V.graph.removed_buffers = set()
493+ mock_V.graph.inplaced_to_remove = set()
494+ sched, kernel, features, nodes, origin_node = self.define_codegen_node_schedule()
495+ NPUTritonScheduling.codegen_node_schedule(
496+ sched,
497+ features,
498+ nodes,
499+ origin_node=origin_node
500+ )
501+ kernel.call_kernel.assert_called_once()
502+ _, kwargs = kernel.call_kernel.call_args
503+ assert kwargs["name"] == "k0"
504+ assert kwargs["origin_node"] is None
505+ 
506+ 
507+ @patch("torch_npu._inductor.codegen.scheduling.is_multi_stream", return_value=True)
508+ def test_codegen_node_multi_stream(
509+ self,
510+ mock_multi_stream
511+ ):
512+ """
513+ case:
514+ 1. multi stream
515+ 2. codegen_node_schedule(..., nodes, node)
516+ """
517+ scheduler = self.define_scheduling()
518+ nodes = [
519+ FakeSchedulerNode(is_reduction=False),
520+ FakeSchedulerNode(is_reduction=True),
521+ ]
522+ fused_node = FakeFusedNode(nodes)
523+ fake_graph = FakeScheduleGraph()
524+ with V.set_graph_handler(fake_graph):
525+ result = scheduler.codegen_node(fused_node)
526+ self.assertEqual(result, "codegen_result")
527+ scheduler.codegen_node_schedule.assert_called_once()
528+ args = scheduler.codegen_node_schedule.call_args[0]
529+ self.assertIsNotNone(args[0])
530+ self.assertEqual(args[1], nodes)
531+ self.assertEqual(args[2], fused_node)
532+ self.assertEqual(len(args), 3)
533+ 
534+ 
535+ @patch("torch_npu._inductor.codegen.scheduling.is_multi_stream", return_value=False)
536+ def test_codegen_node_single_stream(
537+ self,
538+ mock_multi_stream
539+ ):
540+ """
541+ case:
542+ 1. multi stream
543+ 2. codegen_node_schedule(..., nodes)
544+ """
545+ scheduler = self.define_scheduling()
546+ nodes = [
547+ FakeSchedulerNode(is_reduction=False),
548+ FakeSchedulerNode(is_reduction=True),
549+ ]
550+ fused_node = FakeFusedNode(nodes)
551+ fake_graph = FakeScheduleGraph()
552+ with V.set_graph_handler(fake_graph):
553+ result = scheduler.codegen_node(fused_node)
554+ self.assertEqual(result, "codegen_result")
555+ scheduler.codegen_node_schedule.assert_called_once()
556+ args = scheduler.codegen_node_schedule.call_args[0]
557+ self.assertIsNotNone(args[0])
558+ self.assertEqual(args[1], nodes)
559+ self.assertEqual(len(args), 2)
560+ 
561+ 
562+ def test_triton_call_kernel_multi_stream(self):
563+ fake_graph = FakeGraph()
564+ triton_kernel = MagicMock(spec=NPUIndexTritonKernel)
565+ triton_kernel.args = FakeTritonArgs()
566+ triton_kernel.triton_meta = {"meta": 1}
567+ triton_kernel.add_numel_to_call_args = MagicMock()
568+ origin_node = MagicMock()
569+ with patch.object(
570+ torch_npu._inductor.codegen.triton,
571+ "is_multi_stream",
572+ return_value=True,
573+ ):
574+ with V.set_graph_handler(fake_graph):
575+ NPUIndexTritonKernel.call_kernel(
576+ triton_kernel,
577+ "kernel_a",
578+ node=None,
579+ origin_node=origin_node,
580+ )
581+ wrapper = fake_graph.wrapper_code
582+ wrapper.write_triton_header_once.assert_called_once()
583+ triton_kernel.add_numel_to_call_args.assert_called_once()
584+ numel_args = triton_kernel.add_numel_to_call_args.call_args[0]
585+ self.assertEqual(numel_args[0], "kernel_a")
586+ self.assertEqual(numel_args[1], ["a", "b"])
587+ self.assertEqual(numel_args[2], ["int", "int"])
588+ self.assertEqual(
589+ wrapper.generate_workspace_allocation.call_count,
590+ 2,
591+ )
592+ wrapper.generate_workspace_allocation.assert_any_call(
593+ "ws1",
594+ origin_node,
595+ )
596+ wrapper.generate_workspace_allocation.assert_any_call(
597+ "ws2",
598+ origin_node,
599+ )
600+ wrapper.generate_kernel_call.assert_called_once()
601+ self.assertEqual(
602+ wrapper.generate_workspace_deallocation.call_count,
603+ 2,
604+ )
605+ 
606+ def multi_stream_test(
607+ self,
608+ arg0_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1,
609+ arg10_1, arg23_1, arg24_1, arg25_1, arg26_1, arg27_1, arg28_1,
610+ arg29_1, arg30_1, arg31_1, arg32_1, arg33_1, arg34_1,
611+ arg35_1, arg36_1, arg37_1, arg38_1, arg39_1
612+ ):
613+ slice_2 = torch.ops.aten.slice.Tensor(arg0_1, 1, 0, 1)
614+ sum_1 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg2_1, slice_2), [1])
615+ slice_4 = torch.ops.aten.slice.Tensor(arg0_1, 1, 1, 2)
616+ sum_2 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg3_1, slice_4), [1])
617+ slice_6 = torch.ops.aten.slice.Tensor(arg0_1, 1, 2, 3)
618+ sum_3 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg4_1, slice_6), [1])
619+ slice_8 = torch.ops.aten.slice.Tensor(arg0_1, 1, 3, 4)
620+ sum_4 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg5_1, slice_8), [1])
621+ slice_10 = torch.ops.aten.slice.Tensor(arg0_1, 1, 4, 6)
622+ sum_5 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg6_1, slice_10), [1])
623+ slice_12 = torch.ops.aten.slice.Tensor(arg0_1, 1, 6, 7)
624+ sum_6 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg7_1, slice_12), [1])
625+ slice_14 = torch.ops.aten.slice.Tensor(arg0_1, 1, 7, 8)
626+ sum_7 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg8_1, slice_14), [1])
627+ slice_16 = torch.ops.aten.slice.Tensor(arg0_1, 1, 8, 9)
628+ sum_8 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg9_1, slice_16), [1])
629+ slice_18 = torch.ops.aten.slice.Tensor(arg0_1, 1, 9, 10)
630+ sum_9 = torch.ops.aten.sum.dim_IntList(torch.ops.aten.embedding.default(arg10_1, slice_18), [1])
631+ 
632+ cat = torch.ops.aten.cat.default(
633+ [sum_1, sum_2, sum_3, sum_4, sum_5,
634+ sum_6, sum_7, sum_8, sum_9],
635+ 1
636+ )
637+ add_relu = torch.ops.aten.add.Tensor(cat, cat)
638+ 
639+ a = torch.ops.aten.relu.default(torch.ops.aten.mm.default(arg23_1, arg24_1))
640+ a = torch.ops.aten.relu.default(torch.ops.aten.mm.default(a, arg25_1))
641+ a = torch.ops.aten.relu.default(torch.ops.aten.mm.default(a, arg26_1))
642+ a = torch.ops.aten.relu.default(torch.ops.aten.mm.default(a, arg27_1))
643+ mm_4 = torch.ops.aten.relu.default(torch.ops.aten.mm.default(a, arg28_1))
644+ 
645+ b = torch.ops.aten.relu.default(torch.ops.aten.mm.default(arg29_1, arg30_1))
646+ b = torch.ops.aten.relu.default(torch.ops.aten.mm.default(b, arg31_1))
647+ b = torch.ops.aten.relu.default(torch.ops.aten.mm.default(b, arg32_1))
648+ b = torch.ops.aten.relu.default(torch.ops.aten.mm.default(b, arg33_1))
649+ mm_8 = torch.ops.aten.relu.default(torch.ops.aten.mm.default(b, arg34_1))
650+ 
651+ mm_8_t = torch.ops.aten.permute.default(mm_8, [1, 0])
652+ merge = torch.ops.aten.mm.default(mm_4, mm_8_t)
653+ merge = torch.ops.aten.mm.default(merge, arg35_1)
654+ 
655+ add = torch.ops.aten.add.Tensor(merge, add_relu)
656+ mm_15 = torch.ops.aten.mm.default(arg36_1, add)
657+ relu_15 = torch.ops.aten.relu.default(mm_15)
658+ add_15 = torch.ops.aten.add.Tensor(arg37_1, relu_15)
659+ out = torch.ops.aten.addmm.default(
660+ arg39_1,
661+ arg38_1,
662+ add_15,
663+ )
664+ return out
665+ 
666+ @patch("torch_npu._inductor.fx_passes.parallel_scheduler_pass.is_multi_stream", return_value=True)
667+ def test_multi_stream_compile_case(self, mock_multi_stream):
668+ arg0_1 = torch.randint(0, 16, (8, 40), dtype=torch.int64, device="npu")
669+ arg1_1 = torch.randn(16, 8, device="npu")
670+ arg2_1 = torch.randn(16, 8, device="npu")
671+ arg3_1 = torch.randn(16, 8, device="npu")
672+ arg4_1 = torch.randn(16, 8, device="npu")
673+ arg5_1 = torch.randn(16, 8, device="npu")
674+ arg6_1 = torch.randn(16, 8, device="npu")
675+ arg7_1 = torch.randn(16, 8, device="npu")
676+ arg8_1 = torch.randn(16, 8, device="npu")
677+ arg9_1 = torch.randn(16, 8, device="npu")
678+ arg10_1 = torch.randn(16, 8, device="npu")
679+ 
680+ arg23_1 = torch.randn(8, 8, device="npu")
681+ arg24_1 = torch.randn(8, 8, device="npu")
682+ arg25_1 = torch.randn(8, 8, device="npu")
683+ arg26_1 = torch.randn(8, 8, device="npu")
684+ arg27_1 = torch.randn(8, 8, device="npu")
685+ arg28_1 = torch.randn(8, 8, device="npu")
686+ arg29_1 = torch.randn(8, 8, device="npu")
687+ arg30_1 = torch.randn(8, 8, device="npu")
688+ arg31_1 = torch.randn(8, 8, device="npu")
689+ arg32_1 = torch.randn(8, 8, device="npu")
690+ arg33_1 = torch.randn(8, 8, device="npu")
691+ arg34_1 = torch.randn(8, 8, device="npu")
692+ arg35_1 = torch.randn(8, 72, device="npu")
693+ arg36_1 = torch.randn(8, 8, device="npu")
694+ arg37_1 = torch.randn(8, 72, device="npu")
695+ arg38_1 = torch.randn(8, 8, device="npu")
696+ arg39_1 = torch.randn(8, 72, device="npu")
697+ 
698+ std_result = self.multi_stream_test(arg0_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1,
699+ arg10_1, arg23_1, arg24_1, arg25_1, arg26_1, arg27_1, arg28_1, arg29_1,
700+ arg30_1, arg31_1, arg32_1, arg33_1, arg34_1, arg35_1, arg36_1, arg37_1,
701+ arg38_1, arg39_1)
702+ with torch.no_grad():
703+ compiled_op_calc = torch.compile(self.multi_stream_test, backend="inductor")
704+ inductor_result = compiled_op_calc(arg0_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1,
705+ arg10_1, arg23_1, arg24_1, arg25_1, arg26_1, arg27_1, arg28_1, arg29_1,
706+ arg30_1, arg31_1, arg32_1, arg33_1, arg34_1, arg35_1, arg36_1, arg37_1,
707+ arg38_1, arg39_1)
708+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
709+ 
710+ def op_calc(self, arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1):
711+ relu_1 = torch.ops.aten.relu.default(arg0_1)
712+ mul_1 = torch.ops.aten.mm.default(arg1_1, arg2_1)
713+ relu_2 = torch.ops.aten.relu.default(mul_1)
714+ mm_2 = torch.ops.aten.mm.default(relu_1, relu_2)
715+ relu_3 = torch.ops.aten.relu.default(mm_2)
716+ mul_3 = torch.ops.aten.mm.default(relu_3, relu_2)
717+ relu_4 = torch.ops.aten.relu.default(arg3_1)
718+ relu_5 = torch.ops.aten.relu.default(arg4_1)
719+ add_1 = torch.ops.aten.add.Tensor(relu_4, relu_5)
720+ relu_6 = torch.ops.aten.relu.default(arg5_1)
721+ add_2 = torch.ops.aten.add.Tensor(add_1, relu_6)
722+ relu_7 = torch.ops.aten.relu.default(add_2)
723+ silu_1 = torch.ops.aten.silu.default(relu_7)
724+ add_3 = torch.ops.aten.add.Tensor(mul_3, silu_1)
725+ slice_1 = torch.ops.aten.slice.Tensor(add_3, dim=0, start=0, end=128, step=1)
726+ relu_8 = torch.ops.aten.relu.default(slice_1)
727+ slice_2 = torch.ops.aten.slice.Tensor(slice_1, dim=0, start=0, end=128, step=1)
728+ relu_9 = torch.ops.aten.relu.default(slice_1)
729+ add_2 = torch.ops.aten.add.Tensor(relu_8, slice_2)
730+ add_3 = torch.ops.aten.add.Tensor(slice_2, relu_9)
731+ mm_3 = torch.ops.aten.mm.default(add_2, add_3)
732+ relu_10 = torch.ops.aten.relu.default(mm_3)
733+ return relu_10
734+ 
735+ @patch("torch_npu._inductor.fx_passes.parallel_scheduler_pass.is_multi_stream", return_value=False)
736+ def test_single_stream_compile_case(self, mock_multi_stream):
737+ arg0_1 = torch.randn(128, 128, device="npu")
738+ arg1_1 = torch.randn(128, 64, device="npu")
739+ arg2_1 = torch.randn(64, 128, device="npu")
740+ arg3_1 = torch.randn(128, 128, device="npu")
741+ arg4_1 = torch.randn(128, 128, device="npu")
742+ arg5_1 = torch.randn(128, 128, device="npu")
743+ std_result = self.op_calc(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1)
744+ with torch.no_grad():
745+ compiled_op_calc = torch.compile(self.op_calc, backend="inductor")
746+ inductor_result = compiled_op_calc(arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1)
747+ self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3)
748+ 
749+ 
750+instantiate_parametrized_tests(TestMultiStreamPass)
751+ 
752+ 
753+if __name__ == "__main__":
754+ run_tests()