Fixes #256

描述 Description

本 PR 让满足安全切分条件的 hivm.hir.custom 和hivm.hir.custom_macro 参与 TileAndBindSubBlock 的 sub-block 1:2切分。

切分后，CustomOp、CustomMacroOp 及其下游 Store 使用一致的 tile 范围，使 Vector/AIV 计算能够在两个 sub-block 上并行执行。

两个算子共享 HIVM_CustomOp 基类和HIVMStructuredOpInterface。本实现将维度分析和 slice bubble-up 抽象为 Custom-like 通用逻辑；对于 CustomMacroOp 的同步语义，以及当前尚无安全锚点的 leaf 场景，保持保守限制。

本次补充修复以下两个由 Softmax CustomOp E2E 验证发现的问题：

• no_side_effect CustomOp 在 bufferization 后仍会写入声明的 DPS
  output，不能被 canonicalization 或 DCE 当作无写入操作删除。
• Softmax 的结果与输入同 rank，但仍包含 reduction 维。维度分析必须
  根据 iterator_types 和 indexing_map 标记 reduction 维，确保
  TileAndBindSubBlock 只切分 parallel 维。

核心改动

1. Custom-like 维度分析

新增统一的 Custom-like 维度分析逻辑，处理 CustomOp 和CustomMacroOp：

• 未提供完整的 iterator_types 和 indexing_map 时，不建立输入、
  输出和结果之间的可切分维度关系，并保守回退。
• 提供显式 indexing map 时，根据 iterator 类型与 affine dimension
  expression 建立输入、输出和结果之间的维度关系。
• broadcast 常量维不参与普通维度合并。
• 非 Tensor scalar operand 不参与维度分析，切分时保持原 SSA value
  不变。
• 根据 iterator_types 和各 operand/result 的 indexing map 显式标记
  reduction 维。
• reduction 维不会被误选为 parallel 切分维。
• 对于 Softmax 这类输入和结果同 rank 的算子，保留完整 reduction
  范围，只对 parallel 维执行 1:2 切分。
• 对 unranked shaped value 保守回退，避免维度分析断言失败。

2. Slice bubble-up

新增 CustomOpBubbleUpStrategy，根据结果 tensor.extract_slice 计算 Custom-like 算子的 tile-local 输入和输出：

• 要求同时提供 iterator_types 和 indexing_map；任一属性缺失时
  保守回退。
• 支持显式 parallel/reduction iterator。
• 支持 identity 和 broadcast indexing map。
• 支持非 Tensor scalar operand 原值透传。
• 支持多个结果在相同 iteration space 下同步切分。
• 在完成全部合法性检查后才创建新 IR，避免 rewrite 失败后遗留部分
  构造的 slice。
• 重新构造 CustomOp/CustomMacroOp 时显式更新
  operandSegmentSizes。

3. Reduction 轴切分规则

对于包含 reduction iterator 的 CustomOp：

• reduction 维保持完整，不参与 sub-block 1:2 切分；
• parallel 维可以在满足条件时切分；
• 输入、DPS output init、结果和下游 Store 使用一致的 tile 范围。

Softmax 覆盖两种典型布局：

reduce-last:
  iterator_types = [parallel, reduction]
  仅切分 dim0，保持 dim1 完整

reduce-first:
  iterator_types = [reduction, parallel]
  保持 dim0 完整，仅切分 dim1

4. no_side_effect 和 DPS output 语义

no_side_effect 表示 CustomOp 没有未声明的隐藏副作用，不表示其声明的 DPS output 是只读的。

tensor CustomOp 完成 bufferization 后不再依赖 SSA result 保持存活。如果丢失 output write effect，canonicalization 或 DCE 可能删除CustomOp library call，导致 Store 读取未初始化的 output buffer。

本实现对 CustomOp 和 CustomMacroOp 保留保守的 resource-levelRead/Write effect，从而保证：

• bufferized CustomOp 不会被错误删除；
• CustomOp 在下游 Store 之前执行；
• DPS output 与 Store consumer 保持正确顺序；
• no_side_effect CustomOp 的声明输出仍被视为有效写入。

设备端 A/B 验证表明，移除该 effect 后 Softmax CustomOp 会被删除，baseline 输出全部为 0；恢复后精度与性能测试均通过。

5. 保守回退条件

当前仅对能够证明安全的场景执行切分。以下情况会保守回退，保持Custom-like 算子全量执行，并在需要时通过 sub-block 保护避免重复
执行：

• 算子具有未声明副作用。
• 使用 temp buffer 或 extra buffer。
• CustomMacroOp 携带 sync_event_slots 或 sync_related_args。
• 输入包含 memref。
• 使用动态 slice。
• iterator_types 或 indexing map 缺失、不完整。
• indexing map 包含 transpose、重复维度或其他暂不支持的 affine
  expression。
• 待切分 slice 涉及 reduction iterator。
• DPS output init 与对应结果的 shape 或 indexing map 不一致。
• 输入或输出包含无法安全分析的 unranked shaped value。

