描述

engram算子支持训练HOST代码，前向复用现有接口，增加相应参数，反向增加新接口

关联的Issue

https://gitcode.com/cann/ops-transformer/issues/3838

测试

本地验证

文档更新

不涉及

类型标签

• 🐛 Bug修复
• ✨ 新特性
• ⚡ 性能优化
• ♻️ 重构
• 🧪 测试
• 📦 构建/CI
• 🔧 配置变更
• 📝 文档更新
• ⬆️ 依赖升级
• 🔒 安全修复
• 🧹 代码清理
• ❓ 其他，请描述：

PR #8914 代码检视报告

检视概览

• PR URL：https://gitcode.com/cann/ops-transformer/pull/8914
• 算子名：engram_fetch_grad / engram_fetch_train / elastic_buffer
• 代码侧别：Kernel + Tiling + 混合（torch_extension）
• 变更文件数：21
• 文件组数：6（K1, K2, H1, H2, S, E）
• 总条例数：346
• 检视时间：2026-07-21
• 检视模式：大型 PR 检视（21 文件 > 10 阈值）

检视统计

状态	条例数	占比
PASS	341	98.6%
FAIL（发现问题）	5	1.4%
SUSPICIOUS（需关注）	0	0%

逐文件组统计

文件组	侧别	文件数	PASS	FAIL
engram_fetch_grad_kernel_G1 (K1)	Kernel	3	27	0
engram_fetch_train_kernel_G1 (K2)	Kernel	3	37	0
engram_fetch_grad_host_G1 (H1)	Tiling	3	66	2
engram_fetch_train_host_G1 (H2)	Tiling	3	70	2
shared (S)	混合	6	61	1
elastic_buffer (E)	混合	3	80	0
合计	—	21	341	5

系统性风险

SEC-4.1 / TOPK-7 在 H1 和 H2 中均失败：两个 Tiling 侧文件组（engram_fetch_grad 和 engram_fetch_train）均存在 sendCounts/recvLocalEntry（或 sendCountsOut/recvLocalEntryOut）的 dim(0) 值未做非负/范围校验的问题。两个算子的 tiling 代码结构高度相似，校验缺失模式相同，建议全局排查并统一修复。

---

发现问题（HIGH 置信度）

1. [engram_fetch_grad_host_G1] SEC-4.1 外部输入数据需要做合法性校验

• 问题描述：sendCounts 的 dim(0) 和 recvLocalEntry 的 dim(0) 作为外部输入（来自 graph shape）仅校验了维度数（1D），未对 dim(0) 的值做合法性校验，违反 SEC-4.1「外部输入数据需要做合法性校验且确保校验范围正确」。同一函数中 commContext（行 206 GetDim(0) <= 0）和 gradFetched（行 221 numTokens < 0、行 226 hiddenDim <= 0）均对 dim 值做了显式校验，sendCounts/recvLocalEntry 缺失同类校验，形成不一致的安全缺口。

  风险链路 1（sendCounts）：dim(0) 若为负值或超过 UINT32_MAX，行 260 static_cast<uint32_t> 会发生窄化截断/回绕，rankSize 变为巨大错误值 → SetBlockDim 和 workspace 计算均受影响。

  风险链路 2（recvLocalEntry）：dim(0) 赋值给 int64_t totalRecv（行 281），若为负值则保留负数 → wsRecvGrad = totalRecv * hiddenBytes 为负 → wsTotal 求和可能为负 → static_cast<size_t>(wsTotal) 将负数回绕为极大 size_t 值，导致 workspace 申请异常。

