已开启
fix(inductor): optimize DALLE2 attention indexing #45076
伦创建于 19 天前
fix(inductor): optimize DALLE2 attention indexing #45076
已开启
共 2 个文件变更+130-2
| @@ -0,0 +1,81 @@ | |||
| 1 | +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. | ||
| 2 | +# Owner(s): ["module: inductor"] | ||
| 3 | + | ||
| 4 | +from types import SimpleNamespace | ||
| 5 | + | ||
| 6 | +import sympy | ||
| 7 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 8 | + | ||
| 9 | +import torch_npu # noqa: F401 | ||
| 10 | +from torch_npu._inductor.triton_experimental.codegen import triton as npu_triton | ||
| 11 | + | ||
| 12 | + | ||
| 13 | +class _SizeVars: | ||
| 14 | + | ||
| 15 | + def simplify(value): | ||
| 16 | + return sympy.simplify(value) | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + def statically_known_equals(lhs, rhs): | ||
| 20 | + return sympy.simplify(lhs - rhs) == 0 | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + def optimization_hint(value): | ||
| 24 | + return int(value) | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +class TestDalle2CodegenRegressions(TestCase): | ||
| 28 | + def test_fold_dual_decomposition_with_shared_axis(self): | ||
| 29 | + # DALLE2 attention indexes the same 4160 rows as [260, 8, 2] for | ||
| 30 | + # inputs and [260, 16] for output. x0 is shared by both views; x5 | ||
| 31 | + # must be derived from x1/x2 rather than launched as an independent | ||
| 32 | + # axis, which would create a 16x Cartesian grid. | ||
| 33 | + x0, x1, x2 = sympy.symbols("x0 x1 x2", integer=True) | ||
| 34 | + nodes = { | ||
| 35 | + "x0": SimpleNamespace(name="x0", divisor=1, length=260), | ||
| 36 | + "x1": SimpleNamespace(name="x1", divisor=260, length=8), | ||
| 37 | + "x2": SimpleNamespace(name="x2", divisor=2080, length=2), | ||
| 38 | + "x5": SimpleNamespace(name="x5", divisor=260, length=16), | ||
| 39 | + } | ||
| 40 | + tree = SimpleNamespace(nodes=nodes, numel=sympy.Integer(4160)) | ||
| 41 | + kernel = SimpleNamespace() | ||
| 42 | + mapping = {} | ||
| 43 | + matcher = {} | ||
| 44 | + graph = SimpleNamespace(sizevars=_SizeVars()) | ||
| 45 | + | ||
| 46 | + with npu_triton.V.set_graph_handler(graph): | ||
| 47 | + npu_triton.NPUTritonScheduling._fold_dual_decomp( | ||
| 48 | + None, | ||
| 49 | + kernel, | ||
| 50 | + tree, | ||
| 51 | + ([x0, x1, x2], [260, 8, 2]), | ||
| 52 | + mapping, | ||
| 53 | + matcher, | ||
| 54 | + ) | ||
| 55 | + | ||
| 56 | + folded = mapping["x5"] | ||
| 57 | + symbols = {str(symbol): symbol for symbol in folded.free_symbols} | ||
| 58 | + self.assertEqual( | ||
| 59 | + folded.subs( | ||
| 60 | + {symbols["x0"]: 17, symbols["x1"]: 3, symbols["x2"]: 1} | ||
| 61 | + ), | ||
| 62 | + 11, | ||
| 63 | + ) | ||
| 64 | + self.assertNotIn("x5", symbols) | ||
| 65 | + self.assertNotIn("x0", mapping) | ||
| 66 | + self.assertNotIn("x1", mapping) | ||
| 67 | + self.assertNotIn("x2", mapping) | ||
| 68 | + | ||
| 69 | + pattern = "x5 = x5index" | ||
| 70 | + self.assertIn(pattern, matcher) | ||
| 71 | + self.assertNotIn("x5", matcher[pattern].split("=", 1)[1]) | ||
| 72 | + | ||
| 73 | + self.assertTrue(hasattr(kernel, "_npu_addr_text_subs")) | ||
| 74 | + self.assertEqual(len(kernel._npu_addr_text_subs), 1) | ||
| 75 | + source, replacement = next(iter(kernel._npu_addr_text_subs.items())) | ||
| 76 | + self.assertIn("x5", source) | ||
| 77 | + self.assertNotIn("x5", replacement) | ||
| 78 | + | ||
| 79 | + | ||
| 80 | +if __name__ == "__main__": | ||
| 81 | + run_tests() | ||
| @@ -704,8 +704,20 @@ def _npu_apply_promoted_rtree_lines( | |||
| 704 | # Static (fully-resident) node aranges — same innermost level so | 704 | # Static (fully-resident) node aranges — same innermost level so |
| 705 | # the combined mask + reshape below see every node in scope. | 705 | # the combined mask + reshape below see every node in scope. |
| 706 | new_lines.extend(f"{inner_indent}{al}" for al in arange_lines) | 706 | new_lines.extend(f"{inner_indent}{al}" for al in arange_lines) |
| 707 | + # CSE may remove the original flat reduction alias assignment | ||
| 708 | + # even though generated loads still reference that alias. All | ||
| 709 | + # promoted sub-axes are now in scope, so reconstruct aliases | ||
| 710 | + # eagerly instead of waiting for a stale assignment below. | ||
| 711 | + new_lines.extend( | ||
| 712 | + f"{inner_indent}{nm} = {expr}" | ||
| 713 | + for nm, expr in flat_recon.items() | ||
| 714 | + ) | ||
| 707 | else: | 715 | else: |
| 708 | new_lines.extend(f"{indent}{al}" for al in arange_lines) | 716 | new_lines.extend(f"{indent}{al}" for al in arange_lines) |
| 717 | + new_lines.extend( | ||
| 718 | + f"{indent}{nm} = {expr}" | ||
| 719 | + for nm, expr in flat_recon.items() | ||
| 720 | + ) | ||
| 709 | continue | 721 | continue |
| 710 | 722 | ||
| 711 | if inside_rloop: | 723 | if inside_rloop: |
| @@ -726,7 +738,8 @@ def _npu_apply_promoted_rtree_lines( | |||
| 726 | # gone. Must run BEFORE the scaffolding drop below. | 738 | # gone. Must run BEFORE the scaffolding drop below. |
| 727 | _lhs = stripped.split(" = ", 1)[0].strip() if " = " in stripped else None | 739 | _lhs = stripped.split(" = ", 1)[0].strip() if " = " in stripped else None |
| 728 | if _lhs in flat_recon: | 740 | if _lhs in flat_recon: |
| 729 | - new_lines.append(f"{body_indent}{_lhs} = {flat_recon[_lhs]}") | 741 | + # Already emitted unconditionally beside the promoted |
| 742 | + # aranges, including when CSE removed this line entirely. | ||
| 730 | continue | 743 | continue |
| 731 | # Drop r0_index / r0_mask / roffset / rindex scaffolding and | 744 | # Drop r0_index / r0_mask / roffset / rindex scaffolding and |
| 732 | # the per-node mod/div decomposition (aranges replace them). | 745 | # the per-node mod/div decomposition (aranges replace them). |
| @@ -5935,8 +5948,40 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 5935 | 5948 | ||
| 5936 | basis_sorted = _complete_chain(basis_chain) | 5949 | basis_sorted = _complete_chain(basis_chain) |
| 5937 | other_sorted = _complete_chain(other_chain) | 5950 | other_sorted = _complete_chain(other_chain) |
| 5951 | + if basis_sorted and not other_sorted and other_chain: | ||
| 5952 | + # The alternate view may share one or more axes with the basis. | ||
| 5953 | + # Example: [260, 8, 2] and [260, 16] share the stride-1 x0 | ||
| 5954 | + # node. Treating only the disjoint x5 node as the second chain | ||
| 5955 | + # makes it look incomplete and leaves x5 independent, producing a | ||
| 5956 | + # 16x Cartesian grid. Find a complete chain that contains every | ||
| 5957 | + # genuinely alternate node while allowing basis nodes to fill the | ||
| 5958 | + # shared prefix/suffix. | ||
| 5959 | + candidates = list(free_nodes) | ||
| 5960 | + required = {n.name for n in other_chain} | ||
| 5961 | + | ||
| 5962 | + def _find_shared_chain(expected, chain, remaining_required): | ||
| 5963 | + if sizevars.statically_known_equals(expected, tree.numel): | ||
| 5964 | + return chain if not remaining_required else None | ||
| 5965 | + for node in candidates: | ||
| 5966 | + if node in chain or not sizevars.statically_known_equals( | ||
| 5967 | + node.divisor, expected | ||
| 5968 | + ): | ||
| 5969 | + continue | ||
| 5970 | + found = _find_shared_chain( | ||
| 5971 | + node.divisor * node.length, | ||
| 5972 | + chain + [node], | ||
| 5973 | + remaining_required - {node.name}, | ||
| 5974 | + ) | ||
| 5975 | + if found is not None: | ||
| 5976 | + return found | ||
| 5977 | + return None | ||
| 5978 | + | ||
| 5979 | + other_sorted = _find_shared_chain( | ||
| 5980 | + sympy.Integer(1), [], required | ||
| 5981 | + ) | ||
| 5938 | # Only fold a genuine dual decomposition: the basis covers the whole | 5982 | # Only fold a genuine dual decomposition: the basis covers the whole |
| 5939 | - # space on its own AND a second disjoint chain also covers it. | 5983 | + # space on its own AND a second chain (possibly sharing basis axes) also |
| 5984 | + # covers it. | ||
| 5940 | if not (basis_sorted and other_sorted): | 5985 | if not (basis_sorted and other_sorted): |
| 5941 | return | 5986 | return |
| 5942 | 5987 | ||
| @@ -5957,6 +6002,8 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 5957 | 6002 | ||
| 5958 | # Map each secondary node to div/mod of the basis flat index. | 6003 | # Map each secondary node to div/mod of the basis flat index. |
| 5959 | for o in other_sorted: | 6004 | for o in other_sorted: |
| 6005 | + if o.name in basis_names: | ||
| 6006 | + continue | ||
| 5960 | is_top = sizevars.statically_known_equals(o.divisor * o.length, tree.numel) | 6007 | is_top = sizevars.statically_known_equals(o.divisor * o.length, tree.numel) |
| 5961 | if is_top: | 6008 | if is_top: |
| 5962 | node_expr = flat_expr if isinstance(o.divisor, (int, sympy.Integer)) and int(o.divisor) == 1 else FloorDiv(flat_expr, o.divisor) # noqa: B950 | 6009 | node_expr = flat_expr if isinstance(o.divisor, (int, sympy.Integer)) and int(o.divisor) == 1 else FloorDiv(flat_expr, o.divisor) # noqa: B950 |