perf(npu-inductor): permute-gather — 三态分派 + K-tile + 动态 H 兜底 + H>4096 trans-pytorch-AtomGit

已合并

perf(npu-inductor): permute-gather — 三态分派 + K-tile + 动态 H 兜底 + H>4096 trans #44037

huyuchao创建于 8月7日从huyuchao/pytorch： perf/npu-permute-gather-v3 合入到 Ascend/pytorch： master

perf(npu-inductor): permute-gather — 三态分派 + K-tile + 动态 H 兜底 + H>4096 trans #44037

已合并

huyuchao创建于 8月7日

从huyuchao/pytorch perf/npu-permute-gather-v3合入到Ascend/pytorch master

共 6 个文件变更+995-1

Atest/_inductor/test_permute_gather.py+190-0

@@ -0,0 +1,190 @@
	1	+# Owner(s): ["module: tests"]
	2	+import contextlib
	3	+import io
	4	+
	5	+import torch
	6	+from torch._inductor.utils import run_and_get_code
	7	+from torch.testing._internal.common_utils import (
	8	+ run_tests,
	9	+ parametrize,
	10	+ instantiate_parametrized_tests,
	11	+)
	12	+from testutils import TestUtils
	13	+
	14	+import torch_npu # noqa: F401
	15	+from torch_npu._inductor.triton_experimental import config as ncfg
	16	+from torch_npu._inductor.triton_experimental.codegen.triton import NPUTritonKernel
	17	+
	18	+# Rewrite markers asserted in the generated kernel source (run_and_get_code):
	19	+# gather = "tl.gather" (flat DMA + register permute); trans = "tl.make_block_ptr"
	20	+# + "boundary_check" (block_ptr + tl.trans, tails exact). Mode selection and the
	21	+# huge-H trans+K-tile path: see config.permute_gather_mode.
	22	+GATHER_MARKER = "tl.gather"
	23	+TRANS_MARKER = "tl.make_block_ptr"
	24	+TRANS_BC_MARKER = "boundary_check"
	25	+
	26	+
	27	+class TestPermuteGather(TestUtils):
	28	+
	29	+ def setUp(self):
	30	+ super().setUp()
	31	+ torch._dynamo.reset()
	32	+ self._saved_pg = ncfg.enable_permute_gather
	33	+ self._saved_pin_xr = ncfg.pin_xr
	34	+ # In-process compile: forked workers cannot re-init NPU; threads=1 kills the pool.
	35	+ self._saved_threads = torch._inductor.config.compile_threads
	36	+ torch._inductor.config.compile_threads = 1
	37	+
	38	+ def tearDown(self):
	39	+ torch._inductor.config.compile_threads = self._saved_threads
	40	+ ncfg.enable_permute_gather = self._saved_pg
	41	+ ncfg.pin_xr = self._saved_pin_xr
	42	+ torch._dynamo.reset()
	43	+ super().tearDown()
	44	+
	45	+ @staticmethod
	46	+ def _t5(arg1, arg0):
	47	+ # T5 fused pattern: out[b,k,i] = sum_j (arg1[b,k,i,j] + arg0[i,j,k]^T)
	48	+ return (arg1 + arg0.permute(2, 0, 1).unsqueeze(0)).sum(-1)
	49	+
	50	+ def _run(self, h, i=48, j=300, b=2, dynamic=False, enabled=True):
	51	+ """Compile the T5 pattern and return (out, codes, eager_ref)."""
	52	+ ncfg.enable_permute_gather = enabled
	53	+ torch.manual_seed(0)
	54	+ arg1 = torch.randn(b, h, i, j, device="npu")
	55	+ arg0 = torch.randn(i, j, h, device="npu")
	56	+ ref = self._t5(arg1, arg0)
	57	+ if dynamic:
	58	+ torch._dynamo.mark_dynamic(arg1, 1)
	59	+ torch._dynamo.mark_dynamic(arg0, 2)
	60	+ fn = torch.compile(
	61	+ self._t5,
	62	+ options={"npu_backend": "triton_experimental"},
	63	+ dynamic=dynamic,
	64	+ )
	65	+ out, codes = run_and_get_code(fn, arg1, arg0)
	66	+ return out, codes, ref
	67	+
	68	+ @parametrize("h", [12, 32])
	69	+ def test_gather_mode_small_h(self, h):
	70	+ # stride_bytes = h*4 < 256 -> register gather (flat DMA + tl.gather).
	71	+ out, codes, ref = self._run(h)
	72	+ torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
	73	+ self.assertIn(GATHER_MARKER, codes[0])
	74	+ self.assertNotIn(TRANS_MARKER, codes[0])
	75	+ self.assertNotIn("raise ", codes[0]) # P1: no Python raise in any kernel
	76	+
	77	+ @parametrize("h", [64, 72, 1024])
	78	+ def test_trans_mode_large_h(self, h):
	79	+ # stride_bytes >= 256 -> block_ptr + tl.trans; H=72 (tail 8) / H=1024
	80	+ # exercise the int-axis K-tile (tails exact via boundary_check).
	81	+ out, codes, ref = self._run(h)
	82	+ torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
	83	+ self.assertIn(TRANS_MARKER, codes[0])
	84	+ self.assertIn(TRANS_BC_MARKER, codes[0])
	85	+ self.assertNotIn(GATHER_MARKER, codes[0])
	86	+ self.assertNotIn("raise ", codes[0]) # P1: no Python raise in any kernel
	87	+
	88	+ @parametrize("h", [72, 100])
	89	+ def test_dynamic_h_trans_fallback(self, h):
	90	+ # Dynamic (symbolic) H rides trans's shape-generic block_ptr (gather
	91	+ # needs a compile-time-affine stride).
	92	+ out, codes, ref = self._run(h, dynamic=True)
	93	+ torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
	94	+ self.assertIn(TRANS_MARKER, codes[0])
	95	+ self.assertNotIn(GATHER_MARKER, codes[0])
	96	+
	97	+ def test_huge_h_trans_ktile(self):
	98	+ # stride_r > max_xblock gates only gather; trans stays available
	99	+ # (XBLOCK = min(stride_r, ktile)), so H=8192 rides trans+K-tile.
	100	+ out, codes, ref = self._run(8192)
	101	+ torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
	102	+ self.assertIn(TRANS_MARKER, codes[0])
	103	+ self.assertIn(TRANS_BC_MARKER, codes[0])
	104	+ self.assertNotIn(GATHER_MARKER, codes[0])
	105	+
	106	+ def test_disabled_no_rewrite(self):
	107	+ # Flag off (default): no rewrite marker; realize/strided path stays correct.
	108	+ out, codes, ref = self._run(72, enabled=False)
	109	+ torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
	110	+ self.assertNotIn(GATHER_MARKER, codes[0])
	111	+ self.assertNotIn(TRANS_MARKER, codes[0])
	112	+
	113	+ def _with_slot_hook(self, hook):
	114	+ """Run ``hook(geo0)`` on the rewrite's slot sources right before
	115	+ _npu_pg_rewrite_body reads them (injects scheduler-produced layouts the
	116	+ triton_experimental linearize never emits). Records whether the hook
	117	+ fired so an injected test cannot silently degrade to the plain path."""
	118	+ orig = NPUTritonKernel._npu_pg_rewrite_body
	119	+ self._slot_hook_fired = False
	120	+
	121	+ def wrapped(kernel):
	122	+ cands = getattr(kernel, "_npu_pg_candidates", {})
	123	+ if cands:
	124	+ for geo0 in cands.values():
	125	+ hook(geo0)
	126	+ self._slot_hook_fired = True
	127	+ return orig(kernel)
	128	+
	129	+ NPUTritonKernel._npu_pg_rewrite_body = wrapped
	130	+ return orig
	131	+
	132	+ def test_gather_slot_order_r_first(self):
	133	+ # Review: gather pidx is [int, r] flat but reshaped per-slot -- an
	134	+ # R-first layout (r_slot < int_slot) used to transpose silently.
	135	+ # Force the swap; the rewrite must stay exact for any slot order.
	136	+ def swap(geo0):
	137	+ vtd = geo0["int_node"].root.var_tensor_dims
	138	+ r_tree = geo0["r_tree"]
	139	+ int_name = geo0["int_node"].name
	140	+ vtd[int_name], r_tree.tensor_dim = r_tree.tensor_dim, vtd[int_name]
	141	+
	142	+ orig = self._with_slot_hook(swap)
	143	+ try:
	144	+ out, codes, ref = self._run(12)
	145	+ torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
	146	+ self.assertIn(GATHER_MARKER, codes[0])
	147	+ finally:
	148	+ NPUTritonKernel._npu_pg_rewrite_body = orig
	149	+ self.assertTrue(self._slot_hook_fired)
	150	+
	151	+ def test_missing_slot_falls_back_strided(self):
	152	+ # Review: a legitimately-missing slot (vtd key absent / tensor_dim
	153	+ # None) used to raise TypeError; it must fall back to strided.
	154	+ def drop(geo0):
	155	+ geo0["int_node"].root.var_tensor_dims.pop(geo0["int_node"].name, None)
	156	+
	157	+ orig = self._with_slot_hook(drop)
	158	+ try:
	159	+ out, codes, ref = self._run(12)
	160	+ torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
	161	+ self.assertNotIn(GATHER_MARKER, codes[0])
	162	+ finally:
	163	+ NPUTritonKernel._npu_pg_rewrite_body = orig
	164	+ self.assertTrue(self._slot_hook_fired)
	165	+
	166	+ def test_pin_xr_does_not_override_pg(self):
	167	+ # Review: pin_xr (TEMP DIAGNOSTIC) used to return before the PG pin,
	168	+ # compiling the rewritten body under an unvalidated tiling. The PG
	169	+ # marker must win. h=48 (unused elsewhere) so the inductor text-hash
	170	+ # cache cannot mask the pin.
	171	+ ncfg.pin_xr = "256,16"
	172	+ stderr = io.StringIO()
	173	+ try:
	174	+ with contextlib.redirect_stderr(stderr):
	175	+ out, codes, ref = self._run(48)
	176	+ torch.testing.assert_close(out, ref, atol=1e-4, rtol=1e-4)
	177	+ self.assertIn(GATHER_MARKER, codes[0])
	178	+ # Pre-fix the pin_xr tiling never matched the rewrite's pin ->
	179	+ # compile failure -> small-block fallback; the PG pin winning
	180	+ # means no such fallback.
	181	+ self.assertNotIn("All initial configs failed", stderr.getvalue())
	182	+ self.assertNotIn("raise ", codes[0])
	183	+ finally:
	184	+ ncfg.pin_xr = self._saved_pin_xr
	185	+
	186	+
	187	+instantiate_parametrized_tests(TestPermuteGather)
	188	+
	189	+if __name__ == "__main__":
	190	+ run_tests()

