已合并
fix(inductor): preserve scheduler store semantics in linear triton transpose-copy #39857
Xuan Peng创建于 7月2日
fix(inductor): preserve scheduler store semantics in linear triton transpose-copy #39857
已合并
共 4 个文件变更+356-30
| @@ -0,0 +1,82 @@ | |||
| 1 | +import re | ||
| 2 | +from unittest import skip | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | +from torch._inductor.utils import run_and_get_code | ||
| 6 | +from torch.testing._internal.common_utils import run_tests | ||
| 7 | + | ||
| 8 | + | ||
| 9 | +try: | ||
| 10 | + from .testutils import TestUtils | ||
| 11 | +except ImportError: | ||
| 12 | + from testutils import TestUtils | ||
| 13 | + | ||
| 14 | +torch._inductor.config.fx_graph_cache = False | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +class TestLinearTritonStoreSemantics(TestUtils): | ||
| 18 | + _permute_pattern = re.compile(r"\.permute\(") | ||
| 19 | + _axis_replacement_pattern = re.compile(r"\b[trzyx]\d+_\d+(?:_nd)?\b") | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + def transpose_clone_rectangular(x): | ||
| 23 | + y = x.view(-1, 80, 40, 8) | ||
| 24 | + return y.permute(0, 2, 1, 3).clone() | ||
| 25 | + | ||
| 26 | + | ||
| 27 | + def transpose_unary_clone_rectangular(x): | ||
| 28 | + y = x.view(-1, 80, 40, 8) | ||
| 29 | + return y.permute(0, 2, 1, 3).sin().clone() | ||
| 30 | + | ||
| 31 | + | ||
| 32 | + def broadcast_add(x, bias): | ||
| 33 | + return x + bias.view(1, 1, 1, -1) | ||
| 34 | + | ||
| 35 | + def _assert_scheduler_semantic_store_codegen(self, codes): | ||
| 36 | + self.assertTrue(codes) | ||
| 37 | + self.assertFalse( | ||
| 38 | + any(self._permute_pattern.search(code) for code in codes), | ||
| 39 | + msg="expected scheduler-semantic store path without RHS permute", | ||
| 40 | + ) | ||
| 41 | + self.assertFalse( | ||
| 42 | + any(self._axis_replacement_pattern.search(code) for code in codes), | ||
| 43 | + msg="expected scheduler-semantic store path without remapped axis symbols", | ||
| 44 | + ) | ||
| 45 | + | ||
| 46 | + def test_scheduler_semantic_store_path_for_transpose_clone(self): | ||
| 47 | + x = self._generate_tensor((381, 80, 320), "float32") | ||
| 48 | + compiled = torch.compile( | ||
| 49 | + self.transpose_clone_rectangular, backend="inductor", dynamic=False | ||
| 50 | + ) | ||
| 51 | + out, codes = run_and_get_code(compiled, x) | ||
| 52 | + eager = self.transpose_clone_rectangular(x) | ||
| 53 | + torch.testing.assert_close(out, eager, rtol=1e-4, atol=1e-4) | ||
| 54 | + self._assert_scheduler_semantic_store_codegen(codes) | ||
| 55 | + | ||
| 56 | + # AssertionError: assert_size_stride(buf1, (381, 40, 80, 8), (25600, 8, 320, 1), 'torch.ops.aten.sin.default') | ||
| 57 | + | ||
| 58 | + def test_scheduler_semantic_store_path_survives_simple_pointwise(self): | ||
| 59 | + x = self._generate_tensor((381, 80, 320), "float32") | ||
| 60 | + compiled = torch.compile( | ||
| 61 | + self.transpose_unary_clone_rectangular, | ||
| 62 | + backend="inductor", | ||
| 63 | + dynamic=False, | ||
| 64 | + ) | ||
| 65 | + out, codes = run_and_get_code(compiled, x) | ||
| 66 | + eager = self.transpose_unary_clone_rectangular(x) | ||
| 67 | + torch.testing.assert_close(out, eager, rtol=1e-4, atol=1e-4) | ||
| 68 | + self._assert_scheduler_semantic_store_codegen(codes) | ||
| 69 | + | ||
| 70 | + def test_remapped_store_path_still_reachable(self): | ||
| 71 | + x = self._generate_tensor((8, 8, 32, 16), "float32") | ||
| 72 | + bias = self._generate_tensor((16,), "float32") | ||
| 73 | + compiled = torch.compile(self.broadcast_add, backend="inductor", dynamic=False) | ||
| 74 | + out, codes = run_and_get_code(compiled, x, bias) | ||
| 75 | + eager = self.broadcast_add(x, bias) | ||
| 76 | + torch.testing.assert_close(out, eager, rtol=1e-4, atol=1e-4) | ||
| 77 | + self.assertTrue(codes) | ||
| 78 | + self.assertTrue(any("tl.store(" in code for code in codes)) | ||
| 79 | + | ||
| 80 | + | ||
| 81 | +if __name__ == "__main__": | ||
| 82 | + run_tests() | ||
| @@ -20,6 +20,11 @@ class TestPermute(TestUtils): | |||
| 20 | y = a + b | 20 | y = a + b |
| 21 | return y | 21 | return y |
| 22 | 22 | ||
| 23 | + | ||
| 24 | + def transpose_clone_square(x): | ||
| 25 | + y = x.view(-1, 80, 80, 8) | ||
| 26 | + return y.permute(0, 2, 1, 3).clone() | ||
| 27 | + | ||
| 23 | 28 | ||
| 24 | 29 | ||
| 25 | def test_view_cases(self, shape, dtype): | 30 | def test_view_cases(self, shape, dtype): |
| @@ -33,6 +38,15 @@ class TestPermute(TestUtils): | |||
| 33 | 38 | ||
| 34 | self.assertEqual(std_permute, inductor_permute, atol=1e-3, rtol=1e-3) | 39 | self.assertEqual(std_permute, inductor_permute, atol=1e-3, rtol=1e-3) |
| 35 | 40 | ||
| 41 | + def test_transpose_clone_square(self): | ||
| 42 | + x = self._generate_tensor((381, 80, 640), "float32") | ||
| 43 | + eager = self.transpose_clone_square(x) | ||
| 44 | + compiled = torch.compile( | ||
| 45 | + self.transpose_clone_square, backend="inductor", dynamic=False | ||
| 46 | + ) | ||
| 47 | + actual = compiled(x) | ||
| 48 | + torch.testing.assert_close(actual, eager, rtol=1e-4, atol=1e-4) | ||
| 49 | + | ||
| 36 | instantiate_parametrized_tests(TestPermute) | 50 | instantiate_parametrized_tests(TestPermute) |
| 37 | 51 | ||
| 38 | if __name__ == "__main__": | 52 | if __name__ == "__main__": |
| @@ -170,7 +170,7 @@ class IndexAnalysis: | |||
| 170 | index = similar.index(x) | 170 | index = similar.index(x) |
| 171 | self.reshape_sizes[index] = f"{x.name.upper()}BLOCK_SUB" | 171 | self.reshape_sizes[index] = f"{x.name.upper()}BLOCK_SUB" |
| 172 | 172 | ||
| 173 | - def analyze_var_direction(self, nddma=False): | 173 | + def analyze_var_direction(self, nddma=False, materialize_var_directions=True): |
| 174 | if self.var_list == self.gold: | 174 | if self.var_list == self.gold: |
| 175 | return | 175 | return |
| 176 | var_list = self.var_list if len(self.var_list) == len(self.gold) else self.similar | 176 | var_list = self.var_list if len(self.var_list) == len(self.gold) else self.similar |
| @@ -217,13 +217,16 @@ class IndexAnalysis: | |||
| 217 | continue | 217 | continue |
| 218 | self.var_replacements[x] = var_obj | 218 | self.var_replacements[x] = var_obj |
| 219 | self.var_directions[var_obj] = direction_str | 219 | self.var_directions[var_obj] = direction_str |
| 220 | - self.kernel.range_tree_nodes[x].var_directions[var_obj] = direction_str | 220 | + # Store-side semantic checks need a no-side-effect analysis mode so we |
| 221 | + # can inspect replacement directions without polluting final codegen. | ||
| 222 | + if materialize_var_directions: | ||
| 223 | + self.kernel.range_tree_nodes[x].var_directions[var_obj] = direction_str | ||
| 221 | 224 | ||
| 222 | if processed_nddma: | 225 | if processed_nddma: |
| 223 | self.processed_nddma = True | 226 | self.processed_nddma = True |
| 224 | self.need_permute = False | 227 | self.need_permute = False |
| 225 | 228 | ||
| 226 | - def analyze_index(self, nddma=False): | 229 | + def analyze_index(self, nddma=False, materialize_var_directions=True): |
| 227 | if isinstance(self.index, sympy.Integer): | 230 | if isinstance(self.index, sympy.Integer): |
| 228 | return | 231 | return |
| 229 | if not self.kernel.golden_var_list: | 232 | if not self.kernel.golden_var_list: |
| @@ -252,7 +255,9 @@ class IndexAnalysis: | |||
| 252 | pass | 255 | pass |
| 253 | 256 | ||
| 254 | # 4 analyze var direction | 257 | # 4 analyze var direction |
| 255 | - self.analyze_var_direction(nddma) | 258 | + self.analyze_var_direction( |
| 259 | + nddma, materialize_var_directions=materialize_var_directions | ||
| 260 | + ) | ||
| 256 | 261 | ||
| 257 | def generate_statement(self): | 262 | def generate_statement(self): |
| 258 | statement = "" | 263 | statement = "" |
| @@ -2810,39 +2810,66 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 2810 | self, name: str, index: sympy.Expr, value: CSEVariable, mode: StoreMode = None | 2810 | self, name: str, index: sympy.Expr, value: CSEVariable, mode: StoreMode = None |
| 2811 | ) -> None: | 2811 | ) -> None: |
| 2812 | var = self.args.output(name) | 2812 | var = self.args.output(name) |
| 2813 | - index_analyze = IndexAnalysis(self, index, is_store_index=True) | ||
| 2814 | - index_analyze.analyze_index() | ||
| 2815 | - indexing = self.indexing( | ||
| 2816 | - index, | ||
| 2817 | - dense_indexing=True, | ||
| 2818 | - block_ptr=mode is None, | ||
| 2819 | - index_analyze=index_analyze, | ||
| 2820 | - ) | ||
| 2821 | - index_str = indexing.index_str | ||
| 2822 | value_str = f"{value}" | 2813 | value_str = f"{value}" |
| 2823 | - mask_str = indexing.mask_str | 2814 | + selected_indexing = None |
| 2824 | 2815 | ||
| 2825 | - if index_analyze.need_permute: | 2816 | + dryrun_index_analyze = IndexAnalysis(self, index, is_store_index=True) |
| 2826 | - value_str = value_str.replace( | 2817 | + # First inspect the store index without materializing replacement axes. |
| 2827 | - f"{value}", f"{value}{index_analyze.generate_statement()}" | 2818 | + # If the raw scheduler store layout already matches the RHS buffer |
| 2819 | + # layout, preserve the original store semantics and skip remapping. | ||
| 2820 | + # If layout metadata has already been dropped by an earlier update / | ||
| 2821 | + # inplace chain, we conservatively fall back to the remapped store path. | ||
| 2822 | + dryrun_index_analyze.analyze_index(materialize_var_directions=False) | ||
| 2823 | + if getattr(value, "layout_known", False) and dryrun_index_analyze.need_permute: | ||
| 2824 | + raw_indexing = self.indexing( | ||
| 2825 | + index, | ||
| 2826 | + dense_indexing=True, | ||
| 2827 | + block_ptr=False, | ||
| 2828 | + index_analyze=dryrun_index_analyze, | ||
| 2829 | + apply_var_replacements=False, | ||
| 2830 | + materialize_var_directions=False, | ||
| 2828 | ) | 2831 | ) |
| 2832 | + if self._can_preserve_store_semantics_with_value_layout( | ||
| 2833 | + value, raw_indexing, dryrun_index_analyze | ||
| 2834 | + ): | ||
| 2835 | + selected_indexing = raw_indexing | ||
| 2836 | + | ||
| 2837 | + if selected_indexing is None: | ||
| 2838 | + # Fall back to the legacy remapped store path when raw scheduler | ||
| 2839 | + # semantics do not match the current RHS buffer layout. | ||
| 2840 | + index_analyze = IndexAnalysis(self, index, is_store_index=True) | ||
| 2841 | + index_analyze.analyze_index() | ||
| 2842 | + indexing = self.indexing( | ||
| 2843 | + index, | ||
| 2844 | + dense_indexing=True, | ||
| 2845 | + block_ptr=mode is None, | ||
| 2846 | + index_analyze=index_analyze, | ||
| 2847 | + ) | ||
| 2848 | + if index_analyze.need_permute: | ||
| 2849 | + value_str = value_str.replace( | ||
| 2850 | + f"{value}", f"{value}{index_analyze.generate_statement()}" | ||
| 2851 | + ) | ||
| 2852 | + selected_indexing = indexing | ||
| 2853 | + | ||
| 2854 | + index_str = selected_indexing.index_str | ||
| 2855 | + mask_str = selected_indexing.mask_str | ||
| 2829 | 2856 | ||
| 2830 | advance_block_ptr = None | 2857 | advance_block_ptr = None |
| 2831 | - if isinstance(indexing, BlockPtrOptions): | 2858 | + if isinstance(selected_indexing, BlockPtrOptions): |
| 2832 | block_ptr, advance_block_ptr, other = self.codegen_block_ptr( | 2859 | block_ptr, advance_block_ptr, other = self.codegen_block_ptr( |
| 2833 | - name, var, indexing | 2860 | + name, var, selected_indexing |
| 2834 | ) | 2861 | ) |
| 2835 | # block_ptr stores don't do implicit casting | 2862 | # block_ptr stores don't do implicit casting |
| 2836 | line = self.codegen_block_ptr_store_line( | 2863 | line = self.codegen_block_ptr_store_line( |
| 2837 | - name, indexing, block_ptr, value, other | 2864 | + name, selected_indexing, block_ptr, value, other |
| 2838 | ) | 2865 | ) |
| 2839 | elif mode is None: | 2866 | elif mode is None: |
| 2840 | line = f"tl.store({var} + ({index_str}), {value_str}, {mask_str})" | 2867 | line = f"tl.store({var} + ({index_str}), {value_str}, {mask_str})" |
| 2841 | if self.numof_reduction_axis() > 1: | 2868 | if self.numof_reduction_axis() > 1: |
| 2842 | - line = f"tl.store({var} + ({index_str} + tl.arange(0,1) ), {value_str}, {indexing.mask_str})" | 2869 | + line = f"tl.store({var} + ({index_str} + tl.arange(0,1) ), {value_str}, {selected_indexing.mask_str})" |
| 2843 | 2870 | ||
| 2844 | elif mode == "atomic_add": | 2871 | elif mode == "atomic_add": |
| 2845 | - line = f"tl.atomic_add({var} + ({index_str}), {value_str}, {indexing.mask_str})" | 2872 | + line = f"tl.atomic_add({var} + ({index_str}), {value_str}, {selected_indexing.mask_str})" |
| 2846 | else: | 2873 | else: |
| 2847 | raise NotImplementedError(f"store mode={mode}") | 2874 | raise NotImplementedError(f"store mode={mode}") |
| 2848 | 2875 | ||
| @@ -3934,9 +3961,13 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 3934 | 3961 | ||
| 3935 | index_analyze = IndexAnalysis(self, index) | 3962 | index_analyze = IndexAnalysis(self, index) |
| 3936 | nddma_switch = npu_config.nddma_switch | 3963 | nddma_switch = npu_config.nddma_switch |
| 3937 | - index_analyze.analyze_index(nddma=nddma_switch) | ||
| 3938 | indirect_indexing = self.is_indirect_indexing(index) | 3964 | indirect_indexing = self.is_indirect_indexing(index) |
| 3939 | - indexing = self.indexing(index, nddma=nddma_switch, block_ptr=True) | 3965 | + indexing = self.indexing( |
| 3966 | + index, | ||
| 3967 | + nddma=nddma_switch, | ||
| 3968 | + block_ptr=True, | ||
| 3969 | + index_analyze=index_analyze, | ||
| 3970 | + ) | ||
| 3940 | has_rindex = indexing.has_rindex() | 3971 | has_rindex = indexing.has_rindex() |
| 3941 | has_tmpmask = indexing.has_tmpmask() | 3972 | has_tmpmask = indexing.has_tmpmask() |
| 3942 | ep = "" | 3973 | ep = "" |
| @@ -4009,14 +4040,33 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4009 | if not (isinstance(result_var, TritonCSEVariable)): | 4040 | if not (isinstance(result_var, TritonCSEVariable)): |
| 4010 | raise RuntimeError("assert isinstance(result_var, TritonCSEVariable)") | 4041 | raise RuntimeError("assert isinstance(result_var, TritonCSEVariable)") |
| 4011 | result_var.mask_vars = indexing.mask_vars # type: ignore[assignment] | 4042 | result_var.mask_vars = indexing.mask_vars # type: ignore[assignment] |
| 4043 | + # Track the buffer layout produced by the finalized load indexing. | ||
| 4044 | + loaded_layout_axes = ( | ||
| 4045 | + self._infer_layout_axes_from_expr(indexing.index, index_analyze=index_analyze) | ||
| 4046 | + if isinstance(indexing, IndexingOptions) | ||
| 4047 | + else None | ||
| 4048 | + ) | ||
| 4049 | + self._set_layout_axes(result_var, loaded_layout_axes) | ||
| 4012 | 4050 | ||
| 4013 | if append_broadcast and append_broadcast != "[]": | 4051 | if append_broadcast and append_broadcast != "[]": |
| 4014 | line = f"tl.reshape({result_var}, {append_broadcast})" | 4052 | line = f"tl.reshape({result_var}, {append_broadcast})" |
| 4015 | result_var = self.cse.generate(load_buffer, line, dtype=dtype, shape=indexing.expand_shape) | 4053 | result_var = self.cse.generate(load_buffer, line, dtype=dtype, shape=indexing.expand_shape) |
| 4054 | + # Once a scalar/degenerate load is reshaped, we no longer trust the | ||
| 4055 | + # direct axis-to-slot mapping for scheduler-semantic store checks. | ||
| 4056 | + self._mark_layout_unknown(result_var) | ||
| 4016 | # triton can handle broadcast | 4057 | # triton can handle broadcast |
| 4017 | elif index_analyze.need_permute: | 4058 | elif index_analyze.need_permute: |
| 4018 | line = f"{result_var}{index_analyze.generate_statement()}" | 4059 | line = f"{result_var}{index_analyze.generate_statement()}" |
| 4019 | result_var = self.cse.generate(self.loads, line, dtype=dtype, shape=result_var.shape) | 4060 | result_var = self.cse.generate(self.loads, line, dtype=dtype, shape=result_var.shape) |
| 4061 | + if index_analyze.need_reshape or index_analyze.need_broadcast: | ||
| 4062 | + # Mixed reshape/broadcast/permute transforms are conservatively | ||
| 4063 | + # treated as layout-unknown. | ||
| 4064 | + self._mark_layout_unknown(result_var) | ||
| 4065 | + else: | ||
| 4066 | + permuted_layout = self._permute_layout_axes( | ||
| 4067 | + loaded_layout_axes, index_analyze.permute_shape | ||
| 4068 | + ) | ||
| 4069 | + self._set_layout_axes(result_var, permuted_layout) | ||
| 4020 | 4070 | ||
| 4021 | if advance_block_ptr: | 4071 | if advance_block_ptr: |
| 4022 | load_buffer.writeline(advance_block_ptr) | 4072 | load_buffer.writeline(advance_block_ptr) |
| @@ -4026,9 +4076,14 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4026 | 4076 | ||
| 4027 | return result_var | 4077 | return result_var |
| 4028 | 4078 | ||
| 4029 | - # don't call symlify_indexing | 4079 | + # don't call simplify_indexing |
| 4030 | def prepare_indexing( | 4080 | def prepare_indexing( |
| 4031 | - self, index: sympy.Expr, index_analyze, is_index_expr=False, nddma=False | 4081 | + self, |
| 4082 | + index: sympy.Expr, | ||
| 4083 | + index_analyze, | ||
| 4084 | + is_index_expr=False, | ||
| 4085 | + nddma=False, | ||
| 4086 | + materialize_var_directions=True, | ||
| 4032 | ): | 4087 | ): |
| 4033 | index = sympy_subs(index, V.graph.sizevars.precomputed_replacements) | 4088 | index = sympy_subs(index, V.graph.sizevars.precomputed_replacements) |
| 4034 | # if simple replacements didn't get rid of floor/ceil, try full subs | 4089 | # if simple replacements didn't get rid of floor/ceil, try full subs |
| @@ -4054,7 +4109,9 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4054 | ) | 4109 | ) |
| 4055 | 4110 | ||
| 4056 | # to generate range.var_directions for permuted axis | 4111 | # to generate range.var_directions for permuted axis |
| 4057 | - index_analyze.analyze_index(nddma) | 4112 | + index_analyze.analyze_index( |
| 4113 | + nddma, materialize_var_directions=materialize_var_directions | ||
| 4114 | + ) | ||
| 4058 | return self.codegen_indexing(simp_index) | 4115 | return self.codegen_indexing(simp_index) |
| 4059 | 4116 | ||
| 4060 | def replace_index_vars(self, index, index_analyze): | 4117 | def replace_index_vars(self, index, index_analyze): |
| @@ -4088,15 +4145,25 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4088 | index_analyze=None, | 4145 | index_analyze=None, |
| 4089 | is_index_expr=False, | 4146 | is_index_expr=False, |
| 4090 | tma_compatibility_checker: TMACompatibilityChecker | None = None, | 4147 | tma_compatibility_checker: TMACompatibilityChecker | None = None, |
| 4148 | + apply_var_replacements=True, | ||
| 4149 | + materialize_var_directions=True, | ||
| 4091 | ) -> IndexingOptions | BlockPtrOptions: | 4150 | ) -> IndexingOptions | BlockPtrOptions: |
| 4092 | """ | 4151 | """ |
| 4093 | Compute the index and mask to pass to tl.load() or tl.store() | 4152 | Compute the index and mask to pass to tl.load() or tl.store() |
| 4094 | """ | 4153 | """ |
| 4095 | if not index_analyze: | 4154 | if not index_analyze: |
| 4096 | index_analyze = IndexAnalysis(self, index, is_index_expr=is_index_expr) | 4155 | index_analyze = IndexAnalysis(self, index, is_index_expr=is_index_expr) |
| 4097 | - index_analyze.analyze_index(nddma) | 4156 | + index_analyze.analyze_index( |
| 4157 | + nddma, materialize_var_directions=materialize_var_directions | ||
| 4158 | + ) | ||
| 4098 | 4159 | ||
| 4099 | - index = self.prepare_indexing(index, index_analyze, is_index_expr, nddma=nddma) | 4160 | + index = self.prepare_indexing( |
| 4161 | + index, | ||
| 4162 | + index_analyze, | ||
| 4163 | + is_index_expr, | ||
| 4164 | + nddma=nddma, | ||
| 4165 | + materialize_var_directions=materialize_var_directions, | ||
| 4166 | + ) | ||
| 4100 | index_vars = index.free_symbols | 4167 | index_vars = index.free_symbols |
| 4101 | has_rindex = False | 4168 | has_rindex = False |
| 4102 | index = sympy_subs(index, V.graph.sizevars.precomputed_replacements) | 4169 | index = sympy_subs(index, V.graph.sizevars.precomputed_replacements) |
| @@ -4115,7 +4182,8 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4115 | index = sympy_subs(index, replacements) | 4182 | index = sympy_subs(index, replacements) |
| 4116 | 4183 | ||
| 4117 | # if not self.inside_reduction : | 4184 | # if not self.inside_reduction : |
| 4118 | - index = self.replace_index_vars(index, index_analyze) | 4185 | + if apply_var_replacements: |
| 4186 | + index = self.replace_index_vars(index, index_analyze) | ||
| 4119 | index_vars = index.free_symbols | 4187 | index_vars = index.free_symbols |
| 4120 | has_rindex = False | 4188 | has_rindex = False |
| 4121 | 4189 | ||
| @@ -4180,6 +4248,162 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4180 | self.range_tree_nodes[sym].codegen() # type: ignore[index] | 4248 | self.range_tree_nodes[sym].codegen() # type: ignore[index] |
| 4181 | return expr | 4249 | return expr |
| 4182 | 4250 | ||
| 4251 | + | ||
| 4252 | + def _direction_slot_from_str(direction_str: Optional[str]) -> Optional[int]: | ||
| 4253 | + if not direction_str: | ||
| 4254 | + return None | ||
| 4255 | + stripped = direction_str.strip() | ||
| 4256 | + if not (stripped.startswith("[") and stripped.endswith("]")): | ||
| 4257 | + return None | ||
| 4258 | + dims = [x.strip() for x in stripped[1:-1].split(",")] | ||
| 4259 | + slots = [idx for idx, dim in enumerate(dims) if dim == ":"] | ||
| 4260 | + if len(slots) != 1: | ||
| 4261 | + return None | ||
| 4262 | + return slots[0] | ||
| 4263 | + | ||
| 4264 | + def _lookup_symbol_direction( | ||
| 4265 | + self, sym: sympy.Symbol, index_analyze: Optional[IndexAnalysis] = None | ||
| 4266 | + ) -> Optional[str]: | ||
| 4267 | + if index_analyze is not None: | ||
| 4268 | + if sym in index_analyze.var_directions: | ||
| 4269 | + return index_analyze.var_directions[sym] | ||
| 4270 | + if sym in index_analyze.nddma_var_directions: | ||
| 4271 | + return index_analyze.nddma_var_directions[sym] | ||
| 4272 | + if sym in self.range_tree_nodes: | ||
| 4273 | + return self.range_tree_nodes[sym].get_axis_direction() | ||
| 4274 | + if sym in self.range_tree_nodes_removed: | ||
| 4275 | + return self.range_tree_nodes_removed[sym].get_axis_direction() | ||
| 4276 | + for node in itertools.chain( | ||
| 4277 | + self.range_tree_nodes.values(), self.range_tree_nodes_removed.values() | ||
| 4278 | + ): | ||
| 4279 | + if sym in node.var_directions: | ||
| 4280 | + return node.var_directions[sym] | ||
| 4281 | + return None | ||
| 4282 | + | ||
| 4283 | + def _lookup_semantic_axis_name( | ||
| 4284 | + self, sym: sympy.Symbol, index_analyze: Optional[IndexAnalysis] = None | ||
| 4285 | + ) -> Optional[str]: | ||
| 4286 | + if sym in self.range_tree_nodes: | ||
| 4287 | + return self.range_tree_nodes[sym].name | ||
| 4288 | + if sym in self.range_tree_nodes_removed: | ||
| 4289 | + return self.range_tree_nodes_removed[sym].name | ||
| 4290 | + if index_analyze is not None: | ||
| 4291 | + for original, replacement in itertools.chain( | ||
| 4292 | + index_analyze.var_replacements.items(), | ||
| 4293 | + index_analyze.nddma_var_replacements.items(), | ||
| 4294 | + ): | ||
| 4295 | + if replacement == sym: | ||
| 4296 | + return str(original) | ||
| 4297 | + for node in itertools.chain( | ||
| 4298 | + self.range_tree_nodes.values(), self.range_tree_nodes_removed.values() | ||
| 4299 | + ): | ||
| 4300 | + if sym in node.var_directions: | ||
| 4301 | + return node.name | ||
| 4302 | + return None | ||
| 4303 | + | ||
| 4304 | + def _infer_layout_axes_from_expr( | ||
| 4305 | + self, expr: Optional[sympy.Expr], index_analyze: Optional[IndexAnalysis] = None | ||
| 4306 | + ) -> Optional[tuple[str, ...]]: | ||
| 4307 | + if expr is None: | ||
| 4308 | + return None | ||
| 4309 | + # Recover buffer layout from broadcast directions, not from memory | ||
| 4310 | + # stride magnitude. The same address order can still materialize into a | ||
| 4311 | + # different logical axis-to-slot mapping. | ||
| 4312 | + slot_to_axis: dict[int, str] = {} | ||
| 4313 | + found_layout_axis = False | ||
| 4314 | + layout_basis = {str(axis) for axis in (self.golden_var_list or ())} | ||
| 4315 | + for sym in sorted(expr.free_symbols, key=str): | ||
| 4316 | + semantic_axis = self._lookup_semantic_axis_name(sym, index_analyze) | ||
| 4317 | + if semantic_axis is None: | ||
| 4318 | + continue | ||
| 4319 | + if layout_basis and semantic_axis not in layout_basis: | ||
| 4320 | + # This symbol contributes to pointer arithmetic, but it is not | ||
| 4321 | + # part of the dense-layout axis basis for the current kernel. | ||
| 4322 | + # Layout inference must conservatively give up here and let | ||
| 4323 | + # store codegen fall back to the remapped path. | ||
| 4324 | + return None | ||
| 4325 | + direction = self._lookup_symbol_direction(sym, index_analyze) | ||
| 4326 | + if direction is None: | ||
| 4327 | + return None | ||
| 4328 | + slot = self._direction_slot_from_str(direction) | ||
| 4329 | + if slot is None: | ||
| 4330 | + return None | ||
| 4331 | + previous_axis = slot_to_axis.get(slot) | ||
| 4332 | + if previous_axis is not None and previous_axis != semantic_axis: | ||
| 4333 | + return None | ||
| 4334 | + slot_to_axis[slot] = semantic_axis | ||
| 4335 | + found_layout_axis = True | ||
| 4336 | + if not found_layout_axis: | ||
| 4337 | + return None | ||
| 4338 | + return tuple(axis for _, axis in sorted(slot_to_axis.items())) | ||
| 4339 | + | ||
| 4340 | + | ||
| 4341 | + def _permute_layout_axes( | ||
| 4342 | + layout_axes: Optional[tuple[str, ...]], permute_shape: Sequence[int] | ||
| 4343 | + ) -> Optional[tuple[str, ...]]: | ||
| 4344 | + if layout_axes is None: | ||
| 4345 | + return None | ||
| 4346 | + if len(layout_axes) != len(permute_shape): | ||
| 4347 | + return None | ||
| 4348 | + return tuple(layout_axes[idx] for idx in permute_shape) | ||
| 4349 | + | ||
| 4350 | + | ||
| 4351 | + def _mark_layout_unknown(var: TritonCSEVariable) -> None: | ||
| 4352 | + var.layout_axes = None | ||
| 4353 | + var.layout_known = False | ||
| 4354 | + | ||
| 4355 | + | ||
| 4356 | + def _set_layout_axes( | ||
| 4357 | + var: TritonCSEVariable, layout_axes: Optional[tuple[str, ...]] | ||
| 4358 | + ) -> None: | ||
| 4359 | + var.layout_axes = layout_axes | ||
| 4360 | + var.layout_known = layout_axes is not None | ||
| 4361 | + | ||
| 4362 | + def _update_layout_on_args( | ||
| 4363 | + self, | ||
| 4364 | + csevar: TritonCSEVariable, | ||
| 4365 | + args: Sequence[object], | ||
| 4366 | + kwargs: dict[str, object], | ||
| 4367 | + ) -> None: | ||
| 4368 | + # Preserve layout only across simple pointwise expressions where all | ||
| 4369 | + # tensor inputs agree on the same buffer layout. | ||
| 4370 | + # Known limitation: update / inplace lowering chains (for example, | ||
| 4371 | + # permute -> add_ / mutate_to -> clone) may still drop layout metadata | ||
| 4372 | + # here and force store emission to fall back to the remapped path. | ||
| 4373 | + candidate_layouts = [] | ||
| 4374 | + for arg in itertools.chain(args, kwargs.values()): | ||
| 4375 | + if isinstance(arg, TritonCSEVariable): | ||
| 4376 | + if not getattr(arg, "layout_known", False): | ||
| 4377 | + self._mark_layout_unknown(csevar) | ||
| 4378 | + return | ||
| 4379 | + candidate_layouts.append(getattr(arg, "layout_axes", None)) | ||
| 4380 | + if not candidate_layouts: | ||
| 4381 | + self._mark_layout_unknown(csevar) | ||
| 4382 | + return | ||
| 4383 | + first_layout = candidate_layouts[0] | ||
| 4384 | + if first_layout is None or any(layout != first_layout for layout in candidate_layouts[1:]): | ||
| 4385 | + self._mark_layout_unknown(csevar) | ||
| 4386 | + return | ||
| 4387 | + self._set_layout_axes(csevar, first_layout) | ||
| 4388 | + | ||
| 4389 | + def _can_preserve_store_semantics_with_value_layout( | ||
| 4390 | + self, | ||
| 4391 | + value: CSEVariable, | ||
| 4392 | + raw_indexing: IndexingOptions | BlockPtrOptions, | ||
| 4393 | + index_analyze: IndexAnalysis, | ||
| 4394 | + ) -> bool: | ||
| 4395 | + if not isinstance(raw_indexing, IndexingOptions): | ||
| 4396 | + return False | ||
| 4397 | + # Only preserve the raw scheduler store path when the target layout | ||
| 4398 | + # implied by the raw store index matches the RHS buffer layout exactly. | ||
| 4399 | + raw_store_layout = self._infer_layout_axes_from_expr( | ||
| 4400 | + raw_indexing.index, index_analyze=index_analyze | ||
| 4401 | + ) | ||
| 4402 | + return ( | ||
| 4403 | + raw_store_layout is not None | ||
| 4404 | + and raw_store_layout == getattr(value, "layout_axes", None) | ||
| 4405 | + ) | ||
| 4406 | + | ||
| 4183 | # when xindex(16) -> x2:2,x3:8, when new length:16 in , should return (x2,x3) | 4407 | # when xindex(16) -> x2:2,x3:8, when new length:16 in , should return (x2,x3) |
| 4184 | def split_and_set_ranges(self, lengths: Sequence[Sequence[sympy.Expr]]): | 4408 | def split_and_set_ranges(self, lengths: Sequence[Sequence[sympy.Expr]]): |
| 4185 | groups = [rt.numel for rt in self.range_trees] | 4409 | groups = [rt.numel for rt in self.range_trees] |
| @@ -4366,6 +4590,7 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 4366 | output_idx += 1 | 4590 | output_idx += 1 |
| 4367 | 4591 | ||
| 4368 | csevar.update_on_args(name, args, kwargs) | 4592 | csevar.update_on_args(name, args, kwargs) |
| 4593 | + V.kernel._update_layout_on_args(csevar, args, kwargs) | ||
| 4369 | 4594 | ||
| 4370 | return csevar | 4595 | return csevar |
| 4371 | 4596 | ||