这些限制用于保证当前 1:2 切分的执行安全。后续只有在明确 per-sub-block buffer ownership、事件同步和更一般的 indexing semantics 后，才会进一步放宽。

6. Leaf Custom-like 行为

当前不支持没有 Store/Copy 等下游切分锚点的 leafCustomOp/CustomMacroOp 切分。

此类算子保持全量执行；如果不能证明重复执行安全，则限制在sub-block 0 执行。

7. MIX kernel 符号和 sub-block 执行

• 保留 MIX kernel 中用户定义 CustomOp 引用的函数符号，避免 lowering
  或 DCE 阶段错误删除。
• 避免已完成切分的 CustomOp 被继续限制到 sub-block 0。
• 保留 AIV 函数参数，确保 CustomOp 在 MIX Cube/Vector 链路中的
  operand 正确传递。
• 保证 CustomOp、tile-local output 和下游 Store 使用相同的
  sub-block 范围。

测试验证 Testing

1. 主仓 LIT：CustomOp 切分结构

主要测试文件：

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

CustomOp 覆盖：

• 空 iterator_types/indexing_map 的保守回退。
• 显式 parallel identity indexing map。
• parallel/reduction iterator 和 reduction 维保护。
• reduction-first 同 rank 输入/输出。
• broadcast indexing map。
• 非 Tensor scalar operand 原值透传。
• 多结果同步切分。
• 与其他 elementwise op 组成的计算链。
• 多个独立 CustomOp。
• Store 与 CustomOp 链共存。
• leaf、side-effect、temp/extra buffer、memref input、builtin、
  transpose map 和 too-small shape 等保守回退场景。

CustomMacroOp 覆盖：

• 无显式 iterator/indexing map 时的全量执行与 sub-block 0 保护。
• sync_event_slots 场景的全量执行与 sub-block 0 保护。
• 多结果 Macro 的 tile-local 输入、输出和结果。
• 两个 Store 分别消费 #0、#1 结果。

执行命令：

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

结果：

CustomOp/CustomMacroOp LIT: 1/1 passed

2. 主仓 LIT：memory effect 和 bufferization

新增或扩展：

bishengir/test/Dialect/HIVM/IR/custom-op-memory-effects.mlir
bishengir/test/Dialect/HIVM/custom-op-bufferize-inplace.mlir

验证：

• no_side_effect memref CustomOp 不会被 canonicalization 删除。
• tensor CustomOp 在 bufferization 后仍保留 output write。
• lowering 后仍存在 CustomOp library call。
• CustomOp result 与下游 Store 使用同一个 buffer。
• DPS output init 保持 writable 语义。
• CustomOp call 在 Store call 之前执行。

3. 主仓完整回归

执行：

LIT_OPTS="--timeout=30 -j 4" \
  ninja -C build check-bishengir

LIT_OPTS="--timeout=30 -j 4" \
  ninja -C build check-mlir

结果：

check-bishengir:
  Total Discovered Tests: 987
  Passed: 843
  Unsupported: 143
  Expectedly Failed: 1
  Failed: 0

check-mlir:
  Total Discovered Tests: 2895
  Excluded: 1
  Unsupported: 608
  Passed: 2285
  Expectedly Failed: 1
  Failed: 0

端到端精度与性能测试

端到端测试按照社区规范放置在 AscendNPU-IR-DT仓库。

• DT 测试 PR：请填写 AscendNPU-IR-DT PR 链接
• AscendNPU-IR 测试提交：37e32605
• CANN：9.1.0

测试设备和编译后端：

设备	SoC/编译目标	CustomOp bitcode	后端编译器
Ascend 910B2	Ascend910_9382	c220	hivmc
Ascend 950PR	Ascend950PR_9579	c310	hivmc-a5

950 环境使用源码构建的 bishengir-compile，并使用与 CANN 9.1.0 配套的 hivmc-a5 完成设备二进制编译。

1. 通用 CustomOp 精度

通过 Triton-Ascend al.custom 构造真实 Mix Cube/Vector 链路：

tl.dot → al.custom → tl.store

覆盖 7 条 E2E 精度用例：

• 二维 add。
• dim0 broadcast。
• dim1 broadcast。
• 非 Tensor scalar operand。
• dim1 reduction。
• dim0 reduction。
• 双输出 CustomOp。

覆盖 identity、broadcast、parallel/reduction、多输出，以及
f16、bf16、f32。

2. Softmax CustomOp 精度

新增基于数值稳定实现的 Softmax CustomOp：

softmax(x * scale) =
    exp(x * scale - max(x * scale)) /
    sum(exp(x * scale - max(x * scale)))

覆盖以下 shape 和 reduction 轴：

Shape	Reduction axis	类型
64x64	1	reduce-last
128x64	1	reduce-last
256x64	1	reduce-last
128x128	1	reduce-last
64x128	0	reduce-first
64x256	0	reduce-first
128x128	0	reduce-first

同时比较 baseline、tiled 和 PyTorch Reference，并验证沿 reduction轴的概率和接近 1。