Atorch_npu/_inductor/triton_experimental/codegen/permute_gather_rewrite.py+635-0

@@ -0,0 +1,635 @@
	1	+# Copyright (c) 2026, Huawei Technologies Co., Ltd
	2	+#
	3	+# Permute-gather reduction rewrite (ncfg.enable_permute_gather), extracted from
	4	+# codegen/triton.py per review: all rewrite logic lives here (sections:
	5	+# eligibility/geometry/validation, then the gather and trans emissions);
	6	+# NPUTritonKernel keeps thin delegates so call sites and test monkeypatching
	7	+# are unchanged.
	8	+
	9	+import ast
	10	+import logging
	11	+import operator
	12	+import sympy
	13	+
	14	+from torch._inductor.virtualized import V
	15	+
	16	+from .. import config as ncfg
	17	+from .. import device_props
	18	+
	19	+log = logging.getLogger(__name__)
	20	+
	21	+
	22	+def _ast_helpers():
	23	+ # Deferred: .triton imports this module lazily (from its delegates), so a
	24	+ # module-level back-import of its line-parsing helpers would be circular.
	25	+ global _parse_line, _assignment_parts, _parse_tl_load
	26	+ from torch_npu._inductor.triton_experimental.codegen.triton import (
	27	+ _npu_parse_generated_line as _parse_line,
	28	+ _npu_assignment_parts as _assignment_parts,
	29	+ _npu_parse_tl_load_assignment as _parse_tl_load,
	30	+ )
	31	+
	32	+
	33	+# Permute-gather rewrite UB budget (_npu_pg_geometry). The rewrite holds three
	34	+# co-resident tiles of stride_r*R0_BLOCK elements -- flat (input dtype), the
	35	+# int32 gather index, and the gathered result -- and triton caps tensor numel at
	36	+# 2^20. When the full row exceeds either limit the reduction is tiled into
	37	+# multiple r-loop trips (per-chunk load, see _npu_pg_emit). The largest R0_BLOCK
	38	+# that fits both caps is pinned into triton_meta (the rewrite forbids an autotune
	39	+# sweep), so these constants tune the tile at its widest:
	40	+# * ELEM_BYTES: worst-case fp32 for all three tiles (int32 index is always 4B);
	41	+# fp16/fp8 inputs get a conservative (smaller) tile, just more trips.
	42	+# * PIPE: the compiler doubles live buffers for software pipelining (matches
	43	+# _NPU_UB_OVERHEAD_FACTOR in npu_triton_heuristics).
	44	+# * RESERVE: keep a share of UB for the rest of the kernel (other loads, output,
	45	+# masks). Under-estimating only costs extra trips -- it never breaks compile.
	46	+_PG_UB_RESERVE = 0.5
	47	+_PG_UB_TENSORS = 3
	48	+_PG_UB_ELEM_BYTES = 4
	49	+_PG_UB_PIPE = 2
	50	+
	51	+# _npu_pg_eval_expr operator tables (lookup == the if-chain it replaces)
	52	+_PG_EVAL_BINOPS = {
	53	+ ast.Add: operator.add,
	54	+ ast.Sub: operator.sub,
	55	+ ast.Mult: operator.mul,
	56	+ ast.FloorDiv: operator.floordiv,
	57	+ ast.Mod: operator.mod,
	58	+}
	59	+_PG_EVAL_CMPOPS = {
	60	+ ast.Lt: operator.lt,
	61	+ ast.Gt: operator.gt,
	62	+ ast.LtE: operator.le,
	63	+ ast.GtE: operator.ge,
	64	+ ast.Eq: operator.eq,
	65	+ ast.NotEq: operator.ne,
	66	+}
	67	+
	68	+
	69	+def _npu_pg_eval_expr(node, env):
	70	+ """Evaluate a pre_loop constexpr RHS (int arithmetic over ``env`` names).
	71	+
	72	+ A tiny AST walker instead of eval(): the emitted tile-chain lines contain
	73	+ only integer literals, bound names, + - * // %, comparisons, the
	74	+ ``a if cond else b`` tile form and min/max; anything else raises and the
	75	+ caller skips the line.
	76	+ """
	77	+ if isinstance(node, ast.Constant) and isinstance(node.value, int):
	78	+ return node.value
	79	+ if isinstance(node, ast.Name):
	80	+ return env[node.id] # KeyError -> caller skips this line
	81	+ if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
	82	+ return -_npu_pg_eval_expr(node.operand, env)
	83	+ if isinstance(node, ast.BinOp) and type(node.op) in _PG_EVAL_BINOPS:
	84	+ fn = _PG_EVAL_BINOPS[type(node.op)]
	85	+ return fn(_npu_pg_eval_expr(node.left, env), _npu_pg_eval_expr(node.right, env))
	86	+ if (
	87	+ isinstance(node, ast.Compare)
	88	+ and len(node.ops) == 1
	89	+ and type(node.ops[0]) in _PG_EVAL_CMPOPS
	90	+ ):
	91	+ fn = _PG_EVAL_CMPOPS[type(node.ops[0])]
	92	+ return fn(
	93	+ _npu_pg_eval_expr(node.left, env),
	94	+ _npu_pg_eval_expr(node.comparators[0], env),
	95	+ )
	96	+ if isinstance(node, ast.IfExp):
	97	+ branch = node.body if _npu_pg_eval_expr(node.test, env) else node.orelse
	98	+ return _npu_pg_eval_expr(branch, env)
	99	+ if (
	100	+ isinstance(node, ast.Call)
	101	+ and isinstance(node.func, ast.Name)
	102	+ and node.func.id in ("min", "max")
	103	+ and not node.keywords
	104	+ ):
	105	+ vals = [_npu_pg_eval_expr(arg, env) for arg in node.args]
	106	+ return min(vals) if node.func.id == "min" else max(vals)
	107	+ raise ValueError(f"unsupported constexpr expr: {ast.unparse(node)}")
	108	+
	109	+
	110	+# ---- eligibility / geometry / tiling validation (the dispatch half) --------
	111	+
	112	+
	113	+def _npu_pg_record(kernel, result_var, name, index):
	114	+ """Register an eligible strided reduction load for the gather rewrite."""
	115	+ if not ncfg.enable_permute_gather:
	116	+ return
	117	+ if not ncfg.codegen_linearize:
	118	+ return
	119	+ if not kernel.inside_reduction:
	120	+ return
	121	+ if getattr(kernel, "persistent_reduction", False):
	122	+ return
	123	+ var_name = getattr(result_var, "name", None)
	124	+ if not var_name:
	125	+ return
	126	+ try:
	127	+ # Use the buffer dtype (native itemsize). result_var.dtype is the
	128	+ # fp32-promoted codegen dtype (codegen_upcast_to_fp32), which would
	129	+ # dispatch fp16/bf16 permutes at 4B/elt -- same mode only for H%16==0,
	130	+ # and a wrong mode elsewhere (lowering/npu_permute dispatches on the
	131	+ # native 2B itemsize, _elemsize_of).
	132	+ elemsize = V.graph.get_dtype(name).itemsize
	133	+ except Exception:
	134	+ elemsize = 4
	135	+ geo = _npu_pg_geometry(kernel, index, elemsize)
	136	+ if geo is None:
	137	+ log.debug("permute_gather: %s ineligible (geometry/mode gate)", name)
	138	+ return
	139	+ if not hasattr(kernel, "_npu_pg_candidates"):
	140	+ kernel._npu_pg_candidates = {}
	141	+ kernel._npu_pg_candidates[var_name] = geo
	142	+ log.debug("permute_gather: candidate %s mode=%s stride_r=%s",
	143	+ name, geo["mode"], geo["stride_r"])
	144	+
	145	+
	146	+def _npu_pg_geometry(kernel, index, elemsize=4):
	147	+ """Return the rewrite geometry for ``index``, or None if not eligible.
	148	+
	149	+ Single-interior geometry: exactly one reduction leaf with coeff
	150	+ stride_r > 1, and exactly two free leaves -- interior (coeff 1) and row
	151	+ (coeff == SEG = stride_r * r_numel) -- in one x-tree, with the interior
	152	+ node running exactly stride_r elements. Anything else falls back to the
	153	+ strided load (correct, just slow). Mode semantics ("gather" = flat DMA +
	154	+ register tl.gather vs "trans" = block_ptr + tl.trans riding the MTE2
	155	+ read, per-branch caps, gather's static-stride_r requirement): see
	156	+ config.permute_gather_mode and the _npu_pg_emit* docstrings.
	157	+ """
	158	+ leaves = {}
	159	+ size_syms = set()
	160	+ red_sym = None
	161	+ red_leaf = None
	162	+ for sym in index.free_symbols:
	163	+ node = kernel.range_tree_nodes.get(sym)
	164	+ if node is None:
	165	+ # A symbolic size parameter (dynamic dim), e.g. dynamic Sk riding
	166	+ # the row-leaf coeff stride_r * r_numel. Validated against the
	167	+ # symbolic r_numel below; must not appear anywhere else (the
	168	+ # residue check rejects a dynamic base offset). Deferred instead
	169	+ # of rejected so dynamic reduction numels reach the row-leaf
	170	+ # compare and can take the runtime-mask gather path.
	171	+ size_syms.add(sym)
	172	+ continue
	173	+ try:
	174	+ coeff = index.coeff(sym, 1)
	175	+ except Exception:
	176	+ return None
	177	+ if node.root.is_reduction:
	178	+ if red_leaf is not None:
	179	+ return None # >1 reduction leaf -> not single-interior
	180	+ if not isinstance(coeff, (int, sympy.Integer)):
	181	+ # Dynamic H: the reduction stride is symbolic (e.g. ks0 in
	182	+ # x1 + ks0r0_3 + ks0ks1*x0). The gather index must be
	183	+ # compile-time affine in stride_r (vectorized UB permute) and
	184	+ # the flat tile width stride_r*R0_BLOCK constexpr, so a
	185	+ # symbolic stride can only ride the trans mode's shape-generic
	186	+ # block_ptr (mode dispatch below).
	187	+ if not ncfg.permute_gather_dynamic_trans:
	188	+ return None
	189	+ red_leaf = (node, coeff)
	190	+ red_sym = sym
	191	+ continue
	192	+ red_leaf = (node, int(coeff))
	193	+ red_sym = sym
	194	+ else:
	195	+ # free-leaf coeff may be symbolic (e.g. stride_r * dynamic r_numel
	196	+ # on the row leaf); validated against stride_r * r_numel_expr below
	197	+ leaves[sym] = (node, coeff)
	198	+ if red_leaf is None:
	199	+ return None
	200	+ rnode, stride_r = red_leaf
	201	+ if isinstance(stride_r, sympy.Expr) and stride_r.free_symbols:
	202	+ dynamic_h = True
	203	+ mode = "trans"
	204	+ else:
	205	+ stride_r = int(stride_r)
	206	+ dynamic_h = False
	207	+ if stride_r <= 1:
	208	+ return None # unit-stride reduction is already contiguous
	209	+ # Mode dispatch (mirrors lowering.npu_permute): inner strides below
	210	+ # 256B (H<=63 fp32) -> register gather; >= 256B (H>=64 fp32) ->
	211	+ # block_ptr trans, which beats both the old gather (flat tile
	212	+ # overflows UB, measured H=64 R0=128 compile fail) and the realized
	213	+ # transpose (two-kernel 219us) on the MTE2 permutation.
	214	+ mode = ncfg.permute_gather_mode(stride_r, elemsize)
	215	+ if mode is None:
	216	+ # stride_r above the XBLOCK cap (4096): neither rewrite can pin
	217	+ # XBLOCK = stride_r, so keep the strided load -- lowering has
	218	+ # realized the permute (mode None -> realize), making it correct.
	219	+ return None
	220	+ r_len = rnode.length
	221	+ if isinstance(r_len, (int, sympy.Integer)):
	222	+ r_numel = int(r_len)
	223	+ if r_numel <= 0:
	224	+ return None
	225	+ r_numel_expr = sympy.Integer(r_numel)
	226	+ dynamic_r = False
	227	+ else:
	228	+ # Dynamic reduction numel: the runtime masks (flat seg bound in
	229	+ # _npu_pg_emit + the reduction's own r0_mask) keep every output lane
	230	+ # exact for any r_numel, so the rewrite stays sound. r0_block is then
	231	+ # sized from the UB caps alone (can't min against a runtime numel);
	232	+ # since r0_mask is runtime (never the all-True constant-mask fold),
	233	+ # any pow2 R0_BLOCK is safe -- no divisibility requirement.
	234	+ if r_len.free_symbols - size_syms:
	235	+ return None
	236	+ r_numel = None
	237	+ r_numel_expr = r_len
	238	+ dynamic_r = True
	239	+ # Affine residue: removing every linear term must leave a pure constant
	240	+ # base offset (preserved verbatim in the emit); any remaining free
	241	+ # symbol would be a dynamic base offset the flat load cannot reproduce.
	242	+ linear = sum(c * s for s, (n, c) in leaves.items())
	243	+ # stride_r is a Python int (static H) or a sympy Symbol (dynamic H); the
	244	+ # plain multiply covers both, avoiding sympy.Integer(Symbol) (TypeError).
	245	+ linear += stride_r * red_sym
	246	+ residue = index - linear
	247	+ if not isinstance(residue, sympy.Number):
	248	+ return None
	249	+ if len(leaves) != 2:
	250	+ return None
	251	+ int_leaf = row_leaf = None
	252	+ for s, (node, coeff) in leaves.items():
	253	+ if coeff == 1:
	254	+ int_leaf = (node, coeff)
	255	+ elif coeff == stride_r * r_numel_expr:
	256	+ row_leaf = (node, coeff)
	257	+ else:
	258	+ return None # extra / unexpected free stride
	259	+ if int_leaf is None or row_leaf is None:
	260	+ return None
	261	+ int_node, _ = int_leaf
	262	+ row_node, _ = row_leaf
	263	+ if row_node.root is not int_node.root:
	264	+ return None # interior and row must share one x-tree
	265	+ if dynamic_h:
	266	+ # Dynamic H: interior length is the symbolic head axis (== stride_r)
	267	+ if int_node.length != stride_r:
	268	+ return None
	269	+ elif not isinstance(int_node.length, (int, sympy.Integer)) or int(int_node.length) != stride_r:
	270	+ return None
	271	+ r_tree = rnode.root
	272	+ # Register-dim slots (var_tensor_dims) are assigned by _apply_linearize,
	273	+ # which runs AFTER the first body pass emits loads -- so slot resolution is
	274	+ # deferred to _npu_pg_rewrite_body (codegen_kernel time), when they exist.
	275	+ #
	276	+ # Reduction tile for the rewrite: largest R0_BLOCK that keeps the emitted
	277	+ # per-chunk flat/idx/gather tiles (width stride_r * R0_BLOCK, see
	278	+ # _npu_pg_emit) inside triton's 2^20 tensor-numel cap and UB. Larger tile =
	279	+ # fewer r-loop trips, so take the max. At R0_BLOCK == r_numel the loop
	280	+ # runs once and the emission degenerates to the original whole-row burst.
	281	+ #
	282	+ # R0_BLOCK must be a power of two. Upstream's constant-mask optimization
	283	+ # (_has_constant_mask) emits r0_mask = tl.full(..., True) whenever the max
	284	+ # reduction block (TRITON_MAX_BLOCK["R0_"] = 65536) divides r_numel, and
	285	+ # that is only sound if the ACTUAL R0_BLOCK also divides r_numel -- true
	286	+ # for upstream's pow2 blocks, broken by a non-pow2 tile (e.g. 341) which
	287	+ # leaves a partial tail trip whose all-True mask accumulates out-of-bounds
	288	+ # lanes into the reduction. Round the raw cap DOWN to the largest pow2
	289	+ # (cap_ub <= 4096//stride_r <= 2048, so the tile also divides 65536 and
	290	+ # hence every multiple of it).
	291	+ if dynamic_h:
	292	+ # Dynamic-H trans is config-agnostic (block_shape adapts to the swept
	293	+ # (XBLOCK, R0_BLOCK) and boundary_check keeps partial chunks exact),
	294	+ # so no config pin and no r0_block.
	295	+ r0_block = None
	296	+ else:
	297	+ # Effective interior width resident in the tile. Trans chunks the int
	298	+ # axis to permute_gather_ktile (the XBLOCK pin), so the caps must run
	299	+ # on the chunk width -- this is what decouples R0 from H (H=4096:
	300	+ # R0 1 -> 64). Gather's flat tile is stride_r wide by construction
	301	+ # (XBLOCK pin = stride_r, no chunking), so eff == stride_r.
	302	+ eff = stride_r if mode == "gather" else min(stride_r, ncfg.permute_gather_ktile)
	303	+ cap_numel = 1_048_576 // eff
	304	+ budget = int(_PG_UB_RESERVE * device_props.get_npu_ub_size_bytes())
	305	+ cap_ub = max(1, budget // (_PG_UB_TENSORS * _PG_UB_ELEM_BYTES * eff * _PG_UB_PIPE))
	306	+ raw = min(cap_numel, cap_ub) if r_numel is None else min(r_numel, cap_numel, cap_ub)
	307	+ r0_block = 1 << (max(1, raw).bit_length() - 1)
	308	+ return {
	309	+ "mode": mode,
	310	+ "dynamic_h": dynamic_h,
	311	+ "int_node": int_node,
	312	+ "row_node": row_node,
	313	+ "r_tree": r_tree,
	314	+ "rprefix": r_tree.prefix,
	315	+ "rblk": f"{r_tree.prefix.upper()}BLOCK",
	316	+ "stride_r": stride_r,
	317	+ "r_numel": r_numel,
	318	+ "seg": stride_r * r_numel if r_numel is not None else None,
	319	+ "dynamic_r": dynamic_r,
	320	+ "const": int(residue),
	321	+ "r0_block": r0_block,
	322	+ }
	323	+
	324	+
	325	+def _npu_pg_eval_real_blocks(tree, xblock):
	326	+ """Evaluate the emitted pre_loop constexprs under a forced XBLOCK.
	327	+
	328	+ Returns {name: value} over the tile chain (real_block_, _numel,
	329	+ *_blocks) by AST-parsing the pre_loop assignment lines (AnnAssign
	330	+ ``name : tl.constexpr = ...`` normalized like _npu_build_grid_recipe)
	331	+ and walking each RHS with _npu_pg_eval_expr over the names bound so
	332	+ far. Directly reflects the runtime values, immune to ordering-logic
	333	+ drift; a first definition wins, as before.
	334	+ """
	335	+ _ast_helpers()
	336	+ env = {"XBLOCK": int(xblock)}
	337	+ pre = getattr(tree, "pre_loop_code", None)
	338	+ if pre is None:
	339	+ return env
	340	+ for entry in pre._lines:
	341	+ line = entry if isinstance(entry, str) else getattr(entry, "line", None)
	342	+ if not isinstance(line, str):
	343	+ continue
	344	+ parsed = _parse_line(line)
	345	+ if parsed is None:
	346	+ continue
	347	+ _, statement = parsed
	348	+ if isinstance(statement, ast.AnnAssign) and statement.value is not None:
	349	+ statement = ast.Assign(
	350	+ targets=[statement.target], value=statement.value
	351	+ )
	352	+ assignment = _assignment_parts(statement)
	353	+ if assignment is None or assignment[0] in env:
	354	+ continue
	355	+ try:
	356	+ env[assignment[0]] = _npu_pg_eval_expr(assignment[1], env)
	357	+ except Exception:
	358	+ continue
	359	+ return env
	360	+
	361	+
	362	+def _npu_pg_rewrite_body(kernel):
	363	+ """Replace the recorded strided reduction load with DMA + tl.gather, or
	364	+ block_ptr + tl.trans (post codegen_body). Validates the forced tiling
	365	+ against the emitted pre_loop constexprs (gather: XBLOCK=stride_r -> row
	366	+ tile 1 + full int run; trans: XBLOCK=min(stride_r, ktile) -> row tile 1,
	367	+ int chunked), then swaps the load line. Returns the forced config dict
	368	+ for triton_meta, or None to keep strided / leave the sweep free (dyn-H).
	369	+ """
	370	+ _ast_helpers()
	371	+ cands = getattr(kernel, "_npu_pg_candidates", None) or {}
	372	+ if len(cands) != 1:
	373	+ return None # the per-kernel config pin needs exactly one candidate
	374	+ var_name, geo0 = next(iter(cands.items()))
	375	+ # Register-dim slots exist only after _apply_linearize (this hook runs at
	376	+ # codegen_kernel time, so they're present now).
	377	+ vtd = geo0["int_node"].root.var_tensor_dims
	378	+ r_tree = geo0["r_tree"]
	379	+ int_slot = vtd.get(geo0["int_node"].name)
	380	+ row_slot = vtd.get(geo0["row_node"].name)
	381	+ r_slot = getattr(r_tree, "tensor_dim", None)
	382	+ # Slots must be distinct and in range; no relative-order assumption.
	383	+ # The emit derives everything from the slots themselves (trans perm via
	384	+ # sorted slot order, reshape shape filled per-slot), so it stays correct
	385	+ # for any layout -- including the reduction tree taking slot 0, which the
	386	+ # default-backend scheduler patches produce for large H kernels
	387	+ # (range_trees R-first + _npu_repermute_tensor_dims stride order). The
	388	+ # historical int/row-before-r check came from the hard-coded (0,2,1)
	389	+ # trans perm and silently dropped those kernels back to strided loads.
	390	+ slots = (int_slot, row_slot, r_slot)
	391	+ # Type-check before max(): a legitimately-missing slot (vtd key
	392	+ # absent / r_tree.tensor_dim None) must fall back to strided, not
	393	+ # crash codegen with a TypeError on max((None, ...)).
	394	+ if not all(isinstance(s, int) for s in slots):
	395	+ log.debug("permute_gather: %s slot missing -> strided", var_name)
	396	+ return None
	397	+ ndim = max(slots) + 1
	398	+ if not all(0 <= s < ndim for s in slots) or len(set(slots)) != 3:
	399	+ log.debug("permute_gather: %s slots not distinct/in-range -> strided", var_name)
	400	+ return None
	401	+ geo = dict(geo0)
	402	+ geo["int_slot"] = int_slot
	403	+ geo["row_slot"] = row_slot
	404	+ geo["r_slot"] = r_slot
	405	+ geo["ndim"] = ndim
	406	+ tree = geo0["int_node"].root
	407	+ # Static modes pin XBLOCK so the greedy tile chain produces exactly the
	408	+ # tiling the rewrite needs; the rewrite is only valid under that tiling.
	409	+ # Gather pins XBLOCK=stride_r: the flat tile + pidx reshape structurally
	410	+ # need real_block_int == stride_r == int_numel. Trans pins
	411	+ # XBLOCK=min(stride_r, permute_gather_ktile): the interior axis is chunked
	412	+ # across programs (x1_blocks > 1) and the block_ptr/reshape are
	413	+ # chunk-agnostic (real_block_int in both), but the row axis must stay
	414	+ # tile-1 (the reshape's [1, rb_int, rb_row, R] broadcast would break) --
	415	+ # the int axis eating the whole pin budget guarantees that. Dynamic-H
	416	+ # trans is config-agnostic (the block_ptr block_shape adapts and
	417	+ # boundary_check guards partial chunks), so no pin and no tiling invariant.
	418	+ if not geo0["dynamic_h"]:
	419	+ if geo0["mode"] == "gather":
	420	+ pin = geo0["stride_r"]
	421	+ else:
	422	+ pin = min(geo0["stride_r"], ncfg.permute_gather_ktile)
	423	+ env = _npu_pg_eval_real_blocks(tree, pin)
	424	+ rb_row = env.get(f"real_block_{geo0['row_node'].name}")
	425	+ rb_int = env.get(f"real_block_{geo0['int_node'].name}")
	426	+ int_numel = env.get(f"{geo0['int_node'].name}numel")
	427	+ if rb_row != 1 or rb_int != min(int_numel, pin):
	428	+ log.debug("permute_gather: %s tiling mismatch (rb_row=%s rb_int=%s, "
	429	+ "want row=1 int<=%s) -> strided", var_name, rb_row, rb_int, pin)
	430	+ return None # tiling not as forced -> keep strided (correct, slow)
	431	+ # Locate the recorded load structurally (same pattern as
	432	+ # _maybe_rewrite_select_lane_load): match the assignment TARGET, not the
	433	+ # line text -- a load line may carry a trailing ``.to(tl.float32)``
	434	+ # promotion that no ``tl.load(...)$`` regex can delimit, and non-str
	435	+ # body entries (DeferredLine et al.) must pass through untouched.
	436	+ new_lines = []
	437	+ rewritten = False
	438	+ for line in kernel.body._lines:
	439	+ if not rewritten and isinstance(line, str):
	440	+ parsed = _parse_tl_load(line)
	441	+ if parsed is not None and parsed[1] == var_name:
	442	+ indent, _, value_ast, load_ast = parsed
	443	+ emit = _npu_pg_emit(
	444	+ var_name, indent, value_ast, load_ast, geo, mode=geo0["mode"]
	445	+ )
	446	+ if emit:
	447	+ new_lines.extend(emit)
	448	+ rewritten = True
	449	+ continue
	450	+ # unexpected load shape -> keep the strided line below
	451	+ new_lines.append(line)
	452	+ if not rewritten:
	453	+ log.debug("permute_gather: %s load line not found -> strided", var_name)
	454	+ return None
	455	+ kernel.body._lines = new_lines
	456	+ log.debug("permute_gather: rewrote %s (%s)", var_name, geo0["mode"])
	457	+ if geo0["dynamic_h"]:
	458	+ return None # dynamic-H trans: no config pin, autotune sweeps freely
	459	+ if geo0["mode"] == "gather":
	460	+ return {"XBLOCK": geo0["stride_r"], "R0_BLOCK": geo0["r0_block"]}
	461	+ return {
	462	+ "XBLOCK": min(geo0["stride_r"], ncfg.permute_gather_ktile),
	463	+ "R0_BLOCK": geo0["r0_block"],
	464	+ }
	465	+
	466	+
	467	+# ---- emission: "gather" mode (flat contiguous DMA + register tl.gather) ----
	468	+
	469	+
	470	+def _npu_pg_emit(var_name, indent, value_ast, load_ast, geo, mode="gather"):
	471	+ """Build replacement lines for one permute rewrite load.
	472	+
	473	+ ``value_ast``/``load_ast`` come from _npu_parse_tl_load_assignment on the
	474	+ emitted line: the full assignment value (possibly a ``.to(tl.float32)``
	475	+ promotion wrapped around the call) and the ``tl.load`` call node itself.
	476	+ The base pointer, ``other`` and ``eviction_policy`` are read structurally
	477	+ off the call node -- never by regex over the rendered text (an ``other``
	478	+ value containing parens, e.g. ``float('-inf')``, defeats any ``[^)]*``
	479	+ capture). Returns [] to decline the rewrite (keep the strided line).
	480	+
	481	+ ``mode`` selects the primitive: "gather" (flat DMA + register tl.gather)
	482	+ or "trans" (block_ptr [row, r, int] + tl.trans, see _npu_pg_emit_trans).
	483	+ """
	484	+ # args[0] is the emitted ``ptr + (linear index)``; the rewrite rebuilds
	485	+ # the address from the geometry, so only the bare pointer name is needed.
	486	+ if not (
	487	+ isinstance(load_ast.args[0], ast.BinOp)
	488	+ and isinstance(load_ast.args[0].op, ast.Add)
	489	+ and isinstance(load_ast.args[0].left, ast.Name)
	490	+ ):
	491	+ return []
	492	+ ptr = load_ast.args[0].left.id
	493	+ other = "0.0"
	494	+ evict = "evict_last"
	495	+ for kw in load_ast.keywords:
	496	+ if kw.arg == "other":
	497	+ other = ast.unparse(kw.value)
	498	+ elif kw.arg == "eviction_policy":
	499	+ try:
	500	+ evict = ast.literal_eval(kw.value)
	501	+ except (ValueError, TypeError):
	502	+ return []
	503	+ if mode == "trans":
	504	+ return _npu_pg_emit_trans(
	505	+ var_name, indent, ptr, geo, other, value_ast, load_ast
	506	+ )
	507	+ int_n = geo["int_node"].name
	508	+ row_n = geo["row_node"].name
	509	+ rpfx = geo["rprefix"]
	510	+ rblk = geo["rblk"]
	511	+ s = geo["stride_r"]
	512	+ # Row-bound of the source buffer (seg = stride_r * r_numel). Static
	513	+ # r_numel: bake the Python int. Dynamic r_numel: the runtime value
	514	+ # stride_r * r0_numel (the very numel the r-loop and r0_mask iterate),
	515	+ # so the flat mask + gather stay exact for any sequence length -- the
	516	+ # mask-based precision guarantee that lets dynamic Sk ride this path.
	517	+ seg = f"{s} * {rpfx}numel" if geo.get("dynamic_r") else str(geo["seg"])
	518	+ base = f"{ptr} + {geo['const']} + " if geo["const"] else f"{ptr} + "
	519	+ rows = f"{var_name}_pg_rows"
	520	+ flat = f"{var_name}_pg_flat"
	521	+ pr = f"{var_name}_pg_r"
	522	+ pidx = f"{var_name}_pg_idx"
	523	+ pg = f"{var_name}_pg_g"
	524	+ # The var's reshape target must be r-last: the surrounding reduction
	525	+ # code consumes it by axis POSITION, not by register slot -- r0_mask is
	526	+ # [1, 1, 1, R0_BLOCK], the broadcast_to / tl.sum run over the trailing
	527	+ # r axis. Filling ``parts`` per-slot was only right for the int<row<r
	528	+ # layout by coincidence; R-first layouts (r_slot < int_slot) then put
	529	+ # r at slot 0 and either broke the broadcast or silently transposed the
	530	+ # permuted values. Fix the shape as [1, int, row, R0_BLOCK] and build
	531	+ # pidx in the matching [int, r] flat order -- independent of the slots.
	532	+ shape = f"[1, real_block_{int_n}, real_block_{row_n}, {rblk}]"
	533	+ # Per-chunk DMA: this trip of the r-loop loads only its own
	534	+ # stride_rR0_BLOCK stretch of the row, at r0_offsetstride_r past the
	535	+ # row base. At R0_BLOCK == r_numel (single trip, r0_offset == 0) this is
	536	+ # exactly the original whole-row burst. The flat r-bound mask keeps the
	537	+ # tail trip's load lanes inside the row (never reads past seg); the
	538	+ # reduction's own r0_mask still zeroes the corresponding outputs.
	539	+ tile = f"{rpfx}offset * {s} + tl.arange(0, {s} * {rblk})"
	540	+ lines = [
	541	+ f"{indent}{rows} = {row_n}offset + tl.arange(0, real_block_{row_n})",
	542	+ f"{indent}{flat} = tl.load({base}{seg} * {rows}[:, None] + {tile}[None, :], "
	543	+ f"({rows}[:, None] < {row_n}numel) & ({tile} < {seg}), "
	544	+ f"eviction_policy='{evict}', other={other})",
	545	+ f"{indent}{pr} = tl.arange(0, {rblk})[None, :]",
	546	+ ]
	547	+ # pidx's grid is [int, r] (row-major flat = int outer, r inner), which
	548	+ # matches the r-last reshape's flatten exactly. No runtime guard is
	549	+ # emitted: the rewrite only produces this code after its own tiling
	550	+ # check (real_block_row == 1, real_block_int == stride_r) passed at
	551	+ # rewrite time, and Python `raise` is invalid Triton kernel code.
	552	+ lines.append(
	553	+ f"{indent}{pidx} = tl.reshape({int_n}offset + tl.arange(0, real_block_{int_n})[:, None] "
	554	+ f"+ {s} * {pr}, (1, {s} * {rblk})) + tl.full([real_block_{row_n}, 1], 0, tl.int32)"
	555	+ )
	556	+ lines.append(f"{indent}{pg} = tl.reshape(tl.gather({flat}, {pidx}, 1), {shape})")
	557	+ lines.append(
	558	+ _npu_pg_final_line(
	559	+ var_name, indent, value_ast, load_ast, f"tl.where({rpfx}mask, {pg}, {other})"
	560	+ )
	561	+ )
	562	+ return lines
	563	+
	564	+
	565	+def _npu_pg_final_line(var_name, indent, value_ast, load_ast, expr_text):
	566	+ """``var = expr_text`` with the load's surrounding expression preserved:
	567	+ splice the replacement in place of the ``tl.load`` node and unparse, so
	568	+ an fp16/bf16 candidate keeps its ``.to(tl.float32)`` promotion (with no
	569	+ promotion this unparses to exactly ``expr_text``)."""
	570	+ repl = ast.parse(expr_text, mode="eval").body
	571	+
	572	+ class _SwapLoad(ast.NodeTransformer):
	573	+ def visit_Call(self, node):
	574	+ if node is load_ast:
	575	+ return ast.copy_location(repl, node)
	576	+ return self.generic_visit(node)
	577	+
	578	+ value_ast = _SwapLoad().visit(value_ast)
	579	+ ast.fix_missing_locations(value_ast)
	580	+ return f"{indent}{var_name} = {ast.unparse(value_ast)}"
	581	+
	582	+
	583	+# ---- emission: "trans" mode (block_ptr + tl.trans riding the MTE2 read) ----
	584	+
	585	+
	586	+def _npu_pg_emit_trans(var_name, indent, ptr, geo, other, value_ast, load_ast):
	587	+ """block_ptr [row, r, int] + tl.trans + reshape, replacing the strided
	588	+ reduction load. Shape/strides ride the always-present per-axis numel vars
	589	+ (stride_r == int_numel and seg == int_numel * r_numel by the geometry's
	590	+ leaf invariants), so one emission covers static and dynamic shapes;
	591	+ boundary_check + padding_option="zero" keeps tail chunks exact for any
	592	+ constexpr block. The permutation rides the MTE2 read (a UB tile permute)
	593	+ -- no index tensor, so lower UB pressure at large stride_r than gather.
	594	+ """
	595	+ int_n = geo["int_node"].name
	596	+ row_n = geo["row_node"].name
	597	+ rpfx = geo["rprefix"]
	598	+ rblk = geo["rblk"]
	599	+
	600	+ def _numel(node):
	601	+ # Per-node numel args are emitted only for dynamic nodes (triton.py
	602	+ # size-arg loop); bake a static length as the literal int.
	603	+ if isinstance(node.length, (int, sympy.Integer)):
	604	+ return str(int(node.length))
	605	+ return f"{node.name}numel"
	606	+
	607	+ int_nm = _numel(geo["int_node"])
	608	+ row_nm = _numel(geo["row_node"])
	609	+ base = f"{ptr} + {geo['const']}" if geo["const"] else ptr
	610	+ bp = f"{var_name}_pg_bp"
	611	+ t = f"{var_name}_pg_t"
	612	+ g = f"{var_name}_pg_g"
	613	+ parts = ["1"] * geo["ndim"]
	614	+ parts[geo["int_slot"]] = f"real_block_{int_n}"
	615	+ parts[geo["row_slot"]] = f"real_block_{row_n}"
	616	+ parts[geo["r_slot"]] = rblk
	617	+ shape = "[" + ", ".join(parts) + "]"
	618	+ # The loaded tile t is physical (row, r, int); the reshape target is
	619	+ # slot-ascending, so the trans perm must lay the axes in slot order
	620	+ # (a hard-coded (0,2,1) scrambles R-first layouts: measured dynH
	621	+ # diff=109). Derive the perm from the actual slots.
	622	+ pos = {"row": 0, "r": 1, "int": 2}
	623	+ perm = tuple(pos[k] for k in sorted(pos, key=lambda k: geo[f"{k}_slot"]))
	624	+ return [
	625	+ f"{indent}{bp} = tl.make_block_ptr(",
	626	+ f"{indent} {base}, shape=[{row_nm}, {rpfx}numel, {int_nm}], "
	627	+ f"strides=[{int_nm} * {rpfx}numel, {int_nm}, 1], "
	628	+ f"offsets=[{row_n}offset, {rpfx}offset, {int_n}offset], "
	629	+ f"block_shape=[real_block_{row_n}, {rblk}, real_block_{int_n}], order=[2, 1, 0])",
	630	+ f'{indent}{t} = tl.load({bp}, boundary_check=[0, 1, 2], padding_option="zero")',
	631	+ f"{indent}{g} = tl.reshape(tl.trans({t}, {perm}), {shape})",
	632	+ _npu_pg_final_line(
	633	+ var_name, indent, value_ast, load_ast, f"tl.where({rpfx}mask, {g}, {other})"
	634	+ ),
	635	+ ]

Mtorch_npu/_inductor/triton_experimental/codegen/triton.py+24-0

	@@ -1815,6 +1815,7 @@ class NPUTritonKernel(TritonKernel):
1815	name, result_var, index, self._npu_prepared_load_index	1815	name, result_var, index, self._npu_prepared_load_index
1816	)	1816	)
1817	self._record_reduction_load_padinfo(result_var, index)	1817	self._record_reduction_load_padinfo(result_var, index)
		1818	+ self._npu_pg_record(result_var, name, index)
