企业官网建设流程全解析

PyPTO Reshape 全局优化实战指南：外提、合轴与 inplace 消除循环内数据搬运

【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym

本文是 CANN / pypto-gym 仓库中 PyPTO 算子开箱性能调优技能（tune-frontend）的 Reshape 全局优化操作指南。核心目标只有一个：尽量把 reshape 提前到 loop 外面，使用reshape inplace，减少循环体中的 reshape，从而减少数据搬运。读完本文，你将掌握逐 Reshape 系统分析的方法论、四种优化方式（外提 / 合轴 / 删除冗余 / squeeze-unsqueeze 等价替换）的完整代码范式，以及 reshape 前后 vec_tile_shapes 重设等配套约束，可直接用于 GQA Attention、Sparse Flash Attention 等典型算子的开箱性能调优。

为什么 reshape 是开箱性能的关键瓶颈

在 PyPTO 的声明式编程模型里，pypto.reshape并非"零成本"的元数据操作——当 reshape 出现在循环体内部时，每一次循环迭代都可能触发一次完整的数据拷贝与重新布局。尤其对于 Attention 类算子，循环体内往往同时存在 Q/K/V 的 shape 变换、matmul 输出的降维/升维，以及 online softmax 的逐块更新，任何一次多余的 reshape 都会放大为成百上千次重复搬运。

因此，tune-frontend 技能的调优三阶段流程将 Reshape 全局分析列为阶段 A（全局分析）的 A3 制品，并对应优化点 [F-4]（Reshape 全局优化，含 squeeze/unsqueeze → reshape inplace），优先级为 ⭐⭐⭐ P0，见 optimization_catalog.md。该流程强制要求：未完成阶段 A + 阶段 B 的分析，禁止直接进入阶段 C 逐项优化；执行 F-4 时必须加载本文获取完整操作指南，见 tune-frontend/SKILL.md。

第 1 步：逐 Reshape 系统分析

优化前必须对算子中每一个pypto.reshape调用逐个分析，判断其是否必须出现在最内层循环体中，而不是凭直觉猜测优化点。

分析方法：

1. 逐行阅读算子代码，记录每一个pypto.reshape调用的位置（loop 外 / loop 内 / 最内层循环体）；
2. 分析每个 reshape 的输入 tensor 来源（原始输入 / 中间计算结果）和目标 shape；
3. 判断该 reshape 是否一定需要在当前位置执行，还是可以提前到更外层。

SKILL.md 的 A3 小节给出了更工程化的定位手段——直接用 grep 一次性抓取全部相关调用：

grep -nE "pypto\.(reshape|squeeze|unsqueeze)" <算子文件>

拿到调用清单后，按下面的分析表格模板逐行填写（该模板在 SKILL.md 的 A3 小节与本文所依赖的 reshape-global-optimization.md 中均有定义，两处模板略有差异，可合并使用）：

分析时的关键检查项（来自 SKILL.md A3）：

• 源 Shape == 目标 Shape 的冗余 reshape？（应删除）
• 原始输入的 reshape 是否在 loop 外？（应外提）
• 原始输入的 reshape 是否使用了inplace=True？（必须用）
• 原始输入的 squeeze/unsqueeze 是否在 loop 内？（⚠️ squeeze/unsqueeze 不支持 inplace，应替换为等价 reshape 外提）
• 中间结果或输出 tensor 是否使用了inplace=True？（禁止用）
• loop 内的 reshape 是否依赖循环变量？（不依赖则外提）
• reshape 前后是否有对应的set_vec_tile_shapes变更？

⚠️ 特别提醒：optimization_catalog.md 中 [F-4] 的约束明确指出，输出 tensor 不可 inplace reshape，否则会导致切片写入索引断裂（输出全零）。inplace 仅限原始输入（函数参数）。

第 2 步：四种 Reshape 优化方式详解

逐项确认每个可外提 reshape 的优化方式后，按以下四种方式逐一落地。

方式 1：原始输入 reshape 外提（inplace=True）