• 代码片段（mc2/engram_fetch_grad/op_host/op_tiling/engram_fetch_grad_tiling.cpp 行 251-281）：

      // sendCounts: 1D (W,)
      const gert::StorageShape *sendCountsShape = context->GetInputShape(IN_SEND_COUNTS);
      OP_TILING_CHECK(sendCountsShape == nullptr,
          OP_LOGE_WITH_INVALID_INPUT(nodeName, "sendCounts"), return ge::GRAPH_FAILED);
      OP_TILING_CHECK(sendCountsShape->GetStorageShape().GetDimNum() != DIM_ONE,
          OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "sendCounts",
              (std::to_string(sendCountsShape->GetStorageShape().GetDimNum()) + "D").c_str(),
              "The shape dim of sendCounts must be 1D."),
          return ge::GRAPH_FAILED);
      rankSize = static_cast<uint32_t>(sendCountsShape->GetStorageShape().GetDim(0));
  
      // recvCounts: 1D (W,)
      const gert::StorageShape *recvCountsShape = context->GetInputShape(IN_RECV_COUNTS);
      OP_TILING_CHECK(recvCountsShape == nullptr,
          OP_LOGE_WITH_INVALID_INPUT(nodeName, "recvCounts"), return ge::GRAPH_FAILED);
      OP_TILING_CHECK(recvCountsShape->GetStorageShape().GetDimNum() != DIM_ONE,
          OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "recvCounts",
              (std::to_string(recvCountsShape->GetStorageShape().GetDimNum()) + "D").c_str(),
              "The shape dim of recvCounts must be 1D."),
          return ge::GRAPH_FAILED);
  
      // recvLocalEntry: 1D (R,)
      const gert::StorageShape *recvLocalEntryShape = context->GetInputShape(IN_RECV_LOCAL_ENTRY);
      OP_TILING_CHECK(recvLocalEntryShape == nullptr,
          OP_LOGE_WITH_INVALID_INPUT(nodeName, "recvLocalEntry"), return ge::GRAPH_FAILED);
      OP_TILING_CHECK(recvLocalEntryShape->GetStorageShape().GetDimNum() != DIM_ONE,
          OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "recvLocalEntry",
              (std::to_string(recvLocalEntryShape->GetStorageShape().GetDimNum()) + "D").c_str(),
              "The shape dim of recvLocalEntry must be 1D."),
          return ge::GRAPH_FAILED);
      totalRecv = recvLocalEntryShape->GetStorageShape().GetDim(0);
  

• 假设检验证据：

  • (+40%) 规范违反：SEC-4.1 要求外部输入数据需做合法性校验，sendCounts dim(0) 直接 cast uint32_t、recvLocalEntry dim(0) 直接赋值 int64_t，均无值范围校验
  • (+30%) 上下文防御缺失：同函数内 commContext/gradFetched 均有非负校验，唯独 sendCounts/recvLocalEntry 缺失
  • (+20%) PR 归属：新增文件，全部行在 diff 变更范围内
  • (+15%) 数据流风险：rankSize 负值回绕 / totalRecv 负值导致 workspace 回绕为巨大值
  • 自信值 = 105% ≥ 70% → 判定违规

• 修复建议：

  1. 在行 260 之前增加 sendCounts dim(0) 值校验（> 0）
  2. 在行 281 之前增加 recvLocalEntry dim(0) 值校验（>= 0）
  3. 参照同函数内 commContext/gradFetched 的校验模式

---

2. [engram_fetch_grad_host_G1] TOPK-7 融合规则/InferShape/Tiling 外部输入校验

• 问题描述：TOPK-7 要求「Tiling 外部输入使用时必须进行合法性校验」，校验内容包括「Shape 维度是否合法（不为负数、不超过限制）」。sendCounts dim(0) 和 recvLocalEntry dim(0) 作为外部输入被提取后直接使用，未做非负/范围校验。同函数内 commContext dim0（<=0 校验）、gradFetched dim0（<0 校验）、hiddenDim（<=0 校验）均有非负防御，唯独 sendCounts 和 recvLocalEntry 缺失。此问题与 SEC-4.1 同构。

• 代码片段（mc2/engram_fetch_grad/op_host/op_tiling/engram_fetch_grad_tiling.cpp 行 251-281）：

      // sendCounts: 1D (W,)
      const gert::StorageShape *sendCountsShape = context->GetInputShape(IN_SEND_COUNTS);
      OP_TILING_CHECK(sendCountsShape == nullptr,
          OP_LOGE_WITH_INVALID_INPUT(nodeName, "sendCounts"), return ge::GRAPH_FAILED);
      OP_TILING_CHECK(sendCountsShape->GetStorageShape().GetDimNum() != DIM_ONE,
          OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "sendCounts",
              (std::to_string(sendCountsShape->GetStorageShape().GetDimNum()) + "D").c_str(),
              "The shape dim of sendCounts must be 1D."),
          return ge::GRAPH_FAILED);
      rankSize = static_cast<uint32_t>(sendCountsShape->GetStorageShape().GetDim(0));
      // ...（recvLocalEntry 同理）
      totalRecv = recvLocalEntryShape->GetStorageShape().GetDim(0);
  

• 假设检验证据：

  • (+40%) 规范违反：TOPK-7 要求 Shape 维度不为负数
  • (+30%) 上下文防御缺失：同函数内其他 tensor 均有非负校验
  • (+20%) PR 归属：新增文件
  • (+15%) 数据流风险：负值 dim0 经 static_cast 回绕或以负 int64_t 传播至 workspace 计算
  • 自信值 = 105% ≥ 70% → 判定违规

• 修复建议：与 SEC-4.1 同一问题，修复方案一致。

---

