| 【fix】: 解决memset 问题
Co-authored-by: lidaoming1<lidaoming1@huawei.com>
# message auto-generated for no-merge-commit merge:
!3322 merge om2_dts into develop
【fix】: 解决memset 问题
Created-by: lidaoming1
Commit-by: lidaoming1
Merged-by: cann-robot
Description:
# Pull Request
描述
本次变更旨在优化 OM2 代码生成模块的编译期性能,并新增对原子化 TBE Kernel 的支持。核心思路是将运行时动态容器(如
std::vector)替换为编译期确定大小的 C 风格数组及指针,减少生成代码中的 STL
容器构造开销,从而大幅降低生成代码的编译时间。
主要变更
1. 编译期优化 — 去除生成代码中的 std::vector 依赖
- ReportLaunchedOm2Task 签名重构 (stable_part_provider.cc):将入参从 std::vector<Om2TaskIoEntry> /
std::vector<uint64_t> 改为 C 风格指针+长度参数(const Om2TaskIoEntry *inputs, uint64_t
描述
本次变更旨在优化 OM2 代码生成模块的编译期性能,并新增对原子化 TBE Kernel 的支持。核心思路是将运行时动态容器(如
std::vector)替换为编译期确定大小的 C 风格数组及指针,减少生成代码中的 STL
容器构造开销,从而大幅降低生成代码的编译时间。
主要变更
1. 编译期优化 — 去除生成代码中的 std::vector 依赖
- ReportLaunchedOm2Task 签名重构 (stable_part_provider.cc):将入参从 std::vector<Om2TaskIoEntry> /
std::vector<uint64_t> 改为 C 风格指针+长度参数(const Om2TaskIoEntry *inputs, uint64_t
input_num),避免每次调用时构造临时 vector 对象。
- Tensor shape 数组优化 (kernel_task_code_builder.cc):将 Om2Tensor 构建时的 shape 数组从 std::vector<int64_t> 改为
static constexpr int64_t[] 编译期常量数组,消除运行时堆分配。
- AppendReportLaunchedTaskCall 重构 (task_code_builder_util.cc/h):
- 方法签名从返回 ExprRef 改为 Status,将代码直接附加到 items 中
- 生成的输入/输出/workspace 数组使用 C 风格数组声明(如 const Om2TaskIoEntry op0_report_inputs[])
- 返回值用 OM2_CHK_STATUS(ReportLaunchedOm2Task(...)) 包裹,统一错误检查模式
2. 原子化 TBE Kernel 支持
- Kernel 打包 (om2_package_helper.cc):保存 TBE kernel 时,额外检查 ATOMIC_ATTR_TBE_KERNEL_NAME 属性,将关联的原子化
kernel 一并写入 zip 包。
- Kernel 注册 (om2_codegen_model_builder.cc):在 AICore kernel 注册时,同步注册原子化 kernel(sign 为 kernel_name +
"_atomic")。
- Magic 获取 (om2_codegen_utils.cc/h):GetMagic 新增 is_atomic 参数,支持从 atomic_attr_name_magic 属性读取 magic
信息。
- 独立清理任务 (kernel_task_code_builder.cc):移除对 is_separately_clean_task_ 的不支持断言,新增原子化 kernel 的
func_handle_key 解析逻辑。
3. 数据流重构 — const_index 传递链
- AddrSemantic 扩展 (om2_codegen_types.h):新增 std::optional<size_t> const_index 字段,标识常量 tensor 在 constants_
数组中的索引。
- ResolveConstIndex 函数 (om2_model_utils.cc):新增辅助函数,从 const_N 格式的变量名解析出数字索引。
- 地址语义构建 (om2_model_utils.cc):在
ConstructAddrSemanticForInputConst、ResolveInputAddrs、GetRtWeightAddress、GetRtVarAddress 等路径中填充 const_index。
4. 各 TaskCodeBuilder 统一适配
所有 task code builder 不再依赖 is_reused_from_upstream 标志来决定是否生成 BuildOm2Tensor 调用,改为始终生成完整的
tensor 构建代码,并通过 const_index 判断使用常量地址还是设备内存偏移地址。涉及文件:
- dsa_task_code_builder.cc
- kernel_ex_task_code_builder.cc
- cmo_addr_task_code_builder.cc
- memcpy_addr_async_task_code_builder.cc
- memcpy_async_task_code_builder.cc
- label_switch_by_index_task_code_builder.cc
- stream_switch_task_code_builder.cc
5. 任务级变量名隔离
- AssignTaskLocalIoNames (kernel_task_code_builder.cc):为每个 task 的输入/输出分配独立变量名(op<N>_input<i> /
op<N>_output<j>),避免跨 task 的符号名冲突。
- AstBuildContext::Block (ast_build_context.cc/h):新增 Block(std::vector<BodyItem>) 方法,在
LoadAndRunFileCodeGenerator 中用局部作用域包裹每个 task 的代码块。
6. 其他
- 移除 om2_codegen.cc 中未使用的 ge_call_wrapper.h 头文件引入。
- om2_codegen_model_builder.cc 新增多处 kernel 注册日志(RegisterKernel、RegisterAicpu、RegisterCustAicpu 等)。
- memcpy_addr_async_task_code_builder.cc:ResolveKernelName 从 task_def.kernel().kernel_name() 获取 kernel 名称,而非
semantic 成员。
- 若干代码格式化修复。
7. 测试适配
- test_om2.cc:新增 SyncKernelNameForAllModels 辅助函数,将 op_desc 中的 _kernelname 同步到 task_def,多个 ST
用例增加调用。
- om2_codegen_model_builder_unittest.cc、program_generator_unittest.cc、om2_package_helper_unittest.cc:大量 UT
用例扩展与适配。
变更类型
- [x] ✨ 新功能(原子化 TBE Kernel 支持)
- [x] ♻️ 重构(编译期优化、数据流重构)
关联的Issue
N/A
如何测试
1. 编译 GE 项目,确认无编译错误
2. 运行 OM2 相关的 ST 用例:test_om2 中的全部测试用例
3. 运行 OM2 相关的 UT
用例:om2_codegen_model_builder_unittest、om2_codegen_unittest、program_generator_unittest、om2_package_helper_unittest
4. 使用包含 TBE 原子化 kernel 的模型进行端到端 OM2 转换验证
核对清单
- [x] 我的代码遵循了项目的代码风格
- [x] 我已对代码进行了自测
- [x] 我已更新了相关的文档
- [x] 我在标题中使用了合适的类型标签(如:feat:, fix:)
- [x] 我已经详细阅读了贡献指南(CONTRIBUTING.md),并遵守了其中的所有规定
其他信息
本次变更涉及 27 个文件,+1088/-343 行。核心动机是通过将运行时 STL 容器替换为编译期数组/指针,降低 OM2
生成代码的编译开销。同时补齐了原子化 TBE Kernel 场景的支持链路(打包 → 注册 → 代码生成),使 OM2 框架能够正确处理带
atomic_tbe_kernel_name 属性的 AICore 算子。
See merge request: cann/ge!3322 | 23 小时前 |