【合入来源】

如有社区issue，请关联issue链接
请勿携带内部流程信息（需求链接、问题单、内部issue等）

• 需求
• 问题单
• issue/工单 https://gitcode.com/Ascend/pytorch/issues/4366（permute-gather 性能问题）
• 重构优化
• 资料更新

【修改方案】

请描述修改内容的具体实现，涉及哪些组件之间进行交互，可以用1、2、3、...进行罗列
如果是需求或者重构类的PR，需要补充详细设计文档（说明上下游组件关系、时序图、类图、DFX能力等内容）

T5 形状的 permute+归约（out[b,k,i] = sum_j arg1[b,k,i,j] + sum_j arg0[i,j,k]）在默认路径下 permute 物化为独立 HBM 全量转置 + 归约分 kernel，且归约轴非单位 stride 的 strided tl.load 在 Ascend 落标量 gather，DMA 效率低。本 PR 新增 enable_permute_gather（opt-in，默认关）将归约轴 strided load 改写为连续 DMA + 寄存器级置换。组件交互：

1. config.py：permute_gather_mode(stride_r, elemsize) 三态分派——stride 字节 <256 → "gather"（H≤63 fp32）；≥256 → "trans"（H≥64）；gather 分支受 permute_gather_max_xblock(4096) 上限约束（gather pin XBLOCK=stride_r，受 TRITON max_block 限制），超限返回 None 回退 realize。阈值 flags 全部 opt-in（permute_gather_stride_gate_bytes / permute_gather_ktile / permute_gather_dynamic_trans / permute_gather_max_xblock）

2. lowering.py：enable_permute_gather 下 npu_permute 直接构造 ir.PermuteView 零拷贝逻辑 view（绕开 realize 物化）；关闭时走既有 realize 逻辑（基线行为）

3. codegen/triton.py：NPUTritonKernel.load 拦截 eligible 归约 load：

   • gather：整行连续 DMA（L = stride_r·R0_BLOCK 平铺）+ tl.gather 寄存器置换
   • trans：tl.make_block_ptr + tl.trans + boundary_check=[0,1,2] + padding_option="zero"；静态 trans pin XBLOCK=min(stride_r, ktile=64)，int 轴跨程序 chunk（x1_blocks=ceil(H/64)）；trans 分支不受 max_xblock 上限约束（XBLOCK 与 stride_r 解耦），H>4096 同样走 trans+K-tile（H=8192 → 128 个满 chunk，整除无 tail 浪费），realize 仅保留给 gather 超限场景；R0 公式 eff=min(stride_r, ktile) 解耦 R0/H
   • slot 检查放宽为互异且 0<=s<ndim（ndim=max(slots)+1），消除「归约轴在最后」的布局假设（DEFAULT scheduler patch 对大 H kernel 产生 R 树 slot 0 布局时不再误拒）

4. npu_triton_heuristics.py：rewrite 下固定 XBLOCK/R0_BLOCK（real_block_row==1 才合法），autotune 不扫其它配置

5. 评审修复（d308b142b5 / bc0df673bb / 562d97e39f，响应 @rmch 与 robot 评审；其中 562d97e39f 的 autotune winner 缓存键修复后移至独立 PR）：删 kernel 内非法 raise AssertionError（P1，Python raise 非 Triton 合法语句）；slot 预检先类型检查再 max（slot 缺失正常回退 strided，不再 TypeError）；gather pidx/reshape 与 slot 顺序解耦（R-first 布局不再静默重排）；enable_permute_gather 与 pin_xr 同开时 PG marker 优先（避免 body 已 rewrite 而 config 被 pin_xr 决定）；加载缓存 autotune winner 时剥离非 kwarg 键（issue 4367）；新增 4 个回归用例（UT 9→13）

6. 字符串手术 AST 化（925b5372e1，评审遗留整改）：codegen 内三处对生成代码的正则/eval 改写全部替换为结构化 AST 操作（复用本文件既有 _npu_parse_tl_load_assignment / _npu_assignment_parts 基建，与 _maybe_rewrite_select_lane_load 同款模式），删除 PR 自身引入的 import re：

   • _npu_pg_rewrite_body：load 行定位由行文本正则改为按赋值目标名匹配——可穿透 .to(tl.float32) 提升链（贪婪 (.*)\) 无法定界，fp16/bf16 候选原实现会静默错配）；非 str 行（DeferredLine）透传；emit 拒绝时保留原行回退（原实现返回空列表会丢行且仍下发 config pin）
   • _npu_pg_emit：ptr/other/eviction_policy 由 argtext 正则（P3：other=([^)]*) 遇 float('-inf') 等嵌套括号即截断）改为 load AST 节点 args[0]/keywords 结构化读取；末行经节点拼接把 tl.where(...) 换回原 value_ast 再 unparse，保留 .to() dtype 提升
   • _npu_pg_eval_real_blocks：正则+eval 改为 AST 解析（name : tl.constexpr = ... AnnAssign 归一化）+ 新增 _npu_pg_eval_expr 小型求值器（int 算术/比较/IfExp tile 三元式/min-max；不支持的节点照旧跳过，保持「对发射文本求值、免疫排序逻辑漂移」性质）