3. [engram_fetch_train_host_G1] SEC-4.1 外部输入数据需要做合法性校验

• 问题描述：CheckTensorDim 中 sendCountsOut 和 recvLocalEntryOut 的 dim0 值作为外部输入未做范围校验。rankSize（uint32_t）通过 static_cast 从 int64_t 缩窄转换，若 dim0 为负或超过 UINT32_MAX 将静默截断/回绕为巨大值；totalRecv（int64_t）未校验非负，若为负值将导致 wsLocalData = totalRecv * hiddenBytes 为负，进而 wsTotal 为负，再经 static_cast<size_t>(wsTotal) 回绕为约 1.8e19 的巨大 workspace 申请。同函数内 commContextShape dim0 > 0（行 188）、numTokens >= 0（行 202）、numRecvOutShape dim0 == 1（行 272）均已校验，此处校验明显不一致。

• 代码片段（mc2/engram_fetch_train/op_host/op_tiling/engram_fetch_train_tiling.cpp 行 235-262）：

      const gert::StorageShape *sendCountsOutShape = context->GetOutputShape(OUT_SEND_COUNTS);
      OP_TILING_CHECK(sendCountsOutShape == nullptr,
          OP_LOGE_WITH_INVALID_INPUT(nodeName, "sendCountsOut"), return ge::GRAPH_FAILED);
      OP_TILING_CHECK(sendCountsOutShape->GetStorageShape().GetDimNum() != DIM_ONE,
          OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "sendCountsOut",
              (std::to_string(sendCountsOutShape->GetStorageShape().GetDimNum()) + "D").c_str(),
              "The shape dim of sendCountsOut must be 1D."),
          return ge::GRAPH_FAILED);
      rankSize = static_cast<uint32_t>(sendCountsOutShape->GetStorageShape().GetDim(0));
  
      const gert::StorageShape *recvCountsOutShape = context->GetOutputShape(OUT_RECV_COUNTS);
      OP_TILING_CHECK(recvCountsOutShape == nullptr,
          OP_LOGE_WITH_INVALID_INPUT(nodeName, "recvCountsOut"), return ge::GRAPH_FAILED);
      OP_TILING_CHECK(recvCountsOutShape->GetStorageShape().GetDimNum() != DIM_ONE,
          OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "recvCountsOut",
              (std::to_string(recvCountsOutShape->GetStorageShape().GetDimNum()) + "D").c_str(),
              "The shape dim of recvCountsOut must be 1D."),
          return ge::GRAPH_FAILED);
  
      const gert::StorageShape *recvLocalEntryOutShape = context->GetOutputShape(OUT_RECV_LOCAL_ENTRY);
      OP_TILING_CHECK(recvLocalEntryOutShape == nullptr,
          OP_LOGE_WITH_INVALID_INPUT(nodeName, "recvLocalEntryOut"), return ge::GRAPH_FAILED);
      OP_TILING_CHECK(recvLocalEntryOutShape->GetStorageShape().GetDimNum() != DIM_ONE,
          OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(nodeName, "recvLocalEntryOut",
              (std::to_string(recvLocalEntryOutShape->GetStorageShape().GetDimNum()) + "D").c_str(),
              "The shape dim of recvLocalEntryOut must be 1D."),
          return ge::GRAPH_FAILED);
      totalRecv = recvLocalEntryOutShape->GetStorageShape().GetDim(0);
  

• 假设检验证据：

  • (+40%) 规范违反：外部 shape dim 值未做范围/非负校验
  • (+30%) 上下文防御缺失：作用域内仅校验 dim 数，无 dim 值范围校验
  • (+20%) PR 归属：新文件，全量代码均在 diff 范围
  • (+15%) 数据流风险：rankSize 缩窄截断 / totalRecv 负值 → workspace 回绕
  • 自信值 = 105% ≥ 70% → 判定违规

• 修复建议：

  1. 在行 243 前增加 sendCountsOut dim(0) 非负及 <= UINT32_MAX 校验
  2. 在行 262 前增加 recvLocalEntryOut dim(0) >= 0 校验
  3. 参照同函数内 numTokens（行 202-205）的校验模式

---

4. [engram_fetch_train_host_G1] TOPK-7 融合规则/InferShape/Tiling 外部输入校验

• 问题描述：与 SEC-4.1 同一问题。TOPK-7 要求 Tiling 外部输入的 Shape 维度不为负数、不超过限制。sendCountsOut dim0（赋值给 rankSize，缩窄为 uint32_t）和 recvLocalEntryOut dim0（赋值给 totalRecv）均未做非负/范围校验。同函数内 commContextShape dim0 > 0（行 188）、indicesShape numTokens >= 0（行 202）均已校验，存在校验不一致。

