已合并
feat: 完善 IndirectLoad SIMD 和 SIMT 支持 #1542
xiebangrui2025创建于 7月28日
feat: 完善 IndirectLoad SIMD 和 SIMT 支持 #1542
已合并
xiebangrui2025创建于 7月28日
xiebangrui2025
xiebangrui2025成员
7月28日

描述

完善 IndirectLoad 的 SIMD/SIMT 调度与代码生成链路:新增独立 regbase 实现,补充模板角色、轴与 direct-GM 处理,并扩展 Cast、BF16、UINT32 等 E2E 覆盖。同步补齐统一测试脚本中的全部 IndirectLoad E2E target。

变更类型

关联的Issue

无。

如何测试

  1. optimize_ut --gtest_filter='*IndirectLoad*':12/12 通过;Python IndirectLoad UT/ST 各 1/1 通过。
  2. 全部 20 个 IndirectLoad codegen ST 和 20 个 generated kernel E2E 通过;generated tiling/kernel 编译通过,日志无非预期错误。
  3. 目标文件 pre-commit、clang-format、codespell 和 staged OAT 检查通过。

核对清单

其他信息

SIMD 路径使用 Gather offset 表,SIMT 路径使用 direct-GM scalar transform;新增 BF16+INT64、UINT32+INT32 以及 Cast 穿过 IndirectLoad 的数值 E2E。

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 xiebangrui2025 的贡献)
xiebangrui2025xiebangrui2025成员
7月28日 创建了 pull request,commit 48309d33
atomgit-bot
atomgit-bot
7月28日 评论:

变更摘要

此 PR 主要完善了 IndirectLoad 的 SIMD 和 SIMT 调度与代码生成链路。新增了独立的 IndirectLoadSimd 和 IndirectLoadSimt 模板 regbase 实现(以 Gather offset 表和 direct-GM scalar transform 分别驱动两条路径),并将 IsSimtScalarSupported / GenerateSimtScalarExpr 从基类 AscIrCodegen 下沉到 AscIrCodegenV2 并大规模扩展了 Cast、Abs、各类 Bessel/三角函数等 op 的 SIMT 标量支持。调度侧重构了模板角色行为(引入 excludes_tiling_group),SIMD 路径也支持将 input pre 算子移到 IndirectLoad 之后,同时补全了 BF16+INT64、UINT32+INT32、PrecisionCast 穿过等数值组合的 E2E 覆盖。

主要改动

  • 新增独立 regbase 模板实现 indirect_load_simd.h / indirect_load_simt.h:SIMD 路径通过编译期递归 IndirectLoadSimdInnerOffset / IndirectLoadSimdInnerSize 计算 gather offset 表并用 Gather 完成查表取数;SIMT 路径通过 IndirectLoadSimtKernel 模板内核 + IndirectLoadSimtOffset 递归偏移计算,配合用户提供的 IndexTransform / OutputTransform,以 VF_CALL 发射到 direct-GM scalar 管线执行。

  • IsSimtScalarSupported / GenerateSimtScalarExpr 下沉至 AscIrCodegenV2 并大规模扩展:基类 AscIrCodegen 删除了这两个虚函数,移到 AscIrCodegenV2;新增 SimtFloatUnaryAscIrCodegenImplV2 辅助基类,使 Floor、ModifiedBesselI0/I1、BesselJ0/J1/Y0/Y1、Erf/Erfc/Erfcx、Ceil、Cos/Acos/Cosh/Acosh、Sin/Sinh、Tan/Tanh、Asin/Asinh、Atan/Atanh、Ndtr、Erfinv、Round、Expm、Log2/Log10/Log1p、Rsqrt、Lgamma、Trunc 等二十余个 op 继承并获得 SIMT 标量发射能力;Cast 的 IsSimtScalarSupported 扩展了 40+ 对 dtype 组合,Abs 增加了整数/UINT8/BF16 支持,Square/Sign/Isnan/IsFinite/IsInf/LogicalNot/Sigmoid/Reciprocal/Ceil2Int/FloorToInt/RoundToInt/TruncToInt 等也补齐了对应逻辑。

  • 调度侧模板行为与 input-pre 移动重构:TemplateBehavior 新增 excludes_tiling_group 字段,kSimdInputPre 角色改用该字段替代 skips_main_schedule_tiling + skips_api_emit;ApplyIndirectLoadTemplateMerge 增加 merge_tensor_axis 参数,input-inner vectorization 路径不再合并 tensor axis;新增 MoveSimdInputPreNodes / CollectMovableSimdInputPreNodes,使 SIMD 路径也能将可移动的逐元素前处理结点后移到 IndirectLoad 之后;tiling_group.cpp 中 excludes_tiling_group 的结点被跳过不参与 tiling group 生成。

  • IndirectLoad 代码生成路径从内联发射改为调用 regbase 模板函数:reg_indirect_load_api_call.cpp 大幅简化——SIMD GenerateSimd 中移除了原有的 window setup / input window / gather 三段内联代码,改为调用 AscendC::IndirectLoadSimd<...>(...);SIMT GenerateSimt 中移除了自定义 IndirectLoadSimtKernel_* 内核定义与 EmitSimtLaunch,改为调用 AscendC::IndirectLoadSimt<...>(...);IndirectLoadAscIrCodegenImplV2::CalcTmpBufSize 的临时缓冲区计算简化为仅需 offset 表大小。

  • 测试基础设施与 E2E 覆盖扩展:share_graph.h/cc 新增 IndirectLoadInputPreType(None/Relu/Exp2/ReluExp2/ReluRelu/PrecisionCast)和 IndirectLoadOutputPostType(Default/ModifiedBesselK0/DirectStore/CastExp2CastBack)枚举,IndirectLoadStoreFusedGraph 接口扩展为支持 data_type、index_type、输入/输出前后处理类型及静态 shape;统一测试脚本从 6 个 E2E target 扩展到 20 个,覆盖 rank2/3/4、负轴、Relu+Exp2 pre、BF16+INT64、UINT32+INT32、PrecisionCast、ModifiedBesselK0、large tile 等组合。

