已合并
[inductor] fix: remove debug prints for x0 indexing_code #32260
weizhan4创建于 3月24日
[inductor] fix: remove debug prints for x0 indexing_code #32260
已合并
共 3 个文件变更+52-14
| @@ -0,0 +1,31 @@ | |||
| 1 | +import unittest | ||
| 2 | +import torch | ||
| 3 | +from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | ||
| 4 | +from testutils import TestUtils | ||
| 5 | +import torch_npu | ||
| 6 | + | ||
| 7 | + | ||
| 8 | +class TestFactoryReduction(TestUtils): | ||
| 9 | + def op_calc(self, shape, dtype_val): | ||
| 10 | + # Use dtype_val which is actual torch dtype instead of torch.float32 directly | ||
| 11 | + # based on the parameter passed | ||
| 12 | + full_tensor = torch.ops.aten.full.default(shape, 1.0, dtype=dtype_val, layout=torch.strided, device='npu', pin_memory=False) | ||
| 13 | + result = torch.ops.aten.amax.default(full_tensor, [2]) | ||
| 14 | + return result | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + def test_factory_reduction(self, shape, dtype): | ||
| 19 | + # Get actual torch dtype from string | ||
| 20 | + dtype_val = getattr(torch, dtype) | ||
| 21 | + | ||
| 22 | + std_result = self.op_calc(shape, dtype_val) | ||
| 23 | + | ||
| 24 | + compiled_op_calc = torch.compile(self.op_calc, backend="inductor") | ||
| 25 | + inductor_result = compiled_op_calc(shape, dtype_val) | ||
| 26 | + self.assertEqual(std_result, inductor_result, atol=1e-3, rtol=1e-3) | ||
| 27 | + | ||
| 28 | +instantiate_parametrized_tests(TestFactoryReduction) | ||
| 29 | + | ||
| 30 | +if __name__ == "__main__": | ||
| 31 | + run_tests() | ||
| @@ -192,7 +192,7 @@ class IndexAnalysis: | |||
| 192 | 192 | ||
| 193 | def all_tiling_in_var_list(): | 193 | def all_tiling_in_var_list(): |
| 194 | return all([x in self.var_list for x in self.tiling_axis]) | 194 | return all([x in self.var_list for x in self.tiling_axis]) |
| 195 | - # 2 analyze permute shape for full_dim_len index | 195 | + # 2 analyze permute shape for full_dim_len index |
| 196 | 196 | ||
| 197 | if all_tiling_in_var_list(): | 197 | if all_tiling_in_var_list(): |
| 198 | self.similar = self.var_list | 198 | self.similar = self.var_list |
| @@ -325,6 +325,7 @@ class IterationRangesEntryNPUIndex(IterationRangesEntry): | |||
| 325 | if index: | 325 | if index: |
| 326 | self.writeline(index) | 326 | self.writeline(index) |
| 327 | self._codegen_mask() | 327 | self._codegen_mask() |
| 328 | + | ||
| 328 | return self.name | 329 | return self.name |
| 329 | 330 | ||
| 330 | def writeline(self, line): | 331 | def writeline(self, line): |
| @@ -1072,21 +1073,20 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1072 | def find_axis_in_load_store(self, range_val): | 1073 | def find_axis_in_load_store(self, range_val): |
| 1073 | if not range_val: | 1074 | if not range_val: |
| 1074 | return False | 1075 | return False |
| 1076 | + | ||
| 1075 | for line in self.loads._lines: | 1077 | for line in self.loads._lines: |
| 1076 | - if line.find('tl.load') >= 0 and self.is_isolated_symbol(line, range_val): | 1078 | + if self.is_isolated_symbol(line, range_val): |
| 1077 | return True | 1079 | return True |
| 1078 | for line in self.compute._lines: | 1080 | for line in self.compute._lines: |
| 1079 | - if line.find('tl.load') >= 0 and self.is_isolated_symbol(line, range_val): | 1081 | + if self.is_isolated_symbol(line, range_val): |
| 1080 | return True | 1082 | return True |
| 1081 | for line in self.post_loop_store._lines: | 1083 | for line in self.post_loop_store._lines: |
| 1082 | - if isinstance(line, DeferredLine): | 1084 | + str_line = line.line if isinstance(line, DeferredLine) else line |
| 1083 | - line = line.line | 1085 | + if self.is_isolated_symbol(str_line, range_val): |
| 1084 | - if line.find('tl.store') >= 0 and self.is_isolated_symbol(line, range_val): | ||
| 1085 | return True | 1086 | return True |
| 1086 | for line in self.stores._lines: | 1087 | for line in self.stores._lines: |
| 1087 | - if isinstance(line, DeferredLine): | 1088 | + str_line = line.line if isinstance(line, DeferredLine) else line |
| 1088 | - line = line.line | 1089 | + if self.is_isolated_symbol(str_line, range_val): |
| 1089 | - if line.find('tl.store') >= 0 and self.is_isolated_symbol(line, range_val): | ||
| 1090 | return True | 1090 | return True |
| 1091 | return False | 1091 | return False |
| 1092 | 1092 | ||
| @@ -1142,12 +1142,9 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1142 | 1142 | ||
| 1143 | is_last_axis = index == len(self.sorted_axis) - 1 | 1143 | is_last_axis = index == len(self.sorted_axis) - 1 |
| 1144 | indexing_code = getattr(range_val, "indexing_code") | 1144 | indexing_code = getattr(range_val, "indexing_code") |
| 1145 | + | ||
| 1145 | reduction_1d = is_1d_reduction() | 1146 | reduction_1d = is_1d_reduction() |
| 1146 | do_indent = False | 1147 | do_indent = False |
| 1147 | - # do nothing except for writing porintwise | ||
| 1148 | - if len(self.loads._lines) == 0 and len(self.stores._lines) == 0: | ||
| 1149 | - do_indent = False | ||
| 1150 | - indexing_code = None | ||
| 1151 | # tiling axis and last tiling | 1148 | # tiling axis and last tiling |
| 1152 | if range_val.is_tiling_axis and last_tiling: | 1149 | if range_val.is_tiling_axis and last_tiling: |
| 1153 | do_indent = False | 1150 | do_indent = False |
| @@ -1166,7 +1163,8 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1166 | # tiling axis and but not last tiling | 1163 | # tiling axis and but not last tiling |
| 1167 | elif range_val.is_tiling_axis: | 1164 | elif range_val.is_tiling_axis: |
| 1168 | do_indent = False | 1165 | do_indent = False |
| 1169 | - if len(self.loads._lines) == 0 and len(self.stores._lines) == 0: | 1166 | + have_load_store = self.find_axis_in_load_store(range_val) |
| 1167 | + if not have_load_store: | ||
| 1170 | do_indent = False | 1168 | do_indent = False |
| 1171 | indexing_code = None | 1169 | indexing_code = None |
| 1172 | if not range_val.is_no_loop_axis: | 1170 | if not range_val.is_no_loop_axis: |
| @@ -1478,6 +1476,15 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1478 | self.golden_var_list = tuple([x.symbol() for x in self.tiling_axis]) if self.tiling_axis else [] | 1476 | self.golden_var_list = tuple([x.symbol() for x in self.tiling_axis]) if self.tiling_axis else [] |
| 1479 | else: | 1477 | else: |
| 1480 | self.golden_var_list = tuple([x for x in longest if x in self.tiling_axis]) if self.tiling_axis else [] | 1478 | self.golden_var_list = tuple([x for x in longest if x in self.tiling_axis]) if self.tiling_axis else [] |
| 1479 | + | ||
| 1480 | + if self.golden_var_list is not None and self.tiling_axis: | ||
| 1481 | + golden_list = list(self.golden_var_list) | ||
| 1482 | + for x in self.tiling_axis: | ||
| 1483 | + sym = x.symbol() if hasattr(x, 'symbol') else x | ||
| 1484 | + if sym not in golden_list: | ||
| 1485 | + golden_list.append(sym) | ||
| 1486 | + self.golden_var_list = tuple(golden_list) | ||
| 1487 | + | ||
| 1481 | if self.golden_var_list is None: | 1488 | if self.golden_var_list is None: |
| 1482 | raise RuntimeError("assert self.golden_var_list is None") | 1489 | raise RuntimeError("assert self.golden_var_list is None") |
| 1483 | 1490 | ||