7. 评审二轮（ff49a90cde / 73f2144299 / 0c785d501a）：为满足 PR 总行数 ≤1000（现 967），autotune 缓存键修复移至独立 PR（本 PR revert）；_npu_pg_eval_expr 求值器表驱动化、三处重复注释去重、测试 prose 压缩（13 用例与断言不动，232→197 行）——13 个 PG kernel 生成源码 sha256 与整改前逐字节一致；config.py 阈值分节 banner + 三个阈值取值依据导语（256=实测 gather/trans 交叉点+UB 预算而非 DMA 对齐边界、64=256B DMA 效率单元、4096=TRITON max_block）。

8. 评审二轮补齐（89ea3c481c）：全部 _npu_pg_* 实现提取到独立模块 codegen/permute_gather_rewrite.py（section：eligibility/geometry/validation + gather/trans 发射分区；NPUTritonKernel 留薄委托，调用点与测试 monkeypatch 不变）；统一回退 debug 日志（模块 logger：geometry/mode 门、slot 缺失/非互异、tiling 校验失败、load 定位失败、rewrite 成功，均带原因与变量名）；no_raise 断言并入 gather/trans 用例族；PR 总量 996 行。验证：12/12 + 13/13（clean）；11 个唯一 PG kernel 源码 sha256 与迁移前逐字节一致。

设计文档：NPUPermuteGather_upstream_design.md / NPUPermuteGather_v2_design.md（随本 PR 提供，详细设计思路另见 commit message「修复报告」节）

【资料变更】

请确认是否涉及资料变更。如涉及，需要在PR中体现，并简要说明修改内容。如不涉及，需填写“不涉及”

不涉及

【接口变更】

请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及，需要详细说明接口以及对应的变更内容，同时需要在资料中体现。如不涉及，需填写“不涉及”

不涉及（enable_permute_gather 为 Inductor 后端内部配置项；对外 API 不变，生成 kernel 文本属内部产物）

【功能验证】

说明测试场景，测试方法。如果本次测试方式与常规单元测试不同，请详细说明您的测试步骤
新增/变更内容是否已新增/适配UT测试用例看护，并补充测试自验证截图

UT（新增 pytorch/test/_inductor/test_permute_gather.py，12 用例整批通过；no_raise 断言已并入 gather/trans 用例族）：

数值精度：assert_close(atol=1e-4, rtol=1e-4) 全过（H=1024 max_err 3.05e-05、H=8192 max_err 1.53e-05、J=96 bit-exact）。

AST 整改回归（925b5372e1）：test_permute_gather.py 13/13 + test_triton_experimental_enable.py 13/13（source env.sh clean 冷缓存）；13 个 PG kernel 生成源码与整改前 sha256 逐字节一致（纯重构、零行为变化，上表性能数据不受影响）；fp16 端到端补充验证——生成行 tmp1 = tl.where(r0_mask, tmp1_pg_g, 0.0).to(tl.float32) 保留 dtype 提升，compiled 与 fp32 金标 max|diff|=0.0000（eager fp16 自身 3.12e-02，即保 cast 后走 fp32 累加、比 eager 更准；旧正则实现在该场景会丢 cast）。

性能（CANN profiler 设备时长 ground truth，per-launch 各 kernel 均值求和；对比 eager / 基线 enable_permute_gather=False）：

已知倒退（如实）：H=72 trans+K-tile 2.0× 慢于基线（ceil(72/64)=2 程序拆分，tail chunk 仅 8 行按 64 行 tile 读取读放大）；非 32B 对齐 H（65/71/73/129）系统性 1.7–2.0× 慢（block_ptr 非对齐行宽读放大，后续优化方向 H%8≠0 gate 回退）；H=32 gather 14% 慢；墙钟口径下短 kernel 场景（H≤128）被 ~130us/launch host 开销吞掉收益（设备时长收益 ≠ 端到端收益）。

CI：未跑仓 CI 流水线（本环境仅本地 NPU 单机验证）。

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

• 代码注释完备，正确记录错误日志（rewrite 各 gate 回退路径均有注释与文档化实测依据）
• 代码实现进行了返回值、空指针等校验（geometry/rewrite 判定失败一律回退 strided load，正确性保底）
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等（perf(npu-inductor)）
• PR持续集成流水线（CI）执行通过，代码检查无异常（本地 pytest 26 passed（permute_gather 13 + enable 13）；仓 CI 未跑，如实说明）