| 21 | using rnn::RNNCellWeightWrapper; |
| 22 | |
| 23 | void RNNImpl::exec( |
| 24 | _megdnn_tensor_in input, _megdnn_tensor_in hx, |
| 25 | _megdnn_tensor_in flatten_weights, _megdnn_tensor_out output, |
| 26 | _megdnn_tensor_out hy, _megdnn_tensor_out reserve_space, |
| 27 | _megdnn_workspace workspace) { |
| 28 | #if !MGE_BUILD_WITHOUT_NAIVE_EXEC |
| 29 | MIDOUT_BEGIN(megdnn_naive_rnn_fwd) { |
| 30 | auto _param = param(); |
| 31 | size_t D = _param.bidirectional ? 2 : 1; |
| 32 | size_t num_layers = _param.num_layers; |
| 33 | size_t input_size = input.layout.shape[2]; |
| 34 | std::vector<RNNCellWeightWrapper> cells; |
| 35 | size_t used_workspace_size = rnn::get_cells<RNNCellWeightWrapper>( |
| 36 | D, num_layers, input_size, _param.hidden_size, _param.bias, cells, |
| 37 | flatten_weights, workspace); |
| 38 | |
| 39 | Workspace new_workspace( |
| 40 | workspace.raw_ptr + used_workspace_size, |
| 41 | workspace.size - used_workspace_size); |
| 42 | TensorNDArray states, states_new; |
| 43 | states.push_back(hx); |
| 44 | states_new.push_back(hy); |
| 45 | rnn::exec_internal<RNNCellWeightWrapper, RNNCellForward>( |
| 46 | cells, input, states, states_new, output, reserve_space, num_layers, D, |
| 47 | _param.nonlineMode, this->handle(), new_workspace); |
| 48 | } |
| 49 | MIDOUT_END(); |
| 50 | #else |
| 51 | __builtin_trap(); |
| 52 | #endif |
| 53 | } |
| 54 | |
| 55 | size_t RNNImpl::get_workspace_in_bytes( |
| 56 | const TensorLayout& input, const TensorLayout& hx, |
no test coverage detected