| 68 | namespace rnn_cell { |
| 69 | |
| 70 | size_t get_workspace_in_bytes( |
| 71 | const TensorLayout& input, const TensorLayout& weight_ih, |
| 72 | const TensorLayout& /*bias_ih*/, const TensorLayout& hx, |
| 73 | const TensorLayout& weight_hh, const TensorLayout& /*bias_hh*/, |
| 74 | const TensorLayout& dst, Handle* handle) { |
| 75 | auto opr = handle->create_operator<MatrixMulForward>(); |
| 76 | opr->param().transposeB = true; |
| 77 | return dst.span().dist_byte() + |
| 78 | std::max( |
| 79 | opr->get_workspace_in_bytes(hx, weight_hh, dst), |
| 80 | opr->get_workspace_in_bytes(input, weight_ih, dst)); |
| 81 | } |
| 82 | |
| 83 | void exec( |
| 84 | _megdnn_tensor_in input, _megdnn_tensor_in weight_ih, _megdnn_tensor_in bias_ih, |