1818	finally:	1819	finally:
1819	self._npu_capture_prepared_load_index = False	1820	self._npu_capture_prepared_load_index = False
1820	self._npu_prepared_load_index = None	1821	self._npu_prepared_load_index = None
	@@ -1996,6 +1997,18 @@ class NPUTritonKernel(TritonKernel):
1996	"pointer": self.args.input(name),	1997	"pointer": self.args.input(name),
1997	}	1998	}
1998		1999
		2000	+ # -- permute-gather reduction rewrite: thin delegates into
		2001	+ # permute_gather_rewrite.py (kept as methods for the call sites and the
		2002	+ # tests' monkeypatching of _npu_pg_rewrite_body).
		2003	+
		2004	+ def _npu_pg_record(self, result_var, name, index):
		2005	+ from .permute_gather_rewrite import _npu_pg_record as _impl
		2006	+ return _impl(self, result_var, name, index)
		2007	+
		2008	+ def _npu_pg_rewrite_body(self):
		2009	+ from .permute_gather_rewrite import _npu_pg_rewrite_body as _impl
		2010	+ return _impl(self)
		2011	+
1999	def index_to_str(self, index: sympy.Expr) -> str:	2012	def index_to_str(self, index: sympy.Expr) -> str:
2000	# The index carries PyTorch's Max(1, dim) stride clamp from dynamic conv-output	2013	# The index carries PyTorch's Max(1, dim) stride clamp from dynamic conv-output
2001	# layouts (torch.utils._sympy Max, not sympy.Max); the Triton printer has no max	2014	# layouts (torch.utils._sympy Max, not sympy.Max); the Triton printer has no max
	@@ -4033,6 +4046,17 @@ class NPUTritonKernel(TritonKernel):
