已关闭
[Requirement|需求建议]: 新增 block_lstm 算子(TF BlockLSTM/BlockLSTMV2 前向,Ascend 950) #6204
陶贻坤创建于  9月24日关闭于  9月29日
陶贻坤
陶贻坤成员
9月24日 创建

Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.

Backgroud(背景信息)

长尾 RNN 系网络广泛使用 TensorFlow 的 tf.raw_ops.BlockLSTM(V1,gate_order=icfo)与 tf.raw_ops.BlockLSTMV2(V2,gate_order=ifco)作为 LSTM 层的 block 形式前向算子。该算子把整个时间步循环收进一次 kernel launch(block 语义:不逐步暴露中间状态给框架),迁移到 Ascend 950 平台时缺失对应算子,网络无法下沉 NPU 执行。

当前缺口:

  • PyTorch 无同名 block 形式公开算子(torch.nn.LSTM 不暴露逐门输出,无 peephole),无现成对标实现
  • CANN 现有 RNN 族(如 DynamicRNN)为逐步展开路径,与 TF BlockLSTM 语义/接口不匹配

Origin(信息来源)

NA

Benefit / Necessity (价值/作用)

Design(设计方案)

实现方案(已完成)

  • Phase A 投影(cube):x@Wx + h@Wh + b,resident(全驻留)/ wave(M 波 × K 流 × N 板)双路径由片上预算模型(L0A/L0B/L0C/L1/UB)自动择路
  • Phase B 递推(vector):门激活 → cell 递推(含 peephole/cell_clip)→ 7 平面散出
  • 分核:探测核数 + ragged pad-to-divisible(hangSafeCores = aicNum-1 避开满格 grid SyncAll 死锁),实测 B=129 3→26 簇、431_base 时延 15058→9014us(1.67x)
  • 空 tensor 全场景:I==0 走真实计算路径(投影仅 x 侧消失,igates = h@Wh + b,bias 由 vector 侧广播);T/B/H==0 输出全空,GE 图模式剪枝覆盖
  • 工程件:OpDef/InferShape/InferDataType/Tiling/Kernel/TF Plugin/Examples/UT 全链交付,op_list.md 已登记

验证结论

验证项 结果
GEIR 全量 549 条三方交叉校验(NPU vs fp64 golden vs GPU TF) 548/549(唯一失败 305_val 为既有已知问题)
e2e 全量 513 条三方 + 性能(TF Plugin→npu_device→NPU) 512/513(同 305_val,两通路交叉确认)
空 tensor 24 条(I==0 真实计算 15 + B/T/H==0 剪枝 9) 24/24 双通路全 PASS
负向用例 33 条 全部被正确检查拒绝,无极性写反
确定性 --dump full 连跑两次 input/golden/output 逐位一致
ophost UT / kernel UT / examples 8/8、11/11、静态+dynamic(-2) 均 PASS

性能参考(纯 kernel 口径,GEIR)

  • fp32 中位 214us;gate<1e4 中位 92us
  • 分核优化收益:431_base 15058→9014us(1.67x)

交付物

  • 算子代码:rnn/block_lstm/(op_host / op_kernel/arch35 / op_graph / framework / examples / tests)
  • 文档:op_list.md 登记
  • 测试资产:GEIR 549 条 + e2e 513 条用例、golden 插件(fp64 oracle + TF third_party)、双通路测试报告
likedislike
陶贻坤
陶贻坤成员
9月24日 评论:

/assign

likedislike
CANN-robotCANN-robot成员
9月24日 将 eternityk 设为负责人
CANN-robotCANN-robot成员
9月29日 关闭了 issue
CANN-robotCANN-robot成员
9月29日 添加了label:resolved