likedislike
不准确?
atomgit-bot
atomgit-bot
7月28日 评论:

代码审查

审查总结

共审查 23 个变更文件,发现 2 个问题:

优先级 数量
P0 0
P1 0
P2 1
P3 1

整体风险判定:低风险。 主要变更是 IndirectLoad SIMD/SIMT 链路的重构与完善,将内联代码发射迁移到独立的 AscendC regbase 头文件,新增了 Cast/BF16/UINT32 等多类型 E2E 覆盖,并统一了测试脚本中的全部 IndirectLoad target。变更体量大但逻辑自洽,测试覆盖充分(20 个 codegen ST + 20 个 kernel E2E 全部通过)。仅存在一个 P2 的整数收窄转换隐患和一处 P3 的同步模式待确认项,均为非阻塞性建议。

各文件审查结果

文件 结果
autofuse/common/indirect_load_utils.cpp 无问题
autofuse/common/indirect_load_utils.h 无问题
autofuse/inc/graph_metadef/graph/ascendc_ir/ascir_registry.h 无问题
autofuse/optimize/autoschedule/schedule.cpp 无问题
autofuse/optimize/autoschedule/tiling_group.cpp 无问题
autofuse/optimize/task_generator/indirect_load_schedule_case_generator.cpp 无问题
autofuse/tests/framework/share_graph/include/share_graph.h 无问题
autofuse/tests/framework/share_graph/src/share_graph.cc 无问题
autofuse/tests/ut/optimize/task_generator/test_indirect_load_schedule_case_generator.cpp 无问题
autofuse/tests/v35/st/backend_e2e_v2/indirect_load_store_test/CMakeLists.txt 无问题
autofuse/tests/v35/st/backend_e2e_v2/indirect_load_store_test/indirect_load_store_backend_generator.cpp 无问题
autofuse/tests/v35/st/backend_e2e_v2/indirect_load_store_test/test_e2e_indirect_load_store_kernel.cpp 无问题
autofuse/v35/ascendc/api_regbase/CMakeLists.txt 无问题
autofuse/v35/ascendc/api_regbase/indirect_load_simd.h P2:int64→uint32 收窄
autofuse/v35/ascendc/api_regbase/indirect_load_simt.h P3:同步模式确认
autofuse/v35/ascir/ascir_codegen_v2.h 无问题
autofuse/v35/ascir/generator/v2_ascir_codegen_impl.h 无问题
autofuse/v35/codegen/ascendc_reg_base_api_register.cpp 无问题
autofuse/v35/codegen/reg_api_call/reg_indirect_load_api_call.cpp 无问题
autofuse/v35/codegen/reg_api_call/reg_indirect_load_api_call.h 无问题
autofuse/v35/codegen/simt_scalar_call/simt_scalar_emitter.cpp 无问题
autofuse/v35/optimize/partition/vector_func_partitioner.cpp 无问题
scripts/test/run_autofuse_test.sh 无问题

⚠️ 已识别出整体风险,但无法提取行内评论,请参考整体评估。

likedislike
不准确?
CANN-robotCANN-robot成员
7月28日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
7月28日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here。
You can get sig-info at here


PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
repo-cann/graph-autofusion ✅ zhanj, liyuewei, zhang_shengjie (3/2) ✅ zhanj (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

xiebangrui2025, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
此处折叠了41条消息 查看更多
CANN-robotCANN-robot成员
7月29日 删除了label:ci-pipeline-running
CANN-robotCANN-robot成员
7月29日 添加了label:ci-pipeline-passed
zhanj
zhanj成员
7月29日 评论:

/approve

likedislike
CANN-robotCANN-robot成员
7月29日 添加了label:approved
CANN-robotCANN-robot成员
7月29日 合入了pull request