| 75 | namespace lstm_cell { |
| 76 | |
| 77 | size_t get_workspace_in_bytes( |
| 78 | const TensorLayout& input, const TensorLayout& weight_ih, |
| 79 | const TensorLayout& bias_ih, const TensorLayout& hx, |
| 80 | const TensorLayout& weight_hh, const TensorLayout& bias_hh, |
| 81 | const TensorLayout& /*cx*/, const TensorLayout& /*h_new*/, |
| 82 | const TensorLayout& /*c_new*/, const TensorLayout& gates, Handle* handle) { |
| 83 | TensorLayout tmp_layout; |
| 84 | auto opr = handle->create_operator<RNNCellForward>(); |
| 85 | opr->param().nonlineMode = param::RNNCell::NonlineMode::IDENTITY; |
| 86 | opr->deduce_layout(input, weight_ih, bias_ih, hx, weight_hh, bias_hh, tmp_layout); |
| 87 | size_t rnn_cell_need = opr->get_workspace_in_bytes( |
| 88 | input, weight_ih, bias_ih, hx, weight_hh, bias_hh, gates); |
| 89 | size_t lstm_cell_need = 2 * tmp_layout.span().dist_byte(); |
| 90 | return rnn_cell_need > lstm_cell_need ? rnn_cell_need : lstm_cell_need; |
| 91 | } |
| 92 | |
| 93 | void exec( |
| 94 | _megdnn_tensor_in input, _megdnn_tensor_in weight_ih, _megdnn_tensor_in bias_ih, |
no test coverage detected