Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
新增ThnnFusedGruCell算子
对 GRU(Gated Recurrent Unit,门控循环单元)的单个时间步执行门控融合计算:输入两路门控预激活(输入侧 + 隐层侧)、上一步隐状态和可选双 bias,一次算子调用同时产出两个结果——新隐状态 hy 和反向计算复用的中间量 storage。适用于 RNN 逐步推理/训练中的高频单步计算,等价于 PyTorch torch.nn.functional.gru_cell 的 THNN 内核语义。
hy
storage
torch.nn.functional.gru_cell
门控预激活沿最后一维按门序 [r, z, n] 三等分:
input_gates (B, 3H)
hidden_gates (B, 3H)
hx (B, H)
计算过程:
rg=11+e−(gir+ghr+b1r+b2r)rg = \frac{1}{1 + e^{-(gi_r + gh_r + b1_r + b2_r)}} rg=1+e−(gir+ghr+b1r+b2r)1
zg=11+e−(giz+ghz+b1z+b2z)zg = \frac{1}{1 + e^{-(gi_z + gh_z + b1_z + b2_z)}} zg=1+e−(giz+ghz+b1z+b2z)1
ng=tanh(gin+b1n+rg×(ghn+b2n))ng = \tanh(gi_n + b1_n + rg \times (gh_n + b2_n)) ng=tanh(gin+b1n+rg×(ghn+b2n))
hy=ng+zg×(hx−ng)hy = ng + zg \times (hx - ng) hy=ng+zg×(hx−ng)
storage=[rg∣zg∣ng∣hx∣(ghn+b2n)]storage = [rg \mid zg \mid ng \mid hx \mid (gh_n + b2_n)] storage=[rg∣zg∣ng∣hx∣(ghn+b2n)]
语义说明:
rg
zg
ng
tiling 侧按 UB 容量三形态切分:形态 A 沿 H 切(单行 tile)、形态 B 沿 B 切(多行 tile,多核并行)、形态 C 全量单 tile;行级 32B 对齐填充(paddedHidden)消除非对齐 H 的形态退化;微量数据(<160B)回退单核避免 launch-bound 退化。
input_gates
hidden_gates
hx
input_gates.shape[1] == 3 × hx.shape[1]
input_bias
hidden_bias
aclnnThnnFusedGruCell
aclnnThnnFusedGruCellGetWorkspaceSize
op_graph/thnn_fused_gru_cell_proto.h
NA
Thanks for sending an requirement! Please fill in the following template to help quickly solve your problem.
Backgroud(背景信息)
新增ThnnFusedGruCell算子
1. 功能定位
对 GRU(Gated Recurrent Unit,门控循环单元)的单个时间步执行门控融合计算:输入两路门控预激活(输入侧 + 隐层侧)、上一步隐状态和可选双 bias,一次算子调用同时产出两个结果——新隐状态
hy和反向计算复用的中间量storage。适用于 RNN 逐步推理/训练中的高频单步计算,等价于 PyTorchtorch.nn.functional.gru_cell的 THNN 内核语义。hy供下一时间步递推,storage供反向梯度计算复用,避免训练场景重复计算。2. 计算公式
门控预激活沿最后一维按门序 [r, z, n] 三等分:
input_gates (B, 3H)→ gi_r、gi_z、gi_n(各 H 列)hidden_gates (B, 3H)→ gh_r、gh_z、gh_nhx (B, H)为上一步隐状态计算过程:
rg=1+e−(gir+ghr+b1r+b2r)1
zg=1+e−(giz+ghz+b1z+b2z)1
ng=tanh(gin+b1n+rg×(ghn+b2n))
hy=ng+zg×(hx−ng)
storage=[rg∣zg∣ng∣hx∣(ghn+b2n)]
语义说明:
rg重置门、zg更新门:sigmoid 激活,取值 (0, 1)。ng候选状态:tanh 激活,重置门 rg 调制隐层侧候选分量。hy:更新门 zg 在旧状态 hx 与候选状态 ng 之间做凸组合。storage:反向复用中间量,五段 [rg、zg、ng、hx、gh_n+b2_n] 沿最后一维拼接,每段 H 列,故 shape 为 (B, 5H)。3. 输入输出
4. 精度说明
5. 实现要点(kernel 双路径)
tiling 侧按 UB 容量三形态切分:形态 A 沿 H 切(单行 tile)、形态 B 沿 B 切(多行 tile,多核并行)、形态 C 全量单 tile;行级 32B 对齐填充(paddedHidden)消除非对齐 H 的形态退化;微量数据(<160B)回退单核避免 launch-bound 退化。
6. 约束说明
input_gates与hidden_gates的 shape 必须相同且为 (B, 3H),hx的 shape 为 (B, H),需满足input_gates.shape[1] == 3 × hx.shape[1];input_bias与hidden_bias在位时元素个数必须为 3H 且两者相同。7. 调用方式
aclnnThnnFusedGruCell接口调用,两段式:aclnnThnnFusedGruCellGetWorkspaceSize+aclnnThnnFusedGruCellop_graph/thnn_fused_gru_cell_proto.h)构图调用Origin(信息来源)
NA
Benefit / Necessity (价值/作用)
新增ThnnFusedGruCell算子
Design(设计方案)