| 11 | using rnn::LSTMCellWeightWrapper; |
| 12 | |
| 13 | void LSTMImpl::exec( |
| 14 | _megdnn_tensor_in input, _megdnn_tensor_in hx, _megdnn_tensor_in cx, |
| 15 | _megdnn_tensor_in flatten_weights, _megdnn_tensor_out output, |
| 16 | _megdnn_tensor_out hy, _megdnn_tensor_out cy, _megdnn_tensor_out reserve_space, |
| 17 | _megdnn_workspace workspace) { |
| 18 | #if !MGE_BUILD_WITHOUT_NAIVE_EXEC |
| 19 | MIDOUT_BEGIN(megdnn_naive_lstm_fwd) { |
| 20 | auto _param = param(); |
| 21 | size_t D = _param.bidirectional ? 2 : 1; |
| 22 | size_t num_layers = _param.num_layers; |
| 23 | size_t input_size = input.layout.shape[2]; |
| 24 | std::vector<LSTMCellWeightWrapper> cells; |
| 25 | size_t used_workspace_size = rnn::get_cells<LSTMCellWeightWrapper>( |
| 26 | D, num_layers, input_size, _param.hidden_size, _param.bias, cells, |
| 27 | flatten_weights, workspace); |
| 28 | |
| 29 | Workspace new_workspace( |
| 30 | workspace.raw_ptr + used_workspace_size, |
| 31 | workspace.size - used_workspace_size); |
| 32 | TensorNDArray states = {hx, cx}, states_new = {hy, cy}; |
| 33 | rnn::exec_internal<LSTMCellWeightWrapper, LSTMCellForward>( |
| 34 | cells, input, states, states_new, output, reserve_space, num_layers, D, |
| 35 | param::RNNCell::NonlineMode::IDENTITY, this->handle(), new_workspace); |
| 36 | } |
| 37 | MIDOUT_END(); |
| 38 | #else |
| 39 | __builtin_trap(); |
| 40 | #endif |
| 41 | } |
| 42 | |
| 43 | size_t LSTMImpl::get_workspace_in_bytes( |
| 44 | const TensorLayout& input, const TensorLayout& /*hx*/, |