Implements Quantized LSTM op for R.
Also adds support for TENSOR_QUANT8_ASYMM_SIGNED in Test Generator.
Bug: 144841609
Bug: 145916330
Test: NeuralNetworksTest_static
Change-Id: I14b0d284b1945833d532cbaa33c66e4d77afd8b7
diff --git a/common/QuantUtils.cpp b/common/QuantUtils.cpp
new file mode 100644
index 0000000..97b76b7
--- /dev/null
+++ b/common/QuantUtils.cpp
@@ -0,0 +1,193 @@
+#include "QuantUtils.h"
+
+#include <algorithm>
+#include <limits>
+#include <memory>
+
+namespace android {
+namespace nn {
+
+void ApplyLayerNorm(const int16_t* input, const int16_t* layer_norm_weights, const int32_t* bias,
+ int32_t layer_norm_scale_a, int32_t layer_norm_scale_b, int32_t variance_limit,
+ int n_batch, int n_input, int16_t* output) {
+ static const int kOverflowGuard = 1 << 20;
+ for (int i = 0; i < n_batch; ++i) {
+ int64_t sum = 0;
+ int64_t sum_sq = 0;
+ for (int j = 0; j < n_input; ++j) {
+ const int32_t index = i * n_input + j;
+ int32_t val = static_cast<int32_t>(input[index]);
+ sum += val;
+ sum_sq += val * val;
+ }
+ int32_t mean = static_cast<int32_t>(static_cast<int64_t>(sum) * 1024 / n_input);
+ // TODO(jianlijianli): Avoids overflow but only works for POT n_input.
+ int32_t temp = kOverflowGuard / n_input;
+ int64_t variance = sum_sq * temp - static_cast<int64_t>(mean) * static_cast<int64_t>(mean);
+ int32_t variance2 = static_cast<int32_t>(variance / kOverflowGuard);
+ if (variance2 < 1) {
+ variance2 = variance_limit;
+ }
+ int32_t stddev_inverse_a;
+ int stddev_inverse_b;
+ GetInvSqrtQuantizedMultiplierExp(variance2, /*reverse_shift*/ -1, &stddev_inverse_a,
+ &stddev_inverse_b);
+
+ for (int j = 0; j < n_input; ++j) {
+ const int32_t index = i * n_input + j;
+ int32_t val = static_cast<int32_t>(input[index]);
+ int32_t shifted = 1024 * val - mean;
+ int32_t rescaled =
+ MultiplyByQuantizedMultiplier(shifted, stddev_inverse_a, stddev_inverse_b);
+ // TODO(jianlijianli): Saturate this.
+ int64_t val3 = rescaled * layer_norm_weights[j] + bias[j];
+ int32_t val4 = static_cast<int32_t>((val3 > 0 ? val3 + 512 : val3 - 512) / 1024);
+ int32_t val5 = MultiplyByQuantizedMultiplier(val4, layer_norm_scale_a,
+ layer_norm_scale_b + 12);
+ val5 = std::min(std::max(INT16_MIN, val5), INT16_MAX);
+ output[index] = static_cast<int16_t>(val5);
+ }
+ }
+}
+
+void MatrixScalarMultiplyAccumulate(const int8_t* matrix, int32_t scalar, int32_t n_row,
+ int32_t n_col, int32_t* output) {
+ for (int i = 0; i < n_row; ++i) {
+ int32_t row_sum = 0;
+ for (int j = 0; j < n_col; ++j) {
+ row_sum += *matrix++;
+ }
+ output[i] += row_sum * scalar;
+ }
+}
+
+bool PrecomputeZeroPointTimesWeightWithBias(int32_t zero_point, const int8_t* weight_tensor,
+ const Shape& weight_shape, const int32_t* bias_tensor,
+ std::unique_ptr<int32_t[]>* output) {
+ if (weight_tensor == nullptr) {
+ return true;
+ }
+
+ NN_RET_CHECK_EQ(weight_shape.dimensions.size(), 2u);
+ const int row = weight_shape.dimensions[0];
+ const int col = weight_shape.dimensions[1];
+ *output = std::make_unique<int32_t[]>(row);
+ if (bias_tensor == nullptr) {
+ memset(output->get(), 0, row * sizeof(int32_t));
+ } else {
+ memcpy(output->get(), bias_tensor, row * sizeof(int32_t));
+ }
+ if (zero_point != 0) {
+ MatrixScalarMultiplyAccumulate(weight_tensor, zero_point, row, col, output->get());
+ }
+ return true;
+}
+
+void ApplySigmoid(const int16_t* input, int32_t n_batch, int32_t n_input, int16_t* output) {
+ for (int batch = 0; batch < n_batch; ++batch) {
+ for (int c = 0; c < n_input; c++) {
+ using F3 = gemmlowp::FixedPoint<std::int16_t, 3>;
+ using F0 = gemmlowp::FixedPoint<std::int16_t, 0>;
+ const int index = batch * n_input + c;
+ F3 sigmoid_input = F3::FromRaw(input[index]);
+ F0 sigmoid_output = gemmlowp::logistic(sigmoid_input);
+ output[index] = sigmoid_output.raw();
+ }
+ }
+}
+
+void CwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch, int n_input, int shift,
+ int16_t* output) {
+ for (int batch = 0; batch < n_batch; ++batch) {
+ for (int i = 0; i < n_input; ++i) {
+ const int index = batch * n_input + i;
+ const int16_t a = input_1[index];
+ const int16_t b = input_2[index];
+ const int32_t value = static_cast<int32_t>(a) * static_cast<int32_t>(b);
+ output[index] = static_cast<int16_t>(gemmlowp::RoundingDivideByPOT(value, shift));
+ }
+ }
+}
+
+void CwiseMul(const int16_t* input_1, const int16_t* input_2, int32_t multiplier, int32_t shift,
+ int32_t n_batch, int32_t n_input, int32_t output_zp, int8_t* output) {
+ for (int batch = 0; batch < n_batch; ++batch) {
+ for (int i = 0; i < n_input; ++i) {
+ const int index = batch * n_input + i;
+ const int16_t a = input_1[index];
+ const int16_t b = input_2[index];
+ int32_t value = static_cast<int32_t>(a) * static_cast<int32_t>(b);
+ value = MultiplyByQuantizedMultiplier(value, multiplier, shift);
+ value -= output_zp;
+ value = std::min(std::max(-128, value), 127);
+
+ output[index] = static_cast<int8_t>(value);
+ }
+ }
+}
+
+bool CheckedLog2(const float x, int* log2_result) {
+ const float x_log2 = std::log(x) * (1.0f / std::log(2.0f));
+ const float x_log2_rounded = std::round(x_log2);
+ const float x_log2_fracpart = x_log2 - x_log2_rounded;
+
+ *log2_result = static_cast<int>(x_log2_rounded);
+ return std::abs(x_log2_fracpart) < 1e-3;
+}
+
+void CwiseAdd(const int16_t* input_1, const int16_t* input_2, int n_batch, int n_input,
+ int16_t* output) {
+ for (int batch = 0; batch < n_batch; ++batch) {
+ for (int i = 0; i < n_input; ++i) {
+ const int index = batch * n_input + i;
+ int32_t sum = input_1[index] + input_2[index];
+ const int32_t sum_clamped = std::min(INT16_MAX, std::max(INT16_MIN, sum));
+ output[index] = static_cast<int16_t>(sum_clamped);
+ }
+ }
+}
+
+void CwiseClipping(int16_t* input, const int16_t clipping_value, int32_t n_batch, int32_t n_input) {
+ for (int batch = 0; batch < n_batch; ++batch) {
+ for (int i = 0; i < n_input; ++i) {
+ const int index = batch * n_input + i;
+ if (input[index] > clipping_value) {
+ input[index] = clipping_value;
+ }
+ if (input[index] < -clipping_value) {
+ input[index] = -clipping_value;
+ }
+ }
+ }
+}
+
+void CwiseClipping(int8_t* input, const int8_t clipping_value, int32_t n_batch, int32_t n_input) {
+ for (int batch = 0; batch < n_batch; ++batch) {
+ for (int i = 0; i < n_input; ++i) {
+ const int index = batch * n_input + i;
+ if (input[index] > clipping_value) {
+ input[index] = clipping_value;
+ }
+ if (input[index] < -clipping_value) {
+ input[index] = -clipping_value;
+ }
+ }
+ }
+}
+
+void VectorBatchVectorCwiseProductAccumulate(const int16_t* vector, int v_size,
+ const int16_t* batch_vector, int n_batch,
+ int32_t multiplier, int shift, int16_t* result) {
+ for (int b = 0; b < n_batch; b++) {
+ for (int v = 0; v < v_size; v++) {
+ int32_t prod = vector[v] * *batch_vector++;
+ prod = MultiplyByQuantizedMultiplier(prod, multiplier, shift);
+ int32_t output = prod + *result;
+ output = std::max(std::min(32767, output), -32768);
+ *result++ = output;
+ }
+ }
+}
+
+} // namespace nn
+} // namespace android