4033		4046
4034	self.codegen_body()	4047	self.codegen_body()
4035		4048
		4049	+ # Permute-gather rewrite (ncfg.enable_permute_gather): after the body text
		4050	+ # exists, swap the strided reduction load for a contiguous DMA + tl.gather,
		4051	+ # and record the forced (XBLOCK, R0_BLOCK) so reduction() pins exactly the
		4052	+ # config the rewrite was validated against (a sweep would hit other
		4053	+ # (XBLOCK, R0_BLOCK) pairs whose real_block_row != 1 silently miscompile
		4054	+ # the gather). Validation failure -> no marker -> strided load stays.
		4055	+ if triton_codegen_linearize and getattr(self, "_npu_pg_candidates", None):
		4056	+ _pg_marker = self._npu_pg_rewrite_body()
		4057	+ if _pg_marker is not None:
		4058	+ triton_meta["npu_permute_gather"] = _pg_marker
		4059	+
4036	# A5 one-program-per-tile: attach the host-side block-count recipe so the launcher	4060	# A5 one-program-per-tile: attach the host-side block-count recipe so the launcher
4037	# reproduces total_blocks (exact program count) and sizes the grid to it.	4061	# reproduces total_blocks (exact program count) and sizes the grid to it.
4038	# codegen_body() has populated pre_loop_code, so the recipe reads finished	4062	# codegen_body() has populated pre_loop_code, so the recipe reads finished

