| 91 | } |
| 92 | |
| 93 | void exec( |
| 94 | _megdnn_tensor_in input, _megdnn_tensor_in weight_ih, _megdnn_tensor_in bias_ih, |
| 95 | _megdnn_tensor_in hx, _megdnn_tensor_in weight_hh, _megdnn_tensor_in bias_hh, |
| 96 | _megdnn_tensor_in cx, _megdnn_tensor_out h_new, _megdnn_tensor_out c_new, |
| 97 | _megdnn_tensor_out gates, _megdnn_workspace workspace, Handle* handle) { |
| 98 | auto opr = handle->create_operator<RNNCellForward>(); |
| 99 | opr->param().nonlineMode = param::RNNCell::NonlineMode::IDENTITY; |
| 100 | opr->exec(input, weight_ih, bias_ih, hx, weight_hh, bias_hh, gates, workspace); |
| 101 | // activation |
| 102 | size_t batch_size = hx.layout.shape[0]; |
| 103 | size_t hidden_size = hx.layout.shape[1]; |
| 104 | |
| 105 | auto copy_opr = handle->create_operator<TypeCvtForward>(); |
| 106 | TensorND copy_gates{static_cast<void*>(workspace.raw_ptr), gates.layout}; |
| 107 | TensorLayout hidden_layout{TensorShape{hidden_size}, hx.layout.dtype}; |
| 108 | TensorLayout gateinfo_layout{TensorShape{batch_size, hidden_size}, hx.layout.dtype}; |
| 109 | for (size_t i = 0; i < batch_size; i++) { |
| 110 | for (size_t j = 0; j < 4; j++) { |
| 111 | TensorND half_step_states{ |
| 112 | // output |
| 113 | static_cast<uint8_t*>(gates.raw_ptr()) + |
| 114 | (4 * i + j) * hidden_layout.span().dist_byte(), |
| 115 | hidden_layout}; |
| 116 | TensorND half_step_output{ |
| 117 | static_cast<uint8_t*>(copy_gates.raw_ptr()) + |
| 118 | j * gateinfo_layout.span().dist_byte() + |
| 119 | i * hidden_layout.span().dist_byte(), |
| 120 | hidden_layout}; |
| 121 | copy_opr->exec(half_step_states, half_step_output); |
| 122 | } |
| 123 | } |
| 124 | void* workspace_ptr = workspace.raw_ptr + copy_gates.layout.span().dist_byte(); |
| 125 | copy_opr->exec(copy_gates, gates); |
| 126 | |
| 127 | // sigmoid: i f |
| 128 | TensorND tmp{static_cast<void*>(workspace_ptr), copy_gates.layout}; |
| 129 | TensorLayout gates_ifo_layout{ |
| 130 | TensorShape({batch_size, hidden_size * 2}), copy_gates.layout.dtype}; |
| 131 | TensorND gates_ifo_origin{copy_gates.raw_ptr(), gates_ifo_layout}; |
| 132 | TensorND gates_ifo{tmp.raw_ptr(), gates_ifo_layout}; |
| 133 | auto sigmoid = handle->create_operator<ElemwiseForward>(); |
| 134 | sigmoid->param().mode = Elemwise::Param::Mode::SIGMOID; |
| 135 | sigmoid->exec({gates_ifo_origin}, gates_ifo); |
| 136 | // tanh: g |
| 137 | TensorLayout g_layout{ |
| 138 | TensorShape({batch_size, hidden_size}), copy_gates.layout.dtype}; |
| 139 | TensorND g_origin{ |
| 140 | static_cast<char*>(copy_gates.raw_ptr()) + |
| 141 | gates_ifo_layout.span().dist_byte(), |
| 142 | g_layout}; |
| 143 | TensorND g{ |
| 144 | static_cast<char*>(tmp.raw_ptr()) + gates_ifo_layout.span().dist_byte(), |
| 145 | g_layout}; |
| 146 | auto tanh = handle->create_operator<ElemwiseForward>(); |
| 147 | tanh->param().mode = Elemwise::Param::Mode::TANH; |
| 148 | tanh->exec({g_origin}, g); |
| 149 | // sigmoid: o |
| 150 | TensorLayout three_gates_ifo_layout{ |