已开启
[triton_experimental] fixlinearize rank mismatch + fmod #44448
AACAES创建于 8月12日
[triton_experimental] fixlinearize rank mismatch + fmod #44448
已开启
共 5 个文件变更+704-111
| @@ -0,0 +1,100 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | + | ||
| 16 | +"""Guards for the aten.fmod GENERATE_LIST entry (PR #44448). | ||
| 17 | + | ||
| 18 | +fmod was added to the triton_experimental GENERATE_LIST so it stays on the | ||
| 19 | +fused lowering path instead of becoming an extern aten call. The wrapper | ||
| 20 | +must call the fused libdevice lowering (``libdevice.fmod(``) inside the | ||
| 21 | +Triton kernel) and must not emit ``torch.ops.aten.fmod.Tensor`` fallback | ||
| 22 | +calls; results are compared against eager within float rounding. | ||
| 23 | +""" | ||
| 24 | + | ||
| 25 | +import os | ||
| 26 | + | ||
| 27 | +import torch | ||
| 28 | +from torch._inductor.utils import run_and_get_code | ||
| 29 | +from torch.testing._internal.common_utils import run_tests | ||
| 30 | + | ||
| 31 | +import torch_npu # noqa: F401 | ||
| 32 | +from testutils import TestUtils | ||
| 33 | + | ||
| 34 | + | ||
| 35 | +def _use_triton_experimental(): | ||
| 36 | + torch._inductor.config.get_config_copy() | ||
| 37 | + torch._inductor.config.npu_backend = "triton_experimental" | ||
| 38 | + | ||
| 39 | + | ||
| 40 | +class TestTritonExperimentalFmod(TestUtils): | ||
| 41 | + | ||
| 42 | + def setUp(self): | ||
| 43 | + super().setUp() | ||
| 44 | + torch._dynamo.reset() | ||
| 45 | + self._saved_config = getattr(torch._inductor.config, "npu_backend", "default") | ||
| 46 | + self._saved_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND") | ||
| 47 | + | ||
| 48 | + def tearDown(self): | ||
| 49 | + torch._inductor.config.npu_backend = self._saved_config | ||
| 50 | + if self._saved_env is None: | ||
| 51 | + os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None) | ||
| 52 | + else: | ||
| 53 | + os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._saved_env | ||
| 54 | + torch._dynamo.reset() | ||
| 55 | + super().tearDown() | ||
| 56 | + | ||
| 57 | + def test_fmod_stay_fused(self): | ||
| 58 | + _use_triton_experimental() | ||
| 59 | + torch.manual_seed(0) | ||
| 60 | + | ||
| 61 | + def fn(x, y): | ||
| 62 | + return torch.fmod(x, y) * 2.0 + x | ||
| 63 | + | ||
| 64 | + x = torch.randn(64, 128, device="npu") | ||
| 65 | + y = torch.randn(64, 128, device="npu").abs() + 0.5 | ||
| 66 | + expected = fn(x, y) | ||
| 67 | + compiled = torch.compile( | ||
| 68 | + fn, fullgraph=True, options={"npu_backend": "triton_experimental"} | ||
| 69 | + ) | ||
| 70 | + out, codes = run_and_get_code(compiled, x, y) | ||
| 71 | + code = "\n".join(codes) | ||
| 72 | + self.assertIn("libdevice.fmod(", code) | ||
| 73 | + self.assertNotIn(" = torch.ops.aten.fmod.Tensor(", code) | ||
| 74 | + self.assertTrue( | ||
| 75 | + torch.allclose(expected, out, rtol=1e-5, atol=1e-5), | ||
| 76 | + f"max diff {(expected - out).abs().max().item()}", | ||
| 77 | + ) | ||
| 78 | + | ||
| 79 | + def test_fmod_int_stay_fused(self): | ||
| 80 | + # Integer operands exercise the integer fmod lowering. | ||
| 81 | + _use_triton_experimental() | ||
| 82 | + torch.manual_seed(0) | ||
| 83 | + | ||
| 84 | + def fn(x, y): | ||
| 85 | + return torch.fmod(x, y) * 2 - y | ||
| 86 | + | ||
| 87 | + x = torch.randint(-1000, 1000, (64, 64), device="npu") | ||
| 88 | + y = torch.randint(3, 17, (64, 64), device="npu") | ||
| 89 | + expected = fn(x, y) | ||
| 90 | + compiled = torch.compile( | ||
| 91 | + fn, fullgraph=True, options={"npu_backend": "triton_experimental"} | ||
| 92 | + ) | ||
| 93 | + out, codes = run_and_get_code(compiled, x, y) | ||
| 94 | + code = "\n".join(codes) | ||
| 95 | + self.assertNotIn(" = torch.ops.aten.fmod.Tensor(", code) | ||
| 96 | + self.assertEqual(expected, out) | ||
| 97 | + | ||
| 98 | + | ||
| 99 | +if __name__ == "__main__": | ||
| 100 | + run_tests() | ||
| @@ -0,0 +1,208 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +# All rights reserved. | ||
| 3 | +# | ||
| 4 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +# you may not use this file except in compliance with the License. | ||
| 6 | +# You may obtain a copy of the License at | ||
| 7 | +# | ||
| 8 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +# | ||
| 10 | +# Unless required by applicable law or agreed to in writing, software | ||
| 11 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +# See the License for the specific language governing permissions and | ||
| 14 | +# limitations under the License. | ||
| 15 | + | ||
| 16 | +"""Regression guards for the triton_experimental linearize codegen fixes (PR #44448). | ||
| 17 | + | ||
| 18 | +Each test maps to one fix shipped by the PR: | ||
| 19 | + | ||
| 20 | +1. Promoted r-tree output rank: reduction stores fused with pointwise ops | ||
| 21 | + (norm -> div pattern from test_buffer_use_after_remove) used to raise | ||
| 22 | + ``ValueError('Cannot broadcast, rank mismatch')``. | ||
| 23 | +2. ``index_vars_per_node`` initialization order: std/var reduction kernels | ||
| 24 | + used to hit ``AttributeError: 'NPUTritonKernel' object has no attribute | ||
| 25 | + 'index_vars_per_node'``. | ||
| 26 | +3. ``constant()`` shape on 2D pointwise kernels: fill-into-slice mutation | ||
| 27 | + graphs (test_slice_mutation3) used to raise rank mismatch. | ||
| 28 | + | ||
| 29 | +Every graph pattern is taken from the issue's failing cases and compared | ||
| 30 | +against eager outputs. Fused reduction paths differ from eager by float | ||
| 31 | +rounding only, so tolerance-based equality is used where appropriate. | ||
| 32 | +""" | ||
| 33 | + | ||
| 34 | +import os | ||
| 35 | + | ||
| 36 | +import torch | ||
| 37 | +from torch.testing._internal.common_utils import run_tests | ||
| 38 | + | ||
| 39 | +import torch_npu # noqa: F401 | ||
| 40 | +from testutils import TestUtils | ||
| 41 | + | ||
| 42 | +def _use_triton_experimental(): | ||
| 43 | + torch._inductor.config.get_config_copy() | ||
| 44 | + torch._inductor.config.npu_backend = "triton_experimental" | ||
| 45 | + | ||
| 46 | + | ||
| 47 | +class TestLinearizeRegressions(TestUtils): | ||
| 48 | + | ||
| 49 | + def setUp(self): | ||
| 50 | + super().setUp() | ||
| 51 | + torch._dynamo.reset() | ||
| 52 | + self._saved_config = getattr(torch._inductor.config, "npu_backend", "default") | ||
| 53 | + self._saved_env = os.environ.get("TORCHINDUCTOR_NPU_BACKEND") | ||
| 54 | + | ||
| 55 | + def tearDown(self): | ||
| 56 | + torch._inductor.config.npu_backend = self._saved_config | ||
| 57 | + if self._saved_env is None: | ||
| 58 | + os.environ.pop("TORCHINDUCTOR_NPU_BACKEND", None) | ||
| 59 | + else: | ||
| 60 | + os.environ["TORCHINDUCTOR_NPU_BACKEND"] = self._saved_env | ||
| 61 | + torch._dynamo.reset() | ||
| 62 | + super().tearDown() | ||
| 63 | + | ||
| 64 | + # ------------------------------------------------------------------ | ||
| 65 | + # Fix 1: promoted r-tree output rank (test_buffer_use_after_remove) | ||
| 66 | + # ------------------------------------------------------------------ | ||
| 67 | + | ||
| 68 | + def test_reduction_fused_pointwise_broadcast(self): | ||
| 69 | + # norm + div keeps a broadcast div feeding per-component sin*mul; | ||
| 70 | + # before the fix the store rank mismatched the broadcast value rank. | ||
| 71 | + _use_triton_experimental() | ||
| 72 | + torch.manual_seed(0) | ||
| 73 | + | ||
| 74 | + def fn(rotvec): | ||
| 75 | + theta = torch.norm(rotvec, dim=-1) | ||
| 76 | + axis = rotvec / theta[..., None] | ||
| 77 | + return axis * torch.sin(theta)[..., None] | ||
| 78 | + | ||
| 79 | + rotvec = torch.randn(8, 4, 3, device="npu") | ||
| 80 | + expected = fn(rotvec) | ||
| 81 | + compiled = torch.compile( | ||
| 82 | + fn, fullgraph=True, options={"npu_backend": "triton_experimental"} | ||
| 83 | + ) | ||
| 84 | + out = compiled(rotvec) | ||
| 85 | + self.assertTrue( | ||
| 86 | + torch.allclose(expected, out, rtol=2e-6, atol=2e-6), | ||
| 87 | + f"max diff {(expected - out).abs().max().item()}", | ||
| 88 | + ) | ||
| 89 | + | ||
| 90 | + # ------------------------------------------------------------------ | ||
| 91 | + # Fix 2: std / var reductions (index_vars_per_node) | ||
| 92 | + # ------------------------------------------------------------------ | ||
| 93 | + | ||
| 94 | + def test_std_dynamic_shapes(self): | ||
| 95 | + # Same graph family as test_std_dynamic_shapes_npu, which used to hit | ||
| 96 | + # AttributeError on index_vars_per_node. Reductions differ from eager | ||
| 97 | + # by float rounding only. | ||
| 98 | + _use_triton_experimental() | ||
| 99 | + torch.manual_seed(0) | ||
| 100 | + | ||
| 101 | + def fn(x): | ||
| 102 | + return x.std(dim=1, unbiased=True) + 1.0 | ||
| 103 | + | ||
| 104 | + x = torch.randn(16, 128, 32, device="npu") | ||
| 105 | + torch._dynamo.mark_dynamic(x, 0) | ||
| 106 | + expected = fn(x) | ||
| 107 | + compiled = torch.compile( | ||
| 108 | + fn, fullgraph=True, options={"npu_backend": "triton_experimental"} | ||
| 109 | + ) | ||
| 110 | + out = compiled(x) | ||
| 111 | + self.assertTrue( | ||
| 112 | + torch.allclose(expected, out, rtol=1e-5, atol=1e-5), | ||
| 113 | + f"max diff {(expected - out).abs().max().item()}", | ||
| 114 | + ) | ||
| 115 | + | ||
| 116 | + def test_var_mean_multi_output(self): | ||
| 117 | + # var_mean keeps two reduction outputs alive in one graph; both | ||
| 118 | + # stores must use the output rank. | ||
| 119 | + _use_triton_experimental() | ||
| 120 | + torch.manual_seed(0) | ||
| 121 | + | ||
| 122 | + def fn(x): | ||
| 123 | + var, mean = torch.var_mean(x, dim=1) | ||
| 124 | + return var + mean | ||
| 125 | + | ||
| 126 | + x = torch.randn(16, 128, 32, device="npu") | ||
| 127 | + expected = fn(x) | ||
| 128 | + compiled = torch.compile( | ||
| 129 | + fn, fullgraph=True, options={"npu_backend": "triton_experimental"} | ||
| 130 | + ) | ||
| 131 | + out = compiled(x) | ||
| 132 | + self.assertTrue( | ||
| 133 | + torch.allclose(expected, out, rtol=1e-5, atol=1e-5), | ||
| 134 | + f"max diff {(expected - out).abs().max().item()}", | ||
| 135 | + ) | ||
| 136 | + | ||
| 137 | + # ------------------------------------------------------------------ | ||
| 138 | + # Fix 3: fill / slice mutation (test_slice_mutation3) | ||
| 139 | + # ------------------------------------------------------------------ | ||
| 140 | + | ||
| 141 | + def test_slice_mutation3(self): | ||
| 142 | + # Original test_slice_mutation3_npu pattern: in-place fill_ into a | ||
| 143 | + # 2D slice through optimize_assert(compile_fx). Before the fix the | ||
| 144 | + # constant() store broadcast raised rank mismatch ([1] vs [2, 2]). | ||
| 145 | + from torch._inductor.compile_fx import compile_fx | ||
| 146 | + | ||
| 147 | + def fn(a): | ||
| 148 | + a[:2, :2].fill_(10) | ||
| 149 | + | ||
| 150 | + # optimize_assert(compile_fx) has no options entry point, so select | ||
| 151 | + # the backend through the environment variable. | ||
| 152 | + os.environ["TORCHINDUCTOR_NPU_BACKEND"] = "triton_experimental" | ||
| 153 | + opt_fn = torch._dynamo.optimize_assert( | ||
| 154 | + lambda *args, **kwargs: compile_fx(*args, **kwargs) | ||
| 155 | + )(fn) | ||
| 156 | + x1 = torch.randn(8, 8, device="npu") | ||
| 157 | + x2 = x1.clone() | ||
| 158 | + fn(x1) | ||
| 159 | + opt_fn(x2) | ||
| 160 | + self.assertEqual(x1, x2) | ||
| 161 | + | ||
| 162 | + def test_fill_scalar_into_slice(self): | ||
| 163 | + # Functional form of the same graph (clone + slice-assign + consumer): | ||
| 164 | + # fill_(10) into a[:2, :2] lowers to a constant store broadcast to a | ||
| 165 | + # 2D address; before the fix constant() produced a 1D shape and the | ||
| 166 | + # store raised rank mismatch. | ||
| 167 | + _use_triton_experimental() | ||
| 168 | + torch.manual_seed(0) | ||
| 169 | + | ||
| 170 | + def fn(a): | ||
| 171 | + b = a.clone() | ||
| 172 | + b[:2, :2] = 10.0 | ||
| 173 | + return b * 1.5 | ||
| 174 | + | ||
| 175 | + x = torch.randn(8, 8, device="npu") | ||
| 176 | + expected = fn(x) | ||
| 177 | + compiled = torch.compile( | ||
| 178 | + fn, fullgraph=True, options={"npu_backend": "triton_experimental"} | ||
| 179 | + ) | ||
| 180 | + self.assertEqual(expected, compiled(x)) | ||
| 181 | + | ||
| 182 | + def test_slice_mutation_chain(self): | ||
| 183 | + # Two slice writes into one buffer feeding a reduction; guards the | ||
| 184 | + # buffer-use-after-remove store pattern. | ||
| 185 | + _use_triton_experimental() | ||
| 186 | + torch.manual_seed(0) | ||
| 187 | + | ||
| 188 | + def fn(x): | ||
| 189 | + buf = torch.empty_like(x) | ||
| 190 | + buf[:4] = x[:4] * 2.0 | ||
| 191 | + buf[4:] = x[4:] + 1.0 | ||
| 192 | + return buf.sum() | ||
| 193 | + | ||
| 194 | + x = torch.randn(16, 32, device="npu") | ||
| 195 | + expected = fn(x) | ||
| 196 | + compiled = torch.compile( | ||
| 197 | + fn, fullgraph=True, options={"npu_backend": "triton_experimental"} | ||
| 198 | + ) | ||
| 199 | + out = compiled(x) | ||
| 200 | + self.assertTrue( | ||
| 201 | + torch.allclose(expected, out, rtol=1e-5, atol=1e-5), | ||
| 202 | + f"max diff {(expected - out).abs().max().item()}", | ||
| 203 | + ) | ||
| 204 | + | ||
| 205 | + | ||
| 206 | + | ||
| 207 | +if __name__ == "__main__": | ||
| 208 | + run_tests() | ||
| @@ -732,7 +732,20 @@ def _codegen_header_npu_for_tree(kernel, tree, code, outer_blocks=None): | |||
| 732 | header_code.writeline(f"{node.name}offset = 0") | 732 | header_code.writeline(f"{node.name}offset = 0") |
| 733 | else: | 733 | else: |
| 734 | if tree.tensor_dim is not None: | 734 | if tree.tensor_dim is not None: |
| 735 | - size = kernel.indexing_size_str(tree.var_tensor_dims[node.name]) | 735 | + # Use iteration-space rank (_npu_iter_ndim) for the broadcast |
| 736 | + # ndim: slot indices in var_tensor_dims run up to iter_ndim-1 | ||
| 737 | + # (one slot per promoted r-tree sub-node), which can exceed | ||
| 738 | + # triton_tensor_ndim()-1 (output rank). The upstream | ||
| 739 | + # indexing_size_str() uses triton_tensor_ndim() and would | ||
| 740 | + # IndexError for slots >= triton_tensor_ndim(). | ||
| 741 | + _slot = tree.var_tensor_dims[node.name] | ||
| 742 | + if ncfg.codegen_linearize and getattr(kernel, '_linearize_applied', False): | ||
| 743 | + _ndim = kernel._npu_iter_ndim() | ||
| 744 | + else: | ||
| 745 | + _ndim = kernel.triton_tensor_ndim() | ||
| 746 | + _parts = ["None"] * _ndim | ||
| 747 | + _parts[_slot] = ":" | ||
| 748 | + size = "[" + ", ".join(_parts) + "]" | ||
| 736 | line = f"{node.name}offset + tl.arange(0, {arange_upper}){size}" | 749 | line = f"{node.name}offset + tl.arange(0, {arange_upper}){size}" |
| 737 | else: | 750 | else: |
| 738 | line = kernel.iteration_ranges_scalar_code(tree, f"{node.name}offset") | 751 | line = kernel.iteration_ranges_scalar_code(tree, f"{node.name}offset") |
| @@ -258,13 +258,16 @@ def _npu_scalar_odometer_axis_names(kernel, tree): | |||
| 258 | return scalar | 258 | return scalar |
| 259 | 259 | ||
| 260 | 260 | ||
| 261 | -def _npu_rewrite_promoted_rtree_body(kernel, real_sizes, real_ndim): | 261 | +def _npu_rewrite_promoted_rtree_body(kernel, real_sizes, real_ndim, iter_ndim=None): |
| 262 | """Rewrite the assembled reduction body for promoted r-trees. Per tree: turn | 262 | """Rewrite the assembled reduction body for promoted r-trees. Per tree: turn |
| 263 | the flat ``for r0_offset ...`` loop into full-residency aranges (each free | 263 | the flat ``for r0_offset ...`` loop into full-residency aranges (each free |
| 264 | r-node N gets ``N = tl.arange(0, real_block_N)[slot]`` + ``Nmask``, combined | 264 | r-node N gets ``N = tl.arange(0, real_block_N)[slot]`` + ``Nmask``, combined |
| 265 | ``r0_mask``), hoist the real_block_N constexpr defs before the group loop, and | 265 | ``r0_mask``), hoist the real_block_N constexpr defs before the group loop, and |
| 266 | reshape-collapse each ``tl.sum(_acc, dim)`` so the r-slots merge to one axis. | 266 | reshape-collapse each ``tl.sum(_acc, dim)`` so the r-slots merge to one axis. |
| 267 | - Handles only the fully-resident case (guaranteed by the promotability gate).""" | 267 | + Handles only the fully-resident case (guaranteed by the promotability gate). |
| 268 | + iter_ndim is the iteration-space rank (≥ real_ndim): the number of nested | ||
| 269 | + loops / accumulator dims. When the r-tree has more sub-nodes than real_ndim, | ||
| 270 | + iter_ndim > real_ndim.""" | ||
| 268 | sv = V.graph.sizevars | 271 | sv = V.graph.sizevars |
| 269 | for tree in kernel.range_trees: | 272 | for tree in kernel.range_trees: |
| 270 | if not tree.is_reduction: | 273 | if not tree.is_reduction: |
| @@ -278,12 +281,19 @@ def _npu_rewrite_promoted_rtree_body(kernel, real_sizes, real_ndim): | |||
| 278 | 281 | ||
| 279 | # Per-node real_block constexpr defs + arange/mask body lines. | 282 | # Per-node real_block constexpr defs + arange/mask body lines. |
| 280 | # Slot broadcast: ``[None, ..., :, ..., None]`` with ``:`` at the | 283 | # Slot broadcast: ``[None, ..., :, ..., None]`` with ``:`` at the |
| 281 | - # node's tensor_dim slot, one entry per real_ndim. | 284 | + # node's tensor_dim slot, one entry per iter_ndim (the iteration-space |
| 285 | + # rank). All sub-node slots are < iter_ndim (by construction), so the | ||
| 286 | + # arange broadcasts to a iter_ndim-D tensor matching the accumulator | ||
| 287 | + # shape. For the rare slot >= iter_ndim case, collapse to all singletons. | ||
| 288 | + _bcast_ndim = iter_ndim if iter_ndim is not None else real_ndim | ||
| 289 | + | ||
| 282 | def _slot_bcast(slot): | 290 | def _slot_bcast(slot): |
| 283 | - parts = ["None"] * real_ndim | 291 | + parts = ["None"] * _bcast_ndim |
| 284 | - if 0 <= slot < real_ndim: | 292 | + if 0 <= slot < _bcast_ndim: |
| 285 | parts[slot] = ":" | 293 | parts[slot] = ":" |
| 286 | - return "[" + ", ".join(parts) + "]" | 294 | + return "[" + ", ".join(parts) + "]" |
| 295 | + # slot >= _bcast_ndim: collapse to all singletons. | ||
| 296 | + return "[" + ", ".join(parts) + f"].reshape([{', '.join(['1'] * _bcast_ndim)}])" | ||
| 287 | 297 | ||
| 288 | r_slots = sorted(vtd[nm] for nm in ordered_names if nm in vtd) | 298 | r_slots = sorted(vtd[nm] for nm in ordered_names if nm in vtd) |
| 289 | dynamic = kernel._npu_rtree_dynamic.get(prefix, False) | 299 | dynamic = kernel._npu_rtree_dynamic.get(prefix, False) |
| @@ -368,34 +378,80 @@ def _npu_rewrite_promoted_rtree_body(kernel, real_sizes, real_ndim): | |||
| 368 | blk_defs = dyn_defs + static_defs | 378 | blk_defs = dyn_defs + static_defs |
| 369 | combined_mask = " & ".join(mask_terms) | 379 | combined_mask = " & ".join(mask_terms) |
| 370 | 380 | ||
| 371 | - # Collapsed shape for the reshape before tl.sum: fold the contiguous r-slot | 381 | + # Collapsed shape for the reshape before tl.sum: fold ALL r-sub-node |
| 372 | - # run into one product term IN PLACE (row-major reshape reinterprets memory, | 382 | + # slots into one product term (the reduction axis). The internal slot |
| 373 | - # so the fold must sit at the real slot order). r-slots are contiguous (gate) | 383 | + # assignment may give the r-tree MORE slots than real_ndim (one per |
| 374 | - # but may sit anywhere — front or back. Walk slots in order, emit the | 384 | + # sub-node for multi-tile), but the output rank is real_ndim — so all |
| 375 | - # r-product at the first r-slot and each kept token at its own position. | 385 | + # r-slots collapse to a single axis regardless of how many sub-nodes |
| 376 | - first_r = min(r_slots) | 386 | + # they contain. r_tokens use each sub-node's OWN size (dynamic |
| 377 | - last_r = max(r_slots) | 387 | + # ``{nm}_blk`` or static ``real_block_{nm}``), not real_sizes (which |
| 378 | - r_tokens = [real_sizes[s] for s in r_slots] | 388 | + # only has real_ndim entries and would IndexError for slot >= real_ndim). |
| 389 | + r_tokens = [] | ||
| 390 | + for nm in ordered_names: | ||
| 391 | + if dynamic and node_dyn.get(nm, True): | ||
| 392 | + r_tokens.append(f"{nm}_blk") | ||
| 393 | + else: | ||
| 394 | + r_tokens.append(f"real_block_{nm}") | ||
| 379 | r_product_tok = "*".join(r_tokens) if r_tokens else "1" | 395 | r_product_tok = "*".join(r_tokens) if r_tokens else "1" |
| 396 | + # Resize target rank: when a non-reduction tree contributes | ||
| 397 | + # register-tensor slots (var_tensor_dims non-empty), the store address | ||
| 398 | + # (x0index) is broadcast to iter_ndim dimensions, so the value must be | ||
| 399 | + # resized to iter_ndim. When no x-tree has slots (full reduction, or | ||
| 400 | + # all x-tree nodes are scalar-odometer), the store address rank equals | ||
| 401 | + # real_ndim and the value must be resized to real_ndim. | ||
| 402 | + _use_iter_ndim = any( | ||
| 403 | + getattr(t, "var_tensor_dims", {}) | ||
| 404 | + for t in kernel.range_trees | ||
| 405 | + if not t.is_reduction | ||
| 406 | + ) | ||
| 407 | + # The reduction axis sits at the first r-slot position, clamped to | ||
| 408 | + # real_ndim-1 (for full reduction the r-slots may extend past real_ndim, | ||
| 409 | + # e.g. r_slots=[0,1,2] with real_ndim=2 → reduction axis at 0). | ||
| 410 | + first_r = min(r_slots) | ||
| 411 | + collapsed_dim = first_r if first_r < real_ndim else 0 | ||
| 412 | + # Build collapsed_shape: one folded r-product at collapsed_dim, and the | ||
| 413 | + # non-r dimensions (from real_sizes) at the remaining positions. The loop | ||
| 414 | + # must iterate over iter_ndim (not real_ndim) because slot indices run up to | ||
| 415 | + # iter_ndim-1 — the x-tree's slot (e.g. x0 at slot 2 when an r-tree is | ||
| 416 | + # promoted to slots 0,1) sits past real_ndim-1 and would be omitted by | ||
| 417 | + # range(real_ndim), dropping the x-dimension from the reshape target and | ||
| 418 | + # producing a rank mismatch (IndexError in tl.reshape). For full reduction | ||
| 419 | + # (no_x_dim), iter_ndim == real_ndim, so the two are equivalent. | ||
| 380 | collapsed_tokens = [] | 420 | collapsed_tokens = [] |
| 381 | - for s in range(real_ndim): | 421 | + for s in range(iter_ndim if _use_iter_ndim else real_ndim): |
| 382 | - if s < first_r or s > last_r: | 422 | + if s == collapsed_dim: |
| 383 | - collapsed_tokens.append(real_sizes[s]) | ||
| 384 | - elif s == first_r: | ||
| 385 | collapsed_tokens.append(r_product_tok) | 423 | collapsed_tokens.append(r_product_tok) |
| 386 | - # first_r < s <= last_r: r-slot already folded into the product. | 424 | + elif s not in r_slots: |
| 387 | - # Reduction axis after collapse = the position of the folded r-run, | 425 | + collapsed_tokens.append(real_sizes[s]) |
| 388 | - # which is the count of kept slots preceding it. Since r-slots are | 426 | + # s in r_slots but s != collapsed_dim: already folded into the product. |
| 389 | - # contiguous starting at first_r, every slot < first_r is a kept slot. | ||
| 390 | - collapsed_dim = first_r | ||
| 391 | collapsed_shape = "[" + ", ".join(collapsed_tokens) + "]" | 427 | collapsed_shape = "[" + ", ".join(collapsed_tokens) + "]" |
| 392 | - # Post-sum resize: re-insert a singleton at EACH r-slot so the | 428 | + # Post-sum resize: after tl.sum(_, collapsed_dim), the tensor rank is |
| 393 | - # result rank matches the kept-axis stores (x0 index is rank real_ndim). | 429 | + # len(collapsed_tokens) - 1. Insert singletons (None) to bring the rank |
| 394 | - resize_parts = [":"] * real_ndim | 430 | + # back up to the store address rank. The store address rank depends on |
| 395 | - for s in r_slots: | 431 | + # whether an x-tree is present: |
| 396 | - resize_parts[s] = "None" | 432 | + # - When max_slot >= real_ndim (r-tree promoted with extra slots), the |
| 397 | - # The summed tensor has rank (len(keep_tokens)) — re-expand only the | 433 | + # store address (x0index) is broadcast to iter_ndim dimensions, so the |
| 398 | - # r-slots as None; kept slots stay ':'. | 434 | + # value must be resized to iter_ndim to match. |
| 435 | + # - Otherwise (no extra slots), the store address rank = real_ndim, so | ||
| 436 | + # resize to real_ndim. | ||
| 437 | + target_ndim = iter_ndim if _use_iter_ndim else real_ndim | ||
| 438 | + rank_after_sum = len(collapsed_tokens) - 1 | ||
| 439 | + n_singletons = max(0, target_ndim - rank_after_sum) | ||
| 440 | + resize_parts = [":"] * target_ndim | ||
| 441 | + # R-slot positions within target_ndim, in order, take precedence. | ||
| 442 | + r_slots_in_range = sorted(s for s in r_slots if s < target_ndim) | ||
| 443 | + filled = 0 | ||
| 444 | + for s in r_slots_in_range: | ||
| 445 | + if filled < n_singletons: | ||
| 446 | + resize_parts[s] = "None" | ||
| 447 | + filled += 1 | ||
| 448 | + # Remaining singletons (when r-slots exceed target_ndim) fill from left. | ||
| 449 | + for s in range(target_ndim): | ||
| 450 | + if filled >= n_singletons: | ||
| 451 | + break | ||
| 452 | + if resize_parts[s] == ":": | ||
| 453 | + resize_parts[s] = "None" | ||
| 454 | + filled += 1 | ||
| 399 | post_resize = "[" + ", ".join(resize_parts) + "]" | 455 | post_resize = "[" + ", ".join(resize_parts) + "]" |
| 400 | 456 | ||
| 401 | # Split-reconstruction aliases (r0_1 = r0_2 + ks0*r0_3 over promoted | 457 | # Split-reconstruction aliases (r0_1 = r0_2 + ks0*r0_3 over promoted |
| @@ -410,9 +466,13 @@ def _npu_rewrite_promoted_rtree_body(kernel, real_sizes, real_ndim): | |||
| 410 | # r-node's tile token at its slot, and the KEPT axes' real block extent | 466 | # r-node's tile token at its slot, and the KEPT axes' real block extent |
| 411 | # (a live x-axis is XBLOCK, not 1 -- only fully-reduced outputs are 1). | 467 | # (a live x-axis is XBLOCK, not 1 -- only fully-reduced outputs are 1). |
| 412 | # Consumed by the load-ptr re-alignment pass (promoted_rtree_shape_fix). | 468 | # Consumed by the load-ptr re-alignment pass (promoted_rtree_shape_fix). |
| 413 | - slot_block_tokens = [str(real_sizes[s]) for s in range(real_ndim)] | 469 | + # Built over iter_ndim (NOT real_ndim): a promoted r-tree can push the |
| 470 | + # x-tree's register slot past real_ndim-1, and the load index broadcast | ||
| 471 | + # then references that slot -- truncating here drops it from the | ||
| 472 | + # broadcast target and raises "rank mismatch: [.., .., 1], [.., ..]". | ||
| 473 | + slot_block_tokens = [str(real_sizes[s]) for s in range(iter_ndim)] | ||
| 414 | for nm in ordered_names: | 474 | for nm in ordered_names: |
| 415 | - if nm in vtd and 0 <= vtd[nm] < real_ndim: | 475 | + if nm in vtd and 0 <= vtd[nm] < iter_ndim: |
| 416 | is_dyn = dynamic and node_dyn.get(nm, True) | 476 | is_dyn = dynamic and node_dyn.get(nm, True) |
| 417 | slot_block_tokens[vtd[nm]] = f"{nm}_blk" if is_dyn else f"real_block_{nm}" | 477 | slot_block_tokens[vtd[nm]] = f"{nm}_blk" if is_dyn else f"real_block_{nm}" |
| 418 | 478 | ||
| @@ -421,7 +481,7 @@ def _npu_rewrite_promoted_rtree_body(kernel, real_sizes, real_ndim): | |||
| 421 | combined_mask, collapsed_shape, collapsed_dim, post_resize, | 481 | combined_mask, collapsed_shape, collapsed_dim, post_resize, |
| 422 | r_slots, real_ndim, ordered_names, | 482 | r_slots, real_ndim, ordered_names, |
| 423 | dynamic=dynamic, loop_nodes=loop_nodes, slot_bcast=_slot_bcast, | 483 | dynamic=dynamic, loop_nodes=loop_nodes, slot_bcast=_slot_bcast, |
| 424 | - flat_recon=flat_recon, | 484 | + flat_recon=flat_recon, iter_ndim=iter_ndim, |
| 425 | real_sizes=real_sizes, slot_block_tokens=slot_block_tokens, | 485 | real_sizes=real_sizes, slot_block_tokens=slot_block_tokens, |
| 426 | ) | 486 | ) |
| 427 | 487 | ||
| @@ -500,12 +560,62 @@ def _npu_preserve_mixed_leaf_mul_where(lines, leaf_names): | |||
| 500 | return rewritten | 560 | return rewritten |
| 501 | 561 | ||
| 502 | 562 | ||
| 563 | +def _npu_mask_for_load_line(line, prefix, leaf_names, combined_mask): | ||
| 564 | + """Compute the correct r-mask for a tl.load line inside the promoted r-loop. | ||
| 565 | + | ||
| 566 | + The ``combined_mask`` covers ALL r-nodes (e.g. ``r0_2mask & r0_1mask & | ||
| 567 | + r0_0mask``), but a load's index may only use a SUBSET of them. Applying the | ||
| 568 | + full combined mask to a lower-dimensional index causes a Triton broadcast | ||
| 569 | + error (e.g. ``[2,2,3,1]`` vs ``[1,2,3,1]``). | ||
| 570 | + | ||
| 571 | + This function parses the load's index expression to determine which r-nodes | ||
| 572 | + it actually references, then returns a mask containing only those per-node | ||
| 573 | + masks. For non-load lines (e.g. accumulator guards), it falls back to the | ||
| 574 | + full combined_mask. | ||
| 575 | + """ | ||
| 576 | + if "tl.load" not in line: | ||
| 577 | + return combined_mask | ||
| 578 | + # Extract the index expression from: tmp = tl.load(ptr + (INDEX), MASK, ...) | ||
| 579 | + # Delimiter assumptions: the load was emitted by this codegen in the canonical | ||
| 580 | + # ``+ (INDEX),`` form, so the first ``),`` after ``+ (`` closes the index | ||
| 581 | + # parenthesization. If the index itself contains a nested ``( ... ),`` the | ||
| 582 | + # cut lands early — the truncated prefix still tokenizes to a SUBSET of the | ||
| 583 | + # referenced r-nodes, so the derived mask errs on the conservative side | ||
| 584 | + # (extra mask terms), never drops a needed one. Unmatched forms (no | ||
| 585 | + # ``+ (`` / no ``),``) fall back to the full combined_mask, also safe. | ||
| 586 | + ptr_pos = line.find("+ (") | ||
| 587 | + if ptr_pos < 0: | ||
| 588 | + return combined_mask | ||
| 589 | + index_start = ptr_pos + 3 # skip "+ (" | ||
| 590 | + index_end = line.find("),", index_start) | ||
| 591 | + if index_end < 0: | ||
| 592 | + return combined_mask | ||
| 593 | + index_expr = line[index_start:index_end] | ||
| 594 | + # Determine which leaf names (r-nodes) appear in the index. Tokenize to avoid | ||
| 595 | + # false matches (e.g. ``r0_1`` inside ``r0_10``). | ||
| 596 | + tokens = [] | ||
| 597 | + current = [] | ||
| 598 | + for ch in index_expr: | ||
| 599 | + if ch.isalnum() or ch == "_": | ||
| 600 | + current.append(ch) | ||
| 601 | + else: | ||
| 602 | + if current: | ||
| 603 | + tokens.append("".join(current)) | ||
| 604 | + current = [] | ||
| 605 | + if current: | ||
| 606 | + tokens.append("".join(current)) | ||
| 607 | + used = [nm for nm in leaf_names if nm in tokens] | ||
| 608 | + if not used: | ||
| 609 | + return combined_mask | ||
| 610 | + return " & ".join(f"{nm}mask" for nm in used) | ||
| 611 | + | ||
| 612 | + | ||
| 503 | def _npu_apply_promoted_rtree_lines( | 613 | def _npu_apply_promoted_rtree_lines( |
| 504 | kernel, prefix, blk_defs, blk_node_names, arange_lines, | 614 | kernel, prefix, blk_defs, blk_node_names, arange_lines, |
| 505 | combined_mask, collapsed_shape, collapsed_dim, post_resize, | 615 | combined_mask, collapsed_shape, collapsed_dim, post_resize, |
| 506 | r_slots, real_ndim, leaf_names, | 616 | r_slots, real_ndim, leaf_names, |
| 507 | dynamic=False, loop_nodes=None, slot_bcast=None, flat_recon=None, | 617 | dynamic=False, loop_nodes=None, slot_bcast=None, flat_recon=None, |
| 508 | - real_sizes=None, slot_block_tokens=None, | 618 | + iter_ndim=None, real_sizes=None, slot_block_tokens=None, |
| 509 | ): | 619 | ): |
| 510 | """In-place edit of kernel.body._lines for one promoted r-tree. Static mode: | 620 | """In-place edit of kernel.body._lines for one promoted r-tree. Static mode: |
| 511 | replace the flat ``for r0_offset`` loop with loop-free per-node aranges. | 621 | replace the flat ``for r0_offset`` loop with loop-free per-node aranges. |
| @@ -756,13 +866,17 @@ def _npu_apply_promoted_rtree_lines( | |||
| 756 | if len(cur) - len(cur.lstrip()) >= 4: | 866 | if len(cur) - len(cur.lstrip()) >= 4: |
| 757 | rep = rep[4:] | 867 | rep = rep[4:] |
| 758 | # Rewrite the r-mask term inside the deferred line text. | 868 | # Rewrite the r-mask term inside the deferred line text. |
| 869 | + # Per-load: only include masks for r-nodes the index actually uses. | ||
| 870 | + _deferred_mask = _npu_mask_for_load_line( | ||
| 871 | + rep.line, prefix, leaf_names, combined_mask | ||
| 872 | + ) | ||
| 759 | if f"{prefix}mask & xmask" in rep.line: | 873 | if f"{prefix}mask & xmask" in rep.line: |
| 760 | rep = rep._new_line(rep.line.replace( | 874 | rep = rep._new_line(rep.line.replace( |
| 761 | - f"{prefix}mask & xmask", f"({combined_mask}) & xmask" | 875 | + f"{prefix}mask & xmask", f"({_deferred_mask}) & xmask" |
| 762 | )) | 876 | )) |
| 763 | elif f"{prefix}mask" in rep.line: | 877 | elif f"{prefix}mask" in rep.line: |
| 764 | rep = rep._new_line(rep.line.replace( | 878 | rep = rep._new_line(rep.line.replace( |
| 765 | - f"{prefix}mask", f"({combined_mask})" | 879 | + f"{prefix}mask", f"({_deferred_mask})" |
| 766 | )) | 880 | )) |
| 767 | new_lines.append(rep) | 881 | new_lines.append(rep) |
| 768 | else: | 882 | else: |
| @@ -855,14 +969,18 @@ def _npu_apply_promoted_rtree_lines( | |||
| 855 | ): | 969 | ): |
| 856 | continue | 970 | continue |
| 857 | # Rewrite the accumulate guard mask r0_mask -> combined mask. | 971 | # Rewrite the accumulate guard mask r0_mask -> combined mask. |
| 972 | + # Per-load: only include masks for r-nodes the index actually uses. | ||
| 858 | body = stripped | 973 | body = stripped |
| 974 | + _body_mask = _npu_mask_for_load_line( | ||
| 975 | + body, prefix, leaf_names, combined_mask | ||
| 976 | + ) | ||
| 859 | if f"{prefix}mask & xmask" in body: | 977 | if f"{prefix}mask & xmask" in body: |
| 860 | body = body.replace( | 978 | body = body.replace( |
| 861 | f"{prefix}mask & xmask", | 979 | f"{prefix}mask & xmask", |
| 862 | - f"({combined_mask}) & xmask", | 980 | + f"({_body_mask}) & xmask", |
| 863 | ) | 981 | ) |
| 864 | elif f"{prefix}mask" in body: | 982 | elif f"{prefix}mask" in body: |
| 865 | - body = body.replace(prefix + "mask", f"({combined_mask})") | 983 | + body = body.replace(prefix + "mask", f"({_body_mask})") |
| 866 | new_lines.append(f"{body_indent}{body}") | 984 | new_lines.append(f"{body_indent}{body}") |
| 867 | continue | 985 | continue |
| 868 | 986 | ||
| @@ -1557,22 +1675,26 @@ class NPUTritonKernelOverrides(TritonKernelOverrides): | |||
| 1557 | dtype = torch.float32 | 1675 | dtype = torch.float32 |
| 1558 | elif dtype == torch.int64: | 1676 | elif dtype == torch.int64: |
| 1559 | dtype = torch.int32 | 1677 | dtype = torch.int32 |
| 1560 | - # When all tile dimensions are size 1 (degenerate kernel), use scalar | 1678 | + # NPU: always use scalar (0D) constant. The store address may have higher |
| 1561 | - # constant to avoid "Value argument cannot be block type if pointer | 1679 | + # dimensionality than triton_tensor_ndim() reports (e.g. a non-linearized |
| 1562 | - # argument is not a block" error in tl.store. | 1680 | + # 2D pointwise kernel where [None,:]/[:,None] broadcasting yields 2D |
| 1681 | + # indices while triton_tensor_ndim() returns 1). A scalar constant avoids | ||
| 1682 | + # rank mismatch ("Cannot broadcast, rank mismatch: [1], [2, 2]"). | ||
| 1683 | + # | ||
| 1684 | + # Scalar→any-shape implicit broadcasting is standard Triton semantics | ||
| 1685 | + # (the upstream TritonOverrides.constant historically used shape=[]). | ||
| 1686 | + # Upstream later switched to shape=[1]*ndim ONLY because triton-rocm does | ||
| 1687 | + # not support shape=[] (see upstream comment "Cannot use shape=[] as it's | ||
| 1688 | + # not supported by triton-rocm"). NPU Triton supports scalar constants, so | ||
| 1689 | + # we restore the original scalar behavior — this is more robust than | ||
| 1690 | + # matching the store address rank, because the constant's consumers may | ||
| 1691 | + # have varying ranks within the same kernel. | ||
| 1692 | + # | ||
| 1693 | + # ASSUMPTION: NPU Triton compiler supports tl.full([], val) (scalar) and | ||
| 1694 | + # implicitly broadcasts it to any higher-rank shape in tl.store/tl.where. | ||
| 1695 | + # Verified: test_slice_mutation3_npu (2D fill) + reduction accumulators. | ||
| 1563 | from torch._inductor.codegen.triton import TritonOverrides | 1696 | from torch._inductor.codegen.triton import TritonOverrides |
| 1564 | - ndim = V.kernel.triton_tensor_ndim() | 1697 | + return TritonOverrides.constant(value, dtype) |
| 1565 | - if ndim == 0: | ||
| 1566 | - return TritonOverrides.constant(value, dtype) | ||
| 1567 | - # Check if all range tree dimensions have numel=1 | ||
| 1568 | - all_singleton = all( | ||
| 1569 | - tree.numel == 1 | ||
| 1570 | - for tree in V.kernel.range_trees | ||
| 1571 | - if not tree.is_reduction and tree.tensor_dim is not None | ||
| 1572 | - ) | ||
| 1573 | - if all_singleton: | ||
| 1574 | - return TritonOverrides.constant(value, dtype) | ||
| 1575 | - return super().constant(value, dtype) | ||
| 1576 | 1698 | ||
| 1577 | 1699 | ||
| 1578 | def extract_slice(x, offsets, sizes, strides): | 1700 | def extract_slice(x, offsets, sizes, strides): |
| @@ -2007,6 +2129,12 @@ class NPUTritonKernel(TritonKernel): | |||
| 2007 | return self.kexpr(self.rename_indexing(_drop_size_clamp(index))) | 2129 | return self.kexpr(self.rename_indexing(_drop_size_clamp(index))) |
| 2008 | 2130 | ||
| 2009 | def __init__(self, *args, **kwargs): | 2131 | def __init__(self, *args, **kwargs): |
| 2132 | + # Initialize per-node index vars/ranges BEFORE super().__init__() so that | ||
| 2133 | + # any early access (e.g. hasattr checks in triton_tensor_ndim) does not | ||
| 2134 | + # raise AttributeError. Populated in codegen_node_schedule_with_kernel(). | ||
| 2135 | + if triton_codegen_linearize: | ||
| 2136 | + self.index_vars_per_node = [] | ||
| 2137 | + self.var_ranges_per_node = [] | ||
| 2010 | super().__init__(*args, **kwargs) | 2138 | super().__init__(*args, **kwargs) |
| 2011 | self._axis_split_subs: Dict[sympy.Symbol, sympy.Expr] = {} | 2139 | self._axis_split_subs: Dict[sympy.Symbol, sympy.Expr] = {} |
| 2012 | # r-axis cross-core split (OUTER reduction): when True, codegen_kernel | 2140 | # r-axis cross-core split (OUTER reduction): when True, codegen_kernel |
| @@ -3153,6 +3281,20 @@ class NPUTritonKernel(TritonKernel): | |||
| 3153 | 3281 | ||
| 3154 | def triton_tensor_ndim(self): | 3282 | def triton_tensor_ndim(self): |
| 3155 | if triton_codegen_linearize and getattr(self, '_linearize_applied', False): | 3283 | if triton_codegen_linearize and getattr(self, '_linearize_applied', False): |
| 3284 | + # Post-linearize: use the mapping built by _apply_linearize (stored | ||
| 3285 | + # on each tree as tree_node_mapping). Pre-linearize callers (e.g. | ||
| 3286 | + # constant() during body emit) must NOT use this path because | ||
| 3287 | + # index_vars_per_node is only partially populated at that point, | ||
| 3288 | + # which would yield an incomplete mapping and an inflated rank. | ||
| 3289 | + # | ||
| 3290 | + # Returns the OUTPUT rank: each r-tree contributes ONE axis (the | ||
| 3291 | + # reduction axis), regardless of how many internal tiling sub-nodes | ||
| 3292 | + # it is split into. This matches the upstream fallback | ||
| 3293 | + # (``sum(int(tree.tensor_dim is not None)`` = 1 per r-tree) and is | ||
| 3294 | + # the rank expected by reduction_resize / dense_size_str / store | ||
| 3295 | + # address generation. The iteration-space rank (one dim per | ||
| 3296 | + # sub-node for promoted r-trees) is computed separately via | ||
| 3297 | + # _npu_iter_ndim() for accumulator/broadcast shapes. | ||
| 3156 | ndim = 0 | 3298 | ndim = 0 |
| 3157 | for tree in self.range_trees: | 3299 | for tree in self.range_trees: |
| 3158 | if tree.tensor_dim is not None: | 3300 | if tree.tensor_dim is not None: |
| @@ -3165,15 +3307,38 @@ class NPUTritonKernel(TritonKernel): | |||
| 3165 | count += 1 | 3307 | count += 1 |
| 3166 | ndim += count if count else 1 | 3308 | ndim += count if count else 1 |
| 3167 | else: | 3309 | else: |
| 3168 | - # Promoted reduction trees occupy one slot PER surviving | 3310 | + ndim += 1 # each r-tree = 1 output axis |
| 3169 | - # free sub-node (real-block multi-tile), mirroring the | ||
| 3170 | - # x-tree count above; non-promoted r-trees keep 1 slot. | ||
| 3171 | - promoted = getattr(self, "_npu_rtree_promoted", {}) or {} | ||
| 3172 | - slots = promoted.get(tree.prefix) | ||
| 3173 | - ndim += len(slots) if slots else 1 | ||
| 3174 | return ndim | 3311 | return ndim |
| 3175 | return sum(int(tree.tensor_dim is not None) for tree in self.range_trees) | 3312 | return sum(int(tree.tensor_dim is not None) for tree in self.range_trees) |
| 3176 | 3313 | ||
| 3314 | + def _npu_iter_ndim(self): | ||
| 3315 | + """Iteration-space rank: one dimension per sub-node for promoted | ||
| 3316 | + r-trees (matching var_tensor_dims slots). Used for accumulator and | ||
| 3317 | + broadcast shapes, where the iteration space may have more dimensions | ||
| 3318 | + than the output rank. | ||
| 3319 | + | ||
| 3320 | + Example: a reduction over a 16-element axis split into 3 tiles | ||
| 3321 | + (r0_0 len=2, r0_1 len=8, r0_2 len=1, R0_BLOCK=16 → promoted). | ||
| 3322 | + The codegen/store rank is 2 (X plus the collapsed reduction axis), | ||
| 3323 | + while the iteration space uses X plus the promoted reduction slots: | ||
| 3324 | + accumulator shape = [XBLOCK, r0_0_blk, r0_1_blk] (iter_ndim=3) | ||
| 3325 | + post_resize target = [:, None] (real_ndim=2) | ||
| 3326 | + store address = [XBLOCK, 1] (real_ndim=2) | ||
| 3327 | + """ | ||
| 3328 | + if triton_codegen_linearize and getattr(self, '_linearize_applied', False): | ||
| 3329 | + # var_tensor_dims values are slots in one global register-tensor | ||
| 3330 | + # space, so taking each tree's rank and summing double-counts all | ||
| 3331 | + # slots assigned by earlier trees. | ||
| 3332 | + max_slot = -1 | ||
| 3333 | + for tree in self.range_trees: | ||
| 3334 | + if tree.tensor_dim is not None: | ||
| 3335 | + max_slot = max(max_slot, tree.tensor_dim) | ||
| 3336 | + vtd = getattr(tree, "var_tensor_dims", {}) or {} | ||
| 3337 | + if vtd: | ||
| 3338 | + max_slot = max(max_slot, max(vtd.values())) | ||
| 3339 | + return max(self.triton_tensor_ndim(), max_slot + 1) | ||
| 3340 | + return sum(int(tree.tensor_dim is not None) for tree in self.range_trees) | ||
| 3341 | + | ||
| 3177 | def dense_size_str(self): | 3342 | def dense_size_str(self): |
| 3178 | if triton_codegen_linearize: | 3343 | if triton_codegen_linearize: |
| 3179 | ndim = self.triton_tensor_ndim() | 3344 | ndim = self.triton_tensor_ndim() |
| @@ -4211,6 +4376,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 4211 | ): | 4376 | ): |
| 4212 | _npu_rewrite_promoted_rtree_body( | 4377 | _npu_rewrite_promoted_rtree_body( |
| 4213 | self, rw["real_sizes"], rw["real_ndim"], | 4378 | self, rw["real_sizes"], rw["real_ndim"], |
| 4379 | + iter_ndim=rw.get("iter_ndim", rw["real_ndim"]), | ||
| 4214 | ) | 4380 | ) |
| 4215 | self._npu_rtree_rewrite_done = True | 4381 | self._npu_rtree_rewrite_done = True |
| 4216 | self._maybe_rewrite_select_lane_load() | 4382 | self._maybe_rewrite_select_lane_load() |
| @@ -6013,7 +6179,8 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6013 | # contiguous-axis-innermost invariant the divisor sort sets up. | 6179 | # contiguous-axis-innermost invariant the divisor sort sets up. |
| 6014 | kernel._npu_tile_permuted = True | 6180 | kernel._npu_tile_permuted = True |
| 6015 | 6181 | ||
| 6016 | - def _fold_dual_decomp(self, kernel, tree, tree_expr, tree_node_mapping, matcher): | 6182 | + @staticmethod |
| 6183 | + def _fold_dual_decomp(kernel, tree, tree_expr, tree_node_mapping, matcher): | ||
| 6017 | """Collapse a secondary full divisor-chain decomposition onto the basis. | 6184 | """Collapse a secondary full divisor-chain decomposition onto the basis. |
| 6018 | 6185 | ||
| 6019 | See the call site (NPU_FOLD_DUAL_DECOMP) for the data-layout rationale. | 6186 | See the call site (NPU_FOLD_DUAL_DECOMP) for the data-layout rationale. |
| @@ -6217,27 +6384,35 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6217 | recon_nodes.update(mapped) | 6384 | recon_nodes.update(mapped) |
| 6218 | return True | 6385 | return True |
| 6219 | 6386 | ||
| 6220 | - def _apply_linearize(self, kernel, node_schedule): | 6387 | + @staticmethod |
| 6221 | - """ | 6388 | + def _npu_build_tree_node_mapping(kernel): |
| 6222 | - Post-process a kernel for linearize mode: | 6389 | + """Build per-tree node mappings + body-rewrite matcher for linearize mode. |
| 6223 | - - Build tree_node_mapping for multi-dimensional iteration | ||
| 6224 | - - Remap tensor dims to per-node dims | ||
| 6225 | - - Replace codegen_range_tree with NPU version | ||
| 6226 | - - Replace tl.program_id(0) with group-based dispatch | ||
| 6227 | - """ | ||
| 6228 | - if kernel is None: | ||
| 6229 | - return | ||
| 6230 | 6390 | ||
| 6231 | - kernel._npu_rsplit_candidate = False | 6391 | + Runs the SAME mapping passes as _apply_linearize (rank-subrange fold, |
| 6232 | - if _npu_rsplit_outer_applicable(kernel): | 6392 | + flat-node fold, dual-decomp fold, reduction-view fold) without touching |
| 6233 | - try: | 6393 | + the kernel body or node structure (no expand_divmod / collapse_rowmajor). |
| 6234 | - if kernel.features.get_reduction_hint() == ReductionHint.INNER: | 6394 | + Called by _apply_linearize, which writes the result back onto the trees |
| 6235 | - if _npu_rsplit_pick_split_axis(kernel, require_dynamic=False): | 6395 | + (as tree.tree_node_mapping). triton_tensor_ndim() reads those mappings |
| 6236 | - kernel._npu_rsplit_candidate = True | 6396 | + post-linearize — it cannot call this directly because index_vars_per_node |
| 6237 | - else: | 6397 | + is only fully populated after body emit completes. Returns (mappings, matcher). |
| 6238 | - kernel.npu_rsplit_partial = True | 6398 | + """ |
| 6239 | - except Exception: | 6399 | + # Filter r-axes on LOCAL copies: _apply_linearize writes the filtered |
| 6240 | - pass | 6400 | + # lists back onto the kernel, but pre-linearize callers must not mutate |
| 6401 | + # kernel state here. | ||
| 6402 | + index_vars_per_node = [ | ||
| 6403 | + [ | ||
| 6404 | + [item for item in sublist if not str(item).startswith('r')] | ||
| 6405 | + for sublist in nested_list | ||
| 6406 | + ] | ||
| 6407 | + for nested_list in kernel.index_vars_per_node | ||
| 6408 | + ] | ||
| 6409 | + var_ranges_per_node = [ | ||
| 6410 | + [ | ||
| 6411 | + [item2 for item1, item2 in zip(sublist1, sublist2) if not str(item1).startswith('r')] | ||
| 6412 | + for sublist1, sublist2 in zip(nested_list1, nested_list2) | ||
| 6413 | + ] | ||
| 6414 | + for nested_list1, nested_list2 in zip(kernel.index_vars_per_node, kernel.var_ranges_per_node) | ||
| 6415 | + ] | ||
| 6241 | 6416 | ||
| 6242 | def indexer(index_var, var_range): | 6417 | def indexer(index_var, var_range): |
| 6243 | strides = [sympy.Integer(1)] | 6418 | strides = [sympy.Integer(1)] |
| @@ -6254,20 +6429,7 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6254 | 6429 | ||
| 6255 | matcher = {} | 6430 | matcher = {} |
| 6256 | 6431 | ||
| 6257 | - kernel.index_vars_per_node = [ | 6432 | + mappings = {} |
| 6258 | - [ | ||
| 6259 | - [item for item in sublist if not str(item).startswith('r')] | ||
| 6260 | - for sublist in nested_list | ||
| 6261 | - ] | ||
| 6262 | - for nested_list in kernel.index_vars_per_node | ||
| 6263 | - ] | ||
| 6264 | - kernel.var_ranges_per_node = [ | ||
| 6265 | - [ | ||
| 6266 | - [item2 for item1, item2 in zip(sublist1, sublist2) if not str(item1).startswith('r')] | ||
| 6267 | - for sublist1, sublist2 in zip(nested_list1, nested_list2) | ||
| 6268 | - ] | ||
| 6269 | - for nested_list1, nested_list2 in zip(kernel.index_vars_per_node, kernel.var_ranges_per_node) | ||
| 6270 | - ] | ||
| 6271 | 6433 | ||
| 6272 | for i, tree in enumerate(kernel.range_trees): | 6434 | for i, tree in enumerate(kernel.range_trees): |
| 6273 | tree_expr = None | 6435 | tree_expr = None |
| @@ -6279,7 +6441,7 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6279 | 6441 | ||
| 6280 | for id, (var_ranges, index_vars) in enumerate( | 6442 | for id, (var_ranges, index_vars) in enumerate( |
| 6281 | sorted( | 6443 | sorted( |
| 6282 | - zip(kernel.var_ranges_per_node, kernel.index_vars_per_node), | 6444 | + zip(var_ranges_per_node, index_vars_per_node), |
| 6283 | key=lambda pair: len(pair[0][i]) if i < len(pair[0]) else 0, | 6445 | key=lambda pair: len(pair[0][i]) if i < len(pair[0]) else 0, |
| 6284 | reverse=True, | 6446 | reverse=True, |
| 6285 | ) | 6447 | ) |
| @@ -6369,6 +6531,56 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6369 | if sizevars.statically_known_equals(n.length, tree.numel): | 6531 | if sizevars.statically_known_equals(n.length, tree.numel): |
| 6370 | flat_node = n | 6532 | flat_node = n |
| 6371 | break | 6533 | break |
| 6534 | + if flat_node is None: | ||
| 6535 | + # Partial-flat fold: a divisor==1 node whose length is | ||
| 6536 | + # exactly covered by a contiguous divisor-chain of the | ||
| 6537 | + # OTHER free nodes (a flattened view of a sub-space, e.g. | ||
| 6538 | + # x0 len 5304 == x2(26) * x3(204) while x1 strides the | ||
| 6539 | + # outer dim). The basis-entry fold above needs a full-rank | ||
| 6540 | + # entry in var_ranges_per_node, which is only complete | ||
| 6541 | + # AFTER body emit; constant() runs DURING body emit, so | ||
| 6542 | + # the pre-linearize ndim must detect this derived node | ||
| 6543 | + # structurally (divisor chain) instead. | ||
| 6544 | + # | ||
| 6545 | + # KNOWN LIMITATION: handles a SINGLE divisor==1 node covered | ||
| 6546 | + # by a contiguous chain of the other nodes. Does NOT handle: | ||
| 6547 | + # - Multiple divisor==1 nodes sharing the same chain. | ||
| 6548 | + # - A chain that covers multiple divisor==1 nodes. | ||
| 6549 | + # These cases do not arise in current op coverage; if they | ||
| 6550 | + # do, extend this to iterate all candidates and partition | ||
| 6551 | + # the chain among them. | ||
| 6552 | + for cand in free_nodes: | ||
| 6553 | + if not (isinstance(cand.divisor, (int, sympy.Integer)) | ||
| 6554 | + and int(cand.divisor) == 1): | ||
| 6555 | + continue | ||
| 6556 | + if sizevars.statically_known_equals(cand.length, tree.numel): | ||
| 6557 | + continue | ||
| 6558 | + others = [n for n in free_nodes if n is not cand] | ||
| 6559 | + others_sorted = sorted( | ||
| 6560 | + others, key=lambda n: sizevars.optimization_hint(n.divisor) | ||
| 6561 | + ) | ||
| 6562 | + if not others_sorted: | ||
| 6563 | + continue | ||
| 6564 | + if not (isinstance(others_sorted[0].divisor, (int, sympy.Integer)) | ||
| 6565 | + and int(others_sorted[0].divisor) == 1): | ||
| 6566 | + continue | ||
| 6567 | + chain = [] | ||
| 6568 | + expected_div = sympy.Integer(1) | ||
| 6569 | + for n in others_sorted: | ||
| 6570 | + if sizevars.statically_known_equals(n.divisor, expected_div): | ||
| 6571 | + chain.append(n) | ||
| 6572 | + expected_div = n.divisor * n.length | ||
| 6573 | + if sizevars.statically_known_equals(expected_div, cand.length): | ||
| 6574 | + break | ||
| 6575 | + elif (sizevars.optimization_hint(n.divisor) | ||
| 6576 | + != sizevars.optimization_hint(expected_div)): | ||
| 6577 | + break | ||
| 6578 | + if chain and sizevars.statically_known_equals(expected_div, cand.length): | ||
| 6579 | + node_expr = NPUTritonScheduling._flat_node_expr(chain) | ||
| 6580 | + tree_node_mapping[cand.name] = node_expr | ||
| 6581 | + matcher[f"{cand.name} = {cand.name}index"] = f"{cand.name} = {node_expr}" | ||
| 6582 | + break | ||
| 6583 | + | ||
| 6372 | if flat_node is None: | 6584 | if flat_node is None: |
| 6373 | for n in free_nodes: | 6585 | for n in free_nodes: |
| 6374 | if isinstance(n.divisor, (int, sympy.Integer)) and int(n.divisor) == 1: | 6586 | if isinstance(n.divisor, (int, sympy.Integer)) and int(n.divisor) == 1: |
| @@ -6379,9 +6591,9 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6379 | others = [n for n in free_nodes if n is not flat_node] | 6591 | others = [n for n in free_nodes if n is not flat_node] |
| 6380 | others_sorted = sorted(others, key=lambda n: sizevars.optimization_hint(n.divisor)) | 6592 | others_sorted = sorted(others, key=lambda n: sizevars.optimization_hint(n.divisor)) |
| 6381 | if isinstance(others_sorted[0].divisor, (int, sympy.Integer)) and int(others_sorted[0].divisor) == 1: | 6593 | if isinstance(others_sorted[0].divisor, (int, sympy.Integer)) and int(others_sorted[0].divisor) == 1: |
| 6382 | - chain_ok, _ = self._divisor_chain_ok(others_sorted, sizevars) | 6594 | + chain_ok, _ = NPUTritonScheduling._divisor_chain_ok(others_sorted, sizevars) |
| 6383 | if chain_ok: | 6595 | if chain_ok: |
| 6384 | - node_expr = self._flat_node_expr(others_sorted) | 6596 | + node_expr = NPUTritonScheduling._flat_node_expr(others_sorted) |
| 6385 | tree_node_mapping[flat_node.name] = node_expr | 6597 | tree_node_mapping[flat_node.name] = node_expr |
| 6386 | pattern = f"{flat_node.name} = {flat_node.name}index" | 6598 | pattern = f"{flat_node.name} = {flat_node.name}index" |
| 6387 | replacement = f"{flat_node.name} = {node_expr}" | 6599 | replacement = f"{flat_node.name} = {node_expr}" |
| @@ -6397,12 +6609,12 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6397 | and not tree.is_reduction | 6609 | and not tree.is_reduction |
| 6398 | and ncfg.fold_dual_decomp | 6610 | and ncfg.fold_dual_decomp |
| 6399 | ): | 6611 | ): |
| 6400 | - self._fold_dual_decomp( | 6612 | + NPUTritonScheduling._fold_dual_decomp( |
| 6401 | kernel, tree, tree_expr, tree_node_mapping, matcher, | 6613 | kernel, tree, tree_expr, tree_node_mapping, matcher, |
| 6402 | ) | 6614 | ) |
| 6403 | 6615 | ||
| 6404 | if tree.is_reduction and ncfg.fold_flat_rnode: | 6616 | if tree.is_reduction and ncfg.fold_flat_rnode: |
| 6405 | - self._fold_reduction_view_decompositions( | 6617 | + NPUTritonScheduling._fold_reduction_view_decompositions( |
| 6406 | kernel, tree, tree_node_mapping, matcher, | 6618 | kernel, tree, tree_node_mapping, matcher, |
| 6407 | ) | 6619 | ) |
| 6408 | 6620 | ||
| @@ -6434,12 +6646,12 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6434 | if flat_node is not None and others: | 6646 | if flat_node is not None and others: |
| 6435 | others_sorted = sorted(others, key=lambda n: sizevars.optimization_hint(n.divisor)) | 6647 | others_sorted = sorted(others, key=lambda n: sizevars.optimization_hint(n.divisor)) |
| 6436 | if isinstance(others_sorted[0].divisor, (int, sympy.Integer)) and int(others_sorted[0].divisor) == 1: | 6648 | if isinstance(others_sorted[0].divisor, (int, sympy.Integer)) and int(others_sorted[0].divisor) == 1: |
| 6437 | - chain_ok, expected_div = self._divisor_chain_ok(others_sorted, sizevars) | 6649 | + chain_ok, expected_div = NPUTritonScheduling._divisor_chain_ok(others_sorted, sizevars) |
| 6438 | # The decomposition must cover the whole flat space. | 6650 | # The decomposition must cover the whole flat space. |
| 6439 | if chain_ok and not sizevars.statically_known_equals(expected_div, tree.numel): | 6651 | if chain_ok and not sizevars.statically_known_equals(expected_div, tree.numel): |
| 6440 | chain_ok = False | 6652 | chain_ok = False |
| 6441 | if chain_ok: | 6653 | if chain_ok: |
| 6442 | - node_expr = self._flat_node_expr(others_sorted) | 6654 | + node_expr = NPUTritonScheduling._flat_node_expr(others_sorted) |
| 6443 | tree_node_mapping[flat_node.name] = node_expr | 6655 | tree_node_mapping[flat_node.name] = node_expr |
| 6444 | # Flat node is assigned `r0_3 = r0_index` (prefix-level), | 6656 | # Flat node is assigned `r0_3 = r0_index` (prefix-level), |
| 6445 | # so record the sub for codegen_body to rewrite into | 6657 | # so record the sub for codegen_body to rewrite into |
| @@ -6463,7 +6675,53 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6463 | tree._npu_split_recon_nodes = recon_nodes | 6675 | tree._npu_split_recon_nodes = recon_nodes |
| 6464 | recon_nodes.add(flat_node.name) | 6676 | recon_nodes.add(flat_node.name) |
| 6465 | 6677 | ||
| 6466 | - tree.tree_node_mapping = tree_node_mapping | 6678 | + mappings[i] = tree_node_mapping |
| 6679 | + | ||
| 6680 | + return mappings, matcher | ||
| 6681 | + | ||
| 6682 | + | ||
| 6683 | + def _apply_linearize(self, kernel, node_schedule): | ||
| 6684 | + """ | ||
| 6685 | + Post-process a kernel for linearize mode: | ||
| 6686 | + - Build tree_node_mapping for multi-dimensional iteration | ||
| 6687 | + - Remap tensor dims to per-node dims | ||
| 6688 | + - Replace codegen_range_tree with NPU version | ||
| 6689 | + - Replace tl.program_id(0) with group-based dispatch | ||
| 6690 | + """ | ||
| 6691 | + if kernel is None: | ||
| 6692 | + return | ||
| 6693 | + | ||
| 6694 | + kernel._npu_rsplit_candidate = False | ||
| 6695 | + if _npu_rsplit_outer_applicable(kernel): | ||
| 6696 | + try: | ||
| 6697 | + if kernel.features.get_reduction_hint() == ReductionHint.INNER: | ||
| 6698 | + if _npu_rsplit_pick_split_axis(kernel, require_dynamic=False): | ||
| 6699 | + kernel._npu_rsplit_candidate = True | ||
| 6700 | + else: | ||
| 6701 | + kernel.npu_rsplit_partial = True | ||
| 6702 | + except Exception: | ||
| 6703 | + pass | ||
| 6704 | + | ||
| 6705 | + matcher = {} | ||
| 6706 | + | ||
| 6707 | + kernel.index_vars_per_node = [ | ||
| 6708 | + [ | ||
| 6709 | + [item for item in sublist if not str(item).startswith('r')] | ||
| 6710 | + for sublist in nested_list | ||
| 6711 | + ] | ||
| 6712 | + for nested_list in kernel.index_vars_per_node | ||
| 6713 | + ] | ||
| 6714 | + kernel.var_ranges_per_node = [ | ||
| 6715 | + [ | ||
| 6716 | + [item2 for item1, item2 in zip(sublist1, sublist2) if not str(item1).startswith('r')] | ||
| 6717 | + for sublist1, sublist2 in zip(nested_list1, nested_list2) | ||
| 6718 | + ] | ||
| 6719 | + for nested_list1, nested_list2 in zip(kernel.index_vars_per_node, kernel.var_ranges_per_node) | ||
| 6720 | + ] | ||
| 6721 | + | ||
| 6722 | + mappings, matcher = self._npu_build_tree_node_mapping(kernel) | ||
| 6723 | + for i, tree in enumerate(kernel.range_trees): | ||
| 6724 | + tree.tree_node_mapping = mappings[i] | ||
| 6467 | 6725 | ||
| 6468 | # Expand FloorDiv/Mod patterns on single-node trees into sub-nodes. | 6726 | # Expand FloorDiv/Mod patterns on single-node trees into sub-nodes. |
| 6469 | if ncfg.expand_divmod: | 6727 | if ncfg.expand_divmod: |
| @@ -6997,13 +7255,20 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6997 | kernel.body._lines = filtered_lines | 7255 | kernel.body._lines = filtered_lines |
| 6998 | 7256 | ||
| 6999 | # Compute the real linearized dense_size and reduction_dim. | 7257 | # Compute the real linearized dense_size and reduction_dim. |
| 7258 | + # real_ndim = output rank (each r-tree = 1 axis). This is used for | ||
| 7259 | + # reduction_resize / dense_size_str / store address generation. | ||
| 7260 | + # iter_ndim = iteration-space rank (one dim per sub-node for promoted | ||
| 7261 | + # r-trees). This is used for accumulator/broadcast shapes. For full | ||
| 7262 | + # reduction (no_x_dim), iter_ndim == real_ndim. For multi-dim | ||
| 7263 | + # reduction with promoted r-trees, iter_ndim > real_ndim. | ||
| 7000 | real_ndim = kernel.triton_tensor_ndim() | 7264 | real_ndim = kernel.triton_tensor_ndim() |
| 7265 | + iter_ndim = kernel._npu_iter_ndim() | ||
| 7001 | 7266 | ||
| 7002 | # Two real_sizes strategies: permuted kernels (hook moved tensor_dim cross- | 7267 | # Two real_sizes strategies: permuted kernels (hook moved tensor_dim cross- |
| 7003 | # tree, e.g. R→0) index by tree.tensor_dim so slots match old_sizes; un-permuted | 7268 | # tree, e.g. R→0) index by tree.tensor_dim so slots match old_sizes; un-permuted |
| 7004 | # keep the sequential dim_idx walk (tensor_dim-indexing there mis-slots tl.full/ | 7269 | # keep the sequential dim_idx walk (tensor_dim-indexing there mis-slots tl.full/ |
| 7005 | # tl.broadcast_to shapes), so we deviate only when the hook moved a tree. | 7270 | # tl.broadcast_to shapes), so we deviate only when the hook moved a tree. |
| 7006 | - real_sizes = ["1"] * real_ndim | 7271 | + real_sizes = ["1"] * iter_ndim |
| 7007 | permuted = getattr(kernel, "_npu_tile_permuted", False) | 7272 | permuted = getattr(kernel, "_npu_tile_permuted", False) |
| 7008 | 7273 | ||
| 7009 | # Free x-sub-node slot order MUST match var_tensor_dims (divisor-descending); | 7274 | # Free x-sub-node slot order MUST match var_tensor_dims (divisor-descending); |
| @@ -7035,7 +7300,7 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 7035 | if not tree.is_reduction: | 7300 | if not tree.is_reduction: |
| 7036 | node_block_constexpr = getattr(tree, "node_block_constexpr", {}) or {} | 7301 | node_block_constexpr = getattr(tree, "node_block_constexpr", {}) or {} |
| 7037 | for node_slot, node in _ordered_free_nodes(tree, slot): | 7302 | for node_slot, node in _ordered_free_nodes(tree, slot): |
| 7038 | - if node_slot < real_ndim: | 7303 | + if node_slot < iter_ndim: |
| 7039 | real_sizes[node_slot] = node_block_constexpr.get( | 7304 | real_sizes[node_slot] = node_block_constexpr.get( |
| 7040 | node.name, f"real_block_{node.name}" | 7305 | node.name, f"real_block_{node.name}" |
| 7041 | ) | 7306 | ) |
| @@ -7045,16 +7310,20 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 7045 | # its per-node block token (not a single R0_BLOCK). Static | 7310 | # its per-node block token (not a single R0_BLOCK). Static |
| 7046 | # nodes use real_block_{node}; symbolic use the constexpr tile | 7311 | # nodes use real_block_{node}; symbolic use the constexpr tile |
| 7047 | # {node}_blk. Classified per-node so a dynamic tree can hold | 7312 | # {node}_blk. Classified per-node so a dynamic tree can hold |
| 7048 | - # static nodes (ViT mixed seq+batch). | 7313 | + # static nodes (ViT mixed seq+batch). All sub-node slots are |
| 7314 | + # filled (up to iter_ndim) so the accumulator/broadcast shape | ||
| 7315 | + # matches the iteration space, even when the number of | ||
| 7316 | + # sub-nodes exceeds the output rank real_ndim. | ||
| 7049 | _ndyn = getattr(kernel, "_npu_rtree_node_dynamic", {}) | 7317 | _ndyn = getattr(kernel, "_npu_rtree_node_dynamic", {}) |
| 7050 | for node_slot, node in _ordered_free_nodes(tree, slot): | 7318 | for node_slot, node in _ordered_free_nodes(tree, slot): |
| 7051 | - if node_slot < real_ndim: | 7319 | + if node_slot < iter_ndim: |
| 7052 | real_sizes[node_slot] = ( | 7320 | real_sizes[node_slot] = ( |
| 7053 | f"{node.name}_blk" if _ndyn.get(node.name) | 7321 | f"{node.name}_blk" if _ndyn.get(node.name) |
| 7054 | else f"real_block_{node.name}" | 7322 | else f"real_block_{node.name}" |
| 7055 | ) | 7323 | ) |
| 7056 | else: | 7324 | else: |
| 7057 | - real_sizes[slot] = f"{tree.prefix.upper()}BLOCK" | 7325 | + if slot < iter_ndim: |
| 7326 | + real_sizes[slot] = f"{tree.prefix.upper()}BLOCK" | ||
| 7058 | else: | 7327 | else: |
| 7059 | dim_idx = 0 | 7328 | dim_idx = 0 |
| 7060 | for tree in kernel.range_trees: | 7329 | for tree in kernel.range_trees: |
| @@ -7063,7 +7332,7 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 7063 | if not tree.is_reduction: | 7332 | if not tree.is_reduction: |
| 7064 | node_block_constexpr = getattr(tree, "node_block_constexpr", {}) or {} | 7333 | node_block_constexpr = getattr(tree, "node_block_constexpr", {}) or {} |
| 7065 | for node_slot, node in _ordered_free_nodes(tree, dim_idx): | 7334 | for node_slot, node in _ordered_free_nodes(tree, dim_idx): |
| 7066 | - if node_slot < real_ndim: | 7335 | + if node_slot < iter_ndim: |
| 7067 | real_sizes[node_slot] = node_block_constexpr.get( | 7336 | real_sizes[node_slot] = node_block_constexpr.get( |
| 7068 | node.name, f"real_block_{node.name}" | 7337 | node.name, f"real_block_{node.name}" |
| 7069 | ) | 7338 | ) |
| @@ -7072,14 +7341,15 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 7072 | if tree.prefix in getattr(kernel, "_npu_rtree_promoted", {}): | 7341 | if tree.prefix in getattr(kernel, "_npu_rtree_promoted", {}): |
| 7073 | _ndyn = getattr(kernel, "_npu_rtree_node_dynamic", {}) | 7342 | _ndyn = getattr(kernel, "_npu_rtree_node_dynamic", {}) |
| 7074 | for node_slot, node in _ordered_free_nodes(tree, dim_idx): | 7343 | for node_slot, node in _ordered_free_nodes(tree, dim_idx): |
| 7075 | - if node_slot < real_ndim: | 7344 | + if node_slot < iter_ndim: |
| 7076 | real_sizes[node_slot] = ( | 7345 | real_sizes[node_slot] = ( |
| 7077 | f"{node.name}_blk" if _ndyn.get(node.name) | 7346 | f"{node.name}_blk" if _ndyn.get(node.name) |
| 7078 | else f"real_block_{node.name}" | 7347 | else f"real_block_{node.name}" |
| 7079 | ) | 7348 | ) |
| 7080 | dim_idx = max(dim_idx, node_slot + 1) | 7349 | dim_idx = max(dim_idx, node_slot + 1) |
| 7081 | else: | 7350 | else: |
| 7082 | - real_sizes[dim_idx] = f"{tree.prefix.upper()}BLOCK" | 7351 | + if dim_idx < iter_ndim: |
| 7352 | + real_sizes[dim_idx] = f"{tree.prefix.upper()}BLOCK" | ||
| 7083 | dim_idx += 1 | 7353 | dim_idx += 1 |
| 7084 | real_dense_size = f"[{', '.join(real_sizes)}]" | 7354 | real_dense_size = f"[{', '.join(real_sizes)}]" |
| 7085 | # R-tree's actual slot. If the hook pushed it to an outer slot, upstream's | 7355 | # R-tree's actual slot. If the hook pushed it to an outer slot, upstream's |
| @@ -7190,6 +7460,7 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 7190 | kernel._npu_rtree_rewrite_info = { | 7460 | kernel._npu_rtree_rewrite_info = { |
| 7191 | "real_sizes": list(real_sizes), | 7461 | "real_sizes": list(real_sizes), |
| 7192 | "real_ndim": real_ndim, | 7462 | "real_ndim": real_ndim, |
| 7463 | + "iter_ndim": iter_ndim, | ||
| 7193 | } | 7464 | } |
| 7194 | 7465 | ||
| 7195 | def codegen_node_schedule_with_kernel(self, node_schedule, kernel): | 7466 | def codegen_node_schedule_with_kernel(self, node_schedule, kernel): |
| @@ -52,6 +52,7 @@ GENERATE_LIST = [ | |||
| 52 | aten.clamp_max, | 52 | aten.clamp_max, |
| 53 | aten.bitwise_not, | 53 | aten.bitwise_not, |
| 54 | aten.tanh, | 54 | aten.tanh, |
| 55 | + aten.fmod, | ||
| 55 | aten.copy, | 56 | aten.copy, |
| 56 | aten.copy_, | 57 | aten.copy_, |
| 57 | aten.remainder, | 58 | aten.remainder, |