Mtorch_npu/_inductor/triton_experimental/config.py+87-0

	@@ -169,6 +169,13 @@ inject_care_padding: bool = False
169	# Refactor expanded conv-output store strides onto precomputed ks.	169	# Refactor expanded conv-output store strides onto precomputed ks.
170	refactor_clamp_stride: bool = False	170	refactor_clamp_stride: bool = False
171		171
		172	+# =====================================================================
		173	+# Permute-gather strided-reduction rewrite (opt-in: enable_permute_gather).
		174	+# All tuning thresholds live in this section; codegen keeps only the UB-budget
		175	+# formula constants (_PG_UB_*), matching upstream practice of keeping formula
		176	+# internals module-local (e.g. TRITON_MAX_BLOCK in torch/_inductor codegen).
		177	+# =====================================================================
		178	+
172	# Realize a permute+gather into a contiguous buffer at lowering. Default ON: a	179	# Realize a permute+gather into a contiguous buffer at lowering. Default ON: a
173	# non-unit inner stride pushed onto the reduction axis (e.g. T5 fwd softmax with a	180	# non-unit inner stride pushed onto the reduction axis (e.g. T5 fwd softmax with a
174	# relative-position bias, logical [heads,q,k] over [q,k,heads] storage) degrades to	181	# relative-position bias, logical [heads,q,k] over [q,k,heads] storage) degrades to
	@@ -178,6 +185,86 @@ refactor_clamp_stride: bool = False
