你好，认领这个任务

@wcleungaj 你好，请问上次代码审核周会后，功能代码没有改动，只新增了性能测试验证，CI已通过。这次代码审核会议还需要参加么，以及需要汇报什么呢？

@wcleungaj 你好，请问上次代码审核周会后，功能代码没有改动，只新增了性能测试验证，CI已通过。这次代码审核会议还需要参加么，以及需要汇报什么呢？

@sun_teng 你好，你不需要匯報任何東西了，目前需要等待架構師抽空做代碼審視

CustomOp Sub-Block 切分支持评审材料

会议：SIG-AscendNPU IR 双周例会（2026/09/03）

主仓分支：feat-customop-support

初始评审基线：3c0ef9ff（2026/09/02）

功能验证提交：37e32605（feat-customop-support）

性能验证补充日期：2026/09/07

E2E 仓：AscendNPU-IR-DT

E2E 分支：customop-tile-and-bind-e2e-refactor

E2E 精度基线：9239df3；Softmax 性能提交：04e5ac1

做了什么

本任务为 AIV 场景下的 hivm.hir.custom 和hivm.hir.custom_macro 增加 Sub-Block 1:2 切分能力。

原流程从 Store/Copy 创建切片并向 producer 冒泡，但缺少 CustomOp 的切片传播规则。遇到 CustomOp 后，Pass 无法根据算子语义计算 input tile，只能回退或在单一 Sub-Block 上执行完整算子。

本次实现完成了以下功能：

1. 根据 iterator_types 和 indexing_map 分析 Custom-like 算子的维度关系，并将可关联维度加入维度并查集。

2. 从 result slice 反推迭代空间，再投影得到每个 tensor input 的offset/size/stride。

3. 为 input 和 output init 创建切片，重建处理局部 tile 的CustomOp/CustomMacroOp。

4. 支持 identity、broadcast、reduction 输入、标量透传和受限的多结果 CustomOp。

5. 对无法证明安全的情况拒绝切分并回滚，完整算子限制在 Sub-Block 0 执行。

6. 保留非 builtin CustomOp 的用户 symbol，不改变外部 kernel 接口。

示例：原 CustomOp 处理 tensor<16x16xf16>，沿第 0 维切分后，两个 Sub-Block 分别处理一个 tensor<8x16xf16> tile：

Sub-Block 0: offset=[0,0], size=[8,16]
Sub-Block 1: offset=[8,0], size=[8,16]

相对当前本地 master，主仓分支的改动分布如下：

统计基于 git diff master...HEAD：功能代码包括头文件和 C++实现；LIT 的 CHECK、RUN、expected-* 等行虽以 // 开头，但实际是测试断言，因此与说明注释分开计数。删除的 17 行由 13 行代码、3 行注释和 1 个空行组成。E2E 位于独立 DT 仓，不计入上表。

怎么做的

2.1 整体流程

flowchart TD
   A[DimensionAnalyzer 分析维度关系] --> B[选择 Sub-Block 切分维度]
   B --> C[从 Store/Copy 创建 marked result slice]
   C --> D[slice 向 producer 冒泡]
   D --> E{Custom-like 算子满足切分约束?}
   E -- 否 --> F[中断冒泡并回滚]
   F --> G[完整 CustomOp 和 Store/Copy 限制到 Sub-Block 0]
   E -- 是 --> H[result slice 反推迭代空间]
   H --> I[迭代空间投影到各 tensor input]
   I --> J[收集全部 result slices]
   J --> K[创建 input/output-init slices]
   K --> L[重建 tile-local CustomOp]
   L --> M[添加 tiled_op 并替换旧 result slices]
   M --> N[新 input slices 继续向上游冒泡]

2.2 维度分析

DimensionAnalyzer::processCustomOpLike() 处理 CustomOp 和 CustomMacroOp：

• 缺少显式 iterator_types/indexing_map 时，将算子作为不透明节点。

• Custom-like 算子包含 unranked shaped operand 或 result 时，无法建立可靠的轴与 indexing map 对应关系。维度分析将其视为不透明算子；当前TileAndBindSubBlock Pass 会进一步放弃该函数的切分并回退。

• 全 parallel 且为 identity map 时复用普通 parallel 分析路径。

• 其他 map 只沿 parallel iterator 对应的 AffineDimExpr 合并维度。

• non-DPS result 没有对应 output map，分析时跳过，避免错误访问 map。

这一步只负责建立维度关系和候选信息，不负责完整的切分合法性判断。

2.3 合法性检查

CustomOpBubbleUpStrategy::isSupportedOperation() 分派到isSupportedCustomLike()，依次检查：