对原始输入（函数参数）的 reshape，挪到算子入口（所有 loop 之前），并使用inplace=True直接完成 shape 变换，避免在循环体内对同一份数据反复做冗余拷贝：

# ✅ 正确：reshape 挪到算子入口，inplace=True q_grouped = pypto.reshape(query, [num_kv_heads, num_heads_per_group, head_dim], inplace=True) k_cache = pypto.reshape(key_cache, [kv_len, num_kv_heads, head_dim], inplace=True) for i in pypto.loop(num_blocks, ...): # loop 内直接使用已 reshape 的 tensor，无额外搬运 scores = pypto.matmul(q_grouped, k_cache_block, ...) # ❌ 错误：reshape 放在 loop 内部，每次循环重复执行数据拷贝 for i in pypto.loop(num_blocks, ...): q_grouped = pypto.reshape(query, [num_kv_heads, num_heads_per_group, head_dim]) # 冗余搬运

仓库中的真实实践印证了这一范式：在 sparse_flash_attention_quant_impl.py 中，输出 tensorattention_out在进入 batch 循环之前先做了一次pypto.reshape(attention_out, [batch_size_sym * s1_n2_gsym, dn], inplace=True)合轴；lightning_indexer_quant_impl.py 则在算子入口对idx_query、idx_key_cache等输入参数一次性完成inplace=True的 2D 化。而 arctic/sum_lstm.py 对权重/偏置等原始输入使用pypto.reshape(..., [1, hidden_dim], inplace=True)后直接在后续计算中复用——这些都是在循环之前一次性完成 shape 变换的典型写法。

方式 2：高维计算提前合轴（3D/4D → 2D）

如果循环体内的计算超过两维（如 3D/4D），NPU 指令对多维 tensor 处理不友好，性能较差。应在进入循环前对原始输入reshape inplace合轴为 2D，避免循环体内出现 reshape：

# ✅ 正确：循环前合轴为 2D，循环内无 reshape query_2d = pypto.reshape(query, [batch * heads * seq_q, dim], inplace=True) key_2d = pypto.reshape(key, [batch * heads * seq_kv, dim], inplace=True) value_2d = pypto.reshape(value, [batch * heads * seq_kv, dim], inplace=True) for b_idx in pypto.loop(batch, ...): for n_idx in range(heads): q_offset = b_idx * heads * seq_q + n_idx * seq_q + q_start q_block = pypto.view(query_2d, [BLOCK, dim], [q_offset, 0], ...) # ... 计算，循环体中无 reshape

合轴之后循环体内的view取块，offset 计算仍然是直观的线性寻址，完全不损失可读性。这一方式与 F-1（任务粒度）、F-7（外层动态轴切块）协同使用效果最佳。

方式 3：冗余 reshape 删除（source == target）

检查每个 reshape 的源 shape 是否等于目标 shape（常见于代码迭代过程中残留的无效 reshape），直接删除无效调用：

# ❌ 冗余：源 shape 等于目标 shape k_embed = pypto.reshape(k_embed, [8, 128]) # [8,128] → [8,128] # ✅ 删除后直接使用 # k_embed 已经是 [8, 128]，无需 reshape

检查方法：逐行扫描所有pypto.reshape调用，比对源 shape 与目标 shape 是否相同。此检查应在阶段 A3（Reshape 全局分析表）中完成，在表格的"冗余?"列中标记。SKILL.md 的 A3 模板示例中即有k_embed [8,128] → [8,128]被标记为✅ 冗余 → 删除的实例。类似的冗余set_vec_tile_shapes重复调用也应一并合并（见 basic-block-optimization.md 的冗余配置检查）。

方式 4：squeeze/unsqueeze 替换为 reshape inplace 外提

⚠️pypto.squeeze和pypto.unsqueeze不支持inplace=True参数，无法直接原地修改 tensor shape。对原始输入（函数参数）的 squeeze/unsqueeze 操作，如果出现在循环体内部，应替换为等价的pypto.reshape(..., inplace=True)并挪到算子入口（所有 loop 之前），消除循环内的重复数据搬运：

