| 200 | } |
| 201 | |
| 202 | TfLiteStatus EvalHybrid( |
| 203 | const TfLiteTensor* input, const TfLiteTensor* input_weights, |
| 204 | const TfLiteTensor* recurrent_weights, const TfLiteTensor* bias, |
| 205 | const TfLiteSequenceRNNParams* params, TfLiteTensor* input_scratch, |
| 206 | TfLiteTensor* hidden_state_scratch, TfLiteTensor* scaling_factors, |
| 207 | TfLiteTensor* hidden_state, TfLiteTensor* output) { |
| 208 | const bool time_major = params->time_major; |
| 209 | const int batch_size = |
| 210 | (time_major) ? input->dims->data[1] : input->dims->data[0]; |
| 211 | const int max_time = |
| 212 | (time_major) ? input->dims->data[0] : input->dims->data[1]; |
| 213 | const int num_units = input_weights->dims->data[0]; |
| 214 | const int input_size = input->dims->data[2]; |
| 215 | |
| 216 | // Initialize the pointer bias. |
| 217 | const float* bias_ptr = bias->data.f; |
| 218 | |
| 219 | // Initialize input_weights, recurrent_weights, and temporary storage for |
| 220 | // quantized values. |
| 221 | const int8_t* input_weights_ptr; |
| 222 | const int8_t* recurrent_weights_ptr; |
| 223 | int8_t* quantized_input_ptr; |
| 224 | int8_t* quantized_hidden_state_ptr; |
| 225 | if (input_weights->type == kTfLiteUInt8) { |
| 226 | input_weights_ptr = |
| 227 | reinterpret_cast<const int8_t*>(input_weights->data.uint8); |
| 228 | recurrent_weights_ptr = |
| 229 | reinterpret_cast<const int8_t*>(recurrent_weights->data.uint8); |
| 230 | quantized_input_ptr = reinterpret_cast<int8_t*>(input_scratch->data.uint8); |
| 231 | quantized_hidden_state_ptr = |
| 232 | reinterpret_cast<int8_t*>(hidden_state_scratch->data.uint8); |
| 233 | } else { |
| 234 | input_weights_ptr = input_weights->data.int8; |
| 235 | recurrent_weights_ptr = recurrent_weights->data.int8; |
| 236 | quantized_input_ptr = input_scratch->data.int8; |
| 237 | quantized_hidden_state_ptr = hidden_state_scratch->data.int8; |
| 238 | } |
| 239 | |
| 240 | // Get the scale of the quantized weights. |
| 241 | float input_weights_scale = input_weights->params.scale; |
| 242 | float recurrent_weights_scale = recurrent_weights->params.scale; |
| 243 | float* scaling_factors_ptr = scaling_factors->data.f; |
| 244 | |
| 245 | if (time_major) { |
| 246 | // Initialize the pointer to hidden state. |
| 247 | float* hidden_state_ptr_batch = hidden_state->data.f; |
| 248 | // Unroll the sequence and use batch operations for efficiency. |
| 249 | for (int s = 0; s < max_time; s++) { |
| 250 | // Initialize the pointer to input and output. |
| 251 | const float* input_ptr_batch = |
| 252 | input->data.f + s * input_size * batch_size; |
| 253 | float* output_ptr_batch = output->data.f + s * num_units * batch_size; |
| 254 | |
| 255 | kernel_utils::RnnBatchStep( |
| 256 | input_ptr_batch, input_weights_ptr, input_weights_scale, |
| 257 | recurrent_weights_ptr, recurrent_weights_scale, bias_ptr, input_size, |
| 258 | num_units, batch_size, num_units, params->activation, |
| 259 | quantized_input_ptr, quantized_hidden_state_ptr, scaling_factors_ptr, |