| 200 | } // namespace |
| 201 | |
| 202 | void LstmCellCompute::run( |
| 203 | _megdnn_tensor_in input, _megdnn_tensor_in weight_ih, _megdnn_tensor_in bias_ih, |
| 204 | _megdnn_tensor_in hx, _megdnn_tensor_in weight_hh, _megdnn_tensor_in bias_hh, |
| 205 | _megdnn_tensor_in cx, _megdnn_tensor_out h_new, _megdnn_tensor_out c_new, |
| 206 | _megdnn_tensor_out gates, _megdnn_workspace workspace, Handle* handle) { |
| 207 | auto bundle = get_workspace_bundle( |
| 208 | input.layout, weight_ih.layout, bias_ih.layout, hx.layout, weight_hh.layout, |
| 209 | bias_hh.layout, cx.layout, h_new.layout, c_new.layout, gates.layout); |
| 210 | bundle.set(workspace.raw_ptr); |
| 211 | TensorND tmp{static_cast<void*>(bundle.get(0)), gates.layout}; |
| 212 | auto matmul_workspace = |
| 213 | megdnn::Workspace{static_cast<dt_byte*>(bundle.get(1)), bundle.get_size(1)}; |
| 214 | auto opr = handle->create_operator<MatrixMul>(); |
| 215 | opr->param().transposeB = true; |
| 216 | //! the opr will dispatch compute task to device, so record mode |
| 217 | //! performance will not be effect |
| 218 | opr->exec(input, weight_ih, tmp, matmul_workspace); |
| 219 | opr->exec(hx, weight_hh, gates, matmul_workspace); |
| 220 | |
| 221 | //! the optimized post compute, nonlinear(tmp + dst + bias_hx + bias_cx) |
| 222 | if (bias_ih.layout.ndim != 0 && bias_ih.layout.ndim != 0) { |
| 223 | MEGDNN_DISPATCH_CPU_KERN( |
| 224 | static_cast<naive::HandleImpl*>(handle), |
| 225 | rnn_cell_elemwise_compute<true>( |
| 226 | gates, tmp, bias_ih, bias_hh, cx, h_new, c_new)); |
| 227 | } else { |
| 228 | megdnn_assert(bias_ih.layout.ndim == 0 && bias_ih.layout.ndim == 0); |
| 229 | MEGDNN_DISPATCH_CPU_KERN( |
| 230 | static_cast<naive::HandleImpl*>(handle), |
| 231 | rnn_cell_elemwise_compute<false>( |
| 232 | gates, tmp, bias_ih, bias_hh, cx, h_new, c_new)); |
| 233 | } |
| 234 | } |
| 235 | |
| 236 | WorkspaceBundle LstmCellCompute::get_workspace_bundle( |
| 237 | const TensorLayout& input, const TensorLayout& weight_ih, const TensorLayout&, |