已合并
[HFusion] Unify expand_shape tiling with row-major flattening #2041
[HFusion] Unify expand_shape tiling with row-major flattening #2041
已合并
hujiajun创建于 7月29日
hujiajun成员
7月29日

描述 Description

将多分支的 getCollapsedSliceParameters 替换为单一算法:对每个 collapsed 维度的 expanded 组,若存在一个 main dim,其之前仅有 unit tile 或 shape-1 维度、之后仅有全覆盖维度,则该 expanded 空间 上的 tile 可映射为连续的 collapsed slice。collapsed offset 为 expanded offset 的行主序展平;collapsed size 为 main dim 的 tile size 乘以尾部维度乘积。

现支持动态 tile size 和动态 expanded shap,offset 和 size 使用 OpFoldResult 算术并自动常量折叠。

新增 3 个测例:shape-1 前缀、动态 tile size、动态 shape。

变更背景

多 consumer 场景触发:

功能缺陷:合法的 tile pattern 被错误拒绝。被拒后 auto-vectorize-v2 中的 FuseIntoContainingOp 退化为 cloneAndFuse,将未 tiled 的完整 expand_shape 克隆进 loop body。该克隆 op 是循环不变量,被 LICM 提至 loop 外部,导致后续其他 producer 对该 loop 的 FuseIntoContainingOp 操作在 loop 内找不到正确的 user。

性能回退:clone+fuse 的 expand_shape 每次迭代计算完整结果(而非 tiled 子区域),无 tiling 收益。

关联 issue:https://gitcode.com/Ascend/AscendNPU-IR/issues/324

类型 Category

Checklist

likedislike
Pull Request已成功合入, 合并人@gguo
(感谢 hujiajun 的贡献)
Hhujiajun成员
7月29日 创建了 pull request,commit 75d6ea4a
Hhujiajun成员
7月29日 关联了issue:[Feature]: 多consumer实施方案
atomgit-bot
atomgit-bot
7月29日 评论:

变更摘要

该 PR 将 getCollapsedSliceParameters 函数从多分支的 cases 判断重构为统一的行主序展平算法:对每个 collapsed 维度对应的 expanded 组,若满足 [unit || shape==1]* [mainDim] [NoTile]* 模式(即前缀维度均为单位 tile 或 shape-1、后缀维度均为全覆盖),则将 expanded 空间上的 tile 通过行主序展平映射为 collapsed 空间上连续的 slice。同时将 isNoTileisUnitTileisValidTile 等辅助函数简化为接受标量 shape 参数,并引入 OpFoldResult 算术运算以支持动态 tile size 和动态 expanded shape,具备自动常量折叠能力。新增 3 个测例覆盖 shape-1 前缀、动态 tile size 和动态 shape 场景。

主要改动

  • 重构 getCollapsedSliceParameters 为统一行主序展平算法:移除了原函数中 Case 0a/0b/1/2/3/4 的多分支逻辑(约 100 行),替换为基于 mustUnitTile/mustNoTile 的单一模式匹配流程,按 [unit|shape==1]* [mainDim] [NoTile]* 识别可映射的 tile 模式,并通过行主序累乘累加计算 collapsed offset 和 size。

  • 简化辅助函数签名isNoTileisUnitTileisValidTile 的 shape 参数从 ArrayRef<int64_t> 改为 int64_t 标量,移除了内部对数组索引的访问;isUnitTile 返回值统一为布尔表达式,去除了冗余的 if-else 分支。

  • 引入 OpFoldResult 算术运算支持动态值:在新函数中新增 toValuemuladd 三个 lambda,将 OpFoldResult 转换为 Value 并执行 arith::MulIOp/arith::AddIOp,同时自动对两个静态操作数进行常量折叠输出 IndexAttr,使 offset 和 size 的计算支持动态 tile size 与动态 expanded shape。

  • 简化调用方传参ExpandShapeOpTiling 中对 getCollapsedSliceParameters 的调用从传入 associationsexpandedShapecollapsedValue 三个独立参数简化为直接传入 expandedSliceParamsexpandShapeOp,函数内部自行获取所需信息。

  • 新增测试文件 transform-op-tile-expand-shape-advanced.mlir:包含三个 CHECK 测例——shape-1 前缀(@shape1_prefix)、动态 tile size(@dynamic_tile)、动态 expanded shape(@dynamic_shape),均验证 fuse_into_containing_op 能将 expand_shape 成功融合进 loop body 并生成 tensor.extract_slice + tensor.expand_shape

likedislike
不准确?
ascend-robot
ascend-robot成员
7月29日 评论:

libing-pipeline创建中,请等待……

likedislike
atomgit-bot
atomgit-bot
7月29日 评论:

代码审查

我已完成对两个变更文件的全面审查。以下是总结。


审查总结

优先级 数量 说明
P1 1 "all unit" 分支 collapsed offset 被错误硬编码为 0
P3 1 未使用的 structured binding collapsedIdx

已审查的文件:

  • bishengir/lib/Dialect/Tensor/Transforms/TilingInterfaceImpl.cpp — 发现 2 个问题(1×P1 + 1×P3)
  • bishengir/test/Dialect/HFusion/TransformOps/transform-op-tile-expand-shape-advanced.mlir — 无问题(测试用例结构正确,CHECK 模式匹配合理)

整体风险评估:中等。 P1 的 "all unit" 分支 bug 在特定 tile pattern(collapsed group 中所有 expanded dims 的 tile size 均为 1 或 shape 为 1,且存在 shape > 1 的 dim 具有非零 offset)下会导致 fused expand_shape 从错误的 collapsed offset 读取数据。这是一个相对于旧实现的正确性回退(旧代码的 Case 2 正确传递了该 offset)。新增的 3 个测例均未覆盖该触发条件,建议在修复 P1 bug 时同时补充针对该边界情况的回归测试。

类型 数量
🔴 阻塞 1
🟡 建议 0

⛔ 需要修改

likedislike
不准确?
此处折叠了171条消息 查看更多
NNoBBAI成员
7月30日 解决了最后一个问题
lingfeiwu
lingfeiwu成员
7月31日 评论:

/lgtm
/approve

likedislike
ascend-robotascend-robot成员
7月31日 添加了label:approved
gguogguo成员
7月31日 通过测试
gguogguo成员
7月31日 合入了pull request,合并节点 SHA:61a5b15c2106d5aeecb0ef5955227304b23814d7