isSafeToDuplicateCustomLike()
 → 要求 no_side_effect
 → CustomOp 必须非 builtin
 → temp_buffers 和 extra_buffers_info 必须为空
 → CustomMacroOp 的 sync_event_slots 和 sync_related_args 必须为空

hasSliceableCustomInputs()
 → 逐个检查 CustomOp 的 input
 → ranked tensor 可以生成 tensor.extract_slice
 → scalar/index 等非 shaped input 不切分，原值传给新算子
 → memref 无论是否 ranked 都拒绝；unranked tensor 也拒绝

hasSupportedExplicitIndexingMaps()
 → indexing_map 数量必须等于 input 数量 + output 数量；当前不支持带 temp buffer 的 CustomOp。
 → 每个 map 的输入维数必须等于 iterator_types 数量
 → shaped operand 必须 ranked，map 的输出数必须等于 operand rank
 → map 结果只允许按顺序出现、不重复的 dN，或广播常量 0
 → 拒绝 transpose (d1,d0)、重复维 (d0,d0)、非零常量和复杂 affine 表达式

slicesOnlyParallelDimensions()
 → 比较切分前后 result shape，找出 size 实际变小的 result 维度
 → 通过 result map 找到它对应的迭代维 dN
 → 只有 dN 的 iterator type 为 parallel 才允许切分
 → 如果映射为常量/复杂表达式，或对应 reduction iterator，则拒绝

hasAlignedOutputInits()
 → output init 数量必须等于 result 数量，按物理位置 output0/result0 一一对应
 → 所有 output init 和 result 都必须是 ranked tensor
 → 每个 output init 的 shape 和 indexing map 必须与当前被切 result 一致
 → 只有满足这些条件，才能按 result slice 的物理位置安全创建 output-init slice

任何检查失败都不会进入实际 IR 创建阶段。

2.4 result slice 到 input slice

computeCustomInputSliceParams() 先调用recoverIterSpaceFromResultSlice()，通过 result map 恢复迭代空间切片；再调用 projectIterSpaceToOperand() 投影到各 input。

例如：

iterator_types = [parallel, reduction]
result map = affine_map<(d0, d1) -> (d0)>
input map  = affine_map<(d0, d1) -> (d0, d1)>

result slice 为 offset=[8]、size=[8]、stride=[1]，input 类型为tensor<16x32xf16> 时：

result 第0维 -> 迭代维 d0: offset=8, size=8, stride=1
迭代维 d1   -> result 未绑定

input offsets = [8, 0]
input sizes   = [8, 32]
input strides = [1, 1]

未绑定维度可能是 reduction 或 input-only 维度，当前使用 operand 的完整静态范围。常量 0 表示广播维，也使用完整范围。完整范围为动态大小或表达式无法投影时返回失败。

2.5 IR 重建

CustomOpBubbleUpStrategy::execute() 根据 source 类型调用模板化的executeCustomLike()：

executeCustomLike()
   -> computeCustomInputSliceParams()
   -> collectCustomResultSlices()
   -> createCustomInputSlices()
   -> createCustomOutputSlices()
   -> createTiledCustomLike()
   -> 替换并删除旧 result slices
   -> 删除无其他使用的原 CustomOp

所有可能失败的计算都在创建新 IR 前完成，避免失败后留下半完成改写。

createTiledCustomLike()：

• result 类型采用各 resultSlice 的 tile 类型。

• operands 按 inputs、outputs、temp buffers，以及 CustomMacroOp 的 sync-related arguments 顺序组装。

• 重建 operandSegmentSizes。

• 保留原算子的 symbol、map、iterator、pipe 等属性。

• 添加 tiled_op，防止后续逻辑把它误认为未切分算子并限制到 Sub-Block 0。

2.6 失败回退

如果 marked slice 停在未切分的 CustomOp/CustomMacroOp 上，即使广播场景已经关闭 strict mode，也会中断冒泡并回滚。原因是 tiled Store 只能消费对应的局部 producer，不能安全消费一个语义未知的完整 CustomOp result。回滚后，由 LimitUniqueSubBlockIdToStoreCopy 将完整 CustomOp 和 Store/Copy
限制到 Sub-Block 0，避免两个 Sub-Block 重复执行。广播导致的 effectiveStrictMode 调整只在当前函数内生效，不会影响后续函数。

2.7 代码位置

约束边界

需要重点说明的保守假设：

1. Input 之间不要求 shape/map 相同；每个 tensor input 按自身 map 独立生成slice，然后共同传给一个 tiled CustomOp。