178	# no-realize path.	185	# no-realize path.
179	realize_permute_gather: bool = True	186	realize_permute_gather: bool = True
180		187
		188	+# Codegen a contiguous-DMA + tl.gather for a permute that pushes a non-unit stride
		189	+# onto the reduction axis, instead of the strided tl.load (scalar gather on Ascend).
		190	+# Keeps the permute as a zero-copy logical view (no realize buffer) and rewrites the
		191	+# consumer load into a contiguous burst DMA of the whole row, then a register-level
		192	+# gather into the logical tile. Requires a reduction axis with unit-input-coeff interior
		193	+# (single-interior geometry, e.g. bias[Sq,Sk,H] permute(2,0,1) reduce over Sk); other
		194	+# shapes fall back to the strided load. OFF by default; opt-in per compile.
		195	+enable_permute_gather: bool = False
		196	+
		197	+# Why 256: the measured fp32 gather/trans crossover sits between stride_r=24
		198	+# (96B: gather wins) and 64 (256B: gather loses, its flat tile overflows UB at
		199	+# R0=128); 64 fp32 = 256B is also half the 910B2 segment-prefetch granularity
		200	+# knee (512B). NOT a DMA alignment boundary -- the flat DMA is alignment-agnostic.
		201	+#
		202	+# Benefit gate for the permute-gather rewrite, in bytes of transpose granularity
		203	+# (permuted inner stride * elemsize). Below this (stride_bytes < gate) the
		204	+# register gather wins: its flat contiguous DMA is alignment-agnostic and the
		205	+# flat tile stays under UB (H<=63 fp32). At/above the gate the layout goes to
		206	+# the trans mode (block_ptr + tl.trans, tails via boundary_check): a gather flat
		207	+# tile would overflow UB (measured H=64: trans 65.7us vs gather R0=64 93.2us,
		208	+# R0=128 compile fail) and trans beats the realized transpose (two-kernel 219us).
		209	+# Measured crossover (fp32, per-iteration device time): stride_r=24/96B gather
		210	+# 46.7us < strided 66us < eager 58us; stride_r=64/256B gather 99us ~ strided
		211	+# 102us > eager 67us; stride_r=128/512B gather 324us > strided 190us. At the
		212	+# same threshold lowering keeps the permute zero-copy (view for the rewrite) vs
		213	+# realizing it (fast strided fallback).
		214	+permute_gather_stride_gate_bytes: int = 256
		215	+
		216	+# Why 64: 64 fp32 = 256B chunk = the gate's DMA efficiency unit; pow2 and a
		217	+# 32B-multiple (hard constraint below), keeping the XBLOCK pin autotune-legal
		218	+# while boundary_check absorbs tails.
		219	+#
		220	+# Static-trans int-axis chunk width, in ELEMENTS (fp32: 64 = 256B chunk, i.e.
		221	+# the same boundary as the gate above). Static trans pins XBLOCK =
		222	+# min(stride_r, ktile) instead of stride_r, so the interior axis is chunked
		223	+# across programs (x1_blocks = ceil(H/ktile)) through the greedy tile chain +
		224	+# group dispatch -- the same mechanism dynamic-H trans already uses, with tails
		225	+# kept exact by boundary_check. Decouples the R0_BLOCK UB cap from H (R0 stops
		226	+# collapsing as H grows: H=4096 R0 1 -> 64, 300 -> 5 r-trips). Must be a
		227	+# multiple of 32B/elemsize (fp32: 8), else tile_align rounding breaks the
		228	+# forced-tiling eval and the rewrite falls back to strided.
		229	+permute_gather_ktile: int = 64
		230	+
		231	+# Dynamic-H (symbolic reduction stride) fallback: the gather index must be
		232	+# compile-time affine in stride_r (rejected by the geometry's static-int gate), so
		233	+# dynamic H can only ride the trans mode's block_ptr + boundary_check, which is
		234	+# shape-generic. OFF -> the strided-load fallback for dynamic H.
		235	+permute_gather_dynamic_trans: bool = True
		236	+
		237	+# Why 4096: TRITON's max_block -- the same cap family as upstream
		238	+# TRITON_MAX_BLOCK (see npu_triton_config_reduction).
		239	+#
		240	+# Gather-mode XBLOCK cap: the gather rewrite pins XBLOCK = stride_r (the
		241	+# interior axis runs the full head in one tile), and the reduction heuristic
		242	+# caps XBLOCK at TRITON's max_block (4096, see npu_triton_config_reduction). A
		243	+# stride_r above that cannot be gathered -> permute_gather_mode returns None
		244	+# for the gather branch, so lowering realizes the permuted layout (the
		245	+# pre-dispatch behavior) instead of compiling an "XBLOCK too large" kernel for
		246	+# huge H (measured H=16384 fp32). The trans branch is NOT capped: it pins
		247	+# XBLOCK = min(stride_r, permute_gather_ktile), which stays 64 for any H, so
		248	+# huge H goes through trans with the int axis chunked across programs.
		249	+permute_gather_max_xblock: int = 4096
		250	+
		251	+
		252	+def permute_gather_mode(stride_r, elemsize):
		253	+ """Dispatch the permute rewrite for a static permuted inner stride of
		254	+ ``stride_r`` elements of ``elemsize`` bytes: "gather" / "trans" / None
		255	+ (fall back to realize). Per-branch caps and mode semantics: see the flags
		256	+ above. Dynamic (symbolic) strides are dispatched in the codegen geometry,
		257	+ which can see the symbolic reduction coefficient. Reads the live config
		258	+ object (install_config_module moves the typed defaults off module globals
		259	+ into instance attributes)."""
		260	+ from torch_npu._inductor.triton_experimental import config as _cfg
		261	+ stride_bytes = stride_r * elemsize
		262	+ if stride_bytes < _cfg.permute_gather_stride_gate_bytes:
		263	+ if stride_r > _cfg.permute_gather_max_xblock:
		264	+ return None # gather pins XBLOCK = stride_r; capped at max_block
		265	+ return "gather"
		266	+ return "trans"
		267	+
