| 12 | using namespace arm_common; |
| 13 | |
| 14 | void LSTMCellImpl::exec( |
| 15 | _megdnn_tensor_in input, _megdnn_tensor_in weight_ih, _megdnn_tensor_in bias_ih, |
| 16 | _megdnn_tensor_in hx, _megdnn_tensor_in weight_hh, _megdnn_tensor_in bias_hh, |
| 17 | _megdnn_tensor_in cx, _megdnn_tensor_out h_new, _megdnn_tensor_out c_new, |
| 18 | _megdnn_tensor_out gates, _megdnn_workspace workspace) { |
| 19 | //! only float32 and {1, xx} shape bias will be optimized |
| 20 | MIDOUT_BEGIN(megdnn_arm_common_lstm_cell, midout_iv(0)) { |
| 21 | if (!LstmCellCompute::is_optimized( |
| 22 | input.layout, weight_ih.layout, bias_ih.layout, hx.layout, |
| 23 | weight_hh.layout, bias_hh.layout, cx.layout, h_new.layout, |
| 24 | c_new.layout, gates.layout)) { |
| 25 | naive::LSTMCellImpl::exec( |
| 26 | input, weight_ih, bias_ih, hx, weight_hh, bias_hh, cx, h_new, c_new, |
| 27 | gates, workspace); |
| 28 | } else { |
| 29 | LstmCellCompute::run( |
| 30 | input, weight_ih, bias_ih, hx, weight_hh, bias_hh, cx, h_new, c_new, |
| 31 | gates, workspace, handle()); |
| 32 | } |
| 33 | } |
| 34 | MIDOUT_END(); |
| 35 | } |
| 36 | |
| 37 | size_t LSTMCellImpl::get_workspace_in_bytes( |
| 38 | const TensorLayout& input, const TensorLayout& weight_ih, |
no test coverage detected