上述用例已分别在 Ascend 910B2 和 Ascend 950PR 上通过。

3. 简化 Flash Attention

测试流程为：

QK^T → Softmax CustomOp → Store → PV

编译后的 Mix kernel 包含：

QK^T → Softmax CustomOp → Store

最终 PV 在 kernel 外执行，用于验证 Softmax 输出进入 Attention后续计算后的最终精度，不额外引入与本 PR 无关的CustomOp-to-Cube 切分规则。

该用例已分别在 Ascend 910B2 和 Ascend 950PR 上通过。

4. 性能测试

比较同一个 Mix kernel 的两种编译配置：

baseline: enable_auto_bind_sub_block=False
tiled:    enable_auto_bind_sub_block=True

计时范围：

tl.dot → Softmax CustomOp → tl.store

测试方法：

• 5 次预热；
• 默认采集 30 组样本；
• baseline 和 tiled 交替执行；
• 使用 NPU Event 记录设备时间；
• 裁剪最低和最高各 10% 的样本；
• 使用中位数计算 speedup；
• 计时前验证 baseline 和 tiled 精度；
• 单 shape 默认性能下限为 0.98x；
• 几何平均性能下限为 1.05x。

Ascend 910B2

Shape	Axis	Baseline	Tiled	Speedup
128x64	1	3.024 ms	1.516 ms	1.995x
256x64	1	6.043 ms	3.026 ms	1.997x
128x128	1	6.035 ms	3.021 ms	1.997x
64x128	0	2.728 ms	1.368 ms	1.994x
64x256	0	5.451 ms	2.730 ms	1.997x
128x128	0	5.442 ms	2.726 ms	1.997x

性能收益范围为 1.994x–1.997x，几何平均约为 1.996x。

Ascend 950PR

Shape	Axis	Baseline	Tiled	Speedup
128x64	1	0.806 ms	0.404 ms	1.997x
256x64	1	1.611 ms	0.806 ms	1.998x
128x128	1	1.609 ms	0.805 ms	1.998x
64x128	0	0.750 ms	0.376 ms	1.997x
64x256	0	1.499 ms	0.750 ms	1.998x
128x128	0	1.498 ms	0.750 ms	1.998x

性能收益范围为 1.997x–1.998x，几何平均约为 1.998x。

两种设备上的 reduce-first 和 reduce-last 用例均获得接近 2 倍的性能收益，表明 TileAndBindSubBlock 保持 reduction 维完整并切分parallel 维后，两个 Vector sub-block 能够有效并行执行。

5. DT 测试结果

精度测试命令：

python3 -m pytest -v \
  src/modules/cube_vector/test_customop_tile_and_bind.py \
  src/modules/cube_vector/test_customop_softmax_tile_and_bind.py

每种设备覆盖：

• 7 条通用 CustomOp E2E 精度用例；
• 7 条 Softmax 多 shape 精度用例；
• 1 条简化 Flash Attention 精度用例。

精度测试结果：

Ascend 910B2: 15 passed
Ascend 950PR: 15 passed

性能测试命令：

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

性能测试结果：

Ascend 910B2:
  6 个 shape 全部通过
  speedup: 1.994x–1.997x

Ascend 950PR:
  6 个 shape 全部通过
  speedup: 1.997x–1.998x

测试职责边界

主仓 LIT：
  验证 CustomOp/CustomMacroOp 的 1:2 切分 IR；
  验证 reduction 维保持完整、parallel 维被切分；
  验证保守回退和 sub-block 保护；
  验证 no_side_effect CustomOp 的 bufferization 和执行顺序。

DT Pytest：
  验证真实 Mix Cube/Vector 场景的 NPU 最终精度；
  验证 reduce-first 和 reduce-last Softmax；
  验证简化 Flash Attention；
  验证 baseline 与 tiled kernel 的性能收益。

CustomMacroOp 当前没有 Triton 前端接口，因此设备端 E2E 不覆盖
CustomMacroOp；其切分结构、同步限制和回退行为由主仓 LIT 验证。

测试环境说明

Ascend 950PR 测试期间，npu-smi 在测试前后均显示Health: Alarm。设备无其他运行进程，温度保持在 58–59℃，各组计时标准差接近 0，性能结果稳定。该设备健康状态作为测试环境限制记录。

类型 Category

• Bug修复 Bugfix
• 新功能 New feature
• 代码风格更新（格式化，局部变量）Style changes
• 重构（即不是新增功能，也不是修改bug的代码变动）Refactor
• 构建过程或辅助工具的变动 Build-related
• 文档内容更新 Docs

Checklist

• 我的代码遵循这个项目的代码风格 My code follows the project's code style
• 我已经自己测试过我的代码 I have self-tested my code
• 我已经更新了相应的文档 I have updated the relevant documentation
• 我已经根据需要更新了对应的变更日志 I have updated the changelog as needed
• 我已经在标题中正确使用了类型标签（例如：feat:, fix:）I have correctly used type labels in the title (e.g., feat:, fix:)