# ✅ 正确：squeeze 替换为 reshape inplace 外提 # 原代码：pypto.squeeze(query, dim=1) → [B,1,S,D] → [B,S,D] query_3d = pypto.reshape(query, [batch, seq_len, head_dim], inplace=True) # 算子入口，只执行一次 for i in pypto.loop(num_blocks, ...): scores = pypto.matmul(query_3d, key_block, ...) # loop 内直接使用 # ❌ 错误：squeeze 放在 loop 内部，不支持 inplace，每次循环重复搬运 for i in pypto.loop(num_blocks, ...): q_sq = pypto.squeeze(query, dim=1) # 冗余搬运

等价映射规则：

注意：仓库中仍然存在大量pypto.unsqueeze的使用（例如 glm_ffn_shared_expert_quant_impl.py 对 scale 矩阵的pypto.unsqueeze(w13_scale, 0)、quant_matmul_reduce_sum_impl.py 中的升维操作），这些调用如果位于循环体内且作用于原始输入，就是方式 4 的改造对象——替换为等价 reshape 并外提。

约束：

• 仅对原始输入（函数参数）可替换为reshape(inplace=True)，中间结果和输出 tensor 不能 inplace reshape；
• 替换后需检查循环体内所有引用该 tensor 的位置是否使用了正确的 shape；
• 替换后需在 reshape 后、使用前重新设置set_vec_tile_shapes以匹配新 shape。

第 3 步：配套约束——reshape 前后 vec_tile_shapes 重设

reshape 优化不是孤立动作。reshape/view会改变 tensor shape，而vec_tile_shapes必须与当前操作的实际 tensor shape 匹配，因此必须遵循 basic-block-optimization.md 中的重设规则：

1. reshape 前：vec_tile_shapes 按源 tensor（reshape 前）的 shape 设置；
2. reshape 后：在操作 reshape 后的 tensor 之前，重新设置 vec_tile_shapes 按目标 tensor 的 shape。

# ✅ 正确：reshape 前按源 shape 设，reshape 后按目标 shape 重新设 pypto.set_vec_tile_shapes(8, 128) # 源 tensor [8, 128] k_embed_3d = pypto.reshape(k_embed, [8, 1, 128]) pypto.set_vec_tile_shapes(8, 1, 128) # 目标 tensor [8, 1, 128] pypto.assemble(k_embed_3d, [0, pos, 0], cache) # ❌ 错误：reshape 前就按目标 shape 设了 vec_tile pypto.set_vec_tile_shapes(8, 1, 128) # ❌ 此时 tensor 还是 [8, 128] k_embed_3d = pypto.reshape(k_embed, [8, 1, 128]) pypto.assemble(k_embed_3d, [0, pos, 0], cache)

常见遗漏场景：

• reshape 后紧跟assemble：必须在 reshape 后、assemble 前设置匹配目标 shape 的 vec_tile；
• reshape 后紧跟matmul：matmul 由 cube_tile 控制，vec_tile 影响较小，但仍建议按目标 shape 设置；
• 多个连续 reshape：每次 reshape 后都需确认 vec_tile 是否匹配。

仓库中 sparse_flash_attention_quant_impl.py 的 INT8 反量化路径就是这一规则的完整示范：reshape(kn_quant_fp32, [s2_tile * 8, 128])前先设set_vec_tile_shapes(16, 1024)，随后切换到set_vec_tile_shapes(128, 128)执行乘法，再做reshape(kn_fp32, [s2_tile, dn * 2])并重设set_vec_tile_shapes(16, 512)——每次 shape 变化点都显式重设了匹配的 tile。

实战案例：Decode Attention Vector 合轴优化（-6.0%）

本文所依赖文档中引用了一个完整落地案例 vector-axis-merge-softmax.md，该案例的核心正是方式 2（高维合轴）+ 方式 1（inplace 外提）的组合：

场景：GQA decode attention 算子中，softmax 及前后的 vector 操作（mul/amax/sub/exp/sum/div/cast）在 3D shape[8, 4, 2048]下执行，产生 6 个独立 vector 子图，调度开销大。