181	# Route the MASK-COMPOSITE softmax (aten._safe_softmax, produced from	268	# Route the MASK-COMPOSITE softmax (aten._safe_softmax, produced from
182	# transformers-style causal-mask + softmax patterns) through aclnn instead of	269	# transformers-style causal-mask + softmax patterns) through aclnn instead of
183	# Triton fusion: the fused Triton kernel materializes [B,H,S,S] masks and	270	# Triton fusion: the fused Triton kernel materializes [B,H,S,S] masks and

Mtorch_npu/_inductor/triton_experimental/lowering.py+39-1

	@@ -475,6 +475,15 @@ def _permuted_inner_stride(x, dims):
475	return in_stride[last_src]	475	return in_stride[last_src]
476		476
477		477
		478	+def _elemsize_of(x):
		479	+ """Element size in bytes of ``x``, for the permute-gather benefit gate
		480	+ (transpose granularity = inner_stride * elemsize)."""
		481	+ try:
		482	+ return x.get_dtype().itemsize
		483	+ except Exception:
		484	+ return 4
		485	+
		486	+
478	def npu_permute(x, dims):	487	def npu_permute(x, dims):
479	"""NPU permute: realize a permute that pushes a NON-UNIT stride onto the inner	488	"""NPU permute: realize a permute that pushes a NON-UNIT stride onto the inner
480	axis into a contiguous buffer, instead of folding it into the consumer's loads	489	axis into a contiguous buffer, instead of folding it into the consumer's loads
	@@ -484,7 +493,36 @@ def npu_permute(x, dims):