• 代码片段（mc2/engram_fetch_train/op_host/op_tiling/engram_fetch_train_tiling.cpp 行 235-262）：（同 SEC-4.1）

• 假设检验证据：

  • (+40%) 规范违反：Tiling 外部 shape dim 值未做非负/范围校验
  • (+30%) 上下文防御缺失：仅校验 dim 数无 dim 值校验
  • (+20%) PR 归属：新文件全量在 diff 范围
  • (+15%) 数据流风险：rankSize 缩窄截断 / totalRecv 负值 → workspace 回绕
  • 自信值 = 105% ≥ 70% → 判定违规

• 修复建议：与 SEC-4.1 同一问题，修复方案一致。

---

5. [shared] GEN-2.4 禁止 extern 声明引用外部接口

• 问题描述：UT 文件通过 extern "C" 声明引用 aclnn 接口（非 Kernel 入口函数），违反 GEN-2.4。更严重的是，UT 中的声明与 aclnn_engram_fetch_grad.cpp 实现签名严重不一致：

  • 实现签名 14 个参数：6 inputs + numRecv + 3 outputs + numEntriesPerRank + commBufferSize + workspaceSize + executor
  • UT 声明 13 个参数：6 inputs + hiddenSize + numEntriesPerRank + 3 outputs + workspaceSize + executor
  • 差异：实现有 numRecv（第7输入）和 commBufferSize（attr），UT 声明缺这两者；UT 声明多出 hiddenSize（实现中没有）；attr 位置在实现中位于 outputs 之后，在 UT 声明中位于 inputs 之后

  OP_API_UT 框架按 UT 声明的函数指针类型压栈调用，但 C 链接符号解析到实现函数，实现按自己的签名读栈。栈布局完全不匹配，导致 CheckNotNull 把整数当指针校验，后续 Inner 实现解引用无效指针必然段错误或数据错乱。

• 代码片段（mc2/engram_fetch_grad/tests/ut/op_api/test_aclnn_engram_fetch_grad.cpp 行 32-43）：

  extern "C" {
  aclnnStatus aclnnEngramFetchGradGetWorkspaceSize(
      const aclTensor *commContext, const aclTensor *gradFetched,
      const aclTensor *perm, const aclTensor *sendCounts,
      const aclTensor *recvCounts, const aclTensor *recvLocalEntry,
      int32_t hiddenSize, int64_t numEntriesPerRank,
      aclTensor *gradUniqueOut, aclTensor *uniqueLocalEntryOut, aclTensor *numUniqueOut,
      uint64_t *workspaceSize, aclOpExecutor **executor);
  
  aclnnStatus aclnnEngramFetchGrad(void *workspace, uint64_t workspaceSize,
                                        aclOpExecutor *executor, aclrtStream stream);
  }
  

  对照实现（mc2/engram_fetch_grad/op_api/aclnn_engram_fetch_grad.cpp 行 63-82）：

  aclnnStatus aclnnEngramFetchGradGetWorkspaceSize(
      const aclTensor *commContext, const aclTensor *gradFetched,
      const aclTensor *perm, const aclTensor *sendCounts,
      const aclTensor *recvCounts, const aclTensor *recvLocalEntry,
      const aclTensor *numRecv,
      aclTensor *gradUniqueOut, aclTensor *uniqueLocalEntryOut, aclTensor *numUniqueOut,
      int64_t numEntriesPerRank, int64_t commBufferSize,
      uint64_t *workspaceSize, aclOpExecutor **executor)
  

• 假设检验证据：

  • (+40%) 规范违反：GEN-2.4 明确禁止通过 extern 声明引用外部函数接口
  • (+30%) 上下文防御缺失：UT 中无防御，仓库无 aclnn_engram_fetch_grad.h 头文件可引入正确签名
  • (+20%) PR 归属：test_aclnn_engram_fetch_grad.cpp 为新增文件，extern "C" 声明在 diff 变更范围内
  • (+15%) 数据流风险：UT 声明(13参数)与实现(14参数)参数数量、类型、顺序均不同，栈布局不匹配
  • 自信值 = 105% ≥ 70% → 判定违规

• 修复建议：

  1. 修正 UT 中的 extern "C" 声明，使其与实现签名完全一致（补上 numRecv，删除 hiddenSize，补上 commBufferSize，attr 移到 outputs 之后）
  2. 同步修正所有 TEST_F 用例中的 OP_API_UT INPUT/OUTPUT 列表
  3. 更优做法：由 op_host 侧生成 aclnn_engram_fetch_grad.h 对外头文件，UT 通过 #include 引入