2. 多结果不会按 result 分别创建多个 CustomOp。每个 tile 只创建一个仍包含全部 results 的 CustomOp，避免重复调用外部 kernel。

3. 当前只比较各 result slice 的结果类型，没有在此处直接比较全部offset/size/stride；实现依赖 marked slices 由同一次 Sub-Block 切分生成。若要放宽多结果 map，需要显式反推并校验所有 result 的迭代空间切片一致。

4. builtin CustomOp 属于编译器已登记的特殊库调用，可能具有额外 ABI、访存或资源语义；当前不使用通用 CustomOp 规则切分。

测试：LIT 与 E2E

4.1 LIT：验证 Transform 后的 IR

主测试文件：

bishengir/test/Dialect/HIVM/tile-and-bind-sub-block.mlir

LIT 使用 FileCheck 检查 Transform 后的文本 IR：

• CHECK：内容必须按顺序出现。

• CHECK-SAME：内容必须与上一条正向检查出现在同一行。

• CHECK-NOT：内容不能出现在前后两条正向检查界定的范围中。

• CHECK-LABEL：定位一个新的函数检查区域。

主要覆盖：

聚焦命令：

cmake --build build-pinned-check \--target bishengir-opt bishengir-compile --parallel 16

build-pinned-check/bin/llvm-lit -sv \
 bishengir/test/Dialect/HIVM/tile-and-bind-sub-block.mlir

4.2 E2E：验证真实 NPU 精度

测试项目和最终提交：

/workspace/AscendNPU-IR-DT
branch: customop-tile-and-bind-e2e-refactor
precision baseline: 9239df3
Softmax performance: 04e5ac1 [HIVM] test: add tiled CustomOp Softmax benchmark

测试文件：

src/modules/cube_vector/customop_tile_and_bind/customop_test_kernels.cpp
src/modules/cube_vector/customop_tile_and_bind/runtime.py
src/modules/cube_vector/test_customop_tile_and_bind.py
src/modules/cube_vector/test_customop_softmax_tile_and_bind.py
src/modules/cube_vector/test_customop_softmax_tile_and_bind_perf.py

E2E 链路：

C++ CustomOp kernel -> ccec 编译为 bitcode -> al.register_custom_op
   -> Triton JIT -> AscendNPU-IR -> CANN hivmc
   -> Ascend NPU 执行 -> PyTorch reference 精度比较

原有 7 条 Mix Cube/Vector 精度用例覆盖 identity、broadcast、dim0/dim1 reduction、非 Tensor scalar operand 和多输出。它们是通用 CustomOp 切分的 E2E 精度基线，不属于本次 Softmax 性能测试新增用例。

本次新增 8 条 Softmax 相关精度用例：

• 7 组多 shape Softmax CustomOp，覆盖 Reduce-last 和 Reduce-first。
• 1 条基于简单 Flash Attention 场景的 Softmax CustomOp 正确性测试。

Softmax 使用数值稳定实现：先计算 reduce max，再执行 exp、reduce sum 和 归一化。简单 Flash Attention 用例验证打分矩阵经过 CustomOp Softmax 后的
数值结果；它不是完整融合的 Flash Attention 性能基准。

Shape 和 Reduce 轴覆盖如下：

Shape	Reduce axis	类型
64 × 64	1	Reduce-last
128 × 64	1	Reduce-last
256 × 64	1	Reduce-last
128 × 128	1	Reduce-last
64 × 128	0	Reduce-first
64 × 256	0	Reduce-first
128 × 128	0	Reduce-first

Reduce-last 的 iterator 为 [parallel, reduction]，因此只能沿第 0 维切分；
Reduce-first 的 iterator 为 [reduction, parallel]，因此只能沿第 1 维 切分。两类用例共同验证 Pass 根据 iterator 语义保持 Reduce 轴完整，而不是 固定切分某个物理维度。

精度回归命令：

cd /workspace/AscendNPU-IR-DT
source /usr/local/Ascend/cann/set_env.sh
source /workspace/venv/bin/activate

export ASCENDNPU_IR_ROOT=/workspace/AscendNPU-IR
# Ascend 910B2 使用与 CANN 9.1.0 参数对齐的源码编译器入口：
export CUSTOMOP_TILE_AND_BIND_COMPILER=/workspace/customop-toolchain/bishengir-compile
# Ascend 950PR 改为：
# export CUSTOMOP_TILE_AND_BIND_COMPILER=/workspace/AscendNPU-IR/build/bin/bishengir-compile
export TRITON_ALWAYS_COMPILE=1