优化手段：

1. matmul 输出的 3D tensorreshape(inplace)为 2D；
2. 所有 vector 操作在 2D 下执行 +sg_set_scope合图；
3. reshape 回 3D 传给下一个 matmul。

# 优化前（3D，6 个子图） scores_fp32 = pypto.matmul(q_grouped, k_cache, pypto.DT_FP32, b_trans=True) scores_scaled = pypto.mul(scores_fp32, scale) # [8,4,2048] FP32 row_max = pypto.amax(scores_scaled, dim=-1, keepdim=True) # ... softmax 链在 3D 下展开为 6 个独立 vector 子图 # 优化后（2D，1 个合图子图） scores_fp32 = pypto.matmul(q_grouped, k_cache, pypto.DT_FP32, b_trans=True) scores_2d = pypto.reshape(scores_fp32, [32, 2048], inplace=True) pypto.set_vec_tile_shapes(8, 2048) pypto.set_pass_options(sg_set_scope=1) # ... 2D 下的 softmax 链，全部并入同一子图 pypto.set_pass_options(sg_set_scope=-1) attn_weights = pypto.reshape(attn_bf16, [8, 4, 2048], inplace=True) # 传回下一 matmul

实测收益：执行时间 275.44 → 258.98 us（-6.0%），任务数 168 → 137（-18.5%），子图数 6 → 1，精度 Max difference 0.000031 无变化。案例还记录了 4 轮迭代失败分析（详见原案例文件）：归约轴被 vec_tile 切分导致性能回退（+16.6%）、未设 vec_tile 导致编译失败（reduce op 要求尾轴 32B 对齐）、tile 过大超出 UB 容量导致Run pass failed——最终(8, 2048)+sg_set_scope组合胜出。

关键经验（同样适用于本文的 reshape 优化）：

1. 归约轴必须不切分：vec_tile_shapes 第二维应等于实际归约轴长度；
2. UB 容量决定第一维上限：第一维 × 第二维 × dtype字节数 × tensor 总数不能超出 UB（约 128KB 保守估计）；
3. 必须显式设 vec_tile_shapes：合轴后维度变化，不设会导致 reduce op 编译失败；
4. sg_set_scope 合图：将连续 vector 操作合并为单个子图，减少调度开销。

常见失败模式速查

落地检查清单

完成 reshape 全局优化后，对照以下清单逐项确认（对应 SKILL.md 调优检查清单的 [F-4] 项与 A3 关键检查项）：

• 已用grep -nE "pypto\.(reshape|squeeze|unsqueeze)"抓取全部相关调用，并逐一填入 Reshape 全局分析表（源 shape / 目标 shape / 在 loop 内? / 输入类型 / inplace? / 冗余?）
• 原始输入的 reshape 全部外提到所有 loop 之前，且使用inplace=True
• 循环体内计算已尽量合轴为 2D，循环体内无多余 reshape
• 源 shape == 目标 shape 的冗余 reshape 已删除
• 循环体内的 squeeze/unsqueeze 已替换为等价reshape(inplace=True)并外提
• 中间结果与输出 tensor 未使用inplace=True
• 每个 reshape 前后均已重设匹配的set_vec_tile_shapes
• 优化后重新检查受影响的 A1 / A3 / A4 分析表行（SKILL.md 阶段 C 的强制要求：涉及 reshape 移动的结构变更必须更新过期分析结论）

若将本技能与 tune-frontend/SKILL.md 中的 F-1~F-3（任务粒度 / 循环体计算量 / 循环次数）、F-7~F-8（外层切块 / 内层 unroll）、F-9~F-10（Cube/Vector TileShape）等优化点协同使用，即可在算子初始开发阶段获得理想的"开箱性能"。更多同类案例可参考 tune-frontend/cases/README.md 中的案例索引。

【免费下载链接】pypto-gymPyPTO-Gym 是基于 PyPTO 编程框架构建的算子与模型样例仓库项目地址: https://gitcode.com/cann/pypto-gym