in tensorflow/tensorflow/lite/kernels/lstm_eval.cc [459:886]
inline void LstmStepWithAuxInput(
const float* input_ptr_batch, const int8_t* input_to_input_weights_ptr,
float input_to_input_weights_scale,
const int8_t* input_to_forget_weights_ptr,
float input_to_forget_weights_scale,
const int8_t* input_to_cell_weights_ptr, float input_to_cell_weights_scale,
const int8_t* input_to_output_weights_ptr,
float input_to_output_weights_scale, const float* aux_input_ptr_batch,
const int8_t* aux_input_to_input_weights_ptr,
float aux_input_to_input_weights_scale,
const int8_t* aux_input_to_forget_weights_ptr,
float aux_input_to_forget_weights_scale,
const int8_t* aux_input_to_cell_weights_ptr,
float aux_input_to_cell_weights_scale,
const int8_t* aux_input_to_output_weights_ptr,
float aux_input_to_output_weights_scale,
const int8_t* recurrent_to_input_weights_ptr,
float recurrent_to_input_weights_scale,
const int8_t* recurrent_to_forget_weights_ptr,
float recurrent_to_forget_weights_scale,
const int8_t* recurrent_to_cell_weights_ptr,
float recurrent_to_cell_weights_scale,
const int8_t* recurrent_to_output_weights_ptr,
float recurrent_to_output_weights_scale,
const int8_t* cell_to_input_weights_ptr, float cell_to_input_weights_scale,
const int8_t* cell_to_forget_weights_ptr,
float cell_to_forget_weights_scale,
const int8_t* cell_to_output_weights_ptr,
float cell_to_output_weights_scale,
const float* input_layer_norm_coefficients_ptr,
const float* forget_layer_norm_coefficients_ptr,
const float* cell_layer_norm_coefficients_ptr,
const float* output_layer_norm_coefficients_ptr,
const float* input_gate_bias_ptr, const float* forget_gate_bias_ptr,
const float* cell_bias_ptr, const float* output_gate_bias_ptr,
const int8_t* projection_weights_ptr, float projection_weights_scale,
const float* projection_bias_ptr, const TfLiteLSTMParams* params,
int n_batch, int n_cell, int n_input, int n_aux_input, int n_output,
int output_batch_leading_dim, float* input_gate_scratch,
float* forget_gate_scratch, float* cell_scratch, float* output_gate_scratch,
float* scaling_factors, float* product_scaling_factors,
float* recovered_cell_weights, int8_t* quantized_input_ptr_batch,
int8_t* quantized_aux_input_ptr_batch, int8_t* quantized_output_state_ptr,
int8_t* quantized_cell_state_ptr, float* output_state_ptr,
float* cell_state_ptr, float* output_ptr_batch) {
#ifdef GEMMLOWP_PROFILING
gemmlowp::ScopedProfilingLabel label("LstmStepWithAuxInputHybrid");
#endif
// Since we have already checked that weights are all there or none, we
// can check the existence of only one to the get the condition.
const bool use_cifg = (input_to_input_weights_ptr == nullptr);
const bool use_peephole = (cell_to_output_weights_ptr != nullptr);
const bool is_layer_norm_lstm =
(forget_layer_norm_coefficients_ptr != nullptr);
// Initialize scratch buffers with bias.
if (is_layer_norm_lstm) {
if (!use_cifg) {
std::fill_n(input_gate_scratch, n_cell * n_batch, 0.0f);
}
std::fill_n(forget_gate_scratch, n_cell * n_batch, 0.0f);
std::fill_n(cell_scratch, n_cell * n_batch, 0.0f);
std::fill_n(output_gate_scratch, n_cell * n_batch, 0.0f);
} else {
if (!use_cifg) {
tensor_utils::VectorBatchVectorAssign(input_gate_bias_ptr, n_cell,
n_batch, input_gate_scratch);
}
tensor_utils::VectorBatchVectorAssign(forget_gate_bias_ptr, n_cell, n_batch,
forget_gate_scratch);
tensor_utils::VectorBatchVectorAssign(cell_bias_ptr, n_cell, n_batch,
cell_scratch);
tensor_utils::VectorBatchVectorAssign(output_gate_bias_ptr, n_cell, n_batch,
output_gate_scratch);
}
if (!tensor_utils::IsZeroVector(input_ptr_batch, n_batch * n_input)) {
// Save quantization and matmul computation for all zero input.
float unused_min, unused_max;
for (int b = 0; b < n_batch; ++b) {
const int offset = b * n_input;
tensor_utils::SymmetricQuantizeFloats(
input_ptr_batch + offset, n_input, quantized_input_ptr_batch + offset,
&unused_min, &unused_max, &scaling_factors[b]);
}
// For each batch and cell: compute input_weight * input.
if (!use_cifg) {
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * input_to_input_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
input_to_input_weights_ptr, n_cell, n_input,
quantized_input_ptr_batch, product_scaling_factors, n_batch,
input_gate_scratch, /*result_stride=*/1);
}
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * input_to_forget_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
input_to_forget_weights_ptr, n_cell, n_input, quantized_input_ptr_batch,
product_scaling_factors, n_batch, forget_gate_scratch,
/*result_stride=*/1);
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * input_to_cell_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
input_to_cell_weights_ptr, n_cell, n_input, quantized_input_ptr_batch,
product_scaling_factors, n_batch, cell_scratch, /*result_stride=*/1);
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * input_to_output_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
input_to_output_weights_ptr, n_cell, n_input, quantized_input_ptr_batch,
product_scaling_factors, n_batch, output_gate_scratch,
/*result_stride=*/1);
}
if (aux_input_ptr_batch != nullptr &&
!tensor_utils::IsZeroVector(aux_input_ptr_batch, n_batch * n_input)) {
// Save quantization and matmul computation for all zero input.
float unused_min, unused_max;
for (int b = 0; b < n_batch; ++b) {
const int offset = b * n_input;
tensor_utils::SymmetricQuantizeFloats(
aux_input_ptr_batch + offset, n_input,
quantized_aux_input_ptr_batch + offset, &unused_min, &unused_max,
&scaling_factors[b]);
}
// For each batch and cell: compute input_weight * input.
if (!use_cifg) {
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * aux_input_to_input_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
aux_input_to_input_weights_ptr, n_cell, n_input,
quantized_aux_input_ptr_batch, product_scaling_factors, n_batch,
input_gate_scratch, /*result_stride=*/1);
}
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * aux_input_to_forget_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
aux_input_to_forget_weights_ptr, n_cell, n_input,
quantized_aux_input_ptr_batch, product_scaling_factors, n_batch,
forget_gate_scratch, /*result_stride=*/1);
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * aux_input_to_cell_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
aux_input_to_cell_weights_ptr, n_cell, n_input,
quantized_aux_input_ptr_batch, product_scaling_factors, n_batch,
cell_scratch, /*result_stride=*/1);
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * aux_input_to_output_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
aux_input_to_output_weights_ptr, n_cell, n_input,
quantized_aux_input_ptr_batch, product_scaling_factors, n_batch,
output_gate_scratch, /*result_stride=*/1);
}
if (!tensor_utils::IsZeroVector(output_state_ptr, n_batch * n_output)) {
// Save quantization and matmul computation for all zero input.
float unused_min, unused_max;
for (int b = 0; b < n_batch; ++b) {
const int offset = b * n_output;
tensor_utils::SymmetricQuantizeFloats(output_state_ptr + offset, n_output,
quantized_output_state_ptr + offset,
&unused_min, &unused_max,
&scaling_factors[b]);
}
// For each batch and cell: compute recurrent_weight * output_state.
if (!use_cifg) {
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * recurrent_to_input_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
recurrent_to_input_weights_ptr, n_cell, n_output,
quantized_output_state_ptr, product_scaling_factors, n_batch,
input_gate_scratch, /*result_stride=*/1);
}
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * recurrent_to_forget_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
recurrent_to_forget_weights_ptr, n_cell, n_output,
quantized_output_state_ptr, product_scaling_factors, n_batch,
forget_gate_scratch, /*result_stride=*/1);
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * recurrent_to_cell_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
recurrent_to_cell_weights_ptr, n_cell, n_output,
quantized_output_state_ptr, product_scaling_factors, n_batch,
cell_scratch, /*result_stride=*/1);
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * recurrent_to_output_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
recurrent_to_output_weights_ptr, n_cell, n_output,
quantized_output_state_ptr, product_scaling_factors, n_batch,
output_gate_scratch, /*result_stride=*/1);
}
// Save quantization and matmul computation for all zero input.
bool is_cell_state_all_zeros =
tensor_utils::IsZeroVector(cell_state_ptr, n_batch * n_cell);
// For each batch and cell: update input gate.
if (!use_cifg) {
if (use_peephole && !is_cell_state_all_zeros) {
tensor_utils::VectorScalarMultiply(cell_to_input_weights_ptr, n_cell,
cell_to_input_weights_scale,
recovered_cell_weights);
tensor_utils::VectorBatchVectorCwiseProductAccumulate(
recovered_cell_weights, n_cell, cell_state_ptr, n_batch,
input_gate_scratch);
}
if (is_layer_norm_lstm) {
tensor_utils::MeanStddevNormalization(input_gate_scratch,
input_gate_scratch, n_cell, n_batch,
kLayerNormEpsilon);
tensor_utils::VectorBatchVectorCwiseProduct(
input_layer_norm_coefficients_ptr, n_cell, input_gate_scratch,
n_batch, input_gate_scratch);
tensor_utils::VectorBatchVectorAdd(input_gate_bias_ptr, n_cell, n_batch,
input_gate_scratch);
}
ApplyActivationsToVector(input_gate_scratch, n_cell * n_batch,
kTfLiteActSigmoid, input_gate_scratch);
}
// For each batch and cell: update forget gate.
if (use_peephole && !is_cell_state_all_zeros) {
tensor_utils::VectorScalarMultiply(cell_to_forget_weights_ptr, n_cell,
cell_to_forget_weights_scale,
recovered_cell_weights);
tensor_utils::VectorBatchVectorCwiseProductAccumulate(
recovered_cell_weights, n_cell, cell_state_ptr, n_batch,
forget_gate_scratch);
}
if (is_layer_norm_lstm) {
tensor_utils::MeanStddevNormalization(forget_gate_scratch,
forget_gate_scratch, n_cell, n_batch,
kLayerNormEpsilon);
tensor_utils::VectorBatchVectorCwiseProduct(
forget_layer_norm_coefficients_ptr, n_cell, forget_gate_scratch,
n_batch, forget_gate_scratch);
tensor_utils::VectorBatchVectorAdd(forget_gate_bias_ptr, n_cell, n_batch,
forget_gate_scratch);
}
ApplyActivationsToVector(forget_gate_scratch, n_cell * n_batch,
kTfLiteActSigmoid, forget_gate_scratch);
// For each batch and cell: update the cell.
tensor_utils::VectorVectorCwiseProduct(forget_gate_scratch, cell_state_ptr,
n_batch * n_cell, cell_state_ptr);
if (is_layer_norm_lstm) {
tensor_utils::MeanStddevNormalization(cell_scratch, cell_scratch, n_cell,
n_batch, kLayerNormEpsilon);
tensor_utils::VectorBatchVectorCwiseProduct(
cell_layer_norm_coefficients_ptr, n_cell, cell_scratch, n_batch,
cell_scratch);
tensor_utils::VectorBatchVectorAdd(cell_bias_ptr, n_cell, n_batch,
cell_scratch);
}
ApplyActivationsToVector(cell_scratch, n_batch * n_cell, params->activation,
cell_scratch);
if (use_cifg) {
tensor_utils::Sub1Vector(forget_gate_scratch, n_batch * n_cell,
forget_gate_scratch);
tensor_utils::VectorVectorCwiseProductAccumulate(
cell_scratch, forget_gate_scratch, n_batch * n_cell, cell_state_ptr);
} else {
tensor_utils::VectorVectorCwiseProductAccumulate(
cell_scratch, input_gate_scratch, n_batch * n_cell, cell_state_ptr);
}
if (params->cell_clip > 0.0) {
tensor_utils::ClipVector(cell_state_ptr, n_batch * n_cell,
params->cell_clip, cell_state_ptr);
}
is_cell_state_all_zeros =
tensor_utils::IsZeroVector(cell_state_ptr, n_batch * n_cell);
// For each batch and cell: update the output gate.
if (use_peephole && !is_cell_state_all_zeros) {
tensor_utils::VectorScalarMultiply(cell_to_output_weights_ptr, n_cell,
cell_to_output_weights_scale,
recovered_cell_weights);
tensor_utils::VectorBatchVectorCwiseProductAccumulate(
recovered_cell_weights, n_cell, cell_state_ptr, n_batch,
output_gate_scratch);
}
if (is_layer_norm_lstm) {
tensor_utils::MeanStddevNormalization(output_gate_scratch,
output_gate_scratch, n_cell, n_batch,
kLayerNormEpsilon);
tensor_utils::VectorBatchVectorCwiseProduct(
output_layer_norm_coefficients_ptr, n_cell, output_gate_scratch,
n_batch, output_gate_scratch);
tensor_utils::VectorBatchVectorAdd(output_gate_bias_ptr, n_cell, n_batch,
output_gate_scratch);
}
ApplyActivationsToVector(output_gate_scratch, n_batch * n_cell,
kTfLiteActSigmoid, output_gate_scratch);
ApplyActivationsToVector(cell_state_ptr, n_batch * n_cell, params->activation,
cell_scratch);
tensor_utils::VectorVectorCwiseProduct(output_gate_scratch, cell_scratch,
n_batch * n_cell, output_gate_scratch);
const bool use_projection_weight = (projection_weights_ptr != nullptr);
const bool use_projection_bias = (projection_bias_ptr != nullptr);
// For each batch: update the projection and output_state. Note that since
// the output batch rows may not be contiguous (output_batch_leading_dim !=
// n_output), we unroll the batched operations where this is the case.
if (output_batch_leading_dim == n_output) {
if (use_projection_weight) {
if (use_projection_bias) {
tensor_utils::VectorBatchVectorAssign(projection_bias_ptr, n_output,
n_batch, output_ptr_batch);
} else {
std::fill_n(output_ptr_batch, n_batch * n_output, 0.0f);
}
if (!tensor_utils::IsZeroVector(output_gate_scratch, n_batch * n_cell)) {
// Save quantization and matmul computation for all zero input.
float unused_min, unused_max;
for (int b = 0; b < n_batch; ++b) {
const int offset = b * n_cell;
tensor_utils::SymmetricQuantizeFloats(
output_gate_scratch + offset, n_cell,
quantized_cell_state_ptr + offset, &unused_min, &unused_max,
&scaling_factors[b]);
}
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * projection_weights_scale;
}
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
projection_weights_ptr, n_output, n_cell, quantized_cell_state_ptr,
product_scaling_factors, n_batch, output_ptr_batch,
/*result_stride=*/1);
}
if (params->proj_clip > 0.0) {
tensor_utils::ClipVector(output_ptr_batch, n_batch * n_output,
params->proj_clip, output_ptr_batch);
}
} else {
std::copy_n(output_gate_scratch, n_batch * n_output, output_ptr_batch);
}
std::copy_n(output_ptr_batch, n_batch * n_output, output_state_ptr);
} else {
if (use_projection_weight) {
if (use_projection_bias) {
for (int k = 0; k < n_batch; k++) {
std::copy_n(projection_bias_ptr, n_output,
output_ptr_batch + k * output_batch_leading_dim);
}
} else {
for (int k = 0; k < n_batch; k++) {
std::fill_n(output_ptr_batch + k * output_batch_leading_dim, n_output,
0.0f);
}
}
if (!tensor_utils::IsZeroVector(output_gate_scratch, n_batch * n_cell)) {
// Save quantization and matmul computation for all zero input.
float unused_min, unused_max;
for (int b = 0; b < n_batch; ++b) {
const int offset = b * n_cell;
tensor_utils::SymmetricQuantizeFloats(
output_gate_scratch + offset, n_cell,
quantized_cell_state_ptr + offset, &unused_min, &unused_max,
&scaling_factors[b]);
}
for (int b = 0; b < n_batch; ++b) {
product_scaling_factors[b] =
scaling_factors[b] * projection_weights_scale;
}
for (int k = 0; k < n_batch; k++) {
tensor_utils::MatrixBatchVectorMultiplyAccumulate(
projection_weights_ptr, n_output, n_cell,
quantized_cell_state_ptr + k * n_cell,
&product_scaling_factors[k],
/*n_batch=*/1, output_ptr_batch + k * output_batch_leading_dim,
/*result_stride=*/1);
}
}
if (params->proj_clip > 0.0) {
for (int k = 0; k < n_batch; k++) {
tensor_utils::ClipVector(
output_ptr_batch + k * output_batch_leading_dim, n_output,
params->proj_clip,
output_ptr_batch + k * output_batch_leading_dim);
}
}
} else {
for (int k = 0; k < n_batch; k++) {
std::copy_n(output_gate_scratch + k * n_output, n_output,
output_ptr_batch + k * output_batch_leading_dim);
}
}
for (int k = 0; k < n_batch; k++) {
std::copy_n(output_ptr_batch + k * output_batch_leading_dim, n_output,
output_state_ptr + k * n_output);
}
}
}