已合并
fix: iterate _maybe_split_fused_axes to fixpoint (multi-level digit extraction) #45161
AllenGuan创建于 15 天前
fix: iterate _maybe_split_fused_axes to fixpoint (multi-level digit extraction) #45161
已合并
共 1 个文件变更+27-0
| @@ -2602,7 +2602,34 @@ class NPUTritonKernel(TritonKernel): | |||
| 2602 | return index | 2602 | return index |
| 2603 | return sympy_subs(index, replacements) | 2603 | return sympy_subs(index, replacements) |
| 2604 | 2604 | ||
| 2605 | + # Bound for the split fixpoint: a 131072-flat axis needs <=4 rounds | ||
| 2606 | + # (1024 -> 256 -> 16 -> 2); 6 leaves headroom without risking livelock. | ||
| 2607 | + _SPLIT_FIXPOINT_MAX_ROUNDS = 6 | ||
| 2608 | + | ||
| 2605 | def _maybe_split_fused_axes(self, index: sympy.Expr) -> sympy.Expr: | 2609 | def _maybe_split_fused_axes(self, index: sympy.Expr) -> sympy.Expr: |
| 2610 | + """Iterate ``_split_fused_axes_round`` to a fixpoint. | ||
| 2611 | + | ||
| 2612 | + One round splits each detected fused axis once, but a split can expose | ||
| 2613 | + NEW fused symbols inside residual ModularIndexing/FloorDiv bases — | ||
| 2614 | + e.g. a 131072-flat axis splits at c=1024 into batch-outer plus a | ||
| 2615 | + 1024-fused inner whose digits are still embedded as | ||
| 2616 | + ``ModularIndexing(inner, 1, 256)``. Stopping after one round leaves | ||
| 2617 | + those residues un-folded: the base keeps a fused node, ranges cannot | ||
| 2618 | + collapse the mod, and a wrong address formula reaches the printed | ||
| 2619 | + kernel (kernel-17 pass-3 miscompile: ``Mod(4*x0 + x1, 4, 16)`` with | ||
| 2620 | + x1 a 512-fused node survived to output). Iterate until the index stops | ||
| 2621 | + changing (bounded) so it converges to a pure affine expression of | ||
| 2622 | + leaf axes. | ||
| 2623 | + """ | ||
| 2624 | + prev = None | ||
| 2625 | + for _ in range(self._SPLIT_FIXPOINT_MAX_ROUNDS): | ||
| 2626 | + new_index = self._split_fused_axes_round(index) | ||
| 2627 | + if new_index in (index, prev): | ||
| 2628 | + break | ||
| 2629 | + prev, index = index, new_index | ||
| 2630 | + return index | ||
| 2631 | + | ||
| 2632 | + def _split_fused_axes_round(self, index: sympy.Expr) -> sympy.Expr: | ||
| 2606 | """Split a fused iteration axis into outer/inner sub-axes when an index | 2633 | """Split a fused iteration axis into outer/inner sub-axes when an index |
| 2607 | accesses it via FloorDiv(x, c) or ModularIndexing(x, 1, c). | 2634 | accesses it via FloorDiv(x, c) or ModularIndexing(x, 1, c). |
| 2608 | 2635 | ||