484	stride statically != 1, inner length statically > 1, input an unrealized	493	stride statically != 1, inner length statically > 1, input an unrealized
485	producer or plain buffer; else delegate to upstream permute."""	494	producer or plain buffer; else delegate to upstream permute."""
486	result = _upstream_permute(x, dims)	495	result = _upstream_permute(x, dims)
487	- if not ncfg.realize_permute_gather:	496	+ # enable_permute_gather keeps the permute as a zero-copy logical view: the
		497	+ # consumer's reduction load is rewritten (codegen) into a contiguous DMA +
		498	+ # tl.gather. It bypasses the realize path entirely (and the default upstream
		499	+ # fast-path gate below), leaving the view unrealized for the consumer kernel.
		500	+ #
		501	+ # Benefit gate: keep the view only when the rewrite is expected to beat the
		502	+ # realize fallback -- i.e. the transpose granularity (inner stride * elemsize)
		503	+ # is small and the strided/realize path would be DMA-poor. When the gate
		504	+ # rejects (large inner stride, e.g. seg=2M @ stride_r=256), fall through to
		505	+ # the realize path below so the fallback is the fast materialized permute
		506	+ # (measured 434us < eager 709us) rather than the scalar-gather strided load.
		507	+ if ncfg.enable_permute_gather:
		508	+ inner_stride = _permuted_inner_stride(x, list(dims))
		509	+ if not isinstance(inner_stride, (int, sympy.Integer)):
		510	+ # stride unknown (dynamic H): keep the view only when the dynamic-H
		511	+ # trans fallback is on; otherwise fall through to realize so the
		512	+ # permuted layout is materialized instead of a strided scalar-gather
		513	+ # on an unrealized view (codegen geometry returns None in that case).
		514	+ if ncfg.permute_gather_dynamic_trans:
		515	+ return result
		516	+ elif ncfg.permute_gather_mode(int(inner_stride), _elemsize_of(x)):
		517	+ # gather (small inner stride, stride_bytes < 256) or trans (>= 256B,
		518	+ # e.g. H>=64 fp32, int axis chunked to permute_gather_ktile) both keep
		519	+ # the zero-copy view for the codegen rewrite; the geometry picks the
		520	+ # primitive. Realize is no longer the fallback for a large inner
		521	+ # stride -- trans is measured faster there (H=64 65.7us vs two-kernel
		522	+ # realize 219us).
		523	+ return result
		524	+ # else: realize below
		525	+ elif not ncfg.realize_permute_gather:
488	return result	526	return result
489	if not isinstance(x, ir.TensorBox) or not isinstance(result, ir.TensorBox):	527	if not isinstance(x, ir.TensorBox) or not isinstance(result, ir.TensorBox):
490	return result	528	return result

Mtorch_npu/_inductor/triton_experimental/npu_triton_heuristics.py+20-0

	@@ -2438,6 +2438,26 @@ def reduction(
2438	)	2438	)
2439		2439
2440	if len(size_hints) in (2, 3):	2440	if len(size_hints) in (2, 3):
		2441	+ # Permute-gather rewrite (ncfg.enable_permute_gather): the codegen
		2442	+ # rewrote the strided reduction load into a contiguous DMA + tl.gather,
		2443	+ # which is only valid when the greedy tile chain gives the interior axis
		2444	+ # the full XBLOCK run (real_block_row == 1). Pin exactly the config the
		2445	+ # rewrite was validated against; a sweep would run other (XBLOCK, R0_BLOCK)
		2446	+ # pairs whose real_block_row != 1 silently miscompile the gather.
		2447	+ # Checked BEFORE pin_xr: the rewrite's emitted tiling is only valid
		2448	+ # under this forced config, so the PG marker must win over the
		2449	+ # diagnostic pin when both are set.
		2450	+ _pg = triton_meta.get("npu_permute_gather")
		2451	+ if _pg:
		2452	+ return cached_autotune(
		2453	+ size_hints,
		2454	+ [reduction_cfg(int(_pg["XBLOCK"]), int(_pg["R0_BLOCK"]))],
		2455	+ triton_meta=triton_meta,
		2456	+ inductor_meta=inductor_meta,
		2457	+ heuristic_type=HeuristicType.REDUCTION,
		2458	+ filename=filename,
		2459	+ )
		2460	+
2441	# TEMP DIAGNOSTIC: pin a single (XBLOCK, R0_BLOCK) to measure the real	2461	# TEMP DIAGNOSTIC: pin a single (XBLOCK, R0_BLOCK) to measure the real
2442	# kernel at a chosen trip count. Remove after validation.	2462	# kernel at a chosen trip count. Remove after validation.
2443	_pin = ncfg.pin_xr	2463	_pin = ncfg.pin_xr