| 76 | |
| 77 | template <class Cell, typename CellOpr, class States> |
| 78 | void exec_kernel( |
| 79 | SmallVector<Cell>& cells, const TensorNDArray& inputs, |
| 80 | const SmallVector<States>& states_in, SmallVector<States>& states_out, |
| 81 | TensorNDArray& outputs, size_t num_layers, size_t dir_size, Handle* handle, |
| 82 | WorkspaceBundle workspace_bundle) { |
| 83 | megdnn_assert(cells.size() == num_layers * dir_size); |
| 84 | megdnn_assert( |
| 85 | states_in.size() == states_out.size() && |
| 86 | states_in.size() == num_layers * dir_size); |
| 87 | megdnn_assert(outputs.size() == inputs.size()); |
| 88 | //! two tmp state workspace |
| 89 | megdnn_assert(workspace_bundle.nr_workspace() == 4 + States::nr_states()); |
| 90 | |
| 91 | size_t seq_len = inputs.size(); |
| 92 | size_t batch_size = inputs[0].layout.shape[0]; |
| 93 | size_t input_size = inputs[0].layout.shape[1]; |
| 94 | size_t hidden_size = cells[0].m_weight_hh.layout.shape[1]; |
| 95 | |
| 96 | TensorLayout batch_output_layout{ |
| 97 | {hidden_size}, outputs[0].layout.dtype}; // output hy |
| 98 | TensorLayout cell_output_layout{ |
| 99 | {batch_size, hidden_size}, outputs[0].layout.dtype}; // output hy |
| 100 | TensorLayout seq_output_layout{ |
| 101 | {batch_size, dir_size * hidden_size}, outputs[0].layout.dtype}; |
| 102 | TensorLayout cell_first_input_layout{ |
| 103 | {batch_size, input_size}, inputs[0].layout.dtype}; // input |
| 104 | TensorLayout cell_input_layout{ |
| 105 | {batch_size, dir_size * hidden_size}, inputs[0].layout.dtype}; |
| 106 | TensorLayout tmp_output_layout{ |
| 107 | {seq_len, batch_size, dir_size * hidden_size}, outputs[0].layout.dtype}; |
| 108 | |
| 109 | //! workspace get |
| 110 | Workspace cell_workspace( |
| 111 | static_cast<dt_byte*>(workspace_bundle.get(0)), |
| 112 | workspace_bundle.get_size(0) + workspace_bundle.get_size(1)); |
| 113 | auto&& tmp_inputs_1 = split_tensor( |
| 114 | TensorND{workspace_bundle.get(2), tmp_output_layout}, seq_len, |
| 115 | cell_input_layout); |
| 116 | auto&& tmp_outputs_1 = split_tensor( |
| 117 | TensorND{workspace_bundle.get(2), tmp_output_layout}, seq_len, |
| 118 | seq_output_layout); |
| 119 | |
| 120 | auto&& tmp_inputs_2 = split_tensor( |
| 121 | TensorND{workspace_bundle.get(3), tmp_output_layout}, seq_len, |
| 122 | cell_input_layout); |
| 123 | auto&& tmp_outputs_2 = split_tensor( |
| 124 | TensorND{workspace_bundle.get(3), tmp_output_layout}, seq_len, |
| 125 | seq_output_layout); |
| 126 | |
| 127 | using IoPair = std::pair<TensorNDArray, TensorNDArray>; |
| 128 | IoPair io_pair1 = {tmp_inputs_1, tmp_outputs_2}; |
| 129 | IoPair io_pair2 = {tmp_inputs_2, tmp_outputs_1}; |
| 130 | SmallVector<IoPair> io_pairs = {io_pair1, io_pair2}; |
| 131 | |
| 132 | SmallVector<RefPtr> ptr; |
| 133 | for (size_t index = 0; index < States::nr_states(); index++) { |
| 134 | ptr.push_back(workspace_bundle.get(4 + index)); |
| 135 | } |