| 9 | namespace rnn { |
| 10 | |
| 11 | CellWeightsWrapperBase::CellWeightsWrapperBase( |
| 12 | void* weight_ptr, size_t hidden_size, size_t input_size, size_t num_chunks, |
| 13 | bool has_bias, DType dtype, _megdnn_workspace workspace) { |
| 14 | // weight_ih: [gate_hidden_size, input_size] |
| 15 | // weight_hh: [gate_hidden_size, hidden_size] |
| 16 | // bias_ih: [gate_hidden_size] |
| 17 | // bias_hh: [gate_hidden_size] |
| 18 | size_t gate_hidden_size = num_chunks * hidden_size; |
| 19 | TensorLayout weight_ih_layout{{gate_hidden_size, input_size}, dtype}; |
| 20 | TensorLayout weight_hh_layout{{gate_hidden_size, hidden_size}, dtype}; |
| 21 | TensorLayout bias_layout{{gate_hidden_size}, dtype}; |
| 22 | this->_weight_size = 0; |
| 23 | this->weight_ih = TensorND(weight_ptr, weight_ih_layout); |
| 24 | this->_weight_size += weight_ih_layout.span().dist_byte(); |
| 25 | this->weight_hh = TensorND( |
| 26 | static_cast<uint8_t*>(weight_ptr) + this->_weight_size, weight_hh_layout); |
| 27 | this->_weight_size += weight_hh_layout.span().dist_byte(); |
| 28 | if (has_bias) { |
| 29 | this->bias_ih = TensorND( |
| 30 | static_cast<uint8_t*>(weight_ptr) + this->_weight_size, bias_layout); |
| 31 | this->_weight_size += bias_layout.span().dist_byte(); |
| 32 | this->bias_hh = TensorND( |
| 33 | static_cast<uint8_t*>(weight_ptr) + this->_weight_size, bias_layout); |
| 34 | this->_weight_size += bias_layout.span().dist_byte(); |
| 35 | this->_workspace_size = 0; |
| 36 | } else { |
| 37 | this->bias_ih = TensorND(workspace.raw_ptr, bias_layout); |
| 38 | this->bias_hh = TensorND(workspace.raw_ptr, bias_layout); |
| 39 | memset(workspace.raw_ptr, 0, bias_layout.span().dist_byte()); |
| 40 | this->_workspace_size = bias_layout.span().dist_byte(); |
| 41 | } |
| 42 | } |
| 43 | |
| 44 | size_t CellWeightsWrapperBase::weight_size_in_bytes() const { |
| 45 | return this->_weight_size; |