python -m pytest -q -vv \
 src/modules/cube_vector/test_customop_tile_and_bind.py \
 src/modules/cube_vector/test_customop_softmax_tile_and_bind.py

Ascend 910B2 和 Ascend 950PR 实测结果：每种设备上通用 CustomOp 7 条和新增 Softmax 8 条共 15 passed。新增 Softmax 文件包含 7 组多 shape 精度用例和 1 条简单 Flash Attention 正确性用例。

4.3 实际切分验证

精度正确不能单独证明发生了 Sub-Block 切分，因为 Pass 保守回退到 Sub-Block 0 后仍可能得到正确结果。因此除 Pytest 外，还保存并检查了 baseline/tiled 的最终 module.hivm.opt.mlir。

tiled IR 的判定条件包括：

• 存在 hivm.hir.get_sub_block_idx。
• 存在沿目标 parallel 轴生成的 memref.subview。
• CustomOp 输入和 DPS output 使用一致的 tile 范围。
• 已切分计算没有被 limit_sub_block_id0 重新限制到 Sub-Block 0。
• 模块没有 hivm.tile_and_bind_subblock_reverted。

例如 shape=128×128, reduce_axis=0 的 tiled IR 中，CustomOp 输入与输出均 生成 [128, 64] subview。这说明 Reduce 轴第 0 维保持 128 不变，实际沿非 Reduce 的第 1 维执行 1:2 切分。Reduce-last 用例则生成 [64, 128] tile，保持第 1 维 Reduce 范围完整。

4.4 性能测试方法与结果

性能测试严格保持 baseline 与 tiled 的输入、shape、dtype、launch grid、Triton kernel 和 CustomOp bitcode 一致，唯一变量是--enable-auto-bind-sub-block。同时使用独立 compile variant 隔离 JIT缓存，并设置 TRITON_ALWAYS_COMPILE=1，避免复用其他编译器生成的缓存。

计时前先验证 baseline 与 tiled 均非全零并与 PyTorch reference 一致；默认预热 5 次、测量 30 次，baseline/tiled 交替执行，使用 NPU Event 计时，裁剪两端各 10% 样本后取中位数，并报告 std、min 和 max。默认门槛为单个shape 不低于 0.98x，几何平均不低于 1.05x。

两种设备均使用 CANN 9.1.0 完成测试：

设备	SoC/编译目标	CustomOp bitcode	最终编译后端
Ascend 910B2	Ascend910_9382	c220	hivmc 0.3.0
Ascend 950PR	Ascend950PR_9579 / dav-c310-vec	c310	CANN 9.1.0 hivmc-a5 0.3.0

Ascend 950PR 使用源码构建的 bishengir-compile，并使用与
CANN 9.1.0 配套的 hivmc-a5 完成设备二进制编译。
Triton-Ascend wheel 自带的 hivmc-a5 与当前 CANN 后端不是同一
二进制；混用时 baseline Softmax 输出全零。将编译链统一到
CANN 9.1.0 hivmc-a5 后，同一用例和完整回归通过。

Ascend 910B2

Shape	Reduce axis	Baseline	Tiled	Speedup
128 × 64	1	3.024 ms	1.516 ms	1.995x
256 × 64	1	6.043 ms	3.026 ms	1.997x
128 × 128	1	6.035 ms	3.021 ms	1.997x
64 × 128	0	2.728 ms	1.368 ms	1.994x
64 × 256	0	5.451 ms	2.730 ms	1.997x
128 × 128	0	5.442 ms	2.726 ms	1.997x

6 组 shape 的加速比为 1.994x–1.997x，几何平均约 1.996x；各组
std 约为 0.000–0.001 ms。

Ascend 950PR

Shape	Reduce axis	Baseline	Tiled	Speedup
128 × 64	1	0.806 ms	0.404 ms	1.997x
256 × 64	1	1.611 ms	0.806 ms	1.998x
128 × 128	1	1.609 ms	0.805 ms	1.998x
64 × 128	0	0.750 ms	0.376 ms	1.997x
64 × 256	0	1.499 ms	0.750 ms	1.998x
128 × 128	0	1.498 ms	0.750 ms	1.998x

6 组 shape 的加速比为 1.997x–1.998x，几何平均约 1.998x；各组std 为 0.000 ms（按当前输出精度显示）。

两种设备上的 Reduce-first 与 Reduce-last 都获得稳定的、接近1:2 理论并行度的收益。

性能测试命令：

cd /workspace/AscendNPU-IR-DT
source /usr/local/Ascend/cann/set_env.sh
source /workspace/venv/bin/activate

export ASCENDNPU_IR_ROOT=/workspace/AscendNPU-IR
# Ascend 910B2 使用与 CANN 9.1.0 参数对齐的源码编译器入口：
export CUSTOMOP_TILE_AND_BIND_COMPILER=/workspace/customop-toolchain/bishengir-compile
# Ascend 950PR 改为：
# export CUSTOMOP_TILE_AND_BIND_COMPILER=/workspace/AscendNPU-IR/build/bin/bishengir-compile
export TRITON_ALWAYS_COMPILE=1

python -m pytest -q -s \
 src/modules/cube_vector/test_customop_softmax_tile_and_bind_perf.py

4.5 各类测试分别证明什么

LIT
 -> 证明 Pass 生成预期的 tile-local IR 和安全回退结构

E2E precision
 -> 证明生成代码在真实 NPU 上的 shape、dtype 和数值结果正确

Saved final IR
 -> 证明测试运行时确实沿非 Reduce 轴完成切分，而非保守回退

E2E performance
 -> 证明只打开 TileAndBindSubBlock 后具有稳定的实际性能收益

四类证据共同构成闭环。只验证精度不能证明实际发生了切分；只检查 IR 也不能证明设备端数值和性能正确。

4.6 设备覆盖与环境限制

原始需求要求在 Ascend 910 和 Ascend 950 上验证。当前已完成：

• Ascend 910B2：15 passed；6 组性能 case 加速比为
  1.994x–1.997x。
• Ascend 950PR：15 passed；6 组性能 case 加速比为
  1.997x–1.998x。

910B2 上还保存并检查了 baseline/tiled 的最终module.hivm.opt.mlir，确认沿非 Reduce 轴实际生成 1:2 subview，且没有回退到 Sub-Block 0。950PR 上通过完整精度回归和baseline/tiled 性能对比验证 c310/A5 路径。

Ascend 950PR 测试前后 npu-smi 均显示 Health: Alarm。测试时设备无其他运行进程，温度为 58–59℃，各组计时标准差接近 0。该健康状态作为性能测试环境限制记录，不影响本次正确性通过的结论。

总结

本任务已建立 CustomOp/CustomMacroOp 的保守 Sub-Block 切分闭环：

维度分析
 -> 显式语义和安全性检查
 -> result slice 反推迭代空间
 -> 投影 input slices，并按 result slice 创建 output-init slices
 -> 重建 tile-local CustomOp
 -> 成功继续冒泡，失败回滚到 Sub-Block 0

设计原则是“只有能够由显式语义证明安全时才切分”。当前覆盖 identity、broadcast、reduction 输入、标量和受限多结果场景；动态范围、memref、复杂map、builtin、资源/同步语义和叶子 CustomOp 保守回退。

新增 Softmax 测试进一步验证了 reduction 语义：Reduce-first 和 Reduce-last 均保持 Reduce 轴完整，并沿 parallel 轴生成实际的 1:2 subview。Ascend 910B2 和 Ascend 950PR 上的 7 组 Softmax shape 和 1 条简单 Flash Attention 正确性用例全部通过。910B2 上 6 组性能 case 的加速比为 1.994x–1.997x，几何平均约 1.996x；950PR 上为 1.997x–1.998x，几何平均约 1.998x。

因此，已经形成“LIT 结构检查、E2E 精度、最终 IR 切分证据、E2E 性能”的完整验证链，并完成 Ascend 910B2 和 Ascend 950PR 两种设备的真机正确性与性能验证。

@wcleungaj，您好，想跟进一下上上次 SIG 周会汇报后的代码 Review 进展，也想和各位老师进一步确认本次 PR 的合并范围。

目前已完成 CustomOp/CustomMacroOp 的 Sub-Block 切分支持，并在 Ascend 910B2 和 950PR 上完成 Softmax 多 shape 精度与性能验证，相关 CI 也已通过。

当前实现采用保守切分策略，对于能够根据显式语义证明安全的场景进行切分；对于 leaf CustomOp、复杂 indexing map、动态范围等暂未覆盖的场景，则保守回退。

想请教各位老师，当前实现是否可以作为本次 PR 的合并范围？对于尚未覆盖的场景，是否可以在本次 PR 完成 Review 并合并后，再通过后续 PR 逐步补充完善？

如果当前版本还有需要在合并前解决的问题，也请各位老师指出，我这边会继续配合修改。

辛苦各位老师了，感谢！

@sun_teng 我認為當前PR的代碼量已經很多，對於未覆蓋的場景可以在後續的PR中完成
架構師最近比較忙... 我們會儘快跟進社區任務的審視和合入