chore: import upstream snapshot with attribution

tensorflow/lite/kernels/internal/averagepool_quantized_test.cc

@@ -0,0 +1,196 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <functional>
+#include <iterator>
+#include <limits>
+#include <random>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h"
+#include "tensorflow/lite/kernels/internal/test_util.h"
+
+namespace tflite {
+namespace {
+
+// Runs the reference and optimized AveragePool functions and asserts the values
+// are the same.
+void RunOneAveragePoolTest(const PoolParams& params,
+                           const RuntimeShape& input_shape,
+                           const int8_t* input_data,
+                           const RuntimeShape& output_shape) {
+  const int buffer_size = output_shape.FlatSize();
+  std::vector<int8_t> optimized_averagePool_output(buffer_size);
+  std::vector<int8_t> reference_averagePool_output(buffer_size);
+
+  bool reference_success = reference_integer_ops::AveragePool(
+      params, input_shape, input_data, output_shape,
+      reference_averagePool_output.data());
+  bool optimized_success = optimized_integer_ops::AveragePool(
+      params, input_shape, input_data, output_shape,
+      optimized_averagePool_output.data());
+  EXPECT_TRUE(reference_success);
+  EXPECT_TRUE(optimized_success);
+
+  for (int i = 0; i < buffer_size; i++) {
+    EXPECT_TRUE(reference_averagePool_output[i] ==
+                optimized_averagePool_output[i]);
+  }
+}
+
+// Creates random input shape (batch, height, width, depth), then computes
+// output shape based on value of `padding_same`:
+// `padding_same` == true, calculate output with padding == "SAME"
+// `padding_same` == false, calculate output with padding == "VALID"
+// With input/output shapes computed, fills the input data and calls the
+// test function.
+void CreateDataAndRunAveragePool(bool padding_same) {
+  const int batch = UniformRandomInt(1, 2);
+  const int input_depth = UniformRandomInt(1, 700);
+  const int output_depth = input_depth;
+  const int input_width_offset = UniformRandomInt(1, 30);
+  const int input_height_offset = UniformRandomInt(1, 30);
+  const int stride_width = UniformRandomInt(1, 10);
+  const int stride_height = UniformRandomInt(1, 10);
+  const int filter_width = UniformRandomInt(1, 10);
+  const int filter_height = UniformRandomInt(1, 10);
+  const int input_width = input_width_offset + filter_width;
+  const int input_height = input_height_offset + filter_height;
+  const int output_width =
+      padding_same ? (input_width + stride_width - 1) / stride_width
+                   : (input_width - filter_width + stride_width) / stride_width;
+  const int output_height =
+      padding_same
+          ? (input_height + stride_height - 1) / stride_height
+          : (input_height - filter_height + stride_height) / stride_height;
+
+  auto input_shape =
+      RuntimeShape({batch, input_height, input_width, input_depth});
+  auto output_shape =
+      RuntimeShape({batch, output_height, output_width, output_depth});
+  const int buffer_size = input_shape.FlatSize();
+  std::vector<int8_t> input_data(buffer_size);
+  FillRandom(&input_data);
+
+  PoolParams params;
+  params.stride_height = stride_height;
+  params.stride_width = stride_width;
+  params.filter_height = filter_height;
+  params.filter_width = filter_width;
+  params.quantized_activation_min =
+      static_cast<int8_t>(std::numeric_limits<int8_t>::lowest());
+  params.quantized_activation_max =
+      static_cast<int8_t>(std::numeric_limits<int8_t>::max());
+  auto compute_padding = [](int stride, int in_size, int filter_size,
+                            int out_size) {
+    int padding = ((out_size - 1) * stride + filter_size - in_size) / 2;
+    return padding > 0 ? padding : 0;
+  };
+  params.padding_values.width =
+      compute_padding(stride_width, input_width, filter_width, output_width);
+  params.padding_values.height = compute_padding(stride_height, input_height,
+                                                 filter_height, output_height);
+  RunOneAveragePoolTest(params, input_shape, input_data.data(), output_shape);
+}
+
+TEST(TestAveragePool, SymmetricQuantAveragePool) {
+  const int kTestsToRun = 10;
+  for (int i = 0; i < kTestsToRun; i++) {
+    CreateDataAndRunAveragePool(/*padding_same=*/true);
+    CreateDataAndRunAveragePool(/*padding_same=*/false);
+  }
+}
+
+// Creates random input shape (batch, height, width, depth), then computes
+// output shape based on value of `padding_same`:
+// `padding_same` == true, calculate output with padding == "SAME"
+// `padding_same` == false, calculate output with padding == "VALID"
+// With input/output shapes computed, fills the input data and calls the
+// test function.
+void CreateExtremalDataAndRunAveragePool(bool padding_same) {
+  const int batch = UniformRandomInt(1, 2);
+  const int input_depth = UniformRandomInt(1, 700);
+  const int output_depth = input_depth;
+  const int input_width_offset = UniformRandomInt(1, 30);
+  const int input_height_offset = UniformRandomInt(1, 30);
+  const int stride_width = UniformRandomInt(1, 128);
+  const int stride_height = UniformRandomInt(1, 128);
+  const int filter_width = UniformRandomInt(1, 28);
+  const int filter_height = UniformRandomInt(1, 28);
+  if (filter_width * filter_height > 64) {
+    std::cout << "should test 32 version" << std::endl;
+  }
+  const int input_width = input_width_offset + filter_width;
+  const int input_height = input_height_offset + filter_height;
+  const int output_width =
+      padding_same ? (input_width + stride_width - 1) / stride_width
+                   : (input_width - filter_width + stride_width) / stride_width;
+  const int output_height =
+      padding_same
+          ? (input_height + stride_height - 1) / stride_height
+          : (input_height - filter_height + stride_height) / stride_height;
+
+  auto input_shape =
+      RuntimeShape({batch, input_height, input_width, input_depth});
+  auto output_shape =
+      RuntimeShape({batch, output_height, output_width, output_depth});
+
+  PoolParams params;
+  params.stride_height = stride_height;
+  params.stride_width = stride_width;
+  params.filter_height = filter_height;
+  params.filter_width = filter_width;
+  params.quantized_activation_min =
+      static_cast<int8_t>(std::numeric_limits<int8_t>::lowest());
+  params.quantized_activation_max =
+      static_cast<int8_t>(std::numeric_limits<int8_t>::max());
+  auto compute_padding = [](int stride, int in_size, int filter_size,
+                            int out_size) {
+    int padding = ((out_size - 1) * stride + filter_size - in_size) / 2;
+    return padding > 0 ? padding : 0;
+  };
+  params.padding_values.width =
+      compute_padding(stride_width, input_width, filter_width, output_width);
+  params.padding_values.height = compute_padding(stride_height, input_height,
+                                                 filter_height, output_height);
+
+  const int buffer_size = input_shape.FlatSize();
+  std::vector<int8_t> input_data(buffer_size);
+
+  // Test small values
+  int8_t min = std::numeric_limits<int8_t>::min();
+  int8_t max = std::numeric_limits<int8_t>::min() + 10;
+  FillRandom(&input_data, min, max);
+  RunOneAveragePoolTest(params, input_shape, input_data.data(), output_shape);
+
+  // Test large values
+  min = std::numeric_limits<int8_t>::max() - 10;
+  max = std::numeric_limits<int8_t>::max();
+  FillRandom(&input_data, min, max);
+  RunOneAveragePoolTest(params, input_shape, input_data.data(), output_shape);
+}
+
+TEST(TestAveragePool, SymmetricQuantExtremalAveragePool) {
+  CreateExtremalDataAndRunAveragePool(/*padding_same=*/true);
+  CreateExtremalDataAndRunAveragePool(/*padding_same=*/false);
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/batch_to_space_nd_test.cc

@@ -0,0 +1,98 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+
+#include <gtest/gtest.h>
+
+namespace tflite {
+namespace {
+
+// A light wrapper of GetIndexRange which returns a pair of start / end
+// indices.
+std::pair<int, int> GetIndexRange(int spatial_index_dim, int block_shape_dim,
+                                  int input_dim, int output_dim) {
+  int index_start = 0;
+  int index_end = 0;
+  optimized_ops::GetIndexRange(spatial_index_dim, block_shape_dim, input_dim,
+                               output_dim, &index_start, &index_end);
+  return {index_start, index_end};
+}
+
+TEST(BatchToSpaceNDTest, TestIndexRange) {
+  // Simple test case, no cropping.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/3, /*block_shape_dim=*/6,
+                          /*input_dim=*/1, /*output_dim=*/6),
+            std::make_pair(0, 1));
+
+  // No cropping and input_dim > 1.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/2, /*block_shape_dim=*/6,
+                          /*input_dim=*/5, /*output_dim=*/30),
+            std::make_pair(0, 5));
+
+  // With small cropping values (can be either at the beginning or at the end).
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/2,
+                          /*input_dim=*/3, /*output_dim=*/4),
+            std::make_pair(0, 2));
+
+  // With positive cropping values at the beginning.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-2, /*block_shape_dim=*/2,
+                          /*input_dim=*/3, /*output_dim=*/4),
+            std::make_pair(1, 3));
+
+  // Large crop at the beginning.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-30, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/5),
+            std::make_pair(6, 7));
+
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-26, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/5),
+            std::make_pair(6, 7));
+
+  // Large crop at the end.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/5),
+            std::make_pair(0, 1));
+
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/4, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/5),
+            std::make_pair(0, 1));
+
+  // Rounding up incorrectly will fail this test.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/3, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/5),
+            std::make_pair(0, 1));
+
+  // Extreme cropping with output of a single spatial location.
+  // Valid position 1, when large crop at the end.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/0, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/1),
+            std::make_pair(0, 1));
+
+  // Valid position 2, when large crop at the beginning.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-30, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/1),
+            std::make_pair(6, 7));
+
+  // Invalid positions.
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/1, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/1),
+            std::make_pair(0, 0));
+  EXPECT_EQ(GetIndexRange(/*spatial_index_dim=*/-29, /*block_shape_dim=*/5,
+                          /*input_dim=*/7, /*output_dim=*/1),
+            std::make_pair(6, 6));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/common.cc

@@ -0,0 +1,103 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/internal/common.h"
+
+namespace tflite {
+
+// Single-rounding MultiplyByQuantizedMultiplier
+#if TFLITE_SINGLE_ROUNDING
+int32_t MultiplyByQuantizedMultiplier(int32_t x, int32_t quantized_multiplier,
+                                      int shift) {
+  TFLITE_DCHECK(quantized_multiplier >= 0);
+  TFLITE_DCHECK(shift >= -31 && shift <= 30);
+
+  const int64_t total_shift = 31 - shift;
+  const int64_t round = static_cast<int64_t>(1) << (total_shift - 1);
+  int64_t result = x * static_cast<int64_t>(quantized_multiplier) + round;
+  result = result >> total_shift;
+
+  TFLITE_DCHECK(result >= std::numeric_limits<int32_t>::min() &&
+                result <= std::numeric_limits<int32_t>::max());
+  return static_cast<int32_t>(result);
+}
+
+int32_t MultiplyByQuantizedMultiplier(int64_t x, int32_t quantized_multiplier,
+                                      int shift) {
+  // Inputs:
+  // - quantized_multiplier has fixed point at bit 31
+  // - shift is -31 to +7 (negative for right shift)
+  //
+  // Assumptions: The following input ranges are assumed
+  // - quantize_scale>=0  (the usual range is (1<<30) to (1>>31)-1)
+  // - scaling is chosen so final scaled result fits in int32_t
+  // - input x is in the range -(1<<47) <= x < (1<<47)
+  TFLITE_DCHECK(quantized_multiplier >= 0);
+  TFLITE_DCHECK(shift >= -31 && shift < 8);
+  TFLITE_DCHECK(x >= -(static_cast<int64_t>(1) << 47) &&
+                x < (static_cast<int64_t>(1) << 47));
+
+  const int32_t reduced_multiplier =
+      (quantized_multiplier < 0x7FFF0000)
+          ? ((quantized_multiplier + (1 << 15)) >> 16)
+          : 0x7FFF;
+  const int64_t total_shift = 15 - shift;
+  const int64_t round = static_cast<int64_t>(1) << (total_shift - 1);
+  int64_t result = x * static_cast<int64_t>(reduced_multiplier) + round;
+  result = result >> total_shift;
+
+  TFLITE_DCHECK(result >= std::numeric_limits<int32_t>::min() &&
+                result <= std::numeric_limits<int32_t>::max());
+  return static_cast<int32_t>(result);
+}
+// Double-rounding MultiplyByQuantizedMultiplier
+#else
+int32_t MultiplyByQuantizedMultiplier(int32_t x, int32_t quantized_multiplier,
+                                      int shift) {
+  using gemmlowp::RoundingDivideByPOT;
+  using gemmlowp::SaturatingRoundingDoublingHighMul;
+  int left_shift = shift > 0 ? shift : 0;
+  int right_shift = shift > 0 ? 0 : -shift;
+  return RoundingDivideByPOT(SaturatingRoundingDoublingHighMul(
+                                 x * (1 << left_shift), quantized_multiplier),
+                             right_shift);
+}
+
+int32_t MultiplyByQuantizedMultiplier(int64_t x, int32_t quantized_multiplier,
+                                      int shift) {
+  // Inputs:
+  // - quantized_multiplier has fixed point at bit 31
+  // - shift is -31 to +7 (negative for right shift)
+  //
+  // Assumptions: The following input ranges are assumed
+  // - quantize_scale>=0  (the usual range is (1<<30) to (1>>31)-1)
+  // - scaling is chosen so final scaled result fits in int32_t
+  // - input x is in the range -(1<<47) <= x < (1<<47)
+  assert(quantized_multiplier >= 0);
+  assert(shift >= -31 && shift < 8);
+  assert(x >= -(static_cast<int64_t>(1) << 47) &&
+         x < (static_cast<int64_t>(1) << 47));
+
+  int32_t reduced_multiplier = (quantized_multiplier < 0x7FFF0000)
+                                   ? ((quantized_multiplier + (1 << 15)) >> 16)
+                                   : 0x7FFF;
+  int total_shift = 15 - shift;
+  x = (x * (int64_t)reduced_multiplier) + ((int64_t)1 << (total_shift - 1));
+  int32_t result = x >> total_shift;
+  return result;
+}
+#endif  // TFLITE_SINGLE_ROUNDING
+
+}  // namespace tflite

tensorflow/lite/kernels/internal/compatibility.h

@@ -0,0 +1,122 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_
+
+#include <cstdint>
+
+#include "tensorflow/lite/kernels/op_macros.h"
+
+#ifndef TFLITE_DCHECK
+#define TFLITE_DCHECK(condition) (condition) ? (void)0 : TFLITE_ASSERT_FALSE
+#endif
+
+#ifndef TFLITE_DCHECK_EQ
+#define TFLITE_DCHECK_EQ(x, y) ((x) == (y)) ? (void)0 : TFLITE_ASSERT_FALSE
+#endif
+
+#ifndef TFLITE_DCHECK_NE
+#define TFLITE_DCHECK_NE(x, y) ((x) != (y)) ? (void)0 : TFLITE_ASSERT_FALSE
+#endif
+
+#ifndef TFLITE_DCHECK_GE
+#define TFLITE_DCHECK_GE(x, y) ((x) >= (y)) ? (void)0 : TFLITE_ASSERT_FALSE
+#endif
+
+#ifndef TFLITE_DCHECK_GT
+#define TFLITE_DCHECK_GT(x, y) ((x) > (y)) ? (void)0 : TFLITE_ASSERT_FALSE
+#endif
+
+#ifndef TFLITE_DCHECK_LE
+#define TFLITE_DCHECK_LE(x, y) ((x) <= (y)) ? (void)0 : TFLITE_ASSERT_FALSE
+#endif
+
+#ifndef TFLITE_DCHECK_LT
+#define TFLITE_DCHECK_LT(x, y) ((x) < (y)) ? (void)0 : TFLITE_ASSERT_FALSE
+#endif
+
+// TODO(ahentz): Clean up: We should stick to the DCHECK versions.
+#ifndef TFLITE_CHECK
+#define TFLITE_CHECK(condition) (condition) ? (void)0 : TFLITE_ABORT
+#endif
+
+#ifndef TFLITE_CHECK_EQ
+#define TFLITE_CHECK_EQ(x, y) ((x) == (y)) ? (void)0 : TFLITE_ABORT
+#endif
+
+#ifndef TFLITE_CHECK_NE
+#define TFLITE_CHECK_NE(x, y) ((x) != (y)) ? (void)0 : TFLITE_ABORT
+#endif
+
+#ifndef TFLITE_CHECK_GE
+#define TFLITE_CHECK_GE(x, y) ((x) >= (y)) ? (void)0 : TFLITE_ABORT
+#endif
+
+#ifndef TFLITE_CHECK_GT
+#define TFLITE_CHECK_GT(x, y) ((x) > (y)) ? (void)0 : TFLITE_ABORT
+#endif
+
+#ifndef TFLITE_CHECK_LE
+#define TFLITE_CHECK_LE(x, y) ((x) <= (y)) ? (void)0 : TFLITE_ABORT
+#endif
+
+#ifndef TFLITE_CHECK_LT
+#define TFLITE_CHECK_LT(x, y) ((x) < (y)) ? (void)0 : TFLITE_ABORT
+#endif
+
+#ifndef TF_LITE_STATIC_MEMORY
+// TODO(b/162019032): Consider removing these type-aliases.
+using int8 = std::int8_t;
+using uint8 = std::uint8_t;
+using int16 = std::int16_t;
+using uint16 = std::uint16_t;
+using int32 = std::int32_t;
+using uint32 = std::uint32_t;
+#endif  // !defined(TF_LITE_STATIC_MEMORY)
+
+// Allow for cross-compiler usage of function signatures - currently used for
+// specifying named RUY profiler regions in templated methods.
+#if defined(_MSC_VER)
+#define TFLITE_PRETTY_FUNCTION __FUNCSIG__
+#elif defined(__GNUC__)
+#define TFLITE_PRETTY_FUNCTION __PRETTY_FUNCTION__
+#else
+#define TFLITE_PRETTY_FUNCTION __func__
+#endif
+
+// TFLITE_DEPRECATED()
+//
+// Duplicated from absl/base/macros.h to avoid pulling in that library.
+// Marks a deprecated class, struct, enum, function, method and variable
+// declarations. The macro argument is used as a custom diagnostic message (e.g.
+// suggestion of a better alternative).
+//
+// Example:
+//
+//   class TFLITE_DEPRECATED("Use Bar instead") Foo {...};
+//   TFLITE_DEPRECATED("Use Baz instead") void Bar() {...}
+//
+// Every usage of a deprecated entity will trigger a warning when compiled with
+// clang's `-Wdeprecated-declarations` option. This option is turned off by
+// default, but the warnings will be reported by clang-tidy.
+#if defined(__clang__) && __cplusplus >= 201103L
+#define TFLITE_DEPRECATED(message) __attribute__((deprecated(message)))
+#endif
+
+#ifndef TFLITE_DEPRECATED
+#define TFLITE_DEPRECATED(message)
+#endif
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_COMPATIBILITY_H_

tensorflow/lite/kernels/internal/constants.h

@@ -0,0 +1,61 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_CONSTANTS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_CONSTANTS_H_
+
+// Maths constants.
+// The following macros are not always available on all platforms.
+// E.g. MSVC requires additional compile flag to export those.
+#ifndef M_E
+#define M_E 2.7182818284590452354 /* e */
+#endif
+#ifndef M_LOG2E
+#define M_LOG2E 1.4426950408889634074 /* log_2 e */
+#endif
+#ifndef M_LOG10E
+#define M_LOG10E 0.43429448190325182765 /* log_10 e */
+#endif
+#ifndef M_LN2
+#define M_LN2 0.69314718055994530942 /* log_e 2 */
+#endif
+#ifndef M_LN10
+#define M_LN10 2.30258509299404568402 /* log_e 10 */
+#endif
+#ifndef M_PI
+#define M_PI 3.14159265358979323846 /* pi */
+#endif
+#ifndef M_PI_2
+#define M_PI_2 1.57079632679489661923 /* pi/2 */
+#endif
+#ifndef M_PI_4
+#define M_PI_4 0.78539816339744830962 /* pi/4 */
+#endif
+#ifndef M_1_PI
+#define M_1_PI 0.31830988618379067154 /* 1/pi */
+#endif
+#ifndef M_2_PI
+#define M_2_PI 0.63661977236758134308 /* 2/pi */
+#endif
+#ifndef M_2_SQRTPI
+#define M_2_SQRTPI 1.12837916709551257390 /* 2/sqrt(pi) */
+#endif
+#ifndef M_SQRT2
+#define M_SQRT2 1.41421356237309504880 /* sqrt(2) */
+#endif
+#ifndef M_SQRT1_2
+#define M_SQRT1_2 0.70710678118654752440 /* 1/sqrt(2) */
+#endif
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_CONSTANTS_H_

tensorflow/lite/kernels/internal/conv_per_channel_quantized_16x8_test.cc

@@ -0,0 +1,332 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <stdio.h>
+#include <sys/types.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <iterator>
+#include <limits>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/conv.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/conv.h"
+#include "tensorflow/lite/kernels/internal/test_util.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace {
+
+void PickOutputMultiplier(
+    const ConvParams& params, const RuntimeShape& input_shape,
+    const int16_t* input_data, const RuntimeShape& filter_shape,
+    const int8_t* filter_data, const RuntimeShape& bias_shape,
+    const std::int64_t* bias_data, const RuntimeShape& output_shape,
+    float* output_multiplier) {
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int dilation_width_factor = params.dilation_width_factor;
+  const int dilation_height_factor = params.dilation_height_factor;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int input_depth = input_shape.Dims(3);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  const int output_depth = output_shape.Dims(3);
+
+  std::int64_t output_accu_min = std::numeric_limits<std::int64_t>::max();
+  std::int64_t output_accu_max = std::numeric_limits<std::int64_t>::min();
+
+  for (int batch = 0; batch < batches; ++batch) {
+    for (int out_y = 0; out_y < output_height; ++out_y) {
+      for (int out_x = 0; out_x < output_width; ++out_x) {
+        for (int output_channel = 0; output_channel < output_depth;
+             ++output_channel) {
+          const int in_x_origin = (out_x * stride_width) - pad_width;
+          const int in_y_origin = (out_y * stride_height) - pad_height;
+          std::int64_t acc = 0;
+          for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
+            for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
+              for (int in_channel = 0; in_channel < input_depth; ++in_channel) {
+                const int in_x = in_x_origin + dilation_width_factor * filter_x;
+                const int in_y =
+                    in_y_origin + dilation_height_factor * filter_y;
+                // Zero padding by omitting the areas outside the image.
+                const bool is_point_inside_image =
+                    (in_x >= 0) && (in_x < input_width) && (in_y >= 0) &&
+                    (in_y < input_height);
+                if (is_point_inside_image) {
+                  int32_t input_val = input_data[Offset(
+                      input_shape, batch, in_y, in_x, in_channel)];
+                  int32_t filter_val =
+                      filter_data[Offset(filter_shape, output_channel, filter_y,
+                                         filter_x, in_channel)];
+                  acc += static_cast<std::int64_t>(filter_val) *
+                         static_cast<std::int64_t>(input_val);
+                }
+              }
+            }
+          }
+          if (bias_data) {
+            acc += bias_data[output_channel];
+          }
+          output_accu_max = std::max(acc, output_accu_max);
+          output_accu_min = std::min(acc, output_accu_min);
+        }
+      }
+    }
+  }
+
+  // Since int16 ranges from -32768 to 32767, we need to squeeze the accumulator
+  // min/max fit in those ranges correspondingly as much as possible.
+  if (std::abs(output_accu_max) > std::abs(output_accu_min)) {
+    *output_multiplier = 32767.0f / std::abs(output_accu_max);
+  } else {
+    *output_multiplier = 32768.0f / std::abs(output_accu_min);
+  }
+}
+
+void PickReasonableMultiplier(
+    const ConvParams& params, int output_activation_min,
+    int output_activation_max, int output_depth,
+    const RuntimeShape& input_shape_inference, const std::int16_t* input_data,
+    const RuntimeShape& filter_shape_inference, const std::int8_t* filter_data,
+    const RuntimeShape& bias_shape_inference, const std::int64_t* bias_data,
+    const RuntimeShape& output_shape_inference,
+    std::int32_t* output_multiplier_ptr, std::int32_t* output_shift_ptr,
+    std::int16_t* output_data) {
+  float output_multiplier;
+  PickOutputMultiplier(params, input_shape_inference, input_data,
+                       filter_shape_inference, filter_data,
+                       bias_shape_inference, bias_data, output_shape_inference,
+                       &output_multiplier);
+
+  int base_multiplier;
+  int base_shift;
+  QuantizeMultiplier(output_multiplier, &base_multiplier, &base_shift);
+  for (int i = 0; i < output_depth; ++i) {
+    // multipliers typically range in [2^30 ; 2^31 - 1].
+    // Values in [0, 2^30 - 1] are normally unused, but harmless.
+    // Thus a good way to randomize multipliers is to subtract from them
+    // a random value smaller than 2^30 but still significant compared to it.
+    output_multiplier_ptr[i] = base_multiplier - (std::rand() % (1 << 26));
+    output_shift_ptr[i] = base_shift - 1 + (std::rand() % 4);
+  }
+}
+
+bool GenerateValidShapeConfigurations(
+    int filter_width, int filter_height, int dilation_width_factor,
+    int dilation_height_factor, RuntimeShape* input_shape_inference,
+    RuntimeShape* filter_shape_inference, RuntimeShape* output_shape_inference,
+    int* pad_width, int* pad_height, int* stride) {
+  const int batch = UniformRandomInt(1, 3);
+  const int input_depth = 8 * ExponentialRandomPositiveInt(0.9f, 10, 50);
+  const int input_width = UniformRandomInt(5, 50);
+  const int input_height = UniformRandomInt(5, 50);
+  *stride = UniformRandomInt(1, 2);
+  const bool test_pad = UniformRandomInt(0, 1);
+  const auto padding_type = test_pad ? PaddingType::kValid : PaddingType::kSame;
+
+  const int output_depth = 8 * ExponentialRandomPositiveInt(0.9f, 10, 50);
+
+  input_shape_inference->BuildFrom(
+      {batch, input_height, input_width, input_depth});
+
+  filter_shape_inference->BuildFrom(
+      {output_depth, filter_height, filter_width, input_depth});
+
+  EXPECT_TRUE(ComputeConvSizes(
+      *input_shape_inference, output_depth, filter_width, filter_height,
+      *stride, dilation_width_factor, dilation_height_factor, padding_type,
+      output_shape_inference, pad_width, pad_height));
+
+  return true;
+}
+
+void IntToFloat(std::vector<float>* d, std::vector<std::int8_t>* s) {
+  for (unsigned int i = 0; i < s->size(); i++) {
+    d->data()[i] = (float)s->data()[i];
+  }
+}
+
+void IntToFloat(std::vector<float>* d, std::vector<std::int64_t>* s) {
+  for (unsigned int i = 0; i < s->size(); i++) {
+    d->data()[i] = (float)s->data()[i];
+  }
+}
+
+void TryTestOneConvFilter(int test_num) {
+  const int filter_width = UniformRandomInt(2, 5);
+  const int filter_height = UniformRandomInt(2, 5);
+  std::cout << "Test number " << test_num << " (" << filter_width << ","
+            << filter_height << ")\n";
+  // We don't support dilations in the 3x3 filter.
+  const int dilation_width_factor = 1;
+  const int dilation_height_factor = 1;
+
+  const int output_activation_min = -32768;
+  const int output_activation_max = 32767;
+
+  RuntimeShape input_shape_inference;
+  RuntimeShape filter_shape_inference;
+  RuntimeShape output_shape_inference;
+  int pad_width, pad_height;
+  int stride;
+
+  // Keeps trying until we get valid shape/configurations for 3x3 filter case.
+  bool generated_valid_configurations_for_3x3_kernel = false;
+  while (!generated_valid_configurations_for_3x3_kernel) {
+    generated_valid_configurations_for_3x3_kernel =
+        GenerateValidShapeConfigurations(
+            filter_width, filter_height, dilation_width_factor,
+            dilation_height_factor, &input_shape_inference,
+            &filter_shape_inference, &output_shape_inference, &pad_width,
+            &pad_height, &stride);
+  }
+
+  const int output_depth = output_shape_inference.Dims(3);
+
+  RuntimeShape bias_shape_inference({1, 1, 1, output_depth});
+  const int input_buffer_size = input_shape_inference.FlatSize();
+  const int filter_buffer_size = filter_shape_inference.FlatSize();
+  const int output_buffer_size = output_shape_inference.FlatSize();
+  std::vector<std::int16_t> input_data(input_buffer_size);
+  std::vector<std::int8_t> filter_data(filter_buffer_size);
+  std::vector<std::int64_t> bias_data(output_depth);
+
+  if (test_num & 1) {
+    // Use high values samples to give large accumulator
+    FillRandom(&input_data, (std::int16_t)32700, (std::int16_t)32767);
+    FillRandom(&filter_data, (std::int8_t)120, (std::int8_t)127);
+  } else {
+    FillRandom(&input_data);
+    FillRandom(&filter_data);
+  }
+  for (int i = 0; i < output_depth; i++) {
+    bias_data.data()[i] = 0;
+  }
+
+  ConvParams params;
+  params.stride_width = stride;
+  params.stride_height = stride;
+  params.dilation_height_factor = dilation_height_factor;
+  params.dilation_width_factor = dilation_width_factor;
+  params.padding_values.width = pad_width;
+  params.padding_values.height = pad_height;
+  params.weights_offset = 0;
+  params.quantized_activation_min = output_activation_min;
+  params.quantized_activation_max = output_activation_max;
+  params.float_activation_max = (float)(1LL << 40);
+  params.float_activation_min = -params.float_activation_max;
+
+  std::vector<std::int16_t> reference_output_data(output_buffer_size);
+  std::vector<std::int16_t> neon_output_data(output_buffer_size);
+
+  std::vector<std::int32_t> output_multiplier(output_depth);
+  std::vector<std::int32_t> output_shift(output_depth);
+
+  // It's hard to come up with a right multiplier, random guess basically makes
+  // all the results saturated and becomes meaningfulless, so we first use
+  // reference impl to poke the min/max value of the accumulation, then use that
+  // value as a guided suggestion for us to populate meaningful mulitplier &
+  // shift.
+  PickReasonableMultiplier(
+      params, output_activation_min, output_activation_max, output_depth,
+      input_shape_inference, input_data.data(), filter_shape_inference,
+      filter_data.data(), bias_shape_inference, bias_data.data(),
+      output_shape_inference, output_multiplier.data(), output_shift.data(),
+      reference_output_data.data());
+
+  // The following tests compare referene impl and Neon general impl agrees,
+  // and reference impl loosely agrees with fast kernel since they use different
+  // rounding strategy.
+  reference_integer_ops::ConvPerChannel(
+      params, output_multiplier.data(), output_shift.data(),
+      input_shape_inference, input_data.data(), filter_shape_inference,
+      filter_data.data(), bias_shape_inference, bias_data.data(),
+      output_shape_inference, reference_output_data.data());
+
+  std::vector<float> input_data_float(input_buffer_size);
+  std::vector<float> filter_data_float(filter_buffer_size);
+  std::vector<float> bias_data_float(output_depth);
+  std::vector<float> output_data_float(output_buffer_size);
+
+  for (int i = 0; i < input_buffer_size; i++) {
+    input_data_float.data()[i] = (float)(input_data.data()[i]);
+  }
+  IntToFloat(&filter_data_float, &filter_data);
+  IntToFloat(&bias_data_float, &bias_data);
+  RuntimeShape im2col_shape;
+  float im2col_data;
+
+  reference_ops::Conv(params, input_shape_inference, input_data_float.data(),
+                      filter_shape_inference, filter_data_float.data(),
+                      bias_shape_inference, bias_data_float.data(),
+                      output_shape_inference, output_data_float.data(),
+                      im2col_shape, &im2col_data);
+
+  for (int n = 0; n < output_shape_inference.Dims(0); n++) {
+    for (int h = 0; h < output_shape_inference.Dims(1); h++) {
+      for (int w = 0; w < output_shape_inference.Dims(2); w++) {
+        for (int c = 0; c < output_shape_inference.Dims(3); c++) {
+          int offset = Offset(output_shape_inference, n, h, w, c);
+          float float_res = output_data_float.data()[offset];
+          int16_t int16_res = reference_output_data.data()[offset];
+          int32_t output_mul = output_multiplier.data()[c];
+          int shift = output_shift.data()[c];
+          float scale = (float)output_mul / (float)(1ULL << 31);
+          if (shift > 0) scale = scale * (float)(1 << shift);
+          if (shift < 0) scale = scale / (float)(1 << -shift);
+          int ref_res = floor(float_res * scale + 0.5);
+          if (ref_res < output_activation_min) ref_res = output_activation_min;
+          if (ref_res > output_activation_max) ref_res = output_activation_max;
+          int e = (ref_res - int16_res);
+          if (e < 0) e = -e;
+          if (e > 2) {
+            ADD_FAILURE() << "(" << n << ", " << h << ", " << w << ", " << c
+                          << ")"
+                          << " scale=" << output_mul << " shift=" << shift
+                          << " res=" << int16_res
+                          << " float=" << float_res * scale << " (" << float_res
+                          << ", " << scale << ")";
+            EXPECT_TRUE(false);
+          }
+        }
+      }
+    }
+  }
+}
+
+TEST(QuantizedConvPerChannelTest, FastKernelTest) {
+  for (int i = 0; i < 30; ++i) {
+    TryTestOneConvFilter(i);
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/cppmath.h

@@ -0,0 +1,40 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_CPPMATH_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_CPPMATH_H_
+
+#include <cmath>
+
+namespace tflite {
+
+#if defined(TF_LITE_USE_GLOBAL_CMATH_FUNCTIONS) || \
+    (defined(__ANDROID__) && !defined(__NDK_MAJOR__)) || defined(__ZEPHYR__)
+#define TF_LITE_GLOBAL_STD_PREFIX
+#else
+#define TF_LITE_GLOBAL_STD_PREFIX std
+#endif
+
+#define DECLARE_STD_GLOBAL_SWITCH1(tf_name, std_name) \
+  template <class T>                                  \
+  inline T tf_name(const T x) {                       \
+    return TF_LITE_GLOBAL_STD_PREFIX::std_name(x);    \
+  }
+
+DECLARE_STD_GLOBAL_SWITCH1(TfLiteRound, round)
+DECLARE_STD_GLOBAL_SWITCH1(TfLiteExpm1, expm1)
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_CPPMATH_H_

tensorflow/lite/kernels/internal/depthwiseconv_float_test.cc

@@ -0,0 +1,160 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <algorithm>
+#include <cmath>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/test_util.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+#define ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_float.h"
+#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_float.h"
+
+namespace tflite {
+namespace {
+
+// Runs the DepthwiseConv and compares against the reference implementation.
+void TestOneDepthwiseConv(
+    const DepthwiseParams& params, const RuntimeShape& input_shape,
+    const float* input_data, const RuntimeShape& filter_shape,
+    const float* filter_data, const RuntimeShape& bias_shape,
+    const float* bias_data, const RuntimeShape& output_shape) {
+  const int output_buffer_size = output_shape.FlatSize();
+  std::vector<float> output_data(output_buffer_size);
+  std::vector<float> reference_output_data(output_buffer_size);
+  reference_ops::DepthwiseConv(params, input_shape, input_data, filter_shape,
+                               filter_data, bias_shape, bias_data, output_shape,
+                               reference_output_data.data());
+  optimized_ops::DepthwiseConvImpl(
+      params, input_shape, input_data, filter_shape, filter_data, bias_shape,
+      bias_data, output_shape, output_data.data(), CpuFlags(),
+      /*thread_start=*/0,
+      /*thread_end=*/output_shape.Dims(1), /*thread_dim=*/1);
+
+  double sum_abs_diff = 0;
+  float max_abs_val = 0;
+  for (int i = 0; i < output_buffer_size; i++) {
+    sum_abs_diff += std::abs(output_data[i] - reference_output_data[i]);
+    max_abs_val = std::max(max_abs_val, std::abs(reference_output_data[i]));
+  }
+  if (sum_abs_diff != 0.f) {
+    const float mean_diff =
+        static_cast<float>(sum_abs_diff / output_buffer_size);
+    const float relative_error = std::abs(mean_diff) / max_abs_val;
+    ASSERT_LT(relative_error, 1e-5f);
+  }
+}
+
+// This function picks some random DepthwiseConv params, which may or may not
+// be legal. If they're not legal, it returns false. If they're legal,
+// it runs the DepthwiseConv test and returns true. This allows the caller
+// to loop until a test has been run.
+bool TryTestOneDepthwiseConv() {
+  // We have to pick a lot of positive values, where we are particularly
+  // interested in small values because they are most likely to be special
+  // cases in optimized implementations, and secondarily because they allow
+  // tests to run fast, which means we can run more tests and get more
+  // coverage.
+  const int batch = UniformRandomInt(1, 2);
+  const int input_depth = ExponentialRandomPositiveInt(0.9f, 6, 50);
+  const int input_width = ExponentialRandomPositiveInt(0.9f, 20, 200);
+  const int input_height = ExponentialRandomPositiveInt(0.9f, 20, 200);
+  const int filter_width = ExponentialRandomPositiveInt(0.9f, 4, 10);
+  const int filter_height = ExponentialRandomPositiveInt(0.9f, 4, 10);
+  const int depth_multiplier = ExponentialRandomPositiveInt(0.8f, 6, 50);
+  const int stride = ExponentialRandomPositiveInt(0.9f, 3, 8);
+  const int output_depth = input_depth * depth_multiplier;
+  const int dilation_width_factor = RandomElement(std::vector<int>({1, 2, 4}));
+  const int dilation_height_factor = RandomElement(std::vector<int>({1, 2, 4}));
+  float output_activation_min, output_activation_max;
+  FusedActivationFunctionType ac =
+      RandomElement(std::vector<FusedActivationFunctionType>(
+          {FusedActivationFunctionType::kNone,
+           FusedActivationFunctionType::kRelu,
+           FusedActivationFunctionType::kRelu1,
+           FusedActivationFunctionType::kRelu6}));
+  GetActivationMinMax(ac, &output_activation_min, &output_activation_max);
+  // The optimized DepthwiseConv implementation currently uses a fixed-size
+  // accumulator buffer on the stack, with that size. This currently means
+  // that it does not support larger output depths. It CHECK's for it,
+  // so it's safe in the sense that if a larger output depth was encountered,
+  // it would explicitly fail. We just need to adjust our testing to that
+  // constraint.
+  const int kMaxSupportedOutputDepth = 1024;
+  if (output_depth > kMaxSupportedOutputDepth) {
+    return false;
+  }
+  RuntimeShape input_shape_inference(
+      {batch, input_height, input_width, input_depth});
+  RuntimeShape output_shape_inference;
+  int pad_width, pad_height;
+  const auto padding_type =
+      UniformRandomInt(0, 1) ? PaddingType::kSame : PaddingType::kValid;
+  if (!ComputeConvSizes(input_shape_inference, output_depth, filter_width,
+                        filter_height, stride, dilation_width_factor,
+                        dilation_height_factor, padding_type,
+                        &output_shape_inference, &pad_width, &pad_height)) {
+    return false;
+  }
+  RuntimeShape filter_shape_inference(
+      {1, filter_height, filter_width, output_depth});
+  RuntimeShape bias_shape_inference({1, 1, 1, output_depth});
+  const int input_buffer_size = input_shape_inference.FlatSize();
+  const int filter_buffer_size = filter_shape_inference.FlatSize();
+  std::vector<float> input_data(input_buffer_size);
+  std::vector<float> filter_data(filter_buffer_size);
+  std::vector<float> bias_data(output_depth);
+  const float input_amplitude = 1.f;
+  const float filter_amplitude = 1.f;
+  const float bias_amplitude =
+      filter_width * filter_height * input_amplitude * filter_amplitude;
+  FillRandom(&input_data, -input_amplitude, input_amplitude);
+  FillRandom(&filter_data, -filter_amplitude, filter_amplitude);
+  FillRandom(&bias_data, -bias_amplitude, bias_amplitude);
+  DepthwiseParams op_params;
+  op_params.padding_type = PaddingType::kSame;
+  op_params.padding_values.width = pad_width;
+  op_params.padding_values.height = pad_height;
+  op_params.stride_width = stride;
+  op_params.stride_height = stride;
+  op_params.dilation_width_factor = dilation_width_factor;
+  op_params.dilation_height_factor = dilation_height_factor;
+  op_params.depth_multiplier = depth_multiplier;
+  op_params.float_activation_min = output_activation_min;
+  op_params.float_activation_max = output_activation_max;
+  TestOneDepthwiseConv(op_params, input_shape_inference, input_data.data(),
+                       filter_shape_inference, filter_data.data(),
+                       bias_shape_inference, bias_data.data(),
+                       output_shape_inference);
+  return true;
+}
+
+void TestOneDepthwiseConv() {
+  while (!TryTestOneDepthwiseConv()) {
+  }
+}
+
+TEST(TestDepthwiseConv, TestDepthwiseConv) {
+  const int kTestsToRun = 10 * 1000;
+  for (int i = 0; i < kTestsToRun; i++) {
+    TestOneDepthwiseConv();
+  }
+}
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/depthwiseconv_per_channel_quantized_16x8_test.cc

@@ -0,0 +1,320 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <stdio.h>
+#include <sys/types.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <iterator>
+#include <limits>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_float.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/depthwise_conv.h"
+#include "tensorflow/lite/kernels/internal/test_util.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace {
+
+void PickOutputMultiplier(
+    const DepthwiseParams& params, const RuntimeShape& input_shape,
+    const int16_t* input_data, const RuntimeShape& filter_shape,
+    const int8_t* filter_data, const RuntimeShape& bias_shape,
+    const std::int64_t* bias_data, const RuntimeShape& output_shape,
+    float* output_multiplier) {
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int dilation_width_factor = params.dilation_width_factor;
+  const int dilation_height_factor = params.dilation_height_factor;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  const int depth_multiplier = params.depth_multiplier;
+
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int input_depth = input_shape.Dims(3);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+
+  std::int64_t output_accu_min = std::numeric_limits<std::int64_t>::max();
+  std::int64_t output_accu_max = std::numeric_limits<std::int64_t>::min();
+
+  for (int batch = 0; batch < batches; ++batch) {
+    for (int out_y = 0; out_y < output_height; ++out_y) {
+      for (int out_x = 0; out_x < output_width; ++out_x) {
+        for (int in_channel = 0; in_channel < input_depth; ++in_channel) {
+          for (int m = 0; m < depth_multiplier; ++m) {
+            const int output_channel = m + in_channel * depth_multiplier;
+            const int in_x_origin = (out_x * stride_width) - pad_width;
+            const int in_y_origin = (out_y * stride_height) - pad_height;
+            std::int64_t acc = 0;
+            for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
+              for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
+                const int in_x = in_x_origin + dilation_width_factor * filter_x;
+                const int in_y =
+                    in_y_origin + dilation_height_factor * filter_y;
+                // Zero padding by omitting the areas outside the image.
+                const bool is_point_inside_image =
+                    (in_x >= 0) && (in_x < input_width) && (in_y >= 0) &&
+                    (in_y < input_height);
+                if (is_point_inside_image) {
+                  int32_t input_val = input_data[Offset(
+                      input_shape, batch, in_y, in_x, in_channel)];
+                  int32_t filter_val = filter_data[Offset(
+                      filter_shape, 0, filter_y, filter_x, output_channel)];
+                  acc += static_cast<int64_t>(filter_val) *
+                         static_cast<int64_t>(input_val);
+                }
+              }
+            }
+            if (bias_data) {
+              acc += bias_data[output_channel];
+            }
+            output_accu_max = std::max(acc, output_accu_max);
+            output_accu_min = std::min(acc, output_accu_min);
+          }
+        }
+      }
+    }
+  }
+
+  // Since int16 ranges from -32768 to 32767, we need to squeeze the accumulator
+  // min/max fit in those ranges correspondingly as much as possible.
+  if (std::abs(output_accu_max) > std::abs(output_accu_min)) {
+    *output_multiplier = 32767.0f / std::abs(output_accu_max);
+  } else {
+    *output_multiplier = 32768.0f / std::abs(output_accu_min);
+  }
+}
+
+void PickReasonableMultiplier(
+    const DepthwiseParams& params, int output_activation_min,
+    int output_activation_max, int output_depth,
+    const RuntimeShape& input_shape_inference, const std::int16_t* input_data,
+    const RuntimeShape& filter_shape_inference, const std::int8_t* filter_data,
+    const RuntimeShape& bias_shape_inference, const std::int64_t* bias_data,
+    const RuntimeShape& output_shape_inference,
+    std::int32_t* output_multiplier_ptr, std::int32_t* output_shift_ptr,
+    std::int16_t* output_data) {
+  float output_multiplier;
+  PickOutputMultiplier(params, input_shape_inference, input_data,
+                       filter_shape_inference, filter_data,
+                       bias_shape_inference, bias_data, output_shape_inference,
+                       &output_multiplier);
+
+  int base_multiplier;
+  int base_shift;
+  QuantizeMultiplier(output_multiplier, &base_multiplier, &base_shift);
+  for (int i = 0; i < output_depth; ++i) {
+    // multipliers typically range in [2^30 ; 2^31 - 1].
+    // Values in [0, 2^30 - 1] are normally unused, but harmless.
+    // Thus a good way to randomize multipliers is to subtract from them
+    // a random value smaller than 2^30 but still significant compared to it.
+    output_multiplier_ptr[i] = base_multiplier - (std::rand() % (1 << 26));
+    output_shift_ptr[i] = base_shift - 1 + (std::rand() % 4);
+  }
+}
+
+bool GenerateValidShapeConfigurations(
+    int filter_width, int filter_height, int depth_multiplier,
+    int dilation_width_factor, int dilation_height_factor,
+    RuntimeShape* input_shape_inference, RuntimeShape* filter_shape_inference,
+    RuntimeShape* output_shape_inference, int* pad_width, int* pad_height,
+    int* stride) {
+  const int batch = UniformRandomInt(1, 3);
+  const int input_depth = 8 * ExponentialRandomPositiveInt(0.9f, 10, 50);
+  const int input_width = UniformRandomInt(5, 50);
+  const int input_height = UniformRandomInt(5, 50);
+  *stride = UniformRandomInt(1, 2);
+  const bool test_pad = UniformRandomInt(0, 1);
+  const auto padding_type = test_pad ? PaddingType::kValid : PaddingType::kSame;
+
+  const int output_depth = input_depth * depth_multiplier;
+
+  input_shape_inference->BuildFrom(
+      {batch, input_height, input_width, input_depth});
+
+  filter_shape_inference->BuildFrom(
+      {1, filter_height, filter_width, output_depth});
+
+  EXPECT_TRUE(ComputeConvSizes(
+      *input_shape_inference, output_depth, filter_width, filter_height,
+      *stride, dilation_width_factor, dilation_height_factor, padding_type,
+      output_shape_inference, pad_width, pad_height));
+
+  return true;
+}
+
+void IntToFloat(std::vector<float>* d, std::vector<std::int8_t>* s) {
+  for (unsigned int i = 0; i < s->size(); i++) {
+    d->data()[i] = (float)s->data()[i];
+  }
+}
+
+void IntToFloat(std::vector<float>* d, std::vector<std::int64_t>* s) {
+  for (unsigned int i = 0; i < s->size(); i++) {
+    d->data()[i] = (float)s->data()[i];
+  }
+}
+
+void TryTestOneDepthwiseConv3x3Filter() {
+  const int filter_width = 3;
+  const int filter_height = 3;
+  const int depth_multiplier = 1;
+  // We don't support dilations in the 3x3 filter.
+  const int dilation_width_factor = 1;
+  const int dilation_height_factor = 1;
+
+  const int output_activation_min = -32768;
+  const int output_activation_max = 32767;
+
+  RuntimeShape input_shape_inference;
+  RuntimeShape filter_shape_inference;
+  RuntimeShape output_shape_inference;
+  int pad_width, pad_height;
+  int stride;
+
+  // Keeps trying until we get valid shape/configurations for 3x3 filter case.
+  bool generated_valid_configurations_for_3x3_kernel = false;
+  while (!generated_valid_configurations_for_3x3_kernel) {
+    generated_valid_configurations_for_3x3_kernel =
+        GenerateValidShapeConfigurations(
+            filter_width, filter_height, depth_multiplier,
+            dilation_width_factor, dilation_height_factor,
+            &input_shape_inference, &filter_shape_inference,
+            &output_shape_inference, &pad_width, &pad_height, &stride);
+  }
+
+  const int output_depth = output_shape_inference.Dims(3);
+
+  RuntimeShape bias_shape_inference({1, 1, 1, output_depth});
+  const int input_buffer_size = input_shape_inference.FlatSize();
+  const int filter_buffer_size = filter_shape_inference.FlatSize();
+  const int output_buffer_size = output_shape_inference.FlatSize();
+  std::vector<std::int16_t> input_data(input_buffer_size);
+  std::vector<std::int8_t> filter_data(filter_buffer_size);
+  std::vector<std::int64_t> bias_data(output_depth);
+
+  FillRandom(&input_data);
+  FillRandom(&filter_data);
+  for (int i = 0; i < output_depth; i++) {
+    bias_data.data()[i] = 0;
+  }
+
+  DepthwiseParams params;
+  params.stride_width = stride;
+  params.stride_height = stride;
+  params.dilation_height_factor = dilation_height_factor;
+  params.dilation_width_factor = dilation_width_factor;
+  params.padding_values.width = pad_width;
+  params.padding_values.height = pad_height;
+  params.depth_multiplier = depth_multiplier;
+  params.weights_offset = 0;
+  params.quantized_activation_min = output_activation_min;
+  params.quantized_activation_max = output_activation_max;
+  params.float_activation_max = (float)(1LL << 40);
+  params.float_activation_min = -params.float_activation_max;
+
+  std::vector<std::int16_t> reference_output_data(output_buffer_size);
+  std::vector<std::int16_t> neon_output_data(output_buffer_size);
+
+  std::vector<std::int32_t> output_multiplier(output_depth);
+  std::vector<std::int32_t> output_shift(output_depth);
+
+  // It's hard to come up with a right multiplier, random guess basically makes
+  // all the results saturated and becomes meaningfulless, so we first use
+  // reference impl to poke the min/max value of the accumulation, then use that
+  // value as a guided suggestion for us to populate meaningful mulitplier &
+  // shift.
+  PickReasonableMultiplier(
+      params, output_activation_min, output_activation_max, output_depth,
+      input_shape_inference, input_data.data(), filter_shape_inference,
+      filter_data.data(), bias_shape_inference, bias_data.data(),
+      output_shape_inference, output_multiplier.data(), output_shift.data(),
+      reference_output_data.data());
+
+  // The following tests compare reference impl for 16x8 version and
+  // float reference operator.
+  reference_integer_ops::DepthwiseConvPerChannel(
+      params, output_multiplier.data(), output_shift.data(),
+      input_shape_inference, input_data.data(), filter_shape_inference,
+      filter_data.data(), bias_shape_inference, bias_data.data(),
+      output_shape_inference, reference_output_data.data());
+
+  std::vector<float> input_data_float(input_buffer_size);
+  std::vector<float> filter_data_float(filter_buffer_size);
+  std::vector<float> bias_data_float(output_depth);
+  std::vector<float> output_data_float(output_buffer_size);
+
+  for (int i = 0; i < input_buffer_size; i++) {
+    input_data_float.data()[i] = (float)(input_data.data()[i]);
+  }
+  IntToFloat(&filter_data_float, &filter_data);
+  IntToFloat(&bias_data_float, &bias_data);
+
+  reference_ops::DepthwiseConv(
+      params, input_shape_inference, input_data_float.data(),
+      filter_shape_inference, filter_data_float.data(), bias_shape_inference,
+      bias_data_float.data(), output_shape_inference, output_data_float.data());
+
+  for (int n = 0; n < output_shape_inference.Dims(0); n++) {
+    for (int h = 0; h < output_shape_inference.Dims(1); h++) {
+      for (int w = 0; w < output_shape_inference.Dims(2); w++) {
+        for (int c = 0; c < output_shape_inference.Dims(3); c++) {
+          int offset = Offset(output_shape_inference, n, h, w, c);
+          float float_res = output_data_float.data()[offset];
+          int16_t int16_res = reference_output_data.data()[offset];
+          int32_t output_mul = output_multiplier.data()[c];
+          int shift = output_shift.data()[c];
+          float scale = (float)output_mul / (float)(1ULL << 31);
+          if (shift > 0) scale = scale * (float)(1 << shift);
+          if (shift < 0) scale = scale / (float)(1 << -shift);
+          int ref_res = floor(float_res * scale + 0.5);
+          if (ref_res < output_activation_min) ref_res = output_activation_min;
+          if (ref_res > output_activation_max) ref_res = output_activation_max;
+          int e = (ref_res - int16_res);
+          if (e < 0) e = -e;
+          if (e > 1) {
+            printf(
+                "(%d,%d,%d,%d) scale=%08x shift=%d res=%d float=%f (%f,%f)\n",
+                n, h, w, c, output_mul, shift, int16_res, float_res * scale,
+                float_res, scale);
+            EXPECT_TRUE(false);
+          }
+        }
+      }
+    }
+  }
+}
+
+TEST(QuantizedDepthwiseConvPerChannelTest, FastKernelTest) {
+  for (int i = 0; i < 30; ++i) {
+    TryTestOneDepthwiseConv3x3Filter();
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/depthwiseconv_per_channel_quantized_test.cc

@@ -0,0 +1,355 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <sys/types.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <iterator>
+#include <limits>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_3x3_filter_common.h"
+#include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h"
+#include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_3x3_filter.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/depthwise_conv.h"
+#include "tensorflow/lite/kernels/internal/test_util.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace {
+
+void PickOutputMultiplier(
+    const DepthwiseParams& params, const RuntimeShape& input_shape,
+    const int8_t* input_data, const RuntimeShape& filter_shape,
+    const int8_t* filter_data, const RuntimeShape& bias_shape,
+    const int32_t* bias_data, const RuntimeShape& output_shape,
+    float* output_multiplier) {
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int dilation_width_factor = params.dilation_width_factor;
+  const int dilation_height_factor = params.dilation_height_factor;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  const int depth_multiplier = params.depth_multiplier;
+  const int32_t input_offset = params.input_offset;
+
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int input_depth = input_shape.Dims(3);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+
+  int output_accu_min = std::numeric_limits<std::int32_t>::max();
+  int output_accu_max = std::numeric_limits<std::int32_t>::min();
+
+  for (int batch = 0; batch < batches; ++batch) {
+    for (int out_y = 0; out_y < output_height; ++out_y) {
+      for (int out_x = 0; out_x < output_width; ++out_x) {
+        for (int in_channel = 0; in_channel < input_depth; ++in_channel) {
+          for (int m = 0; m < depth_multiplier; ++m) {
+            const int output_channel = m + in_channel * depth_multiplier;
+            const int in_x_origin = (out_x * stride_width) - pad_width;
+            const int in_y_origin = (out_y * stride_height) - pad_height;
+            int32_t acc = 0;
+            for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
+              for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
+                const int in_x = in_x_origin + dilation_width_factor * filter_x;
+                const int in_y =
+                    in_y_origin + dilation_height_factor * filter_y;
+                // Zero padding by omitting the areas outside the image.
+                const bool is_point_inside_image =
+                    (in_x >= 0) && (in_x < input_width) && (in_y >= 0) &&
+                    (in_y < input_height);
+                if (is_point_inside_image) {
+                  int32_t input_val = input_data[Offset(
+                      input_shape, batch, in_y, in_x, in_channel)];
+                  int32_t filter_val = filter_data[Offset(
+                      filter_shape, 0, filter_y, filter_x, output_channel)];
+                  acc += filter_val * (input_val + input_offset);
+                }
+              }
+            }
+            if (bias_data) {
+              acc += bias_data[output_channel];
+            }
+            output_accu_max = std::max(acc, output_accu_max);
+            output_accu_min = std::min(acc, output_accu_min);
+          }
+        }
+      }
+    }
+  }
+
+  // Since int8 ranges from -128 to 127, we need to squeeze the accumulator
+  // min/max fit in those ranges correspondingly as much as possible.
+  if (std::abs(output_accu_max) > std::abs(output_accu_min)) {
+    *output_multiplier = 127.0f / std::abs(output_accu_max);
+  } else {
+    *output_multiplier = 128.0f / std::abs(output_accu_min);
+  }
+}
+
+void PickReasonableMultiplier(
+    const DepthwiseParams& params, int output_activation_min,
+    int output_activation_max, int output_depth,
+    const RuntimeShape& input_shape_inference, const std::int8_t* input_data,
+    const RuntimeShape& filter_shape_inference, const std::int8_t* filter_data,
+    const RuntimeShape& bias_shape_inference, const std::int32_t* bias_data,
+    const RuntimeShape& output_shape_inference,
+    std::int32_t* output_multiplier_ptr, std::int32_t* output_shift_ptr,
+    std::int8_t* output_data) {
+  float output_multiplier;
+  PickOutputMultiplier(params, input_shape_inference, input_data,
+                       filter_shape_inference, filter_data,
+                       bias_shape_inference, bias_data, output_shape_inference,
+                       &output_multiplier);
+
+  int base_multiplier;
+  int base_shift;
+  QuantizeMultiplier(output_multiplier, &base_multiplier, &base_shift);
+  for (int i = 0; i < output_depth; ++i) {
+    // multipliers typically range in [2^30 ; 2^31 - 1].
+    // Values in [0, 2^30 - 1] are normally unused, but harmless.
+    // Thus a good way to randomize multipliers is to subtract from them
+    // a random value smaller than 2^30 but still significant compared to it.
+    output_multiplier_ptr[i] = base_multiplier - (std::rand() % (1 << 26));
+    output_shift_ptr[i] = base_shift - 1 + (std::rand() % 4);
+  }
+}
+
+// The reference implementation & the fast kernel have different rounding
+// mechanism, so we loosely compare the difference.
+void CompareRoundingResults(int flat_size, const int depth_multiplier,
+                            const std::int8_t* reference_result,
+                            const std::int8_t* fast_kernel_result) {
+  std::vector<int> diff(flat_size);
+  std::int64_t sum_diff = 0;
+  std::int64_t sum_abs_diff = 0;
+  for (int i = 0; i < flat_size; i++) {
+    diff[i] = static_cast<int>(fast_kernel_result[i]) -
+              static_cast<int>(reference_result[i]);
+    sum_diff += diff[i];
+    sum_abs_diff += std::abs(diff[i]);
+  }
+  // These stats help understand test failures.
+  std::sort(std::begin(diff), std::end(diff));
+  const int min_diff = diff.front();
+  const int max_diff = diff.back();
+  const int median_diff = diff[diff.size() / 2];
+  const float mean_diff = static_cast<float>(sum_diff) / flat_size;
+  const float mean_abs_diff = static_cast<float>(sum_abs_diff) / flat_size;
+
+  // The tolerance that we apply to means is tight, but we allow for a rounding
+  // difference in one pixel, and loosen by another 1% for float comparison.
+  const float mean_tolerance =
+      std::max(1e-2f, 1.01f / flat_size * std::sqrt(1.f * depth_multiplier));
+  const int diff_mean_tolerance = 256;
+  const int diff_median_tolerance = 225;
+
+  // Normally we should require bit-for-bit exact results. Unfortunately a bug
+  // in the Intel arm_neon_sse.h translation header that we use for x86 tests
+  // causes 1-bit inaccuracy in the vqrdmulh_n_s32 intrinsic, which causes
+  // off-by-1 errors in quantized DepthwiseConv ops. So we have to live with a
+  // few off-by-one errors for now, yet still ensure that no more than a small
+  // minority of values are wrong.
+  EXPECT_LT(std::abs(mean_diff), mean_tolerance);
+  EXPECT_LT(mean_abs_diff, mean_tolerance);
+  EXPECT_LE(std::abs(median_diff), diff_median_tolerance);
+  EXPECT_LE(std::abs(min_diff), diff_mean_tolerance);
+  EXPECT_LE(std::abs(max_diff), diff_mean_tolerance);
+  EXPECT_TRUE(std::abs(mean_diff) < mean_tolerance &&
+              mean_abs_diff < mean_tolerance &&
+              std::abs(median_diff) <= diff_median_tolerance &&
+              std::abs(min_diff) <= diff_mean_tolerance &&
+              std::abs(max_diff) <= diff_mean_tolerance);
+}
+
+bool GenerateValidShapeConfigurations(
+    int filter_width, int filter_height, int depth_multiplier,
+    int dilation_width_factor, int dilation_height_factor,
+    RuntimeShape* input_shape_inference, RuntimeShape* filter_shape_inference,
+    RuntimeShape* output_shape_inference, int* pad_width, int* pad_height,
+    int* stride) {
+  const int batch = UniformRandomInt(1, 3);
+  const int input_depth = 8 * ExponentialRandomPositiveInt(0.9f, 10, 50);
+  const int input_width = UniformRandomInt(5, 50);
+  const int input_height = UniformRandomInt(5, 50);
+  *stride = UniformRandomInt(1, 2);
+  const bool test_pad = UniformRandomInt(0, 1);
+  const auto padding_type = test_pad ? PaddingType::kValid : PaddingType::kSame;
+
+  const int output_depth = input_depth * depth_multiplier;
+
+  input_shape_inference->BuildFrom(
+      {batch, input_height, input_width, input_depth});
+
+  filter_shape_inference->BuildFrom(
+      {1, filter_height, filter_width, output_depth});
+
+  EXPECT_TRUE(ComputeConvSizes(
+      *input_shape_inference, output_depth, filter_width, filter_height,
+      *stride, dilation_width_factor, dilation_height_factor, padding_type,
+      output_shape_inference, pad_width, pad_height));
+
+  // We just care about whether the shape is suitable so we use non-per-channel
+  // case.
+  return optimized_ops::depthwise_conv::Fast3x3FilterKernelSupported<
+      optimized_ops::depthwise_conv::QuantizationType::kNonPerChannelUint8>(
+      *input_shape_inference, *filter_shape_inference, *stride, *stride,
+      dilation_width_factor, dilation_height_factor, *pad_width, *pad_height,
+      depth_multiplier, *output_shape_inference, 0);
+}
+
+void TryTestOneDepthwiseConv3x3Filter() {
+  const int filter_width = 3;
+  const int filter_height = 3;
+  const int depth_multiplier = 1;
+  // We don't support dilations in the 3x3 filter.
+  const int dilation_width_factor = 1;
+  const int dilation_height_factor = 1;
+
+  const int output_activation_min = -128;
+  const int output_activation_max = 127;
+
+  const std::int32_t input_offset = UniformRandomInt(-25, 25);
+  const std::int32_t output_offset = UniformRandomInt(-25, 25);
+
+  RuntimeShape input_shape_inference;
+  RuntimeShape filter_shape_inference;
+  RuntimeShape output_shape_inference;
+  int pad_width, pad_height;
+  int stride;
+
+  // Keeps trying until we get valid shape/configurations for 3x3 filter case.
+  bool generated_valid_configurations_for_3x3_kernel = false;
+  while (!generated_valid_configurations_for_3x3_kernel) {
+    generated_valid_configurations_for_3x3_kernel =
+        GenerateValidShapeConfigurations(
+            filter_width, filter_height, depth_multiplier,
+            dilation_width_factor, dilation_height_factor,
+            &input_shape_inference, &filter_shape_inference,
+            &output_shape_inference, &pad_width, &pad_height, &stride);
+  }
+
+  const int output_depth = output_shape_inference.Dims(3);
+
+  RuntimeShape bias_shape_inference({1, 1, 1, output_depth});
+  const int input_buffer_size = input_shape_inference.FlatSize();
+  const int filter_buffer_size = filter_shape_inference.FlatSize();
+  const int output_buffer_size = output_shape_inference.FlatSize();
+  std::vector<std::int8_t> input_data(input_buffer_size);
+  std::vector<std::int8_t> filter_data(filter_buffer_size);
+  std::vector<std::int32_t> bias_data(output_depth);
+
+  FillRandom(&input_data);
+  FillRandom(&filter_data);
+  FillRandom(&bias_data, -1000, 1000);
+
+  DepthwiseParams params;
+  params.stride_width = stride;
+  params.stride_height = stride;
+  params.dilation_height_factor = dilation_height_factor;
+  params.dilation_width_factor = dilation_width_factor;
+  params.padding_values.width = pad_width;
+  params.padding_values.height = pad_height;
+  params.depth_multiplier = depth_multiplier;
+  params.input_offset = input_offset;
+  params.output_offset = output_offset;
+  params.weights_offset = 0;
+  params.quantized_activation_min = output_activation_min;
+  params.quantized_activation_max = output_activation_max;
+
+  std::vector<std::int8_t> reference_output_data(output_buffer_size);
+  std::vector<std::int8_t> neon_output_data(output_buffer_size);
+
+  std::vector<std::int32_t> output_multiplier(output_depth);
+  std::vector<std::int32_t> output_shift(output_depth);
+
+  // It's hard to come up with a right multiplier, random guess basically makes
+  // all the results saturated and becomes meaningfulless, so we first use
+  // reference impl to poke the min/max value of the accumulation, then use that
+  // value as a guided suggestion for us to populate meaningful multiplier &
+  // shift.
+  PickReasonableMultiplier(
+      params, output_activation_min, output_activation_max, output_depth,
+      input_shape_inference, input_data.data(), filter_shape_inference,
+      filter_data.data(), bias_shape_inference, bias_data.data(),
+      output_shape_inference, output_multiplier.data(), output_shift.data(),
+      reference_output_data.data());
+
+  EXPECT_TRUE(optimized_ops::depthwise_conv::Fast3x3FilterKernelSupported<
+              optimized_ops::depthwise_conv::QuantizationType::kPerChannelInt8>(
+      input_shape_inference, filter_shape_inference, stride, stride,
+      dilation_width_factor, dilation_height_factor, pad_width, pad_height,
+      depth_multiplier, output_shape_inference, 0, output_shift.data()));
+
+  // The following tests compare reference impl and Neon general impl agrees,
+  // and reference impl loosely agrees with fast kernel since they use different
+  // rounding strategy.
+  reference_integer_ops::DepthwiseConvPerChannel(
+      params, output_multiplier.data(), output_shift.data(),
+      input_shape_inference, input_data.data(), filter_shape_inference,
+      filter_data.data(), bias_shape_inference, bias_data.data(),
+      output_shape_inference, reference_output_data.data());
+
+  optimized_integer_ops::depthwise_conv::DepthwiseConvGeneral(
+      params, output_multiplier.data(), output_shift.data(),
+      input_shape_inference, input_data.data(), filter_shape_inference,
+      filter_data.data(), bias_shape_inference, bias_data.data(),
+      output_shape_inference, neon_output_data.data(),
+      /*thread_start=*/0,
+      /*thread_end=*/output_shape_inference.Dims(1), /*thread_dim=*/1);
+
+  // We have changed our rounding strategy to the ARM rounding-right-shift
+  // instruction: breaking tie upward as it's much simpler.
+  // So we allow some difference for the neon output VS. the reference output.
+  CompareRoundingResults(output_buffer_size, depth_multiplier,
+                         reference_output_data.data(), neon_output_data.data());
+
+#if defined(__aarch64__) && !defined(GOOGLE_L4T)
+  std::vector<std::int8_t> fast_kernel_output_data(output_buffer_size);
+  optimized_ops::depthwise_conv::DepthwiseConv3x3FilterPerChannel<
+      DepthwiseConvOutputRounding::kUpward>(
+      params, output_multiplier.data(), output_shift.data(),
+      input_shape_inference, input_data.data(), filter_shape_inference,
+      filter_data.data(), bias_shape_inference, bias_data.data(),
+      output_shape_inference, fast_kernel_output_data.data(),
+      /*thread_start=*/0,
+      /*thread_end=*/output_shape_inference.Dims(1), /*thread_dim=*/1);
+
+  CompareRoundingResults(output_buffer_size, depth_multiplier,
+                         reference_output_data.data(),
+                         fast_kernel_output_data.data());
+#endif
+}
+
+TEST(QuantizedDepthwiseConvPerChannelTest, FastKernelTest) {
+  for (int i = 0; i < 60; ++i) {
+    TryTestOneDepthwiseConv3x3Filter();
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/kernel_utils.cc

@@ -0,0 +1,340 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/kernel_utils.h"
+
+#include <algorithm>
+
+#include "tensorflow/lite/kernels/internal/tensor_utils.h"
+
+namespace tflite {
+namespace kernel_utils {
+
+void RnnBatchStep(const float* input_ptr_batch, const float* input_weights_ptr,
+                  const float* recurrent_weights_ptr, const float* bias_ptr,
+                  int input_size, int num_units, int batch_size,
+                  int output_batch_leading_dim,
+                  TfLiteFusedActivation activation,
+                  float* hidden_state_ptr_batch, float* output_ptr_batch) {
+  RnnBatchStep(input_ptr_batch, input_weights_ptr,
+               /*aux_input_ptr_batch=*/nullptr,
+               /*aux_input_weights_ptr=*/nullptr, recurrent_weights_ptr,
+               bias_ptr, input_size, /*aux_input_size=*/0, num_units,
+               batch_size, output_batch_leading_dim, activation,
+               hidden_state_ptr_batch, output_ptr_batch);
+}
+
+void RnnBatchStep(const float* input_ptr_batch, const float* input_weights_ptr,
+                  const float* aux_input_ptr_batch,
+                  const float* aux_input_weights_ptr,
+                  const float* recurrent_weights_ptr, const float* bias_ptr,
+                  int input_size, int aux_input_size, int num_units,
+                  int batch_size, int output_batch_leading_dim,
+                  TfLiteFusedActivation activation,
+                  float* hidden_state_ptr_batch, float* output_ptr_batch) {
+  // 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 == num_units) {
+    // Output = bias
+    tensor_utils::VectorBatchVectorAssign(bias_ptr, num_units, batch_size,
+                                          output_ptr_batch);
+
+    // Output += input * input_weights
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        input_weights_ptr, num_units, input_size, input_ptr_batch, batch_size,
+        output_ptr_batch);
+
+    // Output += aux_input * aux_input_weights (if they are not empty).
+    if (aux_input_size > 0) {
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          aux_input_weights_ptr, num_units, aux_input_size, aux_input_ptr_batch,
+          batch_size, output_ptr_batch);
+    }
+
+    // Output += recurrent_weights * hidden_state
+    tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+        recurrent_weights_ptr, num_units, num_units, hidden_state_ptr_batch,
+        batch_size, output_ptr_batch);
+
+    // Output = activation(Output) and update hidden_state
+    tensor_utils::ApplyActivationToVector(
+        output_ptr_batch, num_units * batch_size, activation, output_ptr_batch);
+    std::copy_n(output_ptr_batch, num_units * batch_size,
+                hidden_state_ptr_batch);
+  } else {
+    // Output = bias
+    for (int k = 0; k < batch_size; k++) {
+      std::copy_n(bias_ptr, num_units,
+                  output_ptr_batch + k * output_batch_leading_dim);
+    }
+
+    // Output += input * input_weights
+    for (int k = 0; k < batch_size; k++) {
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          input_weights_ptr, num_units, input_size,
+          input_ptr_batch + k * input_size, /*n_batch=*/1,
+          output_ptr_batch + k * output_batch_leading_dim);
+    }
+
+    // Output += aux_input * aux_input_weights (if they are not empty).
+    if (aux_input_size > 0) {
+      for (int k = 0; k < batch_size; k++) {
+        tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+            aux_input_weights_ptr, num_units, aux_input_size,
+            aux_input_ptr_batch + k * aux_input_size,
+            /*n_batch=*/1, output_ptr_batch + k * output_batch_leading_dim);
+      }
+    }
+
+    // Output += recurrent_weights * hidden_state
+    for (int k = 0; k < batch_size; k++) {
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          recurrent_weights_ptr, num_units, num_units,
+          hidden_state_ptr_batch + k * num_units,
+          /*n_batch=*/1, output_ptr_batch + k * output_batch_leading_dim);
+    }
+
+    // Output = activation(Output) and update hidden_state
+    for (int k = 0; k < batch_size; k++) {
+      tensor_utils::ApplyActivationToVector(
+          output_ptr_batch + k * output_batch_leading_dim, num_units,
+          activation, output_ptr_batch + k * output_batch_leading_dim);
+      std::copy_n(output_ptr_batch + k * output_batch_leading_dim, num_units,
+                  hidden_state_ptr_batch + k * num_units);
+    }
+  }
+}
+
+void RnnBatchStep(
+    const float* input_ptr_batch, const int8_t* input_weights_ptr,
+    float input_weights_scale, const int8_t* recurrent_weights_ptr,
+    float recurrent_weights_scale, const float* bias_ptr, int input_size,
+    int num_units, int batch_size, int output_batch_leading_dim,
+    TfLiteFusedActivation activation, int8_t* quantized_input_ptr_batch,
+    int8_t* quantized_hidden_state_ptr_batch, float* scaling_factors,
+    float* hidden_state_ptr_batch, float* output_ptr_batch,
+    bool asymmetric_quantize_inputs, int32_t* zero_points,
+    int32_t* accum_scratch, int32_t* row_sums, bool* compute_row_sums) {
+  RnnBatchStep(input_ptr_batch, input_weights_ptr, input_weights_scale,
+               /*aux_input_ptr_batch=*/nullptr,
+               /*aux_input_weights_ptr=*/nullptr,
+               /*aux_input_weights_scale=*/0.0f, recurrent_weights_ptr,
+               recurrent_weights_scale, bias_ptr, input_size,
+               /*aux_input_size=*/0, num_units, batch_size,
+               output_batch_leading_dim, activation, quantized_input_ptr_batch,
+               /*aux_quantized_input_ptr_batch=*/nullptr,
+               quantized_hidden_state_ptr_batch, scaling_factors,
+               hidden_state_ptr_batch, output_ptr_batch,
+               asymmetric_quantize_inputs, zero_points, accum_scratch, row_sums,
+               compute_row_sums);
+}
+
+void RnnBatchStep(
+    const float* input_ptr_batch, const int8_t* input_weights_ptr,
+    float input_weights_scale, const float* aux_input_ptr_batch,
+    const int8_t* aux_input_weights_ptr, float aux_input_weights_scale,
+    const int8_t* recurrent_weights_ptr, float recurrent_weights_scale,
+    const float* bias_ptr, int input_size, int aux_input_size, int num_units,
+    int batch_size, int output_batch_leading_dim,
+    TfLiteFusedActivation activation, int8_t* quantized_input_ptr_batch,
+    int8_t* aux_quantized_input_ptr_batch,
+    int8_t* quantized_hidden_state_ptr_batch, float* scaling_factors,
+    float* hidden_state_ptr_batch, float* output_ptr_batch,
+    bool asymmetric_quantize_inputs, int32_t* zero_points,
+    int32_t* accum_scratch, int32_t* row_sums, bool* compute_row_sums) {
+  // 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.
+
+  int32_t* input_row_sums = nullptr;
+  int32_t* aux_input_row_sums = nullptr;
+  int32_t* recurrent_row_sums = nullptr;
+  if (asymmetric_quantize_inputs) {
+    input_row_sums = row_sums;
+    aux_input_row_sums = row_sums;
+    if (aux_input_ptr_batch) {
+      aux_input_row_sums += num_units;
+    }
+    recurrent_row_sums = aux_input_row_sums + num_units;
+    if (*compute_row_sums) {
+      tensor_utils::ReductionSumVector(input_weights_ptr, input_row_sums,
+                                       num_units, input_size);
+      if (aux_input_ptr_batch) {
+        tensor_utils::ReductionSumVector(aux_input_weights_ptr,
+                                         aux_input_row_sums, num_units,
+                                         aux_input_size);
+      }
+      tensor_utils::ReductionSumVector(
+          recurrent_weights_ptr, recurrent_row_sums, num_units, num_units);
+      *compute_row_sums = false;
+    }
+  }
+
+  if (output_batch_leading_dim == num_units) {
+    // Output = bias
+    tensor_utils::VectorBatchVectorAssign(bias_ptr, num_units, batch_size,
+                                          output_ptr_batch);
+
+    // Save quantization and matmul computation for all zero input.
+    if (!tensor_utils::IsZeroVector(input_ptr_batch, batch_size * input_size)) {
+      // Quantize input from float to uint8 + quantization params (scaling
+      // factor).
+      tensor_utils::BatchQuantizeFloats(
+          input_ptr_batch, batch_size, input_size, quantized_input_ptr_batch,
+          scaling_factors, zero_points, asymmetric_quantize_inputs);
+      for (int b = 0; b < batch_size; ++b) {
+        scaling_factors[b] *= input_weights_scale;
+      }
+      // Output += input * input_weights
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          input_weights_ptr, num_units, input_size, quantized_input_ptr_batch,
+          scaling_factors, batch_size, output_ptr_batch,
+          /*per_channel_scale=*/nullptr, zero_points, accum_scratch,
+          input_row_sums, compute_row_sums, /*context=*/nullptr);
+    }
+
+    if (aux_input_ptr_batch &&
+        !tensor_utils::IsZeroVector(aux_input_ptr_batch,
+                                    batch_size * aux_input_size)) {
+      tensor_utils::BatchQuantizeFloats(
+          aux_input_ptr_batch, batch_size, aux_input_size,
+          aux_quantized_input_ptr_batch, scaling_factors, zero_points,
+          asymmetric_quantize_inputs);
+      for (int b = 0; b < batch_size; ++b) {
+        scaling_factors[b] *= aux_input_weights_scale;
+      }
+
+      // Output += aux_input * aux_input_weights
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          aux_input_weights_ptr, num_units, aux_input_size,
+          aux_quantized_input_ptr_batch, scaling_factors, batch_size,
+          output_ptr_batch, /*per_channel_scale=*/nullptr, zero_points,
+          accum_scratch, aux_input_row_sums, compute_row_sums,
+          /*context=*/nullptr);
+    }
+
+    // Save quantization and matmul computation for all zero input.
+    if (!tensor_utils::IsZeroVector(hidden_state_ptr_batch,
+                                    batch_size * num_units)) {
+      // Quantize hidden_state
+      tensor_utils::BatchQuantizeFloats(
+          hidden_state_ptr_batch, batch_size, num_units,
+          quantized_hidden_state_ptr_batch, scaling_factors, zero_points,
+          asymmetric_quantize_inputs);
+      for (int b = 0; b < batch_size; ++b) {
+        scaling_factors[b] *= recurrent_weights_scale;
+      }
+
+      // Output += recurrent_weights * hidden_state
+      tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+          recurrent_weights_ptr, num_units, num_units,
+          quantized_hidden_state_ptr_batch, scaling_factors, batch_size,
+          output_ptr_batch, /*per_channel_scale=*/nullptr, zero_points,
+          accum_scratch, recurrent_row_sums, compute_row_sums,
+          /*context=*/nullptr);
+    }
+
+    // Output = activation(Output) and update hidden_state
+    tensor_utils::ApplyActivationToVector(
+        output_ptr_batch, num_units * batch_size, activation, output_ptr_batch);
+    std::copy_n(output_ptr_batch, num_units * batch_size,
+                hidden_state_ptr_batch);
+  } else {
+    // Output = bias
+    for (int k = 0; k < batch_size; k++) {
+      std::copy_n(bias_ptr, num_units,
+                  output_ptr_batch + k * output_batch_leading_dim);
+    }
+
+    // Save quantization and matmul computation for all zero input.
+    if (!tensor_utils::IsZeroVector(input_ptr_batch, batch_size * input_size)) {
+      // Quantize input from float to uint8 + quantization params (scaling
+      // factor).
+      tensor_utils::BatchQuantizeFloats(
+          input_ptr_batch, batch_size, input_size, quantized_input_ptr_batch,
+          scaling_factors, zero_points, asymmetric_quantize_inputs);
+      for (int b = 0; b < batch_size; ++b) {
+        scaling_factors[b] *= input_weights_scale;
+      }
+
+      // Output += input * input_weights
+      for (int k = 0; k < batch_size; k++) {
+        tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+            input_weights_ptr, num_units, input_size,
+            quantized_input_ptr_batch + k * input_size, &scaling_factors[k],
+            /*n_batch=*/1, output_ptr_batch + k * output_batch_leading_dim,
+            /*per_channel_scale=*/nullptr, zero_points + k, accum_scratch,
+            input_row_sums, compute_row_sums, /*context=*/nullptr);
+      }
+    }
+
+    if (aux_input_ptr_batch &&
+        !tensor_utils::IsZeroVector(aux_input_ptr_batch,
+                                    batch_size * aux_input_size)) {
+      tensor_utils::BatchQuantizeFloats(
+          aux_input_ptr_batch, batch_size, aux_input_size,
+          aux_quantized_input_ptr_batch, scaling_factors, zero_points,
+          asymmetric_quantize_inputs);
+      for (int b = 0; b < batch_size; ++b) {
+        scaling_factors[b] *= aux_input_weights_scale;
+      }
+
+      // Output += aux_input * aux_input_weights
+      for (int k = 0; k < batch_size; k++) {
+        tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+            aux_input_weights_ptr, num_units, aux_input_size,
+            aux_quantized_input_ptr_batch + k * aux_input_size,
+            &scaling_factors[k],
+            /*n_batch=*/1, output_ptr_batch + k * output_batch_leading_dim,
+            /*per_channel_scale=*/nullptr, zero_points + k, accum_scratch,
+            aux_input_row_sums, compute_row_sums, /*context=*/nullptr);
+      }
+    }
+
+    // Save quantization and matmul computation for all zero input.
+    if (!tensor_utils::IsZeroVector(hidden_state_ptr_batch,
+                                    batch_size * num_units)) {
+      // Quantize hidden_state
+      tensor_utils::BatchQuantizeFloats(
+          hidden_state_ptr_batch, batch_size, num_units,
+          quantized_hidden_state_ptr_batch, scaling_factors, zero_points,
+          asymmetric_quantize_inputs);
+      for (int b = 0; b < batch_size; ++b) {
+        scaling_factors[b] *= recurrent_weights_scale;
+      }
+
+      // Output += recurrent_weights * hidden_state
+      for (int k = 0; k < batch_size; k++) {
+        tensor_utils::MatrixBatchVectorMultiplyAccumulate(
+            recurrent_weights_ptr, num_units, num_units,
+            quantized_hidden_state_ptr_batch + k * num_units,
+            &scaling_factors[k], /*n_batch=*/1,
+            output_ptr_batch + k * output_batch_leading_dim,
+            /*per_channel_scale=*/nullptr, zero_points + k, accum_scratch,
+            recurrent_row_sums, compute_row_sums, /*context=*/nullptr);
+      }
+    }
+
+    // Output = activation(Output) and update hidden_state
+    for (int k = 0; k < batch_size; k++) {
+      tensor_utils::ApplyActivationToVector(
+          output_ptr_batch + k * output_batch_leading_dim, num_units,
+          activation, output_ptr_batch + k * output_batch_leading_dim);
+      std::copy_n(output_ptr_batch + k * output_batch_leading_dim, num_units,
+                  hidden_state_ptr_batch + k * num_units);
+    }
+  }
+}
+
+}  // namespace kernel_utils
+}  // namespace tflite

tensorflow/lite/kernels/internal/kernel_utils.h

@@ -0,0 +1,93 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_KERNEL_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_KERNEL_UTILS_H_
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+
+namespace tflite {
+namespace kernel_utils {
+
+// Performs an RNN batch inference step for inputs specified by input_ptr_batch.
+// The RNN cell is specified by the pointers to its input and recurrent weights,
+// and biases, along with the input size, number of units, activation.
+//
+// The pointers to the hidden state and the output are updated as a result.
+//
+// The pointers with the suffix "_batch" point to data aligned in batch_major
+// order, and each step processes batch_size many inputs from input_ptr_batch,
+// and updates batch_size many outputs and hidden states.
+//
+// The output_batch_dim is output.shape[-1], i.e. the outermost dimension of the
+// output tensor, and in most cases will be equal to num_units. It is usually
+// not when we want to store the RNN output into a slice of the output tensor,
+// e.g. for bidirectional RNNs with merge_outputs. In this case, the batched
+// operations cannot be used since they assume that the batched outputs are
+// contiguous, and we manually loop over the batched outputs.
+void RnnBatchStep(const float* input_ptr_batch, const float* input_weights_ptr,
+                  const float* recurrent_weights_ptr, const float* bias_ptr,
+                  int input_size, int num_units, int batch_size,
+                  int output_batch_leading_dim,
+                  TfLiteFusedActivation activation,
+                  float* hidden_state_ptr_batch, float* output_ptr_batch);
+
+// Same as above but includes an auxiliary input with the corresponding weights.
+void RnnBatchStep(const float* input_ptr_batch, const float* input_weights_ptr,
+                  const float* aux_input_ptr_batch,
+                  const float* aux_input_weights_ptr,
+                  const float* recurrent_weights_ptr, const float* bias_ptr,
+                  int input_size, int aux_input_size, int num_units,
+                  int batch_size, int output_batch_leading_dim,
+                  TfLiteFusedActivation activation,
+                  float* hidden_state_ptr_batch, float* output_ptr_batch);
+
+// Performs a quantized RNN batch inference step. Same as above, but for
+// quantization purposes, we also pass in quantized_hidden_state_ptr_batch and
+// quantized_input_ptr_batch pointers for temporary storage of the quantized
+// values of hidden_state_ptr_batch and input_ptr_batch, respectively.
+// These temporary storages are expected to be preallocated to the same size as
+// the respective pointers.
+// An additional preallocated temporary storage 'scaling_factors' (of size
+// batch_size) is used to store the scaling factors of the quantization (used
+// for recovery).
+// {input,recurrent}_weights_scale params are used for dequantization/recovery.
+void RnnBatchStep(
+    const float* input_ptr_batch, const int8_t* input_weights_ptr,
+    float input_weights_scale, const int8_t* recurrent_weights_ptr,
+    float recurrent_weights_scale, const float* bias_ptr, int input_size,
+    int num_units, int batch_size, int output_batch_leading_dim,
+    TfLiteFusedActivation activation, int8_t* quantized_input_ptr_batch,
+    int8_t* quantized_hidden_state_ptr_batch, float* scaling_factors,
+    float* hidden_state_ptr_batch, float* output_ptr_batch,
+    bool asymmetric_quantize_inputs, int32_t* zero_points,
+    int32_t* accum_scratch, int32_t* row_sums, bool* compute_row_sums);
+
+void RnnBatchStep(
+    const float* input_ptr_batch, const int8_t* input_weights_ptr,
+    float input_weights_scale, const float* aux_input_ptr_batch,
+    const int8_t* aux_input_weights_ptr, float aux_input_weights_scale,
+    const int8_t* recurrent_weights_ptr, float recurrent_weights_scale,
+    const float* bias_ptr, int input_size, int aux_input_size, int num_units,
+    int batch_size, int output_batch_leading_dim,
+    TfLiteFusedActivation activation, int8_t* quantized_input_ptr_batch,
+    int8_t* aux_quantized_input_ptr_batch,
+    int8_t* quantized_hidden_state_ptr_batch, float* scaling_factors,
+    float* hidden_state_ptr_batch, float* output_ptr_batch,
+    bool asymmetric_quantize_inputs, int32_t* zero_points,
+    int32_t* accum_scratch, int32_t* row_sums, bool* compute_row_sums);
+
+}  // namespace kernel_utils
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_KERNEL_UTILS_H_

tensorflow/lite/kernels/internal/legacy_types.h

@@ -0,0 +1,26 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_LEGACY_TYPES_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_LEGACY_TYPES_H_
+
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+
+// TODO(b/116772710): Insert legacy Dims<> code in here.
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_LEGACY_TYPES_H_

tensorflow/lite/kernels/internal/log_quantized_test.cc

@@ -0,0 +1,312 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <functional>
+#include <iterator>
+#include <limits>
+#include <random>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#define GEMMLOWP_ENABLE_FIXEDPOINT_CONSTANTS_CHECKS
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/string_type.h"
+
+namespace tflite {
+
+class NumberGenerator {
+ public:
+  std::vector<int> RandomIntVector(int n, int min_val, int max_val) {
+    std::vector<int> vec(n);
+    double scale = static_cast<double>(max_val + 1 - min_val) / engine_.max();
+    for (auto& it : vec) {
+      it = min_val + std::floor(engine_() * scale);
+    }
+    return vec;
+  }
+
+  std::mt19937 engine_;
+};
+
+class LogQuantizedTest : public ::testing::Test {
+ public:
+  NumberGenerator generator_;
+};
+
+// input_integer_bits <= 30.  output_integer_bits > 0.
+inline int32_t LogPositiveValuesViaFloat(int32_t input_val,
+                                         int input_integer_bits,
+                                         int output_integer_bits) {
+  const double float_log_sum_of_exps = std::log(
+      static_cast<double>(input_val) * 0.5 / (1 << (30 - input_integer_bits)));
+  static constexpr double min_int =
+      static_cast<double>(std::numeric_limits<int32_t>::min());
+  static constexpr double max_int =
+      static_cast<double>(std::numeric_limits<int32_t>::max());
+  double double_result = tflite::TfLiteRound(float_log_sum_of_exps *
+                                             (1 << (31 - output_integer_bits)));
+  return static_cast<std::int32_t>(
+      std::min(max_int, std::max(min_int, double_result)));
+}
+
+void CheckOutputData(const std::vector<int32_t>& test_output,
+                     const std::vector<int32_t>& reference_output,
+                     const std::vector<int32_t>& test_input,
+                     const string& check_label, int input_integer_bits,
+                     int output_integer_bits, int tolerance) {
+  // In the special case of small input, specifically raw value of 5, a rounding
+  // up leads to difference in the output.  We do not aim to be accurate for
+  // very small input values, and there should be sufficient input fractional
+  // bits that this is a small input.
+  static constexpr double error_from_rounding_up = 0.0224585;
+  const int n = test_output.size();
+  ASSERT_EQ(n, reference_output.size());
+  for (int i = 0; i < n; ++i) {
+    // Adjust tolerance when input <= 5*2^-(31-input_integer_bits).
+    const int adjusted_tolerance =
+        test_input[i] > 5
+            ? tolerance
+            : std::max(tolerance, static_cast<int>(std::ceil(
+                                      error_from_rounding_up *
+                                      (1 << (31 - output_integer_bits)))));
+    ASSERT_LE(std::abs(test_output[i] - reference_output[i]),
+              adjusted_tolerance)
+        << "Failure in \"" << check_label << "\" at i=" << i
+        << ", test_input[i]=" << test_input[i] << "="
+        << static_cast<double>(test_input[i]) / (1 << (31 - input_integer_bits))
+        << ", test_output[i]=" << test_output[i] << "="
+        << static_cast<double>(test_output[i]) /
+               (1 << (31 - output_integer_bits))
+        << ", reference_output[i]=" << reference_output[i] << "="
+        << static_cast<double>(reference_output[i]) /
+               (1 << (31 - output_integer_bits))
+        << ", difference[i]=" << std::abs(reference_output[i] - test_output[i])
+        << "="
+        << static_cast<double>(std::abs(reference_output[i] - test_output[i])) /
+               (1 << (31 - output_integer_bits))
+        << "; tolerance=" << tolerance
+        << ", adj tolerance=" << adjusted_tolerance;
+  }
+}
+
+void RightShiftVector(const std::vector<int32_t>& shifts,
+                      std::vector<int32_t>* vec) {
+  const int n = vec->size();
+  ASSERT_EQ(n, shifts.size());
+  for (int i = 0; i < n; ++i) {
+    vec->at(i) = std::max(1, vec->at(i) >> shifts[i]);
+  }
+}
+
+template <int OutputIntegerBits, int InputIntegerBits>
+void RunSingleTest(const std::vector<int32_t>& test_input,
+                   const string& check_label, int tolerance) {
+  const int n = test_input.size();
+  std::vector<int32_t> float_gen_output(n, 0);
+  std::vector<int32_t> quantized_output(n, 0);
+
+  // Workaround the stupid things that intelligent humans do.
+  // Consequence of __builtin_clz(0u) may equal 31 instead of 32.
+  std::vector<int32_t> fudged_input(n, 0);
+  for (int i = 0; i < n; ++i) {
+    fudged_input[i] = std::max(test_input[i], 2);
+  }
+
+  for (int i = 0; i < n; ++i) {
+    quantized_output[i] =
+        tflite::log_x_for_x_greater_than_or_equal_to_1_impl<OutputIntegerBits,
+                                                            InputIntegerBits>(
+            gemmlowp::FixedPoint<int32_t, InputIntegerBits>::FromRaw(
+                fudged_input[i]))
+            .raw();
+    float_gen_output[i] = LogPositiveValuesViaFloat(
+        fudged_input[i], InputIntegerBits, OutputIntegerBits);
+  }
+  {
+    std::ostringstream label;
+    label << check_label
+          << " / reference vs float-gen / InputIntegerBits=" << InputIntegerBits
+          << ", OutputIntegerBits=" << OutputIntegerBits;
+    CheckOutputData(quantized_output, float_gen_output, test_input, label.str(),
+                    InputIntegerBits, OutputIntegerBits, tolerance);
+  }
+}
+
+template <int OutputIntegerBits>
+void RunSingleTest(const std::vector<int32_t>& test_input,
+                   int input_integer_bits, const string& check_label,
+                   int tolerance) {
+#define INPUT_CASE(K)                                                   \
+  case K:                                                               \
+    return RunSingleTest<OutputIntegerBits, K>(test_input, check_label, \
+                                               tolerance)
+  switch (input_integer_bits) {
+    INPUT_CASE(0);
+    INPUT_CASE(1);
+    INPUT_CASE(2);
+    INPUT_CASE(3);
+    INPUT_CASE(4);
+    INPUT_CASE(5);
+    INPUT_CASE(6);
+    INPUT_CASE(7);
+    INPUT_CASE(8);
+    INPUT_CASE(9);
+    INPUT_CASE(10);
+    INPUT_CASE(11);
+    INPUT_CASE(12);
+    INPUT_CASE(13);
+    INPUT_CASE(14);
+    INPUT_CASE(15);
+    INPUT_CASE(16);
+    INPUT_CASE(17);
+    INPUT_CASE(18);
+    INPUT_CASE(19);
+    INPUT_CASE(20);
+    INPUT_CASE(21);
+    INPUT_CASE(22);
+    INPUT_CASE(23);
+    INPUT_CASE(24);
+    INPUT_CASE(25);
+    INPUT_CASE(26);
+    INPUT_CASE(27);
+    INPUT_CASE(28);
+    INPUT_CASE(29);
+    default:
+      ASSERT_LE(input_integer_bits, 30)
+          << "Input integer bits not handled: " << input_integer_bits;
+  }
+#undef INPUT_CASE
+}
+
+void RunSingleTest(const std::vector<int32_t>& test_input,
+                   int input_integer_bits, int output_integer_bits,
+                   const string& check_label, int tolerance) {
+#define OUTPUT_CASE(K)                                                   \
+  case K:                                                                \
+    return RunSingleTest<K>(test_input, input_integer_bits, check_label, \
+                            tolerance)
+  switch (output_integer_bits) {
+    OUTPUT_CASE(0);
+    OUTPUT_CASE(1);
+    OUTPUT_CASE(2);
+    OUTPUT_CASE(3);
+    OUTPUT_CASE(4);
+    OUTPUT_CASE(5);
+    OUTPUT_CASE(6);
+    OUTPUT_CASE(7);
+    OUTPUT_CASE(8);
+    OUTPUT_CASE(9);
+    OUTPUT_CASE(10);
+    OUTPUT_CASE(11);
+    OUTPUT_CASE(12);
+    OUTPUT_CASE(13);
+    OUTPUT_CASE(14);
+    OUTPUT_CASE(15);
+    OUTPUT_CASE(16);
+    OUTPUT_CASE(17);
+    OUTPUT_CASE(18);
+    OUTPUT_CASE(19);
+    OUTPUT_CASE(20);
+    OUTPUT_CASE(21);
+    OUTPUT_CASE(22);
+    OUTPUT_CASE(23);
+    OUTPUT_CASE(24);
+    OUTPUT_CASE(25);
+    OUTPUT_CASE(26);
+    OUTPUT_CASE(27);
+    OUTPUT_CASE(28);
+    OUTPUT_CASE(29);
+    default:
+      ASSERT_LE(input_integer_bits, 30)
+          << "Input integer bits not handled: " << input_integer_bits;
+  }
+#undef OUTPUT_CASE
+}
+
+void RunUniformTest(int test_size, int input_integer_bits,
+                    int output_integer_bits, const string& check_label,
+                    int tolerance, NumberGenerator* generator) {
+  std::vector<int> test_data = generator->RandomIntVector(
+      test_size, 2, std::numeric_limits<int>::max() - 1);
+  test_data[0] = 2;
+  test_data[1] = 3;
+  test_data[2] = 4;
+  test_data[3] = std::numeric_limits<int32_t>::max() - 2;
+  test_data[4] = std::numeric_limits<int32_t>::max() - 1;
+  test_data[5] = std::numeric_limits<int32_t>::max();
+
+  RunSingleTest(test_data, input_integer_bits, output_integer_bits,
+                check_label + " / uniform test", tolerance);
+}
+
+void RunUniformShiftUniformTest(int test_size, int input_integer_bits,
+                                int output_integer_bits,
+                                const string& check_label, int tolerance,
+                                NumberGenerator* generator) {
+  std::vector<int> test_data = generator->RandomIntVector(
+      test_size, 2, std::numeric_limits<int>::max() - 1);
+  std::vector<int> shifts = generator->RandomIntVector(test_size, 0, 29);
+  RightShiftVector(shifts, &test_data);
+
+  RunSingleTest(test_data, input_integer_bits, output_integer_bits,
+                check_label + " / shifted test", tolerance);
+}
+
+TEST_F(LogQuantizedTest, VariedIntegerBits) {
+  static constexpr int kVariations = 250;
+  static constexpr int kRunSize = 250;
+  static constexpr int kIntegerTolerance = 8;
+  static constexpr double kOutputFloatTolerance = 7.0e-7;
+
+  std::vector<int> input_integer_bits =
+      generator_.RandomIntVector(kVariations, 0, 24);
+  std::vector<int> output_integer_bits =
+      generator_.RandomIntVector(kVariations, 1, 10);
+
+  for (int i = 0; i < kVariations; ++i) {
+    int var_output_integer_bits = output_integer_bits[i];
+    int tolerance =
+        std::max(1.0 * kIntegerTolerance,
+                 (1 << (31 - var_output_integer_bits)) * kOutputFloatTolerance);
+
+    RunUniformTest(kRunSize, input_integer_bits[i], var_output_integer_bits,
+                   "VariedIntegerBits", tolerance, &generator_);
+    RunUniformShiftUniformTest(kRunSize, input_integer_bits[i],
+                               var_output_integer_bits, "VariedIntegerBits",
+                               tolerance, &generator_);
+  }
+}
+
+TEST_F(LogQuantizedTest, SelectedIntegerBits) {
+  static constexpr int kInputBits = 12;
+  static constexpr int kOutputBits = 5;
+  static constexpr int kRunSize = 100000;
+  static constexpr int kIntegerTolerance = 4;
+
+  RunUniformTest(kRunSize, kInputBits, kOutputBits, "SelectedIntegerBits",
+                 kIntegerTolerance, &generator_);
+  RunUniformShiftUniformTest(kRunSize, kInputBits, kOutputBits,
+                             "SelectedIntegerBits", kIntegerTolerance,
+                             &generator_);
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/internal/logsoftmax_quantized_test.cc

@@ -0,0 +1,349 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <functional>
+#include <iterator>
+#include <limits>
+#include <random>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/dequantize.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/log_softmax.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/test_util.h"
+#include "tensorflow/lite/string_type.h"
+
+namespace tflite {
+namespace {
+
+void RunLogSoftmaxFloatReference(const uint8_t* input_data,
+                                 const RuntimeShape& shape_common,
+                                 int32_t input_offset, const double input_scale,
+                                 int stride, float beta,
+                                 uint8_t* reference_output_data) {
+  const int ref_buffer_size = shape_common.FlatSize();
+  std::vector<float> reference_dequant_data(ref_buffer_size);
+  std::vector<float> reference_output_float_data(ref_buffer_size);
+
+  // Reference data generated via Dequant of input into float, and then applying
+  // float LogSoftmax.
+  DequantizationParams dq_params;
+  dq_params.zero_point = input_offset;
+  dq_params.scale = input_scale;
+  reference_ops::Dequantize(dq_params, shape_common, input_data, shape_common,
+                            reference_dequant_data.data());
+  SoftmaxParams sm_params;
+  optimized_ops::LogSoftmax(sm_params, shape_common,
+                            reference_dequant_data.data(), shape_common,
+                            reference_output_float_data.data());
+  // Work with quantized scaling for LogSoftmax, under which 255 represents 0,
+  // and -16 gets nudged up to 0.
+  for (int i = 0; i < ref_buffer_size; i++) {
+    reference_output_data[i] = std::max(
+        0, static_cast<int>(
+               255 + std::round(16.0f * reference_output_float_data[i])));
+  }
+}
+
+// Same as above except for the following change:
+// - input and output data type
+// - Dequnatize function
+// - clamping values
+void RunLogSoftmaxFloatReference(const int8_t* input_data,
+                                 const RuntimeShape& shape_common,
+                                 int32_t input_offset, const double input_scale,
+                                 int stride, float beta,
+                                 int8_t* reference_output_data) {
+  const int ref_buffer_size = shape_common.FlatSize();
+  std::vector<float> reference_dequant_data(ref_buffer_size);
+  std::vector<float> reference_output_float_data(ref_buffer_size);
+
+  // Reference data generated via Dequant of input into float, and then applying
+  // float LogSoftmax.
+  DequantizationParams dq_params;
+  dq_params.zero_point = input_offset;
+  dq_params.scale = input_scale;
+  reference_integer_ops::Dequantize(dq_params, shape_common, input_data,
+                                    shape_common,
+                                    reference_dequant_data.data());
+  SoftmaxParams sm_params;
+  optimized_ops::LogSoftmax(sm_params, shape_common,
+                            reference_dequant_data.data(), shape_common,
+                            reference_output_float_data.data());
+  // Work with quantized scaling for LogSoftmax, under which 255 represents 0,
+  // and -16 gets nudged up to 0.
+  for (int i = 0; i < ref_buffer_size; i++) {
+    reference_output_data[i] = std::max(
+        -128, static_cast<int>(
+                  127 + std::round(16.0f * reference_output_float_data[i])));
+  }
+}
+
+template <typename T>
+void CheckOutputData(const T* test_output, const T* reference_output,
+                     const RuntimeShape& shape_common,
+                     const string& check_label, bool be_exacting) {
+  const int buffer_size = shape_common.FlatSize();
+  // While calculating some metrics in floating point, we work with quantized
+  // scaling.
+  std::vector<int> diff(buffer_size);
+  int64_t sum_diff = 0;
+  int64_t sum_abs_diff = 0;
+  for (int i = 0; i < buffer_size; i++) {
+    diff[i] = static_cast<int>(test_output[i]) - reference_output[i];
+    sum_diff += diff[i];
+    sum_abs_diff += std::abs(diff[i]);
+  }
+  // These stats help understand test failures.
+  std::sort(std::begin(diff), std::end(diff));
+  const int min_diff = diff.front();
+  const int max_diff = diff.back();
+  const int median_diff = diff[diff.size() / 2];
+  const float mean_diff = static_cast<float>(sum_diff) / buffer_size;
+  const float mean_abs_diff = static_cast<float>(sum_abs_diff) / buffer_size;
+  // We either check for bit exactness (against the reference quantized version)
+  // or for general accuracy, allowing off-by-one (against the float reference).
+  if (be_exacting) {
+    ASSERT_TRUE(std::abs(min_diff) == 0 && std::abs(max_diff) == 0)
+        << check_label << ": "
+        << "std::abs(min_diff)=" << std::abs(min_diff)
+        << ", std::abs(max_diff)=" << std::abs(max_diff);
+  } else {
+    // For small numbers of samples, the estimates of the means vary more.
+    // Rather than widen the tolerances, we skip the smaller tests.
+    ASSERT_TRUE(((std::abs(mean_diff) < 2e-2f && mean_abs_diff < 3e-2f) ||
+                 buffer_size < 10000) &&
+                std::abs(median_diff) == 0 && std::abs(min_diff) <= 1 &&
+                std::abs(max_diff) <= 1)
+        << check_label << ": "
+        << "buffer_size=" << buffer_size << ", mean_diff=" << mean_diff
+        << ", mean_abs_diff=" << mean_abs_diff
+        << ", median_diff=" << median_diff << ", min_diff=" << min_diff
+        << ", max_diff=" << max_diff;
+  }
+}
+
+// Runs the LogSoftmax and compares against the float reference implementation
+// and the quantized reference implementation.
+void RunOneLogSoftmaxTest(const uint8_t* input_data,
+                          const RuntimeShape& shape_common,
+                          int32_t input_offset, const double input_scale,
+                          int stride, float beta) {
+  const int buffer_size = shape_common.FlatSize();
+  std::vector<uint8_t> optimized_logsoftmax_output(buffer_size);
+  std::vector<uint8_t> reference_float_logsoftmax_output(buffer_size);
+  std::vector<uint8_t> reference_quant_logsoftmax_output(buffer_size);
+
+  RunLogSoftmaxFloatReference(input_data, shape_common, input_offset,
+                              input_scale, stride, beta,
+                              reference_float_logsoftmax_output.data());
+
+  int32_t input_beta_multiplier;
+  int input_beta_left_shift;
+  int32_t reverse_scaling_divisor;
+  int reverse_scaling_right_shift;
+  static const int kScaledDiffIntegerBits = 5;
+  tflite::PreprocessLogSoftmaxScalingExp(
+      beta, input_scale, kScaledDiffIntegerBits, &input_beta_multiplier,
+      &input_beta_left_shift, &reverse_scaling_divisor,
+      &reverse_scaling_right_shift);
+  reverse_scaling_right_shift *= -1;
+  // diff_min has a negative value, and is used to limit the maximum magnitude
+  // of the diffs, which are <= 0.
+  const int diff_min = -tflite::CalculateInputRadius(kScaledDiffIntegerBits,
+                                                     input_beta_left_shift);
+
+  SoftmaxParams params;
+  float table[256];
+  params.input_multiplier = input_beta_multiplier;
+  params.input_left_shift = input_beta_left_shift;
+  params.reverse_scaling_divisor = reverse_scaling_divisor;
+  params.reverse_scaling_right_shift = reverse_scaling_right_shift;
+  params.diff_min = diff_min;
+
+  params.scale = 1.0f / 16.0f;
+  params.zero_point = 255;
+  params.table = table;
+  optimized_ops::PopulateSoftmaxLookupTable(&params, input_scale, beta);
+  optimized_ops::LogSoftmax(params, input_scale, shape_common, input_data,
+                            shape_common, optimized_logsoftmax_output.data());
+  reference_ops::LogSoftmax(params, shape_common, input_data, shape_common,
+                            reference_quant_logsoftmax_output.data());
+
+  CheckOutputData<uint8_t>(optimized_logsoftmax_output.data(),
+                           reference_float_logsoftmax_output.data(),
+                           shape_common, "Optimized vs float reference", false);
+  CheckOutputData<uint8_t>(optimized_logsoftmax_output.data(),
+                           reference_quant_logsoftmax_output.data(),
+                           shape_common, "Optimized vs quant reference", false);
+  CheckOutputData<uint8_t>(reference_quant_logsoftmax_output.data(),
+                           reference_float_logsoftmax_output.data(),
+                           shape_common, "Quant reference vs float reference",
+                           false);
+}
+
+// Runs the LogSoftmax and compares against the float reference implementation
+// and the int8 quantized reference implementation.
+void RunOneLogSoftmaxTest(const int8_t* input_data,
+                          const RuntimeShape& shape_common,
+                          int32_t input_offset, const double input_scale,
+                          int stride, float beta) {
+  const int buffer_size = shape_common.FlatSize();
+  std::vector<int8_t> quantized_logsoftmax_reference_implementation(
+      buffer_size);
+  std::vector<int8_t> float_logsoftmax_optimized_implementation(buffer_size);
+
+  RunLogSoftmaxFloatReference(input_data, shape_common, input_offset,
+                              input_scale, stride, beta,
+                              float_logsoftmax_optimized_implementation.data());
+
+  int32_t input_beta_multiplier;
+  int input_beta_left_shift;
+  int32_t reverse_scaling_divisor;
+  int reverse_scaling_right_shift;
+  static const int kScaledDiffIntegerBits = 5;
+  tflite::PreprocessLogSoftmaxScalingExp(
+      beta, input_scale, kScaledDiffIntegerBits, &input_beta_multiplier,
+      &input_beta_left_shift, &reverse_scaling_divisor,
+      &reverse_scaling_right_shift);
+  reverse_scaling_right_shift *= -1;
+  // diff_min has a negative value, and is used to limit the maximum magnitude
+  // of the diffs, which are <= 0.
+  const int diff_min = -tflite::CalculateInputRadius(kScaledDiffIntegerBits,
+                                                     input_beta_left_shift);
+
+  const int outer_size =
+      shape_common.Dims(0) * shape_common.Dims(1) * shape_common.Dims(2);
+  const int inner_size = shape_common.Dims(3);
+  reference_integer_ops::LogSoftmax(
+      input_beta_multiplier, input_beta_left_shift, reverse_scaling_divisor,
+      reverse_scaling_right_shift, diff_min, outer_size, inner_size, input_data,
+      quantized_logsoftmax_reference_implementation.data());
+
+  CheckOutputData<int8_t>(quantized_logsoftmax_reference_implementation.data(),
+                          float_logsoftmax_optimized_implementation.data(),
+                          shape_common, "Quant reference vs float reference",
+                          false);
+}
+
+// This function picks some random LogSoftmax params, which are checked for
+// desirability.  If not acceptable, it returns false. If they're OK,
+// it runs the LogSoftmax test and returns true. This allows the caller
+// to loop until a test has been run.
+//
+// Currently we do not reject for any reason.
+template <typename T>
+bool TryOneUniformLogSoftmax() {
+  // We pick mostly positive values, on the whole emphasizing smaller values and
+  // therefore faster tests.  We test a wider range of depths.  In the case of
+  // LogSoftmax, the width and height really just create test repetitions.
+  const int batch = ExponentialRandomPositiveInt(0.9f, 3, 20);
+  const int input_depth = ExponentialRandomPositiveInt(0.75f, 175, 500);
+  const int input_width = ExponentialRandomPositiveInt(0.8f, 20, 200);
+  const int input_height = ExponentialRandomPositiveInt(0.8f, 20, 200);
+  const int stride = ExponentialRandomPositiveInt(0.9f, 3, 8);
+  const double input_scale = std::pow(10.0, UniformRandomFloat(-2.0, 1.0));
+  const int32_t input_offset = UniformRandomInt(-256, 0);
+  static constexpr float beta = 1.0f;
+
+  auto shape_common =
+      RuntimeShape({batch, input_height, input_width, input_depth});
+  const int buffer_size = shape_common.FlatSize();
+
+  std::vector<T> input_data(buffer_size);
+  FillRandom(&input_data);
+  RunOneLogSoftmaxTest(input_data.data(), shape_common, input_offset,
+                       input_scale, stride, beta);
+  return true;
+}
+
+// See TryOneUniformLogSoftmax() for a general description.
+//
+// Tests with "skyscraper" input patterns are included for two reasons. (a)
+// Bimodal distributions are potentially challenging and perhaps more
+// realistic than simple uniform random inputs.  (b) Some implementations of
+// LogSoftmax may adapt as they traverse the depth, and so we test handling of
+// cases where relatively small values are encountered at the beginning and end.
+bool TryOneSkyscraperLogSoftmax(bool small_depth) {
+  // We pick mostly positive values, on the whole emphasizing smaller values and
+  // therefore faster tests.  We test a wider range of depths.  In the case of
+  // LogSoftmax, the width and height really just create test repetitions.
+  const int batch = ExponentialRandomPositiveInt(0.9f, 3, 20);
+  const int input_depth = small_depth
+                              ? ExponentialRandomPositiveInt(0.75f, 40, 500)
+                              : ExponentialRandomPositiveInt(0.75f, 175, 500);
+  const int input_width = ExponentialRandomPositiveInt(0.7f, 20, 200);
+  const int input_height = ExponentialRandomPositiveInt(0.7f, 20, 200);
+  const int stride = ExponentialRandomPositiveInt(0.9f, 3, 8);
+  const double input_scale = std::pow(10.0, UniformRandomFloat(-2.0, 1.0));
+  const int32_t input_offset = UniformRandomInt(-256, 0);
+  static constexpr float beta = 1.0f;
+  // Extra parameters for skyscraper input patterns.
+  const double middle_proportion =
+      ExponentialRandomPositiveFloat(0.65f, 0.1, 1.0);
+  const int middle_min = UniformRandomInt(0, 255);
+  const int sides_max = UniformRandomInt(0, middle_min);
+
+  auto shape_common =
+      RuntimeShape({batch, input_height, input_width, input_depth});
+  const int buffer_size = shape_common.FlatSize();
+
+  std::vector<uint8_t> input_data(buffer_size);
+  FillRandomSkyscraper(&input_data, input_depth, middle_proportion, middle_min,
+                       sides_max);
+  RunOneLogSoftmaxTest(input_data.data(), shape_common, input_offset,
+                       input_scale, stride, beta);
+  return true;
+}
+
+TEST(TestQuantizedLogSoftmax, UniformLogSoftmaxUint8Tests) {
+  const int kTestsToRun = 100;
+  for (int i = 0; i < kTestsToRun; i++) {
+    while (!TryOneUniformLogSoftmax<uint8_t>()) {
+    }
+  }
+}
+
+TEST(TestQuantizedLogSoftmax, UniformLogSoftmaxUint8Int8Tests) {
+  const int kTestsToRun = 100;
+  for (int i = 0; i < kTestsToRun; i++) {
+    while (!TryOneUniformLogSoftmax<int8_t>()) {
+    }
+  }
+}
+
+TEST(TestQuantizedLogSoftmax, SkyscraperLogSoftmaxUint8Tests) {
+  const int kTestsToRun = 100;
+  for (int i = 0; i < kTestsToRun; i++) {
+    while (!TryOneSkyscraperLogSoftmax(false)) {
+    }
+  }
+}
+
+TEST(TestQuantizedLogSoftmax, SmallSkyscraperLogSoftmaxUint8Tests) {
+  const int kTestsToRun = 100;
+  for (int i = 0; i < kTestsToRun; i++) {
+    while (!TryOneSkyscraperLogSoftmax(true)) {
+    }
+  }
+}
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/max.h

@@ -0,0 +1,35 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_MAX_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_MAX_H_
+
+#include <cmath>
+
+namespace tflite {
+
+#if defined(TF_LITE_USE_GLOBAL_MAX) || defined(__ZEPHYR__)
+inline float TfLiteMax(const float& x, const float& y) {
+  return std::max(x, y);
+}
+#else
+template <class T>
+inline T TfLiteMax(const T& x, const T& y) {
+  return std::fmax(x, y);
+}
+#endif
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_MAX_H_

tensorflow/lite/kernels/internal/maxpool_quantized_test.cc

@@ -0,0 +1,117 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <functional>
+#include <iterator>
+#include <limits>
+#include <random>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/pooling.h"
+#include "tensorflow/lite/kernels/internal/test_util.h"
+
+namespace tflite {
+namespace {
+
+// Runs the reference and optimized MaxPool functions and asserts the values
+// are the same.
+void RunOneMaxPoolTest(const PoolParams& params,
+                       const RuntimeShape& input_shape,
+                       const int8_t* input_data,
+                       const RuntimeShape& output_shape) {
+  const int buffer_size = output_shape.FlatSize();
+  std::vector<int8_t> optimized_maxpool_output(buffer_size);
+  std::vector<int8_t> reference_maxpool_output(buffer_size);
+
+  reference_integer_ops::MaxPool(params, input_shape, input_data, output_shape,
+                                 reference_maxpool_output.data());
+  optimized_integer_ops::MaxPool(params, input_shape, input_data, output_shape,
+                                 optimized_maxpool_output.data());
+
+  for (int i = 0; i < buffer_size; i++) {
+    ASSERT_TRUE(reference_maxpool_output[i] == optimized_maxpool_output[i]);
+  }
+}
+
+// Creates random input shape (batch, height, width, depth), then computes
+// output shape based on value of `padding_same`:
+// `padding_same` == true, calculate output with padding == "SAME"
+// `padding_same` == false, calculate output with padding == "VALID"
+// With input/output shapes computed, fills the input data and calls the
+// test function.
+void CreateDataAndRunMaxPool(bool padding_same) {
+  const int batch = UniformRandomInt(1, 2);
+  const int input_depth = UniformRandomInt(1, 700);
+  const int output_depth = input_depth;
+  const int input_width_offset = UniformRandomInt(1, 30);
+  const int input_height_offset = UniformRandomInt(1, 30);
+  const int stride_width = UniformRandomInt(1, 10);
+  const int stride_height = UniformRandomInt(1, 10);
+  const int filter_width = UniformRandomInt(1, 10);
+  const int filter_height = UniformRandomInt(1, 10);
+  const int input_width = input_width_offset + filter_width;
+  const int input_height = input_height_offset + filter_height;
+  const int output_width =
+      padding_same ? (input_width + stride_width - 1) / stride_width
+                   : (input_width - filter_width + stride_width) / stride_width;
+  const int output_height =
+      padding_same
+          ? (input_height + stride_height - 1) / stride_height
+          : (input_height - filter_height + stride_height) / stride_height;
+
+  auto input_shape =
+      RuntimeShape({batch, input_height, input_width, input_depth});
+  auto output_shape =
+      RuntimeShape({batch, output_height, output_width, output_depth});
+  const int buffer_size = input_shape.FlatSize();
+  std::vector<int8_t> input_data(buffer_size);
+  FillRandom(&input_data);
+
+  PoolParams params;
+  params.stride_height = stride_height;
+  params.stride_width = stride_width;
+  params.filter_height = filter_height;
+  params.filter_width = filter_width;
+  params.quantized_activation_min =
+      static_cast<int8_t>(std::numeric_limits<int8_t>::lowest());
+  params.quantized_activation_max =
+      static_cast<int8_t>(std::numeric_limits<int8_t>::max());
+  auto compute_padding = [](int stride, int in_size, int filter_size,
+                            int out_size) {
+    int padding = ((out_size - 1) * stride + filter_size - in_size) / 2;
+    return padding > 0 ? padding : 0;
+  };
+  params.padding_values.width =
+      compute_padding(stride_width, input_width, filter_width, output_width);
+  params.padding_values.height = compute_padding(stride_height, input_height,
+                                                 filter_height, output_height);
+  RunOneMaxPoolTest(params, input_shape, input_data.data(), output_shape);
+}
+
+TEST(TestMaxPool, SymmetricQuantMaxPool) {
+  const int kTestsToRun = 10;
+  for (int i = 0; i < kTestsToRun; i++) {
+    CreateDataAndRunMaxPool(/*padding_same=*/true);
+    CreateDataAndRunMaxPool(/*padding_same=*/false);
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/mfcc.cc

@@ -0,0 +1,65 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <math.h>
+
+#include "tensorflow/lite/kernels/internal/mfcc.h"
+
+namespace tflite {
+namespace internal {
+
+const double kDefaultUpperFrequencyLimit = 4000;
+const double kDefaultLowerFrequencyLimit = 20;
+const double kFilterbankFloor = 1e-12;
+const int kDefaultFilterbankChannelCount = 40;
+const int kDefaultDCTCoefficientCount = 13;
+
+Mfcc::Mfcc()
+    : initialized_(false),
+      lower_frequency_limit_(kDefaultLowerFrequencyLimit),
+      upper_frequency_limit_(kDefaultUpperFrequencyLimit),
+      filterbank_channel_count_(kDefaultFilterbankChannelCount),
+      dct_coefficient_count_(kDefaultDCTCoefficientCount) {}
+
+bool Mfcc::Initialize(int input_length, double input_sample_rate) {
+  bool initialized = mel_filterbank_.Initialize(
+      input_length, input_sample_rate, filterbank_channel_count_,
+      lower_frequency_limit_, upper_frequency_limit_);
+  initialized &=
+      dct_.Initialize(filterbank_channel_count_, dct_coefficient_count_);
+  initialized_ = initialized;
+  return initialized;
+}
+
+void Mfcc::Compute(const std::vector<double>& spectrogram_frame,
+                   std::vector<double>* output) const {
+  if (!initialized_) {
+    // LOG(ERROR) << "Mfcc not initialized.";
+    return;
+  }
+  std::vector<double> working;
+  mel_filterbank_.Compute(spectrogram_frame, &working);
+  for (int i = 0; i < working.size(); ++i) {
+    double val = working[i];
+    if (val < kFilterbankFloor) {
+      val = kFilterbankFloor;
+    }
+    working[i] = log(val);
+  }
+  dct_.Compute(working, output);
+}
+
+}  // namespace internal
+}  // namespace tflite

tensorflow/lite/kernels/internal/mfcc.h

@@ -0,0 +1,78 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Basic class for computing MFCCs from spectrogram slices.
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_H_
+
+#include <vector>
+
+#include "tensorflow/lite/kernels/internal/mfcc_dct.h"
+#include "tensorflow/lite/kernels/internal/mfcc_mel_filterbank.h"
+
+namespace tflite {
+namespace internal {
+
+class Mfcc {
+ public:
+  Mfcc();
+  bool Initialize(int input_length, double input_sample_rate);
+
+  // Input is a single squared-magnitude spectrogram frame. The input spectrum
+  // is converted to linear magnitude and weighted into bands using a
+  // triangular mel filterbank, and a discrete cosine transform (DCT) of the
+  // values is taken. Output is populated with the lowest dct_coefficient_count
+  // of these values.
+  void Compute(const std::vector<double>& spectrogram_frame,
+               std::vector<double>* output) const;
+
+  void set_upper_frequency_limit(double upper_frequency_limit) {
+    // CHECK(!initialized_) << "Set frequency limits before calling
+    // Initialize.";
+    upper_frequency_limit_ = upper_frequency_limit;
+  }
+
+  void set_lower_frequency_limit(double lower_frequency_limit) {
+    // CHECK(!initialized_) << "Set frequency limits before calling
+    // Initialize.";
+    lower_frequency_limit_ = lower_frequency_limit;
+  }
+
+  void set_filterbank_channel_count(int filterbank_channel_count) {
+    /// CHECK(!initialized_) << "Set channel count before calling Initialize.";
+    filterbank_channel_count_ = filterbank_channel_count;
+  }
+
+  void set_dct_coefficient_count(int dct_coefficient_count) {
+    // CHECK(!initialized_) << "Set coefficient count before calling
+    // Initialize.";
+    dct_coefficient_count_ = dct_coefficient_count;
+  }
+
+ private:
+  MfccMelFilterbank mel_filterbank_;
+  MfccDct dct_;
+  bool initialized_;
+  double lower_frequency_limit_;
+  double upper_frequency_limit_;
+  int filterbank_channel_count_;
+  int dct_coefficient_count_;
+};
+
+}  // namespace internal
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_H_

tensorflow/lite/kernels/internal/mfcc_dct.cc

@@ -0,0 +1,78 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/internal/mfcc_dct.h"
+
+#include <math.h>
+
+namespace tflite {
+namespace internal {
+
+MfccDct::MfccDct() : initialized_(false) {}
+
+bool MfccDct::Initialize(int input_length, int coefficient_count) {
+  coefficient_count_ = coefficient_count;
+  input_length_ = input_length;
+
+  if (coefficient_count_ < 1) {
+    return false;
+  }
+
+  if (input_length < 1) {
+    return false;
+  }
+
+  if (coefficient_count_ > input_length_) {
+    return false;
+  }
+
+  cosines_.resize(coefficient_count_);
+  double fnorm = sqrt(2.0 / input_length_);
+  // Some platforms don't have M_PI, so define a local constant here.
+  const double pi = atan(1.0) * 4.0;
+  double arg = pi / input_length_;
+  for (int i = 0; i < coefficient_count_; ++i) {
+    cosines_[i].resize(input_length_);
+    for (int j = 0; j < input_length_; ++j) {
+      cosines_[i][j] = fnorm * cos(i * arg * (j + 0.5));
+    }
+  }
+  initialized_ = true;
+  return true;
+}
+
+void MfccDct::Compute(const std::vector<double> &input,
+                      std::vector<double> *output) const {
+  if (!initialized_) {
+    return;
+  }
+
+  output->resize(coefficient_count_);
+  int length = input.size();
+  if (length > input_length_) {
+    length = input_length_;
+  }
+
+  for (int i = 0; i < coefficient_count_; ++i) {
+    double sum = 0.0;
+    for (int j = 0; j < length; ++j) {
+      sum += cosines_[i][j] * input[j];
+    }
+    (*output)[i] = sum;
+  }
+}
+
+}  // namespace internal
+}  // namespace tflite

tensorflow/lite/kernels/internal/mfcc_dct.h

@@ -0,0 +1,43 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Basic minimal DCT class for MFCC speech processing.
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_DCT_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_DCT_H_
+
+#include <vector>
+
+namespace tflite {
+namespace internal {
+
+class MfccDct {
+ public:
+  MfccDct();
+  bool Initialize(int input_length, int coefficient_count);
+  void Compute(const std::vector<double>& input,
+               std::vector<double>* output) const;
+
+ private:
+  bool initialized_;
+  int coefficient_count_;
+  int input_length_;
+  std::vector<std::vector<double> > cosines_;
+};
+
+}  // namespace internal
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_DCT_H_

tensorflow/lite/kernels/internal/mfcc_mel_filterbank.cc

@@ -0,0 +1,204 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This code resamples the FFT bins, and smooths then with triangle-shaped
+// weights to create a mel-frequency filter bank. For filter i centered at f_i,
+// there is a triangular weighting of the FFT bins that extends from
+// filter f_i-1 (with a value of zero at the left edge of the triangle) to f_i
+// (where the filter value is 1) to f_i+1 (where the filter values returns to
+// zero).
+
+// Note: this code fails if you ask for too many channels.  The algorithm used
+// here assumes that each FFT bin contributes to at most two channels: the
+// right side of a triangle for channel i, and the left side of the triangle
+// for channel i+1.  If you ask for so many channels that some of the
+// resulting mel triangle filters are smaller than a single FFT bin, these
+// channels may end up with no contributing FFT bins.  The resulting mel
+// spectrum output will have some channels that are always zero.
+
+#include "tensorflow/lite/kernels/internal/mfcc_mel_filterbank.h"
+
+#include <math.h>
+
+namespace tflite {
+namespace internal {
+
+MfccMelFilterbank::MfccMelFilterbank() : initialized_(false) {}
+
+bool MfccMelFilterbank::Initialize(int input_length, double input_sample_rate,
+                                   int output_channel_count,
+                                   double lower_frequency_limit,
+                                   double upper_frequency_limit) {
+  num_channels_ = output_channel_count;
+  sample_rate_ = input_sample_rate;
+  input_length_ = input_length;
+
+  if (num_channels_ < 1) {
+    // LOG(ERROR) << "Number of filterbank channels must be positive.";
+    return false;
+  }
+
+  if (sample_rate_ <= 0) {
+    // LOG(ERROR) << "Sample rate must be positive.";
+    return false;
+  }
+
+  if (input_length < 2) {
+    // LOG(ERROR) << "Input length must greater than 1.";
+    return false;
+  }
+
+  if (lower_frequency_limit < 0) {
+    // LOG(ERROR) << "Lower frequency limit must be nonnegative.";
+    return false;
+  }
+
+  if (upper_frequency_limit <= lower_frequency_limit) {
+    /// LOG(ERROR) << "Upper frequency limit must be greater than "
+    //           << "lower frequency limit.";
+    return false;
+  }
+
+  // An extra center frequency is computed at the top to get the upper
+  // limit on the high side of the final triangular filter.
+  center_frequencies_.resize(num_channels_ + 1);
+  const double mel_low = FreqToMel(lower_frequency_limit);
+  const double mel_hi = FreqToMel(upper_frequency_limit);
+  const double mel_span = mel_hi - mel_low;
+  const double mel_spacing = mel_span / static_cast<double>(num_channels_ + 1);
+  for (int i = 0; i < num_channels_ + 1; ++i) {
+    center_frequencies_[i] = mel_low + (mel_spacing * (i + 1));
+  }
+
+  // Always exclude DC; emulate HTK.
+  const double hz_per_sbin =
+      0.5 * sample_rate_ / static_cast<double>(input_length_ - 1);
+  start_index_ = static_cast<int>(1.5 + (lower_frequency_limit / hz_per_sbin));
+  end_index_ = static_cast<int>(upper_frequency_limit / hz_per_sbin);
+
+  // Maps the input spectrum bin indices to filter bank channels/indices. For
+  // each FFT bin, band_mapper tells us which channel this bin contributes to
+  // on the right side of the triangle.  Thus this bin also contributes to the
+  // left side of the next channel's triangle response.
+  band_mapper_.resize(input_length_);
+  int channel = 0;
+  for (int i = 0; i < input_length_; ++i) {
+    double melf = FreqToMel(i * hz_per_sbin);
+    if ((i < start_index_) || (i > end_index_)) {
+      band_mapper_[i] = -2;  // Indicate an unused Fourier coefficient.
+    } else {
+      while ((channel < num_channels_) &&
+             (center_frequencies_[channel] < melf)) {
+        ++channel;
+      }
+      band_mapper_[i] = channel - 1;  // Can be == -1
+    }
+  }
+
+  // Create the weighting functions to taper the band edges.  The contribution
+  // of any one FFT bin is based on its distance along the continuum between two
+  // mel-channel center frequencies.  This bin contributes weights_[i] to the
+  // current channel and 1-weights_[i] to the next channel.
+  weights_.resize(input_length_);
+  for (int i = 0; i < input_length_; ++i) {
+    channel = band_mapper_[i];
+    if ((i < start_index_) || (i > end_index_)) {
+      weights_[i] = 0.0;
+    } else {
+      if (channel >= 0) {
+        weights_[i] =
+            (center_frequencies_[channel + 1] - FreqToMel(i * hz_per_sbin)) /
+            (center_frequencies_[channel + 1] - center_frequencies_[channel]);
+      } else {
+        weights_[i] = (center_frequencies_[0] - FreqToMel(i * hz_per_sbin)) /
+                      (center_frequencies_[0] - mel_low);
+      }
+    }
+  }
+  // Check the sum of FFT bin weights for every mel band to identify
+  // situations where the mel bands are so narrow that they don't get
+  // significant weight on enough (or any) FFT bins -- i.e., too many
+  // mel bands have been requested for the given FFT size.
+  std::vector<int> bad_channels;
+  for (int c = 0; c < num_channels_; ++c) {
+    float band_weights_sum = 0.0;
+    for (int i = 0; i < input_length_; ++i) {
+      if (band_mapper_[i] == c - 1) {
+        band_weights_sum += (1.0 - weights_[i]);
+      } else if (band_mapper_[i] == c) {
+        band_weights_sum += weights_[i];
+      }
+    }
+    // The lowest mel channels have the fewest FFT bins and the lowest
+    // weights sum.  But given that the target gain at the center frequency
+    // is 1.0, if the total sum of weights is 0.5, we're in bad shape.
+    if (band_weights_sum < 0.5) {
+      bad_channels.push_back(c);
+    }
+  }
+  if (!bad_channels.empty()) {
+    /*
+    LOG(ERROR) << "Missing " << bad_channels.size() << " bands "
+               << " starting at " << bad_channels[0]
+               << " in mel-frequency design. "
+               << "Perhaps too many channels or "
+               << "not enough frequency resolution in spectrum. ("
+               << "input_length: " << input_length
+               << " input_sample_rate: " << input_sample_rate
+               << " output_channel_count: " << output_channel_count
+               << " lower_frequency_limit: " << lower_frequency_limit
+               << " upper_frequency_limit: " << upper_frequency_limit;
+               */
+  }
+  initialized_ = true;
+  return true;
+}
+
+// Compute the mel spectrum from the squared-magnitude FFT input by taking the
+// square root, then summing FFT magnitudes under triangular integration windows
+// whose widths increase with frequency.
+void MfccMelFilterbank::Compute(const std::vector<double> &input,
+                                std::vector<double> *output) const {
+  if (!initialized_) {
+    // LOG(ERROR) << "Mel Filterbank not initialized.";
+    return;
+  }
+
+  if (input.size() <= end_index_) {
+    // LOG(ERROR) << "Input too short to compute filterbank";
+    return;
+  }
+
+  // Ensure output is right length and reset all values.
+  output->assign(num_channels_, 0.0);
+
+  for (int i = start_index_; i <= end_index_; i++) {  // For each FFT bin
+    double spec_val = sqrt(input[i]);
+    double weighted = spec_val * weights_[i];
+    int channel = band_mapper_[i];
+    if (channel >= 0)
+      (*output)[channel] += weighted;  // Right side of triangle, downward slope
+    channel++;
+    if (channel < num_channels_)
+      (*output)[channel] += spec_val - weighted;  // Left side of triangle
+  }
+}
+
+double MfccMelFilterbank::FreqToMel(double freq) const {
+  return 1127.0 * log1p(freq / 700.0);
+}
+
+}  // namespace internal
+}  // namespace tflite

tensorflow/lite/kernels/internal/mfcc_mel_filterbank.h

@@ -0,0 +1,63 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Basic class for applying a mel-scale mapping to a power spectrum.
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_MEL_FILTERBANK_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_MEL_FILTERBANK_H_
+
+#include <vector>
+
+namespace tflite {
+namespace internal {
+
+class MfccMelFilterbank {
+ public:
+  MfccMelFilterbank();
+  bool Initialize(int input_length,  // Number of unique FFT bins fftsize/2+1.
+                  double input_sample_rate, int output_channel_count,
+                  double lower_frequency_limit, double upper_frequency_limit);
+
+  // Takes a squared-magnitude spectrogram slice as input, computes a
+  // triangular-mel-weighted linear-magnitude filterbank, and places the result
+  // in output.
+  void Compute(const std::vector<double>& input,
+               std::vector<double>* output) const;
+
+ private:
+  double FreqToMel(double freq) const;
+  bool initialized_;
+  int num_channels_;
+  double sample_rate_;
+  int input_length_;
+  std::vector<double> center_frequencies_;  // In mel, for each mel channel.
+
+  // Each FFT bin b contributes to two triangular mel channels, with
+  // proportion weights_[b] going into mel channel band_mapper_[b], and
+  // proportion (1 - weights_[b]) going into channel band_mapper_[b] + 1.
+  // Thus, weights_ contains the weighting applied to each FFT bin for the
+  // upper-half of the triangular band.
+  std::vector<double> weights_;  // Right-side weight for this fft  bin.
+
+  // FFT bin i contributes to the upper side of mel channel band_mapper_[i]
+  std::vector<int> band_mapper_;
+  int start_index_;  // Lowest FFT bin used to calculate mel spectrum.
+  int end_index_;    // Highest FFT bin used to calculate mel spectrum.
+};
+
+}  // namespace internal
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_MFCC_MEL_FILTERBANK_H_

tensorflow/lite/kernels/internal/min.h

@@ -0,0 +1,35 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_MIN_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_MIN_H_
+
+#include <cmath>
+
+namespace tflite {
+
+#if defined(TF_LITE_USE_GLOBAL_MIN) || defined(__ZEPHYR__)
+inline float TfLiteMin(const float& x, const float& y) {
+  return std::min(x, y);
+}
+#else
+template <class T>
+inline T TfLiteMin(const T& x, const T& y) {
+  return std::fmin(x, y);
+}
+#endif
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_MIN_H_

tensorflow/lite/kernels/internal/non_max_suppression_test.cc

@@ -0,0 +1,325 @@
+/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/reference/non_max_suppression.h"
+
+#include <algorithm>
+#include <cmath>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/test_util.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+constexpr int kNumBoxes = 6;
+
+void InitializeCandidates(std::vector<float>* boxes, std::vector<float>* scores,
+                          bool flip_coordinates = false) {
+  if (!flip_coordinates) {
+    *boxes = {
+        0, 0,    1, 1,     // Box 0
+        0, 0.1,  1, 1.1,   // Box 1
+        0, -0.1, 1, 0.9,   // Box 2
+        0, 10,   1, 11,    // Box 3
+        0, 10.1, 1, 11.1,  // Box 4
+        0, 100,  1, 101    // Box 5
+    };
+  } else {
+    *boxes = {
+        1, 1,     0, 0,     // Box 0
+        0, 0.1,   1, 1.1,   // Box 1
+        0, .9f,   1, -0.1,  // Box 2
+        0, 10,    1, 11,    // Box 3
+        1, 10.1f, 0, 11.1,  // Box 4
+        1, 101,   0, 100    // Box 5
+    };
+  }
+  *scores = {0.9, 0.75, 0.6, 0.95, 0.5, 0.3};
+}
+
+template <typename T>
+void MatchFirstNElements(int num_elements, const std::vector<T>& test_values,
+                         const std::vector<T>& reference_values) {
+  EXPECT_LT(num_elements, test_values.size());
+  EXPECT_EQ(num_elements, reference_values.size());
+
+  for (int i = 0; i < num_elements; ++i) {
+    EXPECT_EQ(test_values[i], reference_values[i]);
+  }
+}
+
+TEST(NonMaxSuppression, TestZeroBoxes) {
+  // Inputs
+  std::vector<float> boxes(1);
+  std::vector<float> scores(1);
+  const float iou_threshold = 0.5;
+  const float score_threshold = 0.4;
+  const int max_output_size = 4;
+
+  // Outputs
+  std::vector<int> selected_indices(6);
+  std::vector<float> selected_scores(6);
+  int num_selected_indices = -1;
+
+  reference_ops::NonMaxSuppression(
+      boxes.data(), /**num_boxes=**/ 0, scores.data(), max_output_size,
+      iou_threshold, score_threshold, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 0);
+}
+
+TEST(NonMaxSuppression, TestSelectFromIdenticalBoxes) {
+  // Inputs
+  std::vector<float> boxes(kNumBoxes * 4);
+  std::vector<float> scores(kNumBoxes);
+  for (int i = 0; i < kNumBoxes; ++i) {
+    boxes[i * 4 + 0] = 0;
+    boxes[i * 4 + 1] = 0;
+    boxes[i * 4 + 2] = 1;
+    boxes[i * 4 + 3] = 1;
+    scores[i] = 0.75;
+  }
+  const float iou_threshold = 0.5;
+  float score_threshold = 0.5;
+  const int max_output_size = kNumBoxes;
+
+  // Outputs
+  std::vector<int> selected_indices(6);
+  std::vector<float> selected_scores(6);
+  int num_selected_indices = -1;
+
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      score_threshold, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 1);
+  MatchFirstNElements(1, selected_scores, {.75});
+
+  score_threshold = 0.95;
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      score_threshold, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 0);
+}
+
+TEST(NonMaxSuppression, TestSelectFromThreeClustersWithZeroScoreThreshold) {
+  // Inputs
+  std::vector<float> boxes;
+  std::vector<float> scores;
+  InitializeCandidates(&boxes, &scores);
+  const float iou_threshold = 0.5;
+  int max_output_size;
+
+  // Outputs
+  std::vector<int> selected_indices(6);
+  std::vector<float> selected_scores(6);
+  int num_selected_indices = -1;
+
+  // Test a large max_output_size.
+  max_output_size = 100;
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      /**score_threshold=**/ 0.0, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 3);
+  MatchFirstNElements(3, selected_indices, {3, 0, 5});
+  MatchFirstNElements(3, selected_scores, {0.95, 0.9, 0.3});
+
+  // Smaller max_output_size.
+  max_output_size = 2;
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      /**score_threshold=**/ 0.0, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, max_output_size);
+  MatchFirstNElements(max_output_size, selected_indices, {3, 0});
+  MatchFirstNElements(max_output_size, selected_scores, {0.95, 0.9});
+
+  // max_output_size = 0.
+  max_output_size = 0;
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      /**score_threshold=**/ 0.0, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 0);
+}
+
+TEST(NonMaxSuppression, TestSelectFromThreeClustersWithScoreThreshold) {
+  // Inputs
+  std::vector<float> boxes;
+  std::vector<float> scores;
+  InitializeCandidates(&boxes, &scores);
+  const float iou_threshold = 0.5;
+  const float score_threshold = 0.4;
+  int max_output_size;
+
+  // Outputs
+  std::vector<int> selected_indices(6);
+  std::vector<float> selected_scores(6);
+  int num_selected_indices = -1;
+
+  // Test a large max_output_size.
+  max_output_size = 100;
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      score_threshold, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 2);
+  MatchFirstNElements(2, selected_indices, {3, 0});
+  MatchFirstNElements(2, selected_scores, {0.95, 0.9});
+
+  // max_output_size = 1.
+  max_output_size = 1;
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      score_threshold, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 1);
+  MatchFirstNElements(1, selected_indices, {3});
+  MatchFirstNElements(1, selected_scores, {0.95});
+}
+
+// This flips the (y1, x1) & (y2, x2) corners for each box. The output should
+// match what we get without flipping.
+TEST(NonMaxSuppression, TestSelectFromThreeClustersWithFlippedCoordinates) {
+  // Inputs
+  std::vector<float> boxes;
+  std::vector<float> scores;
+  InitializeCandidates(&boxes, &scores, /**flipped_coordinates=**/ true);
+  const float iou_threshold = 0.5;
+  const float score_threshold = 0.4;
+  const int max_output_size = 3;
+
+  // Outputs
+  std::vector<int> selected_indices(6);
+  std::vector<float> selected_scores(6);
+  int num_selected_indices = -1;
+
+  // Test a large max_output_size.
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      score_threshold, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 2);
+  MatchFirstNElements(2, selected_indices, {3, 0});
+  MatchFirstNElements(2, selected_scores, {0.95, 0.9});
+
+  // score_threshold = 0.
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      /**score_threshold=**/ 0.0, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 3);
+  MatchFirstNElements(3, selected_indices, {3, 0, 5});
+  MatchFirstNElements(3, selected_scores, {0.95, 0.9, 0.3});
+}
+
+TEST(NonMaxSuppression, TestIoUThresholdBoundaryCases) {
+  // Inputs
+  std::vector<float> boxes;
+  std::vector<float> scores;
+  InitializeCandidates(&boxes, &scores);
+  const float score_threshold = 0.4;
+  const int max_output_size = 4;
+
+  // Outputs
+  std::vector<int> selected_indices(6);
+  std::vector<float> selected_scores(6);
+  int num_selected_indices = -1;
+
+  // IoU threshold is zero. Only one index should get selected.
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size,
+      /**iou_threshold=**/ 0.0, score_threshold, /**sigma=**/ 0.0,
+      selected_indices.data(), selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 1);
+  MatchFirstNElements(1, selected_indices, {3});
+  MatchFirstNElements(1, selected_scores, {0.95});
+
+  // IoU threshold too high. max_output_size number of indices should be
+  // selected.
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size,
+      /**iou_threshold=**/ 0.9999,
+      /**score_threshold=**/ 0.0, /**sigma=**/ 0.0, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, max_output_size);
+  MatchFirstNElements(max_output_size, selected_indices, {3, 0, 1, 2});
+  MatchFirstNElements(max_output_size, selected_scores, {0.95, 0.9, 0.75, 0.6});
+}
+
+TEST(NonMaxSuppression, TestSelectFromThreeClustersWithSoftNMS) {
+  // Inputs
+  std::vector<float> boxes;
+  std::vector<float> scores;
+  InitializeCandidates(&boxes, &scores);
+  const float iou_threshold = 1.0;
+  float score_threshold = 0.0;
+  const float soft_nms_sigma = 0.5;
+  int max_output_size = 6;
+
+  // Outputs
+  std::vector<int> selected_indices(6);
+  std::vector<float> selected_scores(6);
+  int num_selected_indices = -1;
+
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      score_threshold, soft_nms_sigma, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 6);
+  // Box 0 soft-suppresses box 1, but not enough to cause it to fall under
+  // `score_threshold` (which is 0.0) or the score of box 5, so in this test,
+  // box 1 ends up being selected before box 5.
+  EXPECT_THAT(selected_indices, ElementsAreArray({3, 0, 1, 5, 4, 2}));
+  EXPECT_THAT(selected_scores,
+              ElementsAreArray(
+                  ArrayFloatNear({0.95, 0.9, 0.384, 0.3, 0.256, 0.197}, 1e-3)));
+
+  score_threshold = 0.299;
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      score_threshold, soft_nms_sigma, selected_indices.data(),
+      selected_scores.data(), &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 4);
+  MatchFirstNElements(4, selected_indices, {3, 0, 1, 5});
+}
+
+TEST(NonMaxSuppression, TestNullSelectedScoresOutput) {
+  // Inputs
+  std::vector<float> boxes;
+  std::vector<float> scores;
+  InitializeCandidates(&boxes, &scores);
+  const float iou_threshold = 0.5;
+  const float score_threshold = 0.4;
+  int max_output_size;
+
+  // Outputs
+  std::vector<int> selected_indices(6);
+  int num_selected_indices = -1;
+
+  max_output_size = 100;
+  reference_ops::NonMaxSuppression(
+      boxes.data(), kNumBoxes, scores.data(), max_output_size, iou_threshold,
+      score_threshold, /**sigma=**/ 0.0, selected_indices.data(),
+      /**selected_scores=**/ nullptr, &num_selected_indices);
+  EXPECT_EQ(num_selected_indices, 2);
+}
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/opaque_tensor_ctypes.cc

@@ -0,0 +1,38 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/internal/opaque_tensor_ctypes.h"
+
+#include "tensorflow/lite/c/c_api_opaque.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/namespace.h"
+
+namespace tflite {
+namespace TFLITE_CONDITIONAL_NAMESPACE {
+
+RuntimeShape GetTensorShape(const TfLiteOpaqueTensor* tensor) {
+  if (tensor == nullptr) {
+    return RuntimeShape();
+  }
+  const int dims_size = TfLiteOpaqueTensorNumDims(tensor);
+  RuntimeShape shape(dims_size);
+  for (int i = 0; i < dims_size; ++i) {
+    shape.SetDim(i, TfLiteOpaqueTensorDim(tensor, i));
+  }
+  return shape;
+}
+
+}  // namespace TFLITE_CONDITIONAL_NAMESPACE
+}  // namespace tflite

tensorflow/lite/kernels/internal/opaque_tensor_ctypes.h

@@ -0,0 +1,35 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPAQUE_TENSOR_CTYPES_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPAQUE_TENSOR_CTYPES_H_
+
+#include "tensorflow/lite/c/c_api_opaque.h"
+#include "tensorflow/lite/core/macros.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/namespace.h"
+
+namespace tflite {
+namespace TFLITE_CONDITIONAL_NAMESPACE {
+
+/// Returns the dimensions of the given tensor.
+TFLITE_NOINLINE RuntimeShape GetTensorShape(const TfLiteOpaqueTensor* tensor);
+
+}  // namespace TFLITE_CONDITIONAL_NAMESPACE
+
+using ::tflite::TFLITE_CONDITIONAL_NAMESPACE::GetTensorShape;
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPAQUE_TENSOR_CTYPES_H_

tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_common.h

@@ -0,0 +1,47 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_COMMON_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_COMMON_H_
+#include <cstdint>
+
+namespace tflite {
+namespace optimized_4bit {
+
+// Since we need to convert int4 to int8 with shifts, it is faster if we
+// can use unsigned int4, so just subtract zero_point_4bit from all values.
+// Fold input * zero_point into quantization since we need to quantize
+// each input and multiply by zero_point_4bit to convert back to signed int.
+constexpr int zero_point_4bit = -7;
+
+inline int8_t upper(int8_t value) { return value >> 4; }
+
+inline int8_t lower(int8_t value) {
+  uint8_t sign_y = UINT8_C(256) - (value & UINT8_C(8));
+  return (value & UINT8_C(7)) | sign_y;
+}
+
+inline int8_t merge(int8_t upper, int8_t lower) {
+  const auto to_int4 = [](int8_t v) -> uint8_t {
+    int32_t x = v + 7;
+    return static_cast<uint8_t>(x);
+  };
+  return (to_int4(upper) << 4) | to_int4(lower);
+}
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_COMMON_H_

tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference.cc

@@ -0,0 +1,365 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <stdint.h>
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstring>
+#include <vector>
+
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_common.h"
+#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference_impl.h"
+
+namespace tflite {
+namespace optimized_4bit {
+
+void ReferencePackInner(const int8_t* src, uint8_t* box, int src_rows,
+                        int src_cols, int outer_row, int outer_col,
+                        int outer_rows, int outer_cols, int inner_rows,
+                        int inner_cols) {
+  // Create a kernel-specific layout for a packed unit.
+  const int width = inner_rows;
+  const int depth = inner_cols;
+  const int real_depth = depth / 2;
+  const int real_src_cols = src_cols / 2;
+  // Determine row and column of source tensor.
+  const int row = outer_row * inner_rows;
+  const int col = outer_col * inner_cols;
+  // The source width  (rows) and depth (columns).
+  int src_width = std::min(width, src_rows - row);
+  int src_depth = std::min(depth, src_cols - col);
+  int real_col = col / 2;
+  const int8_t* src_data = src + row * real_src_cols + real_col;
+  int real_src_depth = src_depth / 2;
+  // Src is [rows / src_rows, cols / src_depth, src_rows, src_cols]
+  // Reshape and pad to [outer_rows, outer_cols, width, depth]
+  // Interleave values [u1,u2,...,u_depth] to
+  // [u1,u_{depth/2+1},u2,u_{depth/2+2},.., u_{depth/2},u_depth]
+  // So that after shifting, we get [u1,u2...u_{depth/2}] and
+  // [u_{depth/2} + 1, ... u_{depth}].
+  for (int m = 0; m < src_width; ++m) {
+    int i = 0;
+    int k = 0;
+    int half_depth = depth / 2;
+    int half_half_depth = half_depth / 2;
+    for (; i < (real_src_depth & (~(half_depth - 1))); i += half_depth) {
+      for (int j = 0; j < half_half_depth; ++j) {
+        const int8_t v1 = (int8_t)src_data[i + j];
+        int8_t uv1 = upper(v1);
+        int8_t lv1 = lower(v1);
+        const int8_t v2 = (int8_t)src_data[i + j + half_half_depth];
+        int8_t uv2 = upper(v2);
+        int8_t lv2 = lower(v2);
+        box[k] = merge(lv1, lv2);
+        box[k + 1] = merge(uv1, uv2);
+        k += 2;
+      }
+    }
+    // Handle remaining values -- if greater than or equal to
+    // 16 values remaining, do the shuffle.
+    if (i < real_src_depth) {
+      const int remaining = half_half_depth < (real_src_depth - i)
+                                ? half_half_depth
+                                : real_src_depth - i;
+      for (int j = 0; j < remaining; ++j) {
+        const int8_t v1 = (int8_t)src_data[i + j];
+        int8_t uv1 = upper(v1);
+        int8_t lv1 = lower(v1);
+        int8_t uv2 = 0;
+        int8_t lv2 = 0;
+        if ((i + j + half_half_depth) < real_src_depth) {
+          const int8_t v2 = (int8_t)src_data[i + j + half_half_depth];
+          uv2 = upper(v2);
+          lv2 = lower(v2);
+        }
+        box[k] = merge(lv1, lv2);
+        box[k + 1] = merge(uv1, uv2);
+        k += 2;
+      }
+    }
+    box += real_depth;
+    src_data += real_src_cols;
+  }
+}
+
+void ReferencePrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                      int layout_cols, int src_rows, int src_cols, int width,
+                      int depth) {
+  size_t size = layout_rows * layout_cols / 2;
+  memset(dest, static_cast<uint8_t>(0x77), sizeof(uint8_t) * size);
+  int outer_cols = layout_cols / depth;
+  int outer_rows = layout_rows / width;
+  int inner_cols = depth;
+  int inner_rows = width;
+  for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
+    for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+      // Each outer row x outer col contains width x depth, copied
+      // from tensor at the cluster_index.
+      const int cluster_index = outer_row * outer_cols + outer_col;
+      const int real_depth = inner_cols / 2;
+      uint8_t* box = dest + cluster_index * real_depth * inner_rows;
+      ReferencePackInner(tensor, box, src_rows, src_cols, outer_row, outer_col,
+                         outer_rows, outer_cols, inner_rows, inner_cols);
+    }
+  }
+}
+
+void ReferenceBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                      int n_data, int8_t* quantized_data_ptr,
+                                      float* scaling_factors, int width,
+                                      int depth, int32_t* input_offsets) {
+  const int rows = n_batch;
+  const int cols = n_data;
+  // depth is always cols
+  const int layout_rows = (rows + (width - 1)) & ~(width - 1);
+  const int layout_cols = (cols + (depth - 1)) & ~(depth - 1);
+  const int size = layout_rows * layout_cols;
+  int8_t* data = quantized_data_ptr;
+  memset(data, 0, sizeof(int8_t) * size);
+  memset(input_offsets, 0, sizeof(int32_t) * layout_rows);
+  const float* tensor_data = float_data_ptr;
+  // basically, we need to make a new 4D matrix
+  // [rows / width, cols / depth, width, depth] in depth-first
+  const int outer_cols = layout_cols / depth;
+  const int outer_rows = layout_rows / width;
+  float* scaling_factors_ptr = scaling_factors;
+  for (int outer_row = 0; outer_row < outer_rows; outer_row++) {
+    std::vector<float> scale(width);
+    const int row = width * outer_row;
+    scaling_factors_ptr = scaling_factors + row;
+    for (int w = 0; w < width; ++w) {
+      if ((row + w) >= rows) {
+        continue;
+      }
+      const float* start = tensor_data + (row + w) * cols;
+      float scale_denom = 0;
+      for (int c = 0; c < cols; ++c) {
+        scale_denom = std::max(scale_denom, std::abs(*(start++)));
+      }
+      if (scale_denom == 0) {
+        scale_denom = 127.0;
+      }
+      scale[w] = 127.0 / scale_denom;
+      scaling_factors_ptr[w] = scale_denom / 127.0;
+    }
+    for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+      const int col = depth * outer_col;
+      const int src_width = std::min(width, rows - row);
+      const int src_depth = std::min(depth, cols - col);
+      const int cluster_index = outer_row * outer_cols + outer_col;
+      int8_t* box = data + cluster_index * depth * width;
+      for (int w = 0; w < src_width; ++w) {
+        const float* float_data = tensor_data + (row + w) * cols + col;
+        for (int d = 0; d < src_depth; ++d) {
+          int8_t q = static_cast<int8_t>(TfLiteRound(float_data[d] * scale[w]));
+          box[w * depth + d] = q;
+          input_offsets[row + w] += q;
+        }
+      }
+    }
+  }
+  for (int r = 0; r < layout_rows; ++r) {
+    // Multiply the input by zero-point so that we don't have to calculate
+    // later.
+    input_offsets[r] = input_offsets[r] * zero_point_4bit;
+  }
+}
+
+void ReferenceAssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                          const float* batch_scales,
+                                          const float* filter_scales,
+                                          const float* bias_ptr,
+                                          float* output_ptr, int output_depth,
+                                          int batch_size) {
+  if (bias_ptr) {
+    for (int b = 0; b < batch_size; ++b) {
+      const float val = *input_offsets++ * *batch_scales++;
+      const float* filter_scales_ptr = filter_scales;
+      const float* bias_ptr_tmp = bias_ptr;
+      for (int i = 0; i < output_depth; i++) {
+        *output_ptr++ = (val * *filter_scales_ptr++) + *bias_ptr_tmp++;
+      }
+    }
+    return;
+  }
+  for (int b = 0; b < batch_size; ++b) {
+    const float val = *input_offsets++ * *batch_scales++;
+    const float* filter_scales_ptr = filter_scales;
+    for (int i = 0; i < output_depth; i++) {
+      *output_ptr++ = (val * *filter_scales_ptr++);
+    }
+  }
+}
+
+/* Unpack the accumulated scratch buffer by transposing and multiplying
+ * by input and filter scales.
+ * Before, dst contains integer accumulated values with layout:
+ *   [rhs_layout_rows // rhs_width, lhs_layout_rows // lhs_width,
+ *       rhs_width, lhs_width]
+ * Transpose and dequantize to [batch_size, num_units].
+ */
+template <int Depth, int Width>
+void ReferenceUnpack(float* output_ptr, const int32_t* dst, int batch_size,
+                     int num_units, const float* scaling_factors,
+                     const float* filter_scales, int dst_layout_rows,
+                     int dst_layout_cols) {
+  // Width == 1 is when batch size == 1, the most frequent case.
+  // No need to iterate over outer rows.
+  if (Width == 1) {
+    const int outer_rows = dst_layout_rows / Width;
+    const int outer_cols = dst_layout_cols / Depth;
+    const int32_t* dst_ptr = dst;
+    int unit = 0;
+    for (int outer_col = 0; outer_col < outer_cols;
+         ++outer_col, unit += Depth) {
+      float* tmp_output_ptr = output_ptr + unit;
+      int len = num_units - unit < Depth ? num_units - unit : Depth;
+      const float* scaling_factors_ptr = scaling_factors;
+      for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
+        const float scale = *scaling_factors_ptr;
+        const float* filter_scales_ptr = filter_scales + unit;
+        for (int i = 0; i < len; ++i) {
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+        }
+        dst_ptr += (Depth - len);
+        scaling_factors_ptr += Width;
+        tmp_output_ptr += (num_units - len);
+      }
+    }
+    return;
+  }
+  const int outer_rows = dst_layout_rows / Width;
+  const int outer_cols = dst_layout_cols / Depth;
+  for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+    const int unit = outer_col * Depth;
+    const int remaining_units = std::min(num_units - unit, Depth);
+    const int depth_offset = Depth - remaining_units;
+    const int width_offset = num_units - remaining_units;
+    int outer_row = 0;
+    for (; outer_row < outer_rows; ++outer_row) {
+      const int batch = outer_row * Width;
+      const int remaining_width = std::min(batch_size - batch, Width);
+      const int cluster_index = outer_col * outer_rows + outer_row;
+      const int32_t* dst_ptr = dst + cluster_index * Depth * Width;
+      float* tmp_output_ptr = output_ptr + batch * num_units + unit;
+      const float* scale = scaling_factors + batch;
+      int w = remaining_width;
+      for (; w > 0; --w, scale++) {
+        int d = remaining_units;
+        const float* filter_scales_ptr = filter_scales + unit;
+        for (; d > 0; --d) {
+          *tmp_output_ptr++ += *dst_ptr++ * (*scale) * (*filter_scales_ptr++);
+        }
+        dst_ptr += depth_offset;
+        tmp_output_ptr += width_offset;
+      }
+    }
+  }
+}
+
+template <int RowsLeft, int RowsRight, int Cols>
+void ReferenceRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                        int lhs_layout_rows, int lhs_layout_cols,
+                        int rhs_layout_rows, int rhs_layout_cols,
+                        int dst_layout_rows, int dst_layout_cols) {
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  int32_t* elementPtr = dst;
+  const int outer_rows = (clamped_end_row + RowsLeft - 1) / RowsLeft;
+  const int outer_cols = (clamped_end_col + RowsRight - 1) / RowsRight;
+  const int depth = std::min(lhs_layout_cols / Cols, rhs_layout_cols / Cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    int left_index = i * RowsLeft * lhs_layout_cols / 2;
+    const uint8_t* lhs_val_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_val = lhs_val_data;
+      int right_index = j * RowsRight * rhs_layout_cols;
+      const int8_t* rhs_val = rhs + right_index;
+      int32_t accum[RowsLeft * RowsRight];
+      memset(accum, 0, sizeof(int32_t) * RowsLeft * RowsRight);
+      for (int k = 0; k < depth; ++k) {
+        uint8_t lhs_[RowsLeft][Cols];
+        for (int m = 0; m < RowsLeft; ++m) {
+          for (int n = 0; n < Cols / 2; ++n) {
+            uint8_t val = *(lhs_val++);
+            lhs_[m][n] = (val >> 4 & 15);
+            lhs_[m][n + (Cols / 2)] = (val & 15);
+          }
+        }
+        int8_t rhs_[RowsRight][Cols];
+        for (int m = 0; m < RowsRight; ++m) {
+          for (int n = 0; n < Cols; ++n) {
+            rhs_[m][n] = *(rhs_val++);
+          }
+        }
+        for (int r = 0; r < RowsRight; ++r) {
+          for (int l = 0; l < RowsLeft; ++l) {
+            for (int i = 0; i < Cols; ++i) {
+              accum[r * RowsLeft + l] += lhs_[l][i] * rhs_[r][i];
+            }
+          }
+        }
+      }  // end depth
+      for (int r = 0; r < RowsRight; ++r) {
+        for (int l = 0; l < RowsLeft; ++l) {
+          int32_t q = accum[r * RowsLeft + l];
+          *(elementPtr++) = q;
+        }
+      }
+    }
+  }
+}
+
+template void ReferenceUnpack<4, 1>(float* output_ptr, const int32_t* dst,
+                                    int batch_size, int num_units,
+                                    const float* scaling_factors,
+                                    const float* filter_scales,
+                                    int dst_layout_rows, int dst_layout_cols);
+
+template void ReferenceUnpack<4, 2>(float* output_ptr, const int32_t* dst,
+                                    int batch_size, int num_units,
+                                    const float* scaling_factors,
+                                    const float* filter_scales,
+                                    int dst_layout_rows, int dst_layout_cols);
+
+template void ReferenceUnpack<4, 4>(float* output_ptr, const int32_t* dst,
+                                    int batch_size, int num_units,
+                                    const float* scaling_factors,
+                                    const float* filter_scales,
+                                    int dst_layout_rows, int dst_layout_cols);
+
+template void ReferenceRunKernel<4, 1, 32>(
+    const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
+    int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
+    int dst_layout_rows, int dst_layout_cols);
+
+template void ReferenceRunKernel<4, 2, 32>(
+    const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
+    int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
+    int dst_layout_rows, int dst_layout_cols);
+
+template void ReferenceRunKernel<4, 4, 32>(
+    const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
+    int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
+    int dst_layout_rows, int dst_layout_cols);
+
+}  // namespace optimized_4bit
+}  // namespace tflite

tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference.h

@@ -0,0 +1,140 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_H_
+
+#include <cstdint>
+
+#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference_impl.h"
+
+namespace tflite {
+namespace optimized_4bit {
+
+/* Returns the maximum number of rhs rows supported/compiled.
+ *
+ * In general 4 is the most we can do without running out of registers on
+ * aarch64. For x86/aarch64, 4x32 4bit = 64 bytes can be held in cache line
+ * size. This is required to set the packing layout for the rhs.
+ *
+ * For reference, return 1.
+ */
+inline int GetMaxSupportedRows() { return 1; }
+
+/* Pack a 4bit inner_rows x inner_cols array from src.
+ * This is called as an inner function for Prepack.
+ */
+inline void PackInner(const int8_t* src, uint8_t* box, int src_rows,
+                      int src_cols, int outer_row, int outer_col,
+                      int outer_rows, int outer_cols, int inner_rows,
+                      int inner_cols) {
+  ReferencePackInner(src, box, src_rows, src_cols, outer_row, outer_col,
+                     outer_rows, outer_cols, inner_rows, inner_cols);
+}
+
+/* Prepack lhs matrix into dest.
+ * Transform tensor from (src_rows, src_cols) to
+ * (layout_rows / width, layout_cols / depth, width, depth) with possibly
+ * padding, and interleaving values along depth / 2 dimensions.
+ * dest should be aligned and allocated before prepack.
+ */
+inline void Prepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                    int layout_cols, int src_rows, int src_cols, int width,
+                    int depth) {
+  ReferencePrepack(dest, tensor, layout_rows, layout_cols, src_rows, src_cols,
+                   width, depth);
+}
+
+/* Quantize input floats to 8bit and calculate sum of each column.
+ *  Data in float_data_ptr of shape (n_batch x n_data), is quantized and
+ * packed into (n_batch / width, n_data / depth, width, data) into
+ * quantized_data_ptr and input_offsets will contain the product of filter
+ * zero_point and input.
+ */
+inline void BatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                    int n_data, int8_t* quantized_data_ptr,
+                                    float* scaling_factors, int width,
+                                    int depth, int32_t* input_offsets) {
+  ReferenceBatchQuantizeFloats4Bit(float_data_ptr, n_batch, n_data,
+                                   quantized_data_ptr, scaling_factors, width,
+                                   depth, input_offsets);
+}
+
+/* Write bias + input offset * filter_scale to output_ptr.
+ * output_ptr of size (batch_size, output_depth) will have
+ * output_ptr[output_depth * b + o] =
+ *     bias_ptr[o] + input_offsets[b] * batch_scales[b] * filter_scale[o]
+ */
+inline void AssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                        const float* batch_scales,
+                                        float* filter_scales,
+                                        const float* bias_ptr,
+                                        float* output_ptr, int output_depth,
+                                        int batch_size) {
+  ReferenceAssignBiasAndComputeOffsets(input_offsets, batch_scales,
+                                       filter_scales, bias_ptr, output_ptr,
+                                       output_depth, batch_size);
+}
+
+/* Add accumulated integer sums in dst to float output.
+ * output_ptr of size (batch_size, output_depth) will have
+ * output_ptr[b * output_depth + o] = \
+ *   dst[b / dst_layout_rows, o / dst_layout_cols,
+ *       b % dst_layout_rows, o % dst_layout_cols] * scaling_filters[b] *
+ *       filter_scales[o]
+ */
+template <int Depth, int Width>
+void Unpack(float* output_ptr, const int32_t* dst, int batch_size,
+            int num_units, const float* scaling_factors,
+            const float* filter_scales, int dst_layout_rows,
+            int dst_layout_cols) {
+  ReferenceUnpack<Depth, Width>(output_ptr, dst, batch_size, num_units,
+                                scaling_factors, filter_scales, dst_layout_rows,
+                                dst_layout_cols);
+}
+
+/* Computes dst = (lchd,rchd->lr, lhs, rhs)
+ * Where l = lhs_layout_rows, r = rhs_layout_rows,
+ * c = rhs_layout_cols = lhs_layout_cols.
+ */
+template <int RowsLeft, int RowsRight, int Cols>
+void RunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+               int lhs_layout_rows, int lhs_layout_cols, int rhs_layout_rows,
+               int rhs_layout_cols, int dst_layout_rows, int dst_layout_cols) {
+  ReferenceRunKernel<RowsLeft, RowsRight, Cols>(
+      lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
+      rhs_layout_cols, dst_layout_rows, dst_layout_cols);
+}
+
+// Compute sum of lhs * rhs columnwise and write output to output_ptr.
+inline void RunAndUnpack(int rhs_width, const uint8_t* lhs, const int8_t* rhs,
+                         int32_t* dst, int output_depth, int batch_size,
+                         int lhs_layout_rows, int lhs_layout_cols,
+                         int rhs_layout_rows, int rhs_layout_cols,
+                         int dst_layout_rows, int dst_layout_cols,
+                         float* output_ptr, const float* scaling_factors,
+                         const float* filter_scales) {
+  ReferenceRunKernel<4, 1, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
+                               rhs_layout_rows, rhs_layout_cols,
+                               dst_layout_rows, dst_layout_cols);
+  ReferenceUnpack<4, 1>(output_ptr, dst, batch_size, output_depth,
+                        scaling_factors, filter_scales, dst_layout_rows,
+                        dst_layout_cols);
+}
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_H_

tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference_impl.h

@@ -0,0 +1,62 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_IMPL_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_IMPL_H_
+
+#include <stdint.h>
+
+namespace tflite {
+namespace optimized_4bit {
+
+void ReferencePackInner(const int8_t* src, uint8_t* box, int src_rows,
+                        int src_cols, int outer_row, int outer_col,
+                        int outer_rows, int outer_cols, int inner_rows,
+                        int inner_cols);
+
+void ReferencePrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                      int layout_cols, int src_rows, int src_cols, int width,
+                      int depth);
+
+void ReferenceBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                      int n_data, int8_t* quantized_data_ptr,
+                                      float* scaling_factors, int width,
+                                      int depth, int32_t* input_offsets);
+
+void ReferenceAssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                          const float* batch_scales,
+                                          const float* filter_scales,
+                                          const float* bias_ptr,
+                                          float* output_ptr, int output_depth,
+                                          int batch_size);
+
+template <int Depth, int Width>
+extern void ReferenceUnpack(float* output_ptr, const int32_t* dst,
+                            int batch_size, int num_units,
+                            const float* scaling_factors,
+                            const float* filter_scales, int dst_layout_rows,
+                            int dst_layout_cols);
+
+template <int RowsLeft, int RowsRight, int Cols>
+extern void ReferenceRunKernel(const uint8_t* lhs, const int8_t* rhs,
+                               int32_t* dst, int lhs_layout_rows,
+                               int lhs_layout_cols, int rhs_layout_rows,
+                               int rhs_layout_cols, int dst_layout_rows,
+                               int dst_layout_cols);
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_FULLY_CONNECTED_REFERENCE_IMPL_H_

tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected.cc

@@ -0,0 +1,609 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
+#include <arm_neon.h>
+#include <stdint.h>
+
+#include <algorithm>
+#include <cstring>
+#include <vector>
+
+#include "include/cpuinfo.h"
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_common.h"
+#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
+
+namespace tflite {
+namespace optimized_4bit {
+
+#ifndef __aarch64__
+inline int16x8_t vqmovn_high_s32(int16x4_t a_s16x4, int32x4_t b_s32x4) {
+  return vcombine_s16(a_s16x4, vqmovn_s32(b_s32x4));
+}
+
+inline int8x16_t vqmovn_high_s16(int8x8_t a_s8x8, int16x8_t b_s16x8) {
+  return vcombine_s8(a_s8x8, vqmovn_s16(b_s16x8));
+}
+
+inline int32x4_t vpaddq_s32(int32x4_t a, int32x4_t b) {
+  int32x2_t a0 = vpadd_s32(vget_low_s32(a), vget_high_s32(a));
+  int32x2_t b0 = vpadd_s32(vget_low_s32(b), vget_high_s32(b));
+  return vcombine_s32(a0, b0);
+}
+
+inline float vmaxvq_f32(float32x4_t max_f32x4) {
+  float32x2_t max_f32x2 =
+      vmax_f32(vget_low_f32(max_f32x4), vget_high_f32(max_f32x4));
+  max_f32x2 = vpmax_f32(max_f32x2, max_f32x2);
+  return vget_lane_f32(max_f32x2, 0);
+}
+
+inline int32x4_t vcvtaq_s32_f32(float32x4_t a_f32x4) {
+  float32x4_t half = vdupq_n_f32(.5);
+  float32x4_t sign =
+      vcvtq_f32_u32(vshrq_n_u32(vreinterpretq_u32_f32(a_f32x4), 31));
+  float32x4_t add_half = vaddq_f32(a_f32x4, half);
+  float32x4_t round = vsubq_f32(add_half, sign);
+  return vcvtq_s32_f32(round);
+}
+
+void NeonAssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                     const float* batch_scales,
+                                     float* filter_scales,
+                                     const float* bias_ptr, float* output_ptr,
+                                     int output_depth, int batch_size) {
+  if (bias_ptr) {
+    for (int b = 0; b < batch_size; ++b) {
+      const float* filter_scales_ptr = filter_scales;
+      const float* tmp_bias_ptr = bias_ptr;
+      float val = input_offsets[b] * batch_scales[b];
+      int o = output_depth;
+      const float32x4_t v4_f32x4 = vdupq_n_f32(val);
+      for (; o >= 4; o -= 4) {
+        float32x4_t v0_f32x4 = vld1q_f32(filter_scales_ptr);
+        filter_scales_ptr += 4;
+        float32x4_t v5_f32x4 = vld1q_f32(tmp_bias_ptr);
+        tmp_bias_ptr += 4;
+        v5_f32x4 = vmlaq_f32(v5_f32x4, v0_f32x4, v4_f32x4);
+        vst1q_f32(output_ptr, v5_f32x4);
+        output_ptr += 4;
+      }
+      for (; o > 0; --o) {
+        *output_ptr++ = val * (*filter_scales_ptr++) + (*tmp_bias_ptr++);
+      }
+    }
+    return;
+  }
+  for (int b = 0; b < batch_size; ++b) {
+    const float* filter_scales_ptr = filter_scales;
+    float val = input_offsets[b] * batch_scales[b];
+    int o = output_depth;
+    const float32x4_t v4_f32x4 = vdupq_n_f32(val);
+    for (; o >= 4; o -= 4) {
+      float32x4_t v0_f32x4 = vld1q_f32(filter_scales_ptr);
+      filter_scales_ptr += 4;
+      float32x4_t v13_f32x4 = vmulq_f32(v0_f32x4, v4_f32x4);
+      vst1q_f32(output_ptr, v13_f32x4);
+      output_ptr += 4;
+    }
+    for (; o > 0; --o) {
+      *output_ptr++ = val * (*filter_scales_ptr++);
+    }
+  }
+}
+
+#else
+
+void NeonAssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                     const float* batch_scales,
+                                     float* filter_scales,
+                                     const float* bias_ptr, float* output_ptr,
+                                     int output_depth, int batch_size) {
+  if (bias_ptr) {
+    for (int b = 0; b < batch_size; ++b) {
+      const float* filter_scales_ptr = filter_scales;
+      const float* tmp_bias_ptr = bias_ptr;
+      float val = input_offsets[b] * batch_scales[b];
+      int o = output_depth;
+      const float32x4_t v4_f32x4 = vdupq_n_f32(val);
+      for (; o >= 16; o -= 16) {
+        float32x4x4_t v0_to_v3_f32x4x4 = vld1q_f32_x4(filter_scales_ptr);
+        filter_scales_ptr += 16;
+        float32x4x4_t v5_to_v8_f32x4x4 = vld1q_f32_x4(tmp_bias_ptr);
+        tmp_bias_ptr += 16;
+        v5_to_v8_f32x4x4.val[0] = vfmaq_f32(v5_to_v8_f32x4x4.val[0],
+                                            v0_to_v3_f32x4x4.val[0], v4_f32x4);
+        v5_to_v8_f32x4x4.val[1] = vfmaq_f32(v5_to_v8_f32x4x4.val[1],
+                                            v0_to_v3_f32x4x4.val[1], v4_f32x4);
+        v5_to_v8_f32x4x4.val[2] = vfmaq_f32(v5_to_v8_f32x4x4.val[2],
+                                            v0_to_v3_f32x4x4.val[2], v4_f32x4);
+        v5_to_v8_f32x4x4.val[3] = vfmaq_f32(v5_to_v8_f32x4x4.val[3],
+                                            v0_to_v3_f32x4x4.val[3], v4_f32x4);
+        vst1q_f32_x4(output_ptr, v5_to_v8_f32x4x4);
+        output_ptr += 16;
+      }
+      if (o >= 8) {
+        float32x4x2_t v0_to_v1_f32x4x2 = vld1q_f32_x2(filter_scales_ptr);
+        filter_scales_ptr += 8;
+        float32x4x2_t v5_to_v6_f32x4x2 = vld1q_f32_x2(tmp_bias_ptr);
+        tmp_bias_ptr += 8;
+        v5_to_v6_f32x4x2.val[0] = vfmaq_f32(v5_to_v6_f32x4x2.val[0],
+                                            v0_to_v1_f32x4x2.val[0], v4_f32x4);
+        v5_to_v6_f32x4x2.val[1] = vfmaq_f32(v5_to_v6_f32x4x2.val[1],
+                                            v0_to_v1_f32x4x2.val[1], v4_f32x4);
+        vst1q_f32_x2(output_ptr, v5_to_v6_f32x4x2);
+        output_ptr += 8;
+        o -= 8;
+      }
+      if (o >= 4) {
+        float32x4_t v0_f32x4 = vld1q_f32(filter_scales_ptr);
+        filter_scales_ptr += 4;
+        float32x4_t v5_f32x4 = vld1q_f32(tmp_bias_ptr);
+        tmp_bias_ptr += 4;
+        v5_f32x4 = vfmaq_f32(v5_f32x4, v0_f32x4, v4_f32x4);
+        vst1q_f32(output_ptr, v5_f32x4);
+        output_ptr += 4;
+        o -= 4;
+      }
+      for (; o > 0; --o) {
+        *output_ptr++ = val * (*filter_scales_ptr++) + (*tmp_bias_ptr++);
+      }
+    }
+    return;
+  }
+  for (int b = 0; b < batch_size; ++b) {
+    const float* filter_scales_ptr = filter_scales;
+    float val = input_offsets[b] * batch_scales[b];
+    int o = output_depth;
+    const float32x4_t v4_f32x4 = vdupq_n_f32(val);
+    for (; o >= 16; o -= 16) {
+      float32x4x4_t v0_to_v3_f32x4x4 = vld1q_f32_x4(filter_scales_ptr);
+      filter_scales_ptr += 16;
+      float32x4x4_t v13_to_v16_f32x4x4;
+      v13_to_v16_f32x4x4.val[0] = vmulq_f32(v0_to_v3_f32x4x4.val[0], v4_f32x4);
+      v13_to_v16_f32x4x4.val[1] = vmulq_f32(v0_to_v3_f32x4x4.val[1], v4_f32x4);
+      v13_to_v16_f32x4x4.val[2] = vmulq_f32(v0_to_v3_f32x4x4.val[2], v4_f32x4);
+      v13_to_v16_f32x4x4.val[3] = vmulq_f32(v0_to_v3_f32x4x4.val[3], v4_f32x4);
+      vst1q_f32_x4(output_ptr, v13_to_v16_f32x4x4);
+      output_ptr += 16;
+    }
+    if (o >= 8) {
+      float32x4x2_t v0_to_v1_f32x4x2 = vld1q_f32_x2(filter_scales_ptr);
+      filter_scales_ptr += 8;
+      float32x4x2_t v11_to_v12_f32x4x2;
+      v11_to_v12_f32x4x2.val[0] = vmulq_f32(v0_to_v1_f32x4x2.val[0], v4_f32x4);
+      v11_to_v12_f32x4x2.val[1] = vmulq_f32(v0_to_v1_f32x4x2.val[1], v4_f32x4);
+      vst1q_f32_x2(output_ptr, v11_to_v12_f32x4x2);
+      output_ptr += 8;
+      o -= 8;
+    }
+    if (o >= 4) {
+      float32x4_t v0_f32x4 = vld1q_f32(filter_scales_ptr);
+      filter_scales_ptr += 4;
+      float32x4_t v13_f32x4 = vmulq_f32(v0_f32x4, v4_f32x4);
+      vst1q_f32(output_ptr, v13_f32x4);
+      output_ptr += 4;
+      o -= 4;
+    }
+    for (; o > 0; --o) {
+      *output_ptr++ = val * (*filter_scales_ptr++);
+    }
+  }
+}
+#endif
+
+void NeonPackInner(const int8_t* src, uint8_t* box, int src_rows, int src_cols,
+                   int outer_row, int outer_col, int outer_rows, int outer_cols,
+                   int inner_rows, int inner_cols) {
+  // create a kernel-specific layout and store it into packed
+  const int width = inner_rows;
+  const int depth = inner_cols;
+  const int real_depth = depth / 2;
+  const int real_src_cols = src_cols / 2;
+  // which virtual row
+  const int row = outer_row * inner_rows;
+  const int col = outer_col * inner_cols;
+  int src_width = std::min(width, src_rows - row);
+  int src_depth = std::min(depth, src_cols - col);
+  int real_col = col / 2;
+  const int8_t* src_data = src + row * real_src_cols + real_col;
+  int real_src_depth = src_depth / 2;
+  const int8x16_t seven = vdupq_n_s8(7);
+  const int8x8_t seven8 = vdup_n_s8(7);
+  for (int m = 0; m < src_width; ++m) {
+    int i = 0;
+    int k = 0;
+    for (; i < (real_src_depth & (~15)); i += 16) {
+      int8x16_t values_16x8 = vld1q_s8(src_data + i);
+      int8x16_t uv1 = vshrq_n_s8(values_16x8, 4);
+      int8x16_t lv1 = vshlq_n_s8(values_16x8, 4);
+      uv1 = vaddq_s8(uv1, seven);
+      lv1 = vshrq_n_s8(lv1, 4);
+      lv1 = vaddq_s8(lv1, seven);
+      int8x8_t iuvl = vget_low_s8(uv1);
+      int8x8_t iuvh = vget_high_s8(uv1);
+      int8x8_t ilvl = vget_low_s8(lv1);
+      int8x8_t ilvh = vget_high_s8(lv1);
+      uint8x8_t uvl = vshl_n_u8(vreinterpret_u8_s8(iuvl), 4);
+      uint8x8_t lvl = vshl_n_u8(vreinterpret_u8_s8(ilvl), 4);
+      uint8x8_t uv = vorr_u8(uvl, vreinterpret_u8_s8(iuvh));
+      uint8x8_t lv = vorr_u8(lvl, vreinterpret_u8_s8(ilvh));
+      uint8x8x2_t zipped = vzip_u8(lv, uv);
+      uint8x16_t combined = vcombine_u8(zipped.val[0], zipped.val[1]);
+      vst1q_u8(box + k, combined);
+      k += 16;
+    }
+    // If exactly 16 values remaining, use fast path
+    if (real_src_depth == (real_src_depth & (~7))) {
+      for (; i < (real_src_depth & (~7)); i += 8) {
+        int8x8_t values_8x8 = vld1_s8(src_data + i);
+        int8x8_t uv1 = vshr_n_s8(values_8x8, 4);
+        int8x8_t lv1 = vshl_n_s8(values_8x8, 4);
+        uv1 = vadd_s8(uv1, seven8);
+        lv1 = vshr_n_s8(lv1, 4);
+        lv1 = vadd_s8(lv1, seven8);
+        uint8x8_t uvl = vshl_n_u8(vreinterpret_u8_s8(uv1), 4);
+        uint8x8_t lvl = vshl_n_u8(vreinterpret_u8_s8(lv1), 4);
+        uint8x8x2_t zipped = vzip_u8(lvl, uvl);
+        uint8x16_t combined = vcombine_u8(zipped.val[0], zipped.val[1]);
+        vst1q_u8(box + k, combined);
+        k += 16;
+      }
+    }
+    // Handle remaining values -- if greater than 16 values,
+    // shuffle.
+    if (i < real_src_depth) {
+      int remaining = 8;
+      remaining =
+          remaining < (real_src_depth - i) ? remaining : real_src_depth - i;
+      for (int j = 0; j < remaining; ++j) {
+        const int8_t v1 = (int8_t)src_data[i + j];
+        int8_t uv1 = upper(v1);
+        int8_t lv1 = lower(v1);
+        int8_t uv2 = 0;
+        int8_t lv2 = 0;
+        if ((i + j + 8) < real_src_depth) {
+          const int8_t v2 = (int8_t)src_data[i + j + 8];
+          uv2 = upper(v2);
+          lv2 = lower(v2);
+        }
+        box[k] = merge(lv1, lv2);
+        box[k + 1] = merge(uv1, uv2);
+        k += 2;
+      }
+    }
+    box += real_depth;
+    src_data += real_src_cols;
+  }
+}
+
+void NeonPrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                 int layout_cols, int src_rows, int src_cols, int width,
+                 int depth) {
+  // depth is always cols
+  size_t size = layout_rows * layout_cols / 2;
+  memset(dest, static_cast<uint8_t>(119), sizeof(uint8_t) * size);
+  // basically, we need to make a new 4D matrix
+  // [rows / width, cols / depth, width, depth] in depth-first
+  int outer_cols = layout_cols / depth;
+  int outer_rows = layout_rows / width;
+  int inner_cols = depth;
+  int inner_rows = width;
+  for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
+    for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+      const int cluster_index = outer_row * outer_cols + outer_col;
+      const int real_depth = inner_cols / 2;
+      uint8_t* box = dest + cluster_index * real_depth * inner_rows;
+      NeonPackInner(tensor, box, src_rows, src_cols, outer_row, outer_col,
+                    outer_rows, outer_cols, inner_rows, inner_cols);
+    }
+  }
+}
+
+void NeonBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                 int n_data, int8_t* quantized_data_ptr,
+                                 float* scaling_factors, int width, int depth,
+                                 int32_t* input_offsets) {
+  const int rows = n_batch;
+  const int cols = n_data;
+  // depth is alpways cols
+  const int layout_rows = (rows + (width - 1)) & ~(width - 1);
+  const int layout_cols = (cols + (depth - 1)) & ~(depth - 1);
+  const int size = layout_rows * layout_cols;
+  int8_t* data = quantized_data_ptr;
+  memset(data, 0, sizeof(int8_t) * size);
+  memset(input_offsets, 0, sizeof(int32_t) * layout_rows);
+  const float* tensor_data = float_data_ptr;
+  // basically, we need to make a new 4D matrix
+  // [rows / width, cols / depth, width, depth] in depth-first
+  const int outer_cols = layout_cols / depth;
+  const int outer_rows = layout_rows / width;
+  float* scaling_factors_ptr = scaling_factors;
+  for (int outer_row = 0; outer_row < outer_rows; outer_row++) {
+    std::vector<float> scale(width);
+    const int row = width * outer_row;
+    scaling_factors_ptr = scaling_factors + row;
+    for (int w = 0; w < width; ++w) {
+      if ((row + w) >= rows) {
+        continue;
+      }
+      int c = 0;
+      const float* start = tensor_data + (row + w) * cols;
+      float32x4_t v1_f32x4 = vdupq_n_f32(0);
+      for (; c < (cols & ~3); c += 4) {
+        float32x4_t v0_f32x4 = vld1q_f32(start);
+        v0_f32x4 = vabsq_f32(v0_f32x4);
+        start += 4;
+        v1_f32x4 = vmaxq_f32(v0_f32x4, v1_f32x4);
+      }
+      float scale_denom = vmaxvq_f32(v1_f32x4);
+      for (; c < cols; ++c) {
+        scale_denom = std::max(scale_denom, std::abs(*(start++)));
+      }
+      if (scale_denom == 0) {
+        scale_denom = 127.0;
+      }
+      scale[w] = 127.0 / scale_denom;
+      scaling_factors_ptr[w] = scale_denom / 127.0;
+    }
+    for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+      const int col = depth * outer_col;
+      const int src_width = std::min(width, rows - row);
+      const int src_depth = std::min(depth, cols - col);
+      const int cluster_index = outer_row * outer_cols + outer_col;
+      int8_t* box = data + cluster_index * depth * width;
+      __builtin_prefetch(box, 1, 3);
+      for (int w = 0; w < src_width; ++w) {
+        const float scale_w = scale[w];
+        const float* float_data = tensor_data + (row + w) * cols + col;
+        __builtin_prefetch(float_data, 0, 3);
+        int32_t* input_offsets_ptr = input_offsets + (row + w);
+        __builtin_prefetch(input_offsets_ptr, 1, 3);
+        int8_t* x0 = box + w * depth;
+        const float* x1 = float_data;
+        int16x8_t v12_s16x8 = vdupq_n_s16(0);
+        int32x4_t v13_s32x4 = vdupq_n_s32(0);
+        size_t run_depth = 0;
+        float32x4_t v0_f32x4 = vdupq_n_f32(scale_w);
+#ifdef __aarch64__
+        for (; run_depth < (src_depth & ~15); run_depth += 16) {
+          const float32x4x4_t v1_f32x4x4 = vld1q_f32_x4(x1);
+          x1 += 16;
+          const float32x4_t v5_f32x4 = vmulq_f32(v1_f32x4x4.val[0], v0_f32x4);
+          const float32x4_t v6_f32x4 = vmulq_f32(v1_f32x4x4.val[1], v0_f32x4);
+          const float32x4_t v7_f32x4 = vmulq_f32(v1_f32x4x4.val[2], v0_f32x4);
+          const float32x4_t v8_f32x4 = vmulq_f32(v1_f32x4x4.val[3], v0_f32x4);
+          const int32x4_t v5_s32x4 = vcvtaq_s32_f32(v5_f32x4);
+          const int32x4_t v6_s32x4 = vcvtaq_s32_f32(v6_f32x4);
+          const int32x4_t v7_s32x4 = vcvtaq_s32_f32(v7_f32x4);
+          const int32x4_t v8_s32x4 = vcvtaq_s32_f32(v8_f32x4);
+          const int16x4_t v9_low_s16x4 = vqmovn_s32(v5_s32x4);
+          const int16x8_t v9_s16x8 = vqmovn_high_s32(v9_low_s16x4, v6_s32x4);
+          const int16x4_t v10_low_s16x4 = vqmovn_s32(v7_s32x4);
+          const int16x8_t v10_s16x8 = vqmovn_high_s32(v10_low_s16x4, v8_s32x4);
+          const int8x8_t v11_low_s8x8 = vqmovn_s16(v9_s16x8);
+          const int8x16_t v11_s8x16 = vqmovn_high_s16(v11_low_s8x8, v10_s16x8);
+          v12_s16x8 = vaddq_s16(v12_s16x8, v9_s16x8);
+          v12_s16x8 = vaddq_s16(v12_s16x8, v10_s16x8);
+          vst1q_s8(x0, v11_s8x16);
+          x0 += 16;
+        }
+        for (; run_depth < (src_depth & ~7); run_depth += 8) {
+          const float32x4x2_t v1_f32x4x2 = vld1q_f32_x2(x1);
+          x1 += 8;
+          const float32x4_t v5_f32x4 = vmulq_f32(v1_f32x4x2.val[0], v0_f32x4);
+          const float32x4_t v6_f32x4 = vmulq_f32(v1_f32x4x2.val[1], v0_f32x4);
+          const int32x4_t v5_s32x4 = vcvtaq_s32_f32(v5_f32x4);
+          const int32x4_t v6_s32x4 = vcvtaq_s32_f32(v6_f32x4);
+          const int16x4_t v9_low_s16x4 = vqmovn_s32(v5_s32x4);
+          const int16x8_t v9_s16x8 = vqmovn_high_s32(v9_low_s16x4, v6_s32x4);
+          const int8x8_t v11_low_s8x8 = vqmovn_s16(v9_s16x8);
+          vst1_s8(x0, v11_low_s8x8);
+          x0 += 8;
+          v12_s16x8 = vaddq_s16(v12_s16x8, v9_s16x8);
+        }
+#else
+        for (; run_depth < (src_depth & ~7); run_depth += 8) {
+          const float32x4_t v1_f32x4_0 = vld1q_f32(x1);
+          x1 += 4;
+          const float32x4_t v1_f32x4_1 = vld1q_f32(x1);
+          x1 += 4;
+          const float32x4_t v5_f32x4 = vmulq_f32(v1_f32x4_0, v0_f32x4);
+          const float32x4_t v6_f32x4 = vmulq_f32(v1_f32x4_1, v0_f32x4);
+          const int32x4_t v5_s32x4 = vcvtaq_s32_f32(v5_f32x4);
+          const int32x4_t v6_s32x4 = vcvtaq_s32_f32(v6_f32x4);
+          const int16x4_t v9_low_s16x4 = vqmovn_s32(v5_s32x4);
+          const int16x8_t v9_s16x8 = vqmovn_high_s32(v9_low_s16x4, v6_s32x4);
+          const int8x8_t v11_low_s8x8 = vqmovn_s16(v9_s16x8);
+          vst1_s8(x0, v11_low_s8x8);
+          x0 += 8;
+          v12_s16x8 = vaddq_s16(v12_s16x8, v9_s16x8);
+        }
+#endif
+        int32_t row_sum = 0;
+        if (run_depth > 0) {
+          v13_s32x4 = vpadalq_s16(v13_s32x4, v12_s16x8);
+          v13_s32x4 = vpaddq_s32(v13_s32x4, v13_s32x4);
+          v13_s32x4 = vpaddq_s32(v13_s32x4, v13_s32x4);
+          row_sum += vgetq_lane_s32(v13_s32x4, 0);
+        }
+        for (; run_depth < src_depth; run_depth++) {
+          const float f = *x1++;
+          const int8_t q =
+              static_cast<int8_t>(::tflite::TfLiteRound(f * scale_w));
+          *x0++ = q;
+          row_sum += q;
+        }
+        input_offsets[row + w] += row_sum;
+      }
+    }
+  }
+  for (int r = 0; r < layout_rows; ++r) {
+    input_offsets[r] = input_offsets[r] * zero_point_4bit;
+  }
+}
+
+void NeonAssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                     const float* batch_scales,
+                                     float* filter_scales,
+                                     const float* bias_ptr, float* output_ptr,
+                                     int output_depth, int batch_size);
+
+template <int Depth, int Width>
+void NeonUnpack(float* output_ptr, const int32_t* dst, int batch_size,
+                int num_units, const float* scaling_factors,
+                const float* filter_scales, int dst_layout_rows,
+                int dst_layout_cols) {
+  if (Width == 1) {
+    const int outer_rows = dst_layout_rows / Width;
+    const int outer_cols = dst_layout_cols / Depth;
+    const int32_t* dst_ptr = dst;
+    int unit = 0;
+    for (int outer_col = 0; outer_col < outer_cols;
+         ++outer_col, unit += Depth) {
+      float* tmp_output_ptr = output_ptr + unit;
+      int len = num_units - unit < Depth ? num_units - unit : Depth;
+      int cond = len & ~3;
+      const float* scaling_factors_ptr = scaling_factors;
+      for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
+        const float scale = *scaling_factors_ptr;
+        const float* filter_scales_ptr = filter_scales + unit;
+        int i = 0;
+        for (; i < cond; i += 4) {
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+        }
+        for (; i < len; ++i) {
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+        }
+        dst_ptr += (Depth - len);
+        scaling_factors_ptr += Width;
+        tmp_output_ptr += (num_units - len);
+      }
+    }
+    return;
+  }
+  const int outer_rows = dst_layout_rows / Width;
+  const int outer_cols = dst_layout_cols / Depth;
+  for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+    const int unit = outer_col * Depth;
+    const int remaining_units = std::min(num_units - unit, Depth);
+    const int depth_offset = Depth - remaining_units;
+    const int width_offset = num_units - remaining_units;
+    int outer_row = 0;
+    for (; outer_row < outer_rows; ++outer_row) {
+      const int batch = outer_row * Width;
+      const int remaining_width = std::min(batch_size - batch, Width);
+      const int cluster_index = outer_col * outer_rows + outer_row;
+      const int32_t* dst_ptr = dst + cluster_index * Depth * Width;
+      float* tmp_output_ptr = output_ptr + batch * num_units + unit;
+      const float* scale = scaling_factors + batch;
+      int w = remaining_width;
+      for (; w > 0; --w, scale++) {
+        int d = remaining_units;
+        const float* filter_scales_ptr = filter_scales + unit;
+        float32x4_t v3_f32x4 = vld1q_dup_f32(scale);
+        for (; d > 3; d -= 4) {
+          int32x4_t v0_s32x4 = vld1q_s32(dst_ptr);
+          dst_ptr += 4;
+          float32x4_t v1_f32x4 = vcvtq_f32_s32(v0_s32x4);
+          float32x4_t v2_f32x4 = vld1q_f32(filter_scales_ptr);
+          filter_scales_ptr += 4;
+          float32x4_t v5_f32x4 = vmulq_f32(v1_f32x4, v3_f32x4);
+          float32x4_t v6_f32x4 = vmulq_f32(v5_f32x4, v2_f32x4);
+          float32x4_t v7_f32x4 = vld1q_f32(tmp_output_ptr);
+          float32x4_t v8_f32x4 = vaddq_f32(v6_f32x4, v7_f32x4);
+          vst1q_f32(tmp_output_ptr, v8_f32x4);
+          tmp_output_ptr += 4;
+        }
+        for (; d > 0; --d) {
+          *tmp_output_ptr++ += *dst_ptr++ * (*scale) * (*filter_scales_ptr++);
+        }
+        dst_ptr += depth_offset;
+        tmp_output_ptr += width_offset;
+      }
+    }
+  }
+}
+
+inline bool HasSDot() {
+  // CPUInfo already guards against double init
+  if (!cpuinfo_initialize()) {
+    // If we failed to init CPUInfo, assume ARM v8.2a-dotprod is not supported.
+    return false;
+  };
+  return cpuinfo_has_arm_neon_dot();
+}
+
+template <int RowsLeft, int RowsRight, int Cols>
+void NeonRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                   int lhs_layout_rows, int lhs_layout_cols,
+                   int rhs_layout_rows, int rhs_layout_cols,
+                   int dst_layout_rows, int dst_layout_cols) {
+  if (HasSDot()) {
+    NeonRunKernelSDot<RowsLeft, RowsRight, Cols>(
+        lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
+        rhs_layout_cols, dst_layout_rows, dst_layout_cols);
+    return;
+  }
+  NeonRunKernelNoSDot<RowsLeft, RowsRight, Cols>(
+      lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
+      rhs_layout_cols, dst_layout_rows, dst_layout_cols);
+}
+
+template void NeonUnpack<4, 1>(float* output_ptr, const int32_t* dst,
+                               int batch_size, int num_units,
+                               const float* scaling_factors,
+                               const float* filter_scales, int dst_layout_rows,
+                               int dst_layout_cols);
+
+template void NeonRunKernel<4, 1, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                      int32_t* dst, int lhs_layout_rows,
+                                      int lhs_layout_cols, int rhs_layout_rows,
+                                      int rhs_layout_cols, int dst_layout_rows,
+                                      int dst_layout_cols);
+
+#ifdef __aarch64__
+template void NeonUnpack<4, 2>(float* output_ptr, const int32_t* dst,
+                               int batch_size, int num_units,
+                               const float* scaling_factors,
+                               const float* filter_scales, int dst_layout_rows,
+                               int dst_layout_cols);
+
+template void NeonRunKernel<4, 2, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                      int32_t* dst, int lhs_layout_rows,
+                                      int lhs_layout_cols, int rhs_layout_rows,
+                                      int rhs_layout_cols, int dst_layout_rows,
+                                      int dst_layout_cols);
+
+template void NeonUnpack<4, 4>(float* output_ptr, const int32_t* dst,
+                               int batch_size, int num_units,
+                               const float* scaling_factors,
+                               const float* filter_scales, int dst_layout_rows,
+                               int dst_layout_cols);
+
+template void NeonRunKernel<4, 4, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                      int32_t* dst, int lhs_layout_rows,
+                                      int lhs_layout_cols, int rhs_layout_rows,
+                                      int rhs_layout_cols, int dst_layout_rows,
+                                      int dst_layout_cols);
+#endif
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+#endif  // defined(FC_4BIT_NEON)...

tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected.h

@@ -0,0 +1,130 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_H_
+#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
+#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
+
+namespace tflite {
+namespace optimized_4bit {
+
+// Maximum RowsRight compiled RunKernel implementations.
+inline int GetMaxSupportedRows() {
+#ifdef __aarch64__
+  return 4;
+#else
+  return 1;
+#endif
+}
+
+// Pack a 4bit inner_rows x inner_cols array from src.
+inline void PackInner(const int8_t* src, uint8_t* box, int src_rows,
+                      int src_cols, int outer_row, int outer_col,
+                      int outer_rows, int outer_cols, int inner_rows,
+                      int inner_cols) {
+  NeonPackInner(src, box, src_rows, src_cols, outer_row, outer_col, outer_rows,
+                outer_cols, inner_rows, inner_cols);
+}
+
+// Prepack lhs matrix into dest.
+inline void Prepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                    int layout_cols, int src_rows, int src_cols, int width,
+                    int depth) {
+  NeonPrepack(dest, tensor, layout_rows, layout_cols, src_rows, src_cols, width,
+              depth);
+}
+
+// Quantize input floats to 8bit and calculate sum of each column.
+inline void BatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                    int n_data, int8_t* quantized_data_ptr,
+                                    float* scaling_factors, int width,
+                                    int depth, int32_t* input_offsets) {
+  NeonBatchQuantizeFloats4Bit(float_data_ptr, n_batch, n_data,
+                              quantized_data_ptr, scaling_factors, width, depth,
+                              input_offsets);
+}
+
+// Write bias + input offset * filter_scale to output_ptr.
+inline void AssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                        const float* batch_scales,
+                                        float* filter_scales,
+                                        const float* bias_ptr,
+                                        float* output_ptr, int output_depth,
+                                        int batch_size) {
+  NeonAssignBiasAndComputeOffsets(input_offsets, batch_scales, filter_scales,
+                                  bias_ptr, output_ptr, output_depth,
+                                  batch_size);
+}
+
+// Add accumulated integer sums in dst to float output.
+template <int Depth, int Width>
+void Unpack(float* output_ptr, const int32_t* dst, int batch_size,
+            int num_units, const float* scaling_factors,
+            const float* filter_scales, int dst_layout_rows,
+            int dst_layout_cols) {
+  NeonUnpack<Depth, Width>(output_ptr, dst, batch_size, num_units,
+                           scaling_factors, filter_scales, dst_layout_rows,
+                           dst_layout_cols);
+}
+
+// Compute sum of lhs * rhs columnwise.
+template <int RowsLeft, int RowsRight, int Cols>
+void RunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+               int lhs_layout_rows, int lhs_layout_cols, int rhs_layout_rows,
+               int rhs_layout_cols, int dst_layout_rows, int dst_layout_cols) {
+  NeonRunKernel<RowsLeft, RowsRight, Cols>(
+      lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
+      rhs_layout_cols, dst_layout_rows, dst_layout_cols);
+}
+
+// Compute sum of lhs * rhs columnwise and write output to output_ptr.
+inline void RunAndUnpack(int rhs_width, const uint8_t* lhs, const int8_t* rhs,
+                         int32_t* dst, int output_depth, int batch_size,
+                         int lhs_layout_rows, int lhs_layout_cols,
+                         int rhs_layout_rows, int rhs_layout_cols,
+                         int dst_layout_rows, int dst_layout_cols,
+                         float* output_ptr, const float* scaling_factors,
+                         const float* filter_scales) {
+#ifdef __aarch64__
+  if (rhs_width >= 4) {
+    NeonRunKernel<4, 4, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
+                            rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
+                            dst_layout_cols);
+    NeonUnpack<4, 4>(output_ptr, dst, batch_size, output_depth, scaling_factors,
+                     filter_scales, dst_layout_rows, dst_layout_cols);
+    return;
+  }
+  if (rhs_width >= 2) {
+    NeonRunKernel<4, 2, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
+                            rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
+                            dst_layout_cols);
+    NeonUnpack<4, 2>(output_ptr, dst, batch_size, output_depth, scaling_factors,
+                     filter_scales, dst_layout_rows, dst_layout_cols);
+    return;
+  }
+#endif
+  NeonRunKernel<4, 1, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
+                          rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
+                          dst_layout_cols);
+  NeonUnpack<4, 1>(output_ptr, dst, batch_size, output_depth, scaling_factors,
+                   filter_scales, dst_layout_rows, dst_layout_cols);
+}
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // defined(FC_4BIT_NEON)...
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_H_

tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_aarch64_nosdot.cc

@@ -0,0 +1,607 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
+#include <arm_neon.h>
+#include <stdint.h>
+
+#include <algorithm>
+#include <vector>
+
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
+
+namespace tflite {
+namespace optimized_4bit {
+
+template <int RowsLeft, int RowsRight, int Cols>
+void NeonRunKernelNoSDot(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                         int lhs_layout_rows, int lhs_layout_cols,
+                         int rhs_layout_rows, int rhs_layout_cols,
+                         int dst_layout_rows, int dst_layout_cols) {}
+
+template <>
+void NeonRunKernelNoSDot<4, 1, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                   int32_t* dst, int lhs_layout_rows,
+                                   int lhs_layout_cols, int rhs_layout_rows,
+                                   int rhs_layout_cols, int dst_layout_rows,
+                                   int dst_layout_cols) {
+  const int rows_left = 4;
+  const int rows_right = 1;
+  const int cols = 32;
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
+  const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
+  const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    int left_index = i * rows_left * lhs_layout_cols / 2;
+    const uint8_t* lhs_ptr_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_ptr = lhs_ptr_data;
+      int right_index = j * rows_right * rhs_layout_cols;
+      const int8_t* rhs_ptr = rhs + right_index;
+      int run_depth = depth;
+      asm(R"asm(
+          movi v24.16b, #15
+          ld1 {v4.16b}, [%[lhs_ptr]], #16
+          movi v16.4s, #0
+          movi v17.4s, #0
+          ld1 {v5.16b}, [%[lhs_ptr]], #16
+          movi v18.4s, #0
+          movi v19.4s, #0
+          ld1 {v6.16b}, [%[lhs_ptr]], #16
+          and v8.16b, v4.16b, v24.16b
+          and v9.16b, v5.16b, v24.16b
+          ld1 {v7.16b}, [%[lhs_ptr]], #16
+          ushr v12.16b, v4.16b, #4
+          ushr v13.16b, v5.16b, #4
+          ld1 {v0.16b}, [%[rhs_ptr]], #16
+          and v10.16b, v6.16b, v24.16b
+          and v11.16b, v7.16b, v24.16b
+          ld1 {v1.16b}, [%[rhs_ptr]], #16
+          ushr v14.16b, v6.16b, #4
+          ushr v15.16b, v7.16b, #4
+          subs %w[run_depth], %w[run_depth], #1
+          b.ls 1f /* skip loop */
+            0: /* loop start */
+            ld1 {v4.16b}, [%[lhs_ptr]], #16
+            smull v20.8h, v12.8b, v0.8b
+            smull v21.8h, v13.8b, v0.8b
+            smull v22.8h, v14.8b, v0.8b
+            ld1 {v5.16b}, [%[lhs_ptr]], #16
+            smull v23.8h, v15.8b, v0.8b
+            smlal v20.8h, v8.8b, v1.8b
+            smlal v21.8h, v9.8b, v1.8b
+            ld1 {v6.16b}, [%[lhs_ptr]], #16
+            smlal v22.8h, v10.8b, v1.8b
+            smlal v23.8h, v11.8b, v1.8b
+            smlal2 v20.8h, v12.16b, v0.16b
+            ld1 {v7.16b}, [%[lhs_ptr]], #16
+            smlal2 v21.8h, v13.16b, v0.16b
+            smlal2 v22.8h, v14.16b, v0.16b
+            smlal2 v23.8h, v15.16b, v0.16b
+            smlal2 v20.8h, v8.16b, v1.16b
+            smlal2 v21.8h, v9.16b, v1.16b
+            smlal2 v22.8h, v10.16b, v1.16b
+            ld1 {v0.16b}, [%[rhs_ptr]], #16
+            smlal2 v23.8h, v11.16b, v1.16b
+            sadalp v16.4s, v20.8h
+            sadalp v17.4s, v21.8h
+            sadalp v18.4s, v22.8h
+            sadalp v19.4s, v23.8h
+            ld1 {v1.16b}, [%[rhs_ptr]], #16
+            and v8.16b, v4.16b, v24.16b
+            and v9.16b, v5.16b, v24.16b
+            ushr v12.16b, v4.16b, #4
+            ushr v13.16b, v5.16b, #4
+            and v10.16b, v6.16b, v24.16b
+            and v11.16b, v7.16b, v24.16b
+            ushr v14.16b, v6.16b, #4
+            ushr v15.16b, v7.16b, #4
+            subs %w[run_depth], %w[run_depth], #1
+            b.hi 0b /* loop branch */
+          1: /* loop end */
+          smull v20.8h, v12.8b, v0.8b
+          smull v21.8h, v13.8b, v0.8b
+          smull v22.8h, v14.8b, v0.8b
+          smull v23.8h, v15.8b, v0.8b
+          smlal v20.8h, v8.8b, v1.8b
+          smlal v21.8h, v9.8b, v1.8b
+          smlal v22.8h, v10.8b, v1.8b
+          smlal v23.8h, v11.8b, v1.8b
+          smlal2 v20.8h, v12.16b, v0.16b
+          smlal2 v21.8h, v13.16b, v0.16b
+          smlal2 v22.8h, v14.16b, v0.16b
+          smlal2 v23.8h, v15.16b, v0.16b
+          smlal2 v20.8h, v8.16b, v1.16b
+          smlal2 v21.8h, v9.16b, v1.16b
+          smlal2 v22.8h, v10.16b, v1.16b
+          smlal2 v23.8h, v11.16b, v1.16b
+          sadalp v16.4s, v20.8h
+          sadalp v17.4s, v21.8h
+          sadalp v18.4s, v22.8h
+          sadalp v19.4s, v23.8h
+          addp v4.4s, v16.4s, v17.4s
+          addp v5.4s, v18.4s, v19.4s
+          addp v6.4s, v4.4s, v5.4s
+          st1 {v6.4s}, [%[dst]], #16
+          )asm"
+          : [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
+            [run_depth] "+r"(run_depth)
+          :
+          : "cc", "memory", "v0", "v1", "v4", "v5", "v6", "v7", "v8", "v9",
+            "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18",
+            "v19", "v20", "v21", "v22", "v23", "v24");
+    }
+  }
+}
+
+template <>
+void NeonRunKernelNoSDot<4, 2, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                   int32_t* dst, int lhs_layout_rows,
+                                   int lhs_layout_cols, int rhs_layout_rows,
+                                   int rhs_layout_cols, int dst_layout_rows,
+                                   int dst_layout_cols) {
+  const int rows_left = 4;
+  const int rows_right = 2;
+  const int cols = 32;
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
+  const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
+  const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    int left_index = i * rows_left * lhs_layout_cols / 2;
+    const uint8_t* lhs_ptr_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_ptr = lhs_ptr_data;
+      int right_index = j * rows_right * rhs_layout_cols;
+      const int8_t* rhs_ptr = rhs + right_index;
+      int run_depth = depth;
+      asm(R"asm(
+          ld1 {v4.16b}, [%[lhs_ptr]], #16
+          movi v31.16b, #15
+          movi v16.4s, #0
+          movi v17.4s, #0
+          ld1 {v5.16b}, [%[lhs_ptr]], #16
+          movi v18.4s, #0
+          movi v19.4s, #0
+          ld1 {v6.16b}, [%[lhs_ptr]], #16
+          and v8.16b, v4.16b, v31.16b
+          and v9.16b, v5.16b, v31.16b
+          ld1 {v7.16b}, [%[lhs_ptr]], #16
+          ushr v12.16b, v4.16b, #4
+          ushr v13.16b, v5.16b, #4
+          movi v24.4s, #0
+          ld1 {v0.16b}, [%[rhs_ptr]], #16
+          movi v25.4s, #0
+          movi v26.4s, #0
+          ld1 {v1.16b}, [%[rhs_ptr]], #16
+          movi v27.4s, #0
+          and v10.16b, v6.16b, v31.16b
+          ld1 {v2.16b}, [%[rhs_ptr]], #16
+          and v11.16b, v7.16b, v31.16b
+          ushr v14.16b, v6.16b, #4
+          ld1 {v3.16b}, [%[rhs_ptr]], #16
+          ushr v15.16b, v7.16b, #4
+          subs %w[run_depth], %w[run_depth], #1
+          b.ls 1f /* skip loop */
+            0: /* loop start */
+            smull v20.8h, v12.8b, v0.8b
+            smull v21.8h, v13.8b, v0.8b
+            smull v22.8h, v14.8b, v0.8b
+            ld1 {v4.16b}, [%[lhs_ptr]], #16
+            smull v23.8h, v15.8b, v0.8b
+            smlal v20.8h, v8.8b, v1.8b
+            smlal v21.8h, v9.8b, v1.8b
+            ld1 {v5.16b}, [%[lhs_ptr]], #16
+            smlal v22.8h, v10.8b, v1.8b
+            smlal v23.8h, v11.8b, v1.8b
+            smlal2 v20.8h, v12.16b, v0.16b
+            ld1 {v6.16b}, [%[lhs_ptr]], #16
+            smlal2 v21.8h, v13.16b, v0.16b
+            smlal2 v22.8h, v14.16b, v0.16b
+            smlal2 v23.8h, v15.16b, v0.16b
+            ld1 {v7.16b}, [%[lhs_ptr]], #16
+            smlal2 v20.8h, v8.16b, v1.16b
+            smlal2 v21.8h, v9.16b, v1.16b
+            smlal2 v22.8h, v10.16b, v1.16b
+            smlal2 v23.8h, v11.16b, v1.16b
+            ld1 {v0.16b}, [%[rhs_ptr]], #16
+            sadalp v16.4s, v20.8h
+            sadalp v17.4s, v21.8h
+            sadalp v18.4s, v22.8h
+            sadalp v19.4s, v23.8h
+            ld1 {v1.16b}, [%[rhs_ptr]], #16
+            smull v28.8h, v12.8b, v2.8b
+            smull v29.8h, v13.8b, v2.8b
+            smull v30.8h, v14.8b, v2.8b
+            smull v20.8h, v15.8b, v2.8b
+            smlal v28.8h, v8.8b, v3.8b
+            smlal v29.8h, v9.8b, v3.8b
+            smlal v30.8h, v10.8b, v3.8b
+            smlal v20.8h, v11.8b, v3.8b
+            smlal2 v28.8h, v12.16b, v2.16b
+            smlal2 v29.8h, v13.16b, v2.16b
+            smlal2 v30.8h, v14.16b, v2.16b
+            smlal2 v20.8h, v15.16b, v2.16b
+            smlal2 v28.8h, v8.16b, v3.16b
+            smlal2 v29.8h, v9.16b, v3.16b
+            smlal2 v30.8h, v10.16b, v3.16b
+            smlal2 v20.8h, v11.16b, v3.16b
+            ld1 {v2.16b}, [%[rhs_ptr]], #16
+            sadalp v24.4s, v28.8h
+            sadalp v25.4s, v29.8h
+            sadalp v26.4s, v30.8h
+            sadalp v27.4s, v20.8h
+            ld1 {v3.16b}, [%[rhs_ptr]], #16
+            and v8.16b, v4.16b, v31.16b
+            and v9.16b, v5.16b, v31.16b
+            ushr v12.16b, v4.16b, #4
+            ushr v13.16b, v5.16b, #4
+            and v10.16b, v6.16b, v31.16b
+            and v11.16b, v7.16b, v31.16b
+            ushr v14.16b, v6.16b, #4
+            ushr v15.16b, v7.16b, #4
+            subs %w[run_depth], %w[run_depth], #1
+            b.hi 0b /* loop branch */
+          1: /* loop end */
+          smull v20.8h, v12.8b, v0.8b
+          smull v21.8h, v13.8b, v0.8b
+          smull v22.8h, v14.8b, v0.8b
+          smull v23.8h, v15.8b, v0.8b
+          smlal v20.8h, v8.8b, v1.8b
+          smlal v21.8h, v9.8b, v1.8b
+          smlal v22.8h, v10.8b, v1.8b
+          smlal v23.8h, v11.8b, v1.8b
+          smlal2 v20.8h, v12.16b, v0.16b
+          smlal2 v21.8h, v13.16b, v0.16b
+          smlal2 v22.8h, v14.16b, v0.16b
+          smlal2 v23.8h, v15.16b, v0.16b
+          smlal2 v20.8h, v8.16b, v1.16b
+          smlal2 v21.8h, v9.16b, v1.16b
+          smlal2 v22.8h, v10.16b, v1.16b
+          smlal2 v23.8h, v11.16b, v1.16b
+          smull v28.8h, v12.8b, v2.8b
+          smull v29.8h, v13.8b, v2.8b
+          smull v30.8h, v14.8b, v2.8b
+          smull v31.8h, v15.8b, v2.8b
+          smlal v28.8h, v8.8b, v3.8b
+          smlal v29.8h, v9.8b, v3.8b
+          smlal v30.8h, v10.8b, v3.8b
+          smlal v31.8h, v11.8b, v3.8b
+          smlal2 v28.8h, v12.16b, v2.16b
+          smlal2 v29.8h, v13.16b, v2.16b
+          smlal2 v30.8h, v14.16b, v2.16b
+          smlal2 v31.8h, v15.16b, v2.16b
+          smlal2 v28.8h, v8.16b, v3.16b
+          smlal2 v29.8h, v9.16b, v3.16b
+          smlal2 v30.8h, v10.16b, v3.16b
+          smlal2 v31.8h, v11.16b, v3.16b
+          sadalp v16.4s, v20.8h
+          sadalp v17.4s, v21.8h
+          sadalp v18.4s, v22.8h
+          sadalp v19.4s, v23.8h
+          sadalp v24.4s, v28.8h
+          sadalp v25.4s, v29.8h
+          sadalp v26.4s, v30.8h
+          sadalp v27.4s, v31.8h
+          addp v4.4s, v16.4s, v17.4s
+          addp v5.4s, v18.4s, v19.4s
+          addp v8.4s, v24.4s, v25.4s
+          addp v9.4s, v26.4s, v27.4s
+          addp v6.4s, v4.4s, v5.4s
+          addp v7.4s, v8.4s, v9.4s
+          st1 {v6.4s, v7.4s}, [%[dst]], #32
+          )asm"
+          : [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
+            [run_depth] "+r"(run_depth)
+          :
+          : "cc", "memory", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
+            "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
+            "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26",
+            "v27", "v28", "v29", "v30", "v31");
+    }
+  }
+}
+
+template <>
+void NeonRunKernelNoSDot<4, 4, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                   int32_t* dst, int lhs_layout_rows,
+                                   int lhs_layout_cols, int rhs_layout_rows,
+                                   int rhs_layout_cols, int dst_layout_rows,
+                                   int dst_layout_cols) {
+  const int rows_left = 4;
+  const int rows_right = 4;
+  const int cols = 32;
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
+  const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
+  const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    int left_index = i * rows_left * lhs_layout_cols / 2;
+    const uint8_t* lhs_ptr_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_ptr = lhs_ptr_data;
+      int right_index = j * rows_right * rhs_layout_cols;
+      const int8_t* rhs_ptr = rhs + right_index;
+      int run_depth = depth;
+      asm(R"asm(
+          movi v3.16b, #15
+          ld1 {v4.16b}, [%[lhs_ptr]], #16
+          movi v16.4s, #0
+          movi v17.4s, #0
+          movi v18.4s, #0
+          movi v19.4s, #0
+          ld1 {v5.16b}, [%[lhs_ptr]], #16
+          movi v20.4s, #0
+          movi v21.4s, #0
+          movi v22.4s, #0
+          ld1 {v6.16b}, [%[lhs_ptr]], #16
+          movi v23.4s, #0
+          movi v24.4s, #0
+          movi v25.4s, #0
+          ld1 {v7.16b}, [%[lhs_ptr]], #16
+          movi v26.4s, #0
+          movi v27.4s, #0
+          movi v28.4s, #0
+          movi v29.4s, #0
+          ld1 {v0.16b}, [%[rhs_ptr]], #16
+          movi v30.4s, #0
+          movi v31.4s, #0
+          ld1 {v1.16b}, [%[rhs_ptr]], #16
+          and v8.16b, v4.16b, v3.16b
+          and v9.16b, v5.16b, v3.16b
+          and v10.16b, v6.16b, v3.16b
+          and v11.16b, v7.16b, v3.16b
+          ushr v12.16b, v4.16b, #4
+          ushr v13.16b, v5.16b, #4
+          ushr v14.16b, v6.16b, #4
+          ushr v15.16b, v7.16b, #4
+          subs %w[run_depth], %w[run_depth], #1
+          b.ls 1f /* skip loop */
+            0: /* loop start */
+            smull v4.8h, v12.8b, v0.8b
+            smull v5.8h, v13.8b, v0.8b
+            smull v6.8h, v14.8b, v0.8b
+            smull v7.8h, v15.8b, v0.8b
+            smlal2 v4.8h, v12.16b, v0.16b
+            smlal2 v5.8h, v13.16b, v0.16b
+            smlal2 v6.8h, v14.16b, v0.16b
+            smlal2 v7.8h, v15.16b, v0.16b
+            ld1 {v2.16b}, [%[rhs_ptr]], #16
+            smlal v4.8h, v8.8b, v1.8b
+            smlal v5.8h, v9.8b, v1.8b
+            smlal v6.8h, v10.8b, v1.8b
+            smlal v7.8h, v11.8b, v1.8b
+            smlal2 v4.8h, v8.16b, v1.16b
+            smlal2 v5.8h, v9.16b, v1.16b
+            smlal2 v6.8h, v10.16b, v1.16b
+            smlal2 v7.8h, v11.16b, v1.16b
+            ld1 {v0.16b}, [%[rhs_ptr]], #16
+            sadalp v16.4s, v4.8h
+            sadalp v17.4s, v5.8h
+            sadalp v18.4s, v6.8h
+            sadalp v19.4s, v7.8h
+            smull v4.8h, v12.8b, v2.8b
+            smull v5.8h, v13.8b, v2.8b
+            smull v6.8h, v14.8b, v2.8b
+            smull v7.8h, v15.8b, v2.8b
+            smlal2 v4.8h, v12.16b, v2.16b
+            smlal2 v5.8h, v13.16b, v2.16b
+            smlal2 v6.8h, v14.16b, v2.16b
+            smlal2 v7.8h, v15.16b, v2.16b
+            ld1 {v1.16b}, [%[rhs_ptr]], #16
+            smlal v4.8h, v8.8b, v0.8b
+            smlal v5.8h, v9.8b, v0.8b
+            smlal v6.8h, v10.8b, v0.8b
+            smlal v7.8h, v11.8b, v0.8b
+            smlal2 v4.8h, v8.16b, v0.16b
+            smlal2 v5.8h, v9.16b, v0.16b
+            smlal2 v6.8h, v10.16b, v0.16b
+            smlal2 v7.8h, v11.16b, v0.16b
+            ld1 {v2.16b}, [%[rhs_ptr]], #16
+            sadalp v20.4s, v4.8h
+            sadalp v21.4s, v5.8h
+            sadalp v22.4s, v6.8h
+            sadalp v23.4s, v7.8h
+            smull v4.8h, v12.8b, v1.8b
+            smull v5.8h, v13.8b, v1.8b
+            smull v6.8h, v14.8b, v1.8b
+            smull v7.8h, v15.8b, v1.8b
+            smlal2 v4.8h, v12.16b, v1.16b
+            smlal2 v5.8h, v13.16b, v1.16b
+            smlal2 v6.8h, v14.16b, v1.16b
+            smlal2 v7.8h, v15.16b, v1.16b
+            ld1 {v0.16b}, [%[rhs_ptr]], #16
+            smlal v4.8h, v8.8b, v2.8b
+            smlal v5.8h, v9.8b, v2.8b
+            smlal v6.8h, v10.8b, v2.8b
+            smlal v7.8h, v11.8b, v2.8b
+            smlal2 v4.8h, v8.16b, v2.16b
+            smlal2 v5.8h, v9.16b, v2.16b
+            smlal2 v6.8h, v10.16b, v2.16b
+            smlal2 v7.8h, v11.16b, v2.16b
+            ld1 {v1.16b}, [%[rhs_ptr]], #16
+            sadalp v24.4s, v4.8h
+            sadalp v25.4s, v5.8h
+            sadalp v26.4s, v6.8h
+            sadalp v27.4s, v7.8h
+            smull v4.8h, v12.8b, v0.8b
+            smull v5.8h, v13.8b, v0.8b
+            smull v6.8h, v14.8b, v0.8b
+            smull v7.8h, v15.8b, v0.8b
+            smlal2 v4.8h, v12.16b, v0.16b
+            smlal2 v5.8h, v13.16b, v0.16b
+            smlal2 v6.8h, v14.16b, v0.16b
+            smlal2 v7.8h, v15.16b, v0.16b
+            ld1 {v12.16b}, [%[lhs_ptr]], #16
+            smlal v4.8h, v8.8b, v1.8b
+            smlal v5.8h, v9.8b, v1.8b
+            smlal v6.8h, v10.8b, v1.8b
+            smlal v7.8h, v11.8b, v1.8b
+            ld1 {v13.16b}, [%[lhs_ptr]], #16
+            smlal2 v4.8h, v8.16b, v1.16b
+            smlal2 v5.8h, v9.16b, v1.16b
+            smlal2 v6.8h, v10.16b, v1.16b
+            smlal2 v7.8h, v11.16b, v1.16b
+            ld1 {v14.16b}, [%[lhs_ptr]], #16
+            sadalp v28.4s, v4.8h
+            sadalp v29.4s, v5.8h
+            sadalp v30.4s, v6.8h
+            sadalp v31.4s, v7.8h
+            ld1 {v15.16b}, [%[lhs_ptr]], #16
+            and v8.16b, v12.16b, v3.16b
+            and v9.16b, v13.16b, v3.16b
+            and v10.16b, v14.16b, v3.16b
+            ld1 {v0.16b}, [%[rhs_ptr]], #16
+            and v11.16b, v15.16b, v3.16b
+            ushr v12.16b, v12.16b, #4
+            ushr v13.16b, v13.16b, #4
+            ld1 {v1.16b}, [%[rhs_ptr]], #16
+            ushr v14.16b, v14.16b, #4
+            ushr v15.16b, v15.16b, #4
+            subs %w[run_depth], %w[run_depth], #1
+            b.hi 0b /* loop branch */
+          1: /* loop end */
+          smull v4.8h, v12.8b, v0.8b
+          smull v5.8h, v13.8b, v0.8b
+          smull v6.8h, v14.8b, v0.8b
+          smull v7.8h, v15.8b, v0.8b
+          smlal2 v4.8h, v12.16b, v0.16b
+          smlal2 v5.8h, v13.16b, v0.16b
+          smlal2 v6.8h, v14.16b, v0.16b
+          smlal2 v7.8h, v15.16b, v0.16b
+          ld1 {v2.16b}, [%[rhs_ptr]], #16
+          smlal v4.8h, v8.8b, v1.8b
+          smlal v5.8h, v9.8b, v1.8b
+          smlal v6.8h, v10.8b, v1.8b
+          smlal v7.8h, v11.8b, v1.8b
+          smlal2 v4.8h, v8.16b, v1.16b
+          smlal2 v5.8h, v9.16b, v1.16b
+          smlal2 v6.8h, v10.16b, v1.16b
+          smlal2 v7.8h, v11.16b, v1.16b
+          ld1 {v0.16b}, [%[rhs_ptr]], #16
+          sadalp v16.4s, v4.8h
+          sadalp v17.4s, v5.8h
+          sadalp v18.4s, v6.8h
+          sadalp v19.4s, v7.8h
+          smull v4.8h, v12.8b, v2.8b
+          smull v5.8h, v13.8b, v2.8b
+          smull v6.8h, v14.8b, v2.8b
+          smull v7.8h, v15.8b, v2.8b
+          smlal2 v4.8h, v12.16b, v2.16b
+          smlal2 v5.8h, v13.16b, v2.16b
+          smlal2 v6.8h, v14.16b, v2.16b
+          smlal2 v7.8h, v15.16b, v2.16b
+          ld1 {v1.16b}, [%[rhs_ptr]], #16
+          smlal v4.8h, v8.8b, v0.8b
+          smlal v5.8h, v9.8b, v0.8b
+          smlal v6.8h, v10.8b, v0.8b
+          smlal v7.8h, v11.8b, v0.8b
+          smlal2 v4.8h, v8.16b, v0.16b
+          smlal2 v5.8h, v9.16b, v0.16b
+          smlal2 v6.8h, v10.16b, v0.16b
+          smlal2 v7.8h, v11.16b, v0.16b
+          ld1 {v2.16b}, [%[rhs_ptr]], #16
+          sadalp v20.4s, v4.8h
+          sadalp v21.4s, v5.8h
+          sadalp v22.4s, v6.8h
+          sadalp v23.4s, v7.8h
+          smull v4.8h, v12.8b, v1.8b
+          smull v5.8h, v13.8b, v1.8b
+          smull v6.8h, v14.8b, v1.8b
+          smull v7.8h, v15.8b, v1.8b
+          smlal2 v4.8h, v12.16b, v1.16b
+          smlal2 v5.8h, v13.16b, v1.16b
+          smlal2 v6.8h, v14.16b, v1.16b
+          smlal2 v7.8h, v15.16b, v1.16b
+          ld1 {v0.16b}, [%[rhs_ptr]], #16
+          smlal v4.8h, v8.8b, v2.8b
+          smlal v5.8h, v9.8b, v2.8b
+          smlal v6.8h, v10.8b, v2.8b
+          smlal v7.8h, v11.8b, v2.8b
+          smlal2 v4.8h, v8.16b, v2.16b
+          smlal2 v5.8h, v9.16b, v2.16b
+          smlal2 v6.8h, v10.16b, v2.16b
+          smlal2 v7.8h, v11.16b, v2.16b
+          ld1 {v1.16b}, [%[rhs_ptr]], #16
+          sadalp v24.4s, v4.8h
+          sadalp v25.4s, v5.8h
+          sadalp v26.4s, v6.8h
+          sadalp v27.4s, v7.8h
+          smull v4.8h, v12.8b, v0.8b
+          smull v5.8h, v13.8b, v0.8b
+          smull v6.8h, v14.8b, v0.8b
+          smull v7.8h, v15.8b, v0.8b
+          smlal2 v4.8h, v12.16b, v0.16b
+          smlal2 v5.8h, v13.16b, v0.16b
+          smlal2 v6.8h, v14.16b, v0.16b
+          smlal2 v7.8h, v15.16b, v0.16b
+          smlal v4.8h, v8.8b, v1.8b
+          smlal v5.8h, v9.8b, v1.8b
+          smlal v6.8h, v10.8b, v1.8b
+          smlal v7.8h, v11.8b, v1.8b
+          smlal2 v4.8h, v8.16b, v1.16b
+          smlal2 v5.8h, v9.16b, v1.16b
+          smlal2 v6.8h, v10.16b, v1.16b
+          smlal2 v7.8h, v11.16b, v1.16b
+          sadalp v28.4s, v4.8h
+          sadalp v29.4s, v5.8h
+          sadalp v30.4s, v6.8h
+          sadalp v31.4s, v7.8h
+          addp v14.4s, v16.4s, v17.4s
+          addp v15.4s, v18.4s, v19.4s
+          addp v12.4s, v20.4s, v21.4s
+          addp v13.4s, v22.4s, v23.4s
+          addp v10.4s, v24.4s, v25.4s
+          addp v11.4s, v26.4s, v27.4s
+          addp v8.4s, v28.4s, v29.4s
+          addp v9.4s, v30.4s, v31.4s
+          addp v4.4s, v14.4s, v15.4s
+          addp v5.4s, v12.4s, v13.4s
+          addp v6.4s, v10.4s, v11.4s
+          addp v7.4s, v8.4s, v9.4s
+          st1 {v4.4s, v5.4s, v6.4s, v7.4s}, [%[dst]], #64
+          )asm"
+          : [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
+            [run_depth] "+r"(run_depth)
+          :
+          : "cc", "memory", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
+            "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
+            "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26",
+            "v27", "v28", "v29", "v30", "v31");
+    }
+  }
+}
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+#endif  // defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) ||
+        // defined(__ARM_NEON))

tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_aarch64_sdot.cc

@@ -0,0 +1,418 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
+#include <stdint.h>
+
+#include <algorithm>
+#include <vector>
+
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
+
+#define DOTPROD_ATTRIBUTE __attribute__((target("dotprod")))
+
+namespace tflite {
+namespace optimized_4bit {
+
+template <int RowsLeft, int RowsRight, int Cols>
+void NeonRunKernelSDot(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                       int lhs_layout_rows, int lhs_layout_cols,
+                       int rhs_layout_rows, int rhs_layout_cols,
+                       int dst_layout_rows, int dst_layout_cols);
+
+template <>
+DOTPROD_ATTRIBUTE void NeonRunKernelSDot<4, 1, 32>(
+    const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
+    int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
+    int dst_layout_rows, int dst_layout_cols) {
+  const int rows_left = 4;
+  const int rows_right = 1;
+  const int cols = 32;
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
+  const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
+  const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    int left_index = i * rows_left * lhs_layout_cols / 2;
+    const uint8_t* lhs_ptr_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_ptr = lhs_ptr_data;
+      int right_index = j * rows_right * rhs_layout_cols;
+      const int8_t* rhs_ptr = rhs + right_index;
+      int run_depth = depth;
+      asm(R"asm(
+          movi v24.16b, #15
+          ld1 {v4.16b}, [%[lhs_ptr]], #16
+          movi v16.4s, #0
+          movi v17.4s, #0
+          ld1 {v5.16b}, [%[lhs_ptr]], #16
+          movi v18.4s, #0
+          movi v19.4s, #0
+          ld1 {v6.16b, v7.16b}, [%[lhs_ptr]], #32
+          and v8.16b, v4.16b, v24.16b
+          and v9.16b, v5.16b, v24.16b
+          ushr v12.16b, v4.16b, #4
+          ushr v13.16b, v5.16b, #4
+          ld1 {v0.16b, v1.16b}, [%[rhs_ptr]], #32
+          and v10.16b, v6.16b, v24.16b
+          and v11.16b, v7.16b, v24.16b
+          ushr v14.16b, v6.16b, #4
+          ushr v15.16b, v7.16b, #4
+          subs %w[run_depth], %w[run_depth], #1
+          b.ls 1f /* skip loop */
+            0: /* loop start */
+            ld1 {v4.16b, v5.16b, v6.16b, v7.16b}, [%[lhs_ptr]], #64
+            sdot v16.4s, v8.16b, v1.16b
+            sdot v17.4s, v9.16b, v1.16b
+            sdot v18.4s, v10.16b, v1.16b
+            sdot v19.4s, v11.16b, v1.16b
+            sdot v16.4s, v12.16b, v0.16b
+            sdot v17.4s, v13.16b, v0.16b
+            sdot v18.4s, v14.16b, v0.16b
+            sdot v19.4s, v15.16b, v0.16b
+            and v8.16b, v4.16b, v24.16b
+            and v9.16b, v5.16b, v24.16b
+            ushr v12.16b, v4.16b, #4
+            ushr v13.16b, v5.16b, #4
+            ld1 {v0.16b, v1.16b}, [%[rhs_ptr]], #32
+            and v10.16b, v6.16b, v24.16b
+            and v11.16b, v7.16b, v24.16b
+            ushr v14.16b, v6.16b, #4
+            ushr v15.16b, v7.16b, #4
+            subs %w[run_depth], %w[run_depth], #1
+            b.hi 0b /* loop branch */
+          1: /* loop end */
+          sdot v16.4s, v8.16b, v1.16b
+          sdot v17.4s, v9.16b, v1.16b
+          sdot v18.4s, v10.16b, v1.16b
+          sdot v19.4s, v11.16b, v1.16b
+          sdot v16.4s, v12.16b, v0.16b
+          sdot v17.4s, v13.16b, v0.16b
+          sdot v18.4s, v14.16b, v0.16b
+          sdot v19.4s, v15.16b, v0.16b
+          addp v4.4s, v16.4s, v17.4s
+          addp v5.4s, v18.4s, v19.4s
+          addp v6.4s, v4.4s, v5.4s
+          st1 {v6.4s}, [%[dst]], #16
+          )asm"
+          : [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
+            [run_depth] "+r"(run_depth)
+          :
+          : "cc", "memory", "v0", "v1", "v4", "v5", "v6", "v7", "v8", "v9",
+            "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18",
+            "v19", "v24");
+    }
+  }
+}
+
+// Note: NeonRunKernelSDot<4, 2, 32> does not mutate registers v25-v31 in its
+// inline assembly block, so they are intentionally omitted from the clobber
+// list to avoid redundant register preservation.
+template <>
+DOTPROD_ATTRIBUTE void NeonRunKernelSDot<4, 2, 32>(
+    const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
+    int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
+    int dst_layout_rows, int dst_layout_cols) {
+  const int rows_left = 4;
+  const int rows_right = 2;
+  const int cols = 32;
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
+  const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
+  const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    int left_index = i * rows_left * lhs_layout_cols / 2;
+    const uint8_t* lhs_ptr_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_ptr = lhs_ptr_data;
+      int right_index = j * rows_right * rhs_layout_cols;
+      const int8_t* rhs_ptr = rhs + right_index;
+      int run_depth = depth;
+      asm(R"asm(
+          movi v24.16b, #15
+          ld1 {v4.16b}, [%[lhs_ptr]], #16
+          movi v16.4s, #0
+          movi v17.4s, #0
+          ld1 {v5.16b}, [%[lhs_ptr]], #16
+          movi v18.4s, #0
+          movi v19.4s, #0
+          ld1 {v6.16b, v7.16b}, [%[lhs_ptr]], #32
+          movi v20.4s, #0
+          movi v21.4s, #0
+          and v8.16b, v4.16b, v24.16b
+          and v9.16b, v5.16b, v24.16b
+          movi v22.4s, #0
+          movi v23.4s, #0
+          ushr v12.16b, v4.16b, #4
+          ushr v13.16b, v5.16b, #4
+          ld1 {v0.16b, v1.16b, v2.16b, v3.16b}, [%[rhs_ptr]], #64
+          and v10.16b, v6.16b, v24.16b
+          and v11.16b, v7.16b, v24.16b
+          ushr v14.16b, v6.16b, #4
+          ushr v15.16b, v7.16b, #4
+          subs %w[run_depth], %w[run_depth], #1
+          b.ls 1f /* skip loop */
+            0: /* loop start */
+            ld1 {v4.16b, v5.16b, v6.16b, v7.16b}, [%[lhs_ptr]], #64
+            sdot v16.4s, v12.16b, v0.16b
+            sdot v17.4s, v13.16b, v0.16b
+            sdot v18.4s, v14.16b, v0.16b
+            sdot v19.4s, v15.16b, v0.16b
+            sdot v16.4s, v8.16b, v1.16b
+            sdot v17.4s, v9.16b, v1.16b
+            sdot v18.4s, v10.16b, v1.16b
+            sdot v19.4s, v11.16b, v1.16b
+            sdot v20.4s, v12.16b, v2.16b
+            sdot v21.4s, v13.16b, v2.16b
+            sdot v22.4s, v14.16b, v2.16b
+            sdot v23.4s, v15.16b, v2.16b
+            sdot v20.4s, v8.16b, v3.16b
+            sdot v21.4s, v9.16b, v3.16b
+            sdot v22.4s, v10.16b, v3.16b
+            sdot v23.4s, v11.16b, v3.16b
+            and v8.16b, v4.16b, v24.16b
+            and v9.16b, v5.16b, v24.16b
+            ushr v12.16b, v4.16b, #4
+            ld1 {v0.16b, v1.16b, v2.16b, v3.16b}, [%[rhs_ptr]], #64
+            ushr v13.16b, v5.16b, #4
+            and v10.16b, v6.16b, v24.16b
+            and v11.16b, v7.16b, v24.16b
+            ushr v14.16b, v6.16b, #4
+            ushr v15.16b, v7.16b, #4
+            subs %w[run_depth], %w[run_depth], #1
+            b.hi 0b /* loop branch */
+          1: /* loop end */
+          sdot v16.4s, v12.16b, v0.16b
+          sdot v17.4s, v13.16b, v0.16b
+          sdot v18.4s, v14.16b, v0.16b
+          sdot v19.4s, v15.16b, v0.16b
+          sdot v16.4s, v8.16b, v1.16b
+          sdot v17.4s, v9.16b, v1.16b
+          sdot v18.4s, v10.16b, v1.16b
+          sdot v19.4s, v11.16b, v1.16b
+          sdot v20.4s, v12.16b, v2.16b
+          sdot v21.4s, v13.16b, v2.16b
+          sdot v22.4s, v14.16b, v2.16b
+          sdot v23.4s, v15.16b, v2.16b
+          sdot v20.4s, v8.16b, v3.16b
+          sdot v21.4s, v9.16b, v3.16b
+          sdot v22.4s, v10.16b, v3.16b
+          sdot v23.4s, v11.16b, v3.16b
+          addp v4.4s, v16.4s, v17.4s
+          addp v5.4s, v18.4s, v19.4s
+          addp v8.4s, v20.4s, v21.4s
+          addp v9.4s, v22.4s, v23.4s
+          addp v6.4s, v4.4s, v5.4s
+          addp v7.4s, v8.4s, v9.4s
+          st1 {v6.4s, v7.4s}, [%[dst]], #32
+          )asm"
+          : [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
+            [run_depth] "+r"(run_depth)
+          :
+          : "cc", "memory", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
+            "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
+            "v18", "v19", "v20", "v21", "v22", "v23", "v24");
+    }
+  }
+}
+
+template <>
+DOTPROD_ATTRIBUTE void NeonRunKernelSDot<4, 4, 32>(
+    const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
+    int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
+    int dst_layout_rows, int dst_layout_cols) {
+  const int rows_left = 4;
+  const int rows_right = 4;
+  const int cols = 32;
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
+  const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
+  const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    int left_index = i * rows_left * lhs_layout_cols / 2;
+    const uint8_t* lhs_ptr_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_ptr = lhs_ptr_data;
+      int right_index = j * rows_right * rhs_layout_cols;
+      const int8_t* rhs_ptr = rhs + right_index;
+      int run_depth = depth;
+      asm(R"asm(
+          movi v3.16b, #15
+          ld1 {v4.16b}, [%[lhs_ptr]], #16
+          movi v16.4s, #0
+          movi v17.4s, #0
+          ld1 {v5.16b}, [%[lhs_ptr]], #16
+          movi v18.4s, #0
+          movi v19.4s, #0
+          ld1 {v6.16b, v7.16b}, [%[lhs_ptr]], #32
+          and v8.16b, v4.16b, v3.16b
+          and v9.16b, v5.16b, v3.16b
+          movi v20.4s, #0
+          movi v21.4s, #0
+          movi v22.4s, #0
+          movi v23.4s, #0
+          movi v24.4s, #0
+          movi v25.4s, #0
+          movi v26.4s, #0
+          movi v27.4s, #0
+          movi v28.4s, #0
+          movi v29.4s, #0
+          movi v30.4s, #0
+          movi v31.4s, #0
+          ushr v12.16b, v4.16b, #4
+          ushr v13.16b, v5.16b, #4
+          ld1 {v0.16b, v1.16b}, [%[rhs_ptr]], #32
+          and v10.16b, v6.16b, v3.16b
+          and v11.16b, v7.16b, v3.16b
+          ushr v14.16b, v6.16b, #4
+          ushr v15.16b, v7.16b, #4
+          subs %w[run_depth], %w[run_depth], #1
+          b.ls 1f /* skip loop */
+            0: /* loop start */
+            ld1 {v4.16b, v5.16b, v6.16b, v7.16b}, [%[lhs_ptr]], #64
+            sdot v16.4s, v12.16b, v0.16b
+            sdot v17.4s, v13.16b, v0.16b
+            sdot v18.4s, v14.16b, v0.16b
+            sdot v19.4s, v15.16b, v0.16b
+            ld1 {v2.16b}, [%[rhs_ptr]], #16
+            sdot v16.4s, v8.16b, v1.16b
+            sdot v17.4s, v9.16b, v1.16b
+            sdot v18.4s, v10.16b, v1.16b
+            sdot v19.4s, v11.16b, v1.16b
+            ld1 {v0.16b}, [%[rhs_ptr]], #16
+            sdot v20.4s, v12.16b, v2.16b
+            sdot v21.4s, v13.16b, v2.16b
+            sdot v22.4s, v14.16b, v2.16b
+            sdot v23.4s, v15.16b, v2.16b
+            ld1 {v1.16b}, [%[rhs_ptr]], #16
+            sdot v20.4s, v8.16b, v0.16b
+            sdot v21.4s, v9.16b, v0.16b
+            sdot v22.4s, v10.16b, v0.16b
+            sdot v23.4s, v11.16b, v0.16b
+            ld1 {v2.16b}, [%[rhs_ptr]], #16
+            sdot v24.4s, v12.16b, v1.16b
+            sdot v25.4s, v13.16b, v1.16b
+            sdot v26.4s, v14.16b, v1.16b
+            sdot v27.4s, v15.16b, v1.16b
+            ld1 {v0.16b}, [%[rhs_ptr]], #16
+            sdot v24.4s, v8.16b, v2.16b
+            sdot v25.4s, v9.16b, v2.16b
+            sdot v26.4s, v10.16b, v2.16b
+            sdot v27.4s, v11.16b, v2.16b
+            ld1 {v1.16b}, [%[rhs_ptr]], #16
+            sdot v28.4s, v12.16b, v0.16b
+            sdot v29.4s, v13.16b, v0.16b
+            sdot v30.4s, v14.16b, v0.16b
+            sdot v31.4s, v15.16b, v0.16b
+            sdot v28.4s, v8.16b, v1.16b
+            sdot v29.4s, v9.16b, v1.16b
+            sdot v30.4s, v10.16b, v1.16b
+            sdot v31.4s, v11.16b, v1.16b
+            ld1 {v0.16b}, [%[rhs_ptr]], #16
+            and v8.16b, v4.16b, v3.16b
+            and v9.16b, v5.16b, v3.16b
+            ushr v12.16b, v4.16b, #4
+            ushr v13.16b, v5.16b, #4
+            ld1 {v1.16b}, [%[rhs_ptr]], #16
+            and v10.16b, v6.16b, v3.16b
+            and v11.16b, v7.16b, v3.16b
+            ushr v14.16b, v6.16b, #4
+            ushr v15.16b, v7.16b, #4
+            subs %w[run_depth], %w[run_depth], #1
+            b.hi 0b /* loop branch */
+          1: /* loop end */
+          sdot v16.4s, v12.16b, v0.16b
+          sdot v17.4s, v13.16b, v0.16b
+          sdot v18.4s, v14.16b, v0.16b
+          sdot v19.4s, v15.16b, v0.16b
+          ld1 {v2.16b}, [%[rhs_ptr]], #16
+          sdot v16.4s, v8.16b, v1.16b
+          sdot v17.4s, v9.16b, v1.16b
+          sdot v18.4s, v10.16b, v1.16b
+          sdot v19.4s, v11.16b, v1.16b
+          ld1 {v0.16b}, [%[rhs_ptr]], #16
+          sdot v20.4s, v12.16b, v2.16b
+          sdot v21.4s, v13.16b, v2.16b
+          sdot v22.4s, v14.16b, v2.16b
+          sdot v23.4s, v15.16b, v2.16b
+          ld1 {v1.16b}, [%[rhs_ptr]], #16
+          sdot v20.4s, v8.16b, v0.16b
+          sdot v21.4s, v9.16b, v0.16b
+          sdot v22.4s, v10.16b, v0.16b
+          sdot v23.4s, v11.16b, v0.16b
+          ld1 {v2.16b}, [%[rhs_ptr]], #16
+          sdot v24.4s, v12.16b, v1.16b
+          sdot v25.4s, v13.16b, v1.16b
+          sdot v26.4s, v14.16b, v1.16b
+          sdot v27.4s, v15.16b, v1.16b
+          ld1 {v0.16b}, [%[rhs_ptr]], #16
+          sdot v24.4s, v8.16b, v2.16b
+          sdot v25.4s, v9.16b, v2.16b
+          sdot v26.4s, v10.16b, v2.16b
+          sdot v27.4s, v11.16b, v2.16b
+          ld1 {v1.16b}, [%[rhs_ptr]], #16
+          sdot v28.4s, v12.16b, v0.16b
+          sdot v29.4s, v13.16b, v0.16b
+          sdot v30.4s, v14.16b, v0.16b
+          sdot v31.4s, v15.16b, v0.16b
+          sdot v28.4s, v8.16b, v1.16b
+          sdot v29.4s, v9.16b, v1.16b
+          sdot v30.4s, v10.16b, v1.16b
+          sdot v31.4s, v11.16b, v1.16b
+          addp v14.4s, v16.4s, v17.4s
+          addp v15.4s, v18.4s, v19.4s
+          addp v12.4s, v20.4s, v21.4s
+          addp v13.4s, v22.4s, v23.4s
+          addp v10.4s, v24.4s, v25.4s
+          addp v11.4s, v26.4s, v27.4s
+          addp v8.4s, v28.4s, v29.4s
+          addp v9.4s, v30.4s, v31.4s
+          addp v4.4s, v14.4s, v15.4s
+          addp v5.4s, v12.4s, v13.4s
+          addp v6.4s, v10.4s, v11.4s
+          addp v7.4s, v8.4s, v9.4s
+          st1 {v4.4s, v5.4s, v6.4s, v7.4s}, [%[dst]], #64
+          )asm"
+          : [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
+            [run_depth] "+r"(run_depth)
+          :
+          : "cc", "memory", "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7",
+            "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17",
+            "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", "v26",
+            "v27", "v28", "v29", "v30", "v31");
+    }
+  }
+}
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+#endif  // defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) ||
+        // defined(__ARM_NEON))

tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_arm32_nosdot.cc

@@ -0,0 +1,159 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
+
+#include <stdint.h>
+
+#include <algorithm>
+
+#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
+
+namespace tflite {
+namespace optimized_4bit {
+
+template <int RowsLeft, int RowsRight, int Cols>
+void NeonRunKernelNoSDot(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                         int lhs_layout_rows, int lhs_layout_cols,
+                         int rhs_layout_rows, int rhs_layout_cols,
+                         int dst_layout_rows, int dst_layout_cols);
+
+template <>
+void NeonRunKernelNoSDot<4, 1, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                   int32_t* dst, int lhs_layout_rows,
+                                   int lhs_layout_cols, int rhs_layout_rows,
+                                   int rhs_layout_cols, int dst_layout_rows,
+                                   int dst_layout_cols) {
+  const int rows_left = 4;
+  const int rows_right = 1;
+  const int cols = 32;
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
+  const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
+  const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    const int left_index = i * rows_left * lhs_layout_cols / 2;
+    const uint8_t* lhs_ptr_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_ptr = lhs_ptr_data;
+      const int right_index = j * rows_right * rhs_layout_cols;
+      const int8_t* rhs_ptr = rhs + right_index;
+      int run_depth = depth;
+      asm(R"asm(
+          vmov.i8 q14, #15
+          vld1.8 {q4}, [%[lhs_ptr]]!
+          vmov.i32 q0, #0
+          vmov.i32 q1, #0
+          vld1.8 {q5}, [%[lhs_ptr]]!
+          vmov.i32 q2, #0
+          vmov.i32 q3, #0
+          vld1.8 {q6}, [%[lhs_ptr]]!
+          vand q8, q4, q14
+          vand q9, q5, q14
+          vld1.8 {q7}, [%[lhs_ptr]]!
+          vshr.u8 q4, q4, #4
+          vshr.u8 q5, q5, #4
+          vld1.8 {d24, d25}, [%[rhs_ptr]]!
+          vand q10, q6, q14
+          vand q11, q7, q14
+          vld1.8 {d26, d27}, [%[rhs_ptr]]!
+          vshr.u8 q6, q6, #4
+          vshr.u8 q7, q7, #4
+          subs %[run_depth], %[run_depth], #1
+          bls 1f /* skip loop */
+            0: /* loop start */
+            vmull.s8 q15, d8, d24
+            vmlal.s8 q15, d9, d25
+            vmlal.s8 q15, d16, d26
+            vmlal.s8 q15, d17, d27
+            vpadal.s16 q0, q15
+            vmull.s8 q15, d10, d24
+            vmlal.s8 q15, d11, d25
+            vmlal.s8 q15, d18, d26
+            vmlal.s8 q15, d19, d27
+            vpadal.s16 q1, q15
+            vld1.8 {q4, q5}, [%[lhs_ptr]]!
+            vmull.s8 q15, d12, d24
+            vmlal.s8 q15, d13, d25
+            vmlal.s8 q15, d20, d26
+            vmlal.s8 q15, d21, d27
+            vpadal.s16 q2, q15
+            vmull.s8 q15, d14, d24
+            vmlal.s8 q15, d15, d25
+            vmlal.s8 q15, d22, d26
+            vmlal.s8 q15, d23, d27
+            vpadal.s16 q3, q15
+            vld1.8 {q6, q7}, [%[lhs_ptr]]!
+            vld1.8 {d24, d25, d26, d27}, [%[rhs_ptr]]!
+            vand q8, q4, q14
+            vand q9, q5, q14
+            vand q10, q6, q14
+            vand q11, q7, q14
+            vshr.u8 q4, q4, #4
+            vshr.u8 q5, q5, #4
+            vshr.u8 q6, q6, #4
+            vshr.u8 q7, q7, #4
+            subs %[run_depth], %[run_depth], #1
+            bhi 0b /* loop branch */
+          1: /* loop end */
+          vmull.s8 q15, d8, d24
+          vmlal.s8 q15, d9, d25
+          vmlal.s8 q15, d16, d26
+          vmlal.s8 q15, d17, d27
+          vpadal.s16 q0, q15
+          vmull.s8 q15, d10, d24
+          vmlal.s8 q15, d11, d25
+          vmlal.s8 q15, d18, d26
+          vmlal.s8 q15, d19, d27
+          vpadal.s16 q1, q15
+          vmull.s8 q15, d12, d24
+          vmlal.s8 q15, d13, d25
+          vmlal.s8 q15, d20, d26
+          vmlal.s8 q15, d21, d27
+          vpadal.s16 q2, q15
+          vmull.s8 q15, d14, d24
+          vmlal.s8 q15, d15, d25
+          vmlal.s8 q15, d22, d26
+          vmlal.s8 q15, d23, d27
+          vpadal.s16 q3, q15
+          vpadd.i32 d0, d0, d1
+          vpadd.i32 d1, d2, d3
+          vpadd.i32 d2, d4, d5
+          vpadd.i32 d3, d6, d7
+          vpadd.i32 d4, d0, d1
+          vpadd.i32 d5, d2, d3
+          vst1.32 {d4, d5}, [%[dst]]!
+          )asm"
+          : [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
+            [run_depth] "+r"(run_depth)
+          :
+          : "cc", "memory", "d0", "d1", "d2", "d3", "d4", "d5", "d6", "d7",
+            "d8", "d9", "d10", "d11", "d12", "d13", "d14", "d15", "d16", "d17",
+            "d18", "d19", "d20", "d21", "d22", "d23", "d24", "d25", "d26",
+            "d27", "q0", "q1", "q2", "q3", "q4", "q5", "q6", "q7", "q8", "q9",
+            "q10", "q11", "q14", "q15");
+    }
+  }
+}
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+#endif  // defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) ||
+        // defined(__ARM_NEON))

tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_arm32_sdot.cc

@@ -0,0 +1,134 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
+
+#include <stdint.h>
+
+#include <algorithm>
+
+#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h"
+
+#define DOTPROD_ATTRIBUTE __attribute__((target("dotprod")))
+
+namespace tflite {
+namespace optimized_4bit {
+
+template <int RowsLeft, int RowsRight, int Cols>
+void NeonRunKernelSDot(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                       int lhs_layout_rows, int lhs_layout_cols,
+                       int rhs_layout_rows, int rhs_layout_cols,
+                       int dst_layout_rows, int dst_layout_cols);
+
+template <>
+DOTPROD_ATTRIBUTE void NeonRunKernelSDot<4, 1, 32>(
+    const uint8_t* lhs, const int8_t* rhs, int32_t* dst, int lhs_layout_rows,
+    int lhs_layout_cols, int rhs_layout_rows, int rhs_layout_cols,
+    int dst_layout_rows, int dst_layout_cols) {
+  const int rows_left = 4;
+  const int rows_right = 1;
+  const int cols = 32;
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  const int outer_rows = (clamped_end_row + rows_left - 1) / rows_left;
+  const int outer_cols = (clamped_end_col + rows_right - 1) / rows_right;
+  const int depth = std::min(lhs_layout_cols / cols, rhs_layout_cols / cols);
+  for (int i = start_row; i < outer_rows; ++i) {
+    const int left_index = i * rows_left * lhs_layout_cols / 2;
+    const uint8_t* lhs_ptr_data = lhs + left_index;
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_ptr = lhs_ptr_data;
+      const int right_index = j * rows_right * rhs_layout_cols;
+      const int8_t* rhs_ptr = rhs + right_index;
+      int run_depth = depth;
+      asm(R"asm(
+          vmov.i8 q14, #15
+          vld1.8 {q4}, [%[lhs_ptr]]!
+          vmov.i32 q0, #0
+          vmov.i32 q1, #0
+          vld1.8 {q5}, [%[lhs_ptr]]!
+          vmov.i32 q2, #0
+          vmov.i32 q3, #0
+          vld1.8 {q6}, [%[lhs_ptr]]!
+          vand q8, q4, q14
+          vand q9, q5, q14
+          vld1.8 {q7}, [%[lhs_ptr]]!
+          vshr.u8 q4, q4, #4
+          vshr.u8 q5, q5, #4
+          vld1.8 {q12}, [%[rhs_ptr]]!
+          vand q10, q6, q14
+          vand q11, q7, q14
+          vld1.8 {q13}, [%[rhs_ptr]]!
+          vshr.u8 q6, q6, #4
+          vshr.u8 q7, q7, #4
+          subs %[run_depth], %[run_depth], #1
+          bls 1f /* skip loop */
+            0: /* loop start */
+            vsdot.s8 q0, q4, q12
+            vsdot.s8 q0, q8, q13
+            vsdot.s8 q1, q5, q12
+            vsdot.s8 q1, q9, q13
+            vld1.8 {q4, q5}, [%[lhs_ptr]]!
+            vsdot.s8 q2, q6, q12
+            vsdot.s8 q2, q10, q13
+            vsdot.s8 q3, q7, q12
+            vsdot.s8 q3, q11, q13
+            vld1.8 {q6, q7}, [%[lhs_ptr]]!
+            vld1.8 {q12, q13}, [%[rhs_ptr]]!
+            vand q8, q4, q14
+            vand q9, q5, q14
+            vand q10, q6, q14
+            vand q11, q7, q14
+            vshr.u8 q4, q4, #4
+            vshr.u8 q5, q5, #4
+            vshr.u8 q6, q6, #4
+            vshr.u8 q7, q7, #4
+            subs %[run_depth], %[run_depth], #1
+            bhi 0b /* loop branch */
+          1: /* loop end */
+          vsdot.s8 q0, q4, q12
+          vsdot.s8 q0, q8, q13
+          vsdot.s8 q1, q5, q12
+          vsdot.s8 q1, q9, q13
+          vsdot.s8 q2, q6, q12
+          vsdot.s8 q2, q10, q13
+          vsdot.s8 q3, q7, q12
+          vsdot.s8 q3, q11, q13
+          vpadd.i32 d0, d0, d1
+          vpadd.i32 d1, d2, d3
+          vpadd.i32 d2, d4, d5
+          vpadd.i32 d3, d6, d7
+          vpadd.i32 d4, d0, d1
+          vpadd.i32 d5, d2, d3
+          vst1.32 {d4, d5}, [%[dst]]!
+          )asm"
+          : [lhs_ptr] "+r"(lhs_ptr), [rhs_ptr] "+r"(rhs_ptr), [dst] "+r"(dst),
+            [run_depth] "+r"(run_depth)
+          :
+          : "cc", "memory", "d0", "d1", "d2", "d3", "d4", "d5", "d6", "d7",
+            "q0", "q1", "q2", "q3", "q4", "q5", "q6", "q7", "q8", "q9", "q10",
+            "q11", "q12", "q13", "q14");
+    }
+  }
+}
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+#endif  // defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) ||
+        // defined(__ARM_NEON))

tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected_impl.h

@@ -0,0 +1,78 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_IMPL_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_IMPL_H_
+#if defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
+#include <stdint.h>
+
+#if !defined(EIGEN_MAX_ALIGN_BYTES) && !defined(__aarch64__)
+#define EIGEN_MAX_ALIGN_BYTES 32
+#elif !defined(EIGEN_MAX_ALIGN_BYTES)
+#define EIGEN_MAX_ALIGN_BYTES 64
+#endif
+
+namespace tflite {
+namespace optimized_4bit {
+
+void NeonPackInner(const int8_t* src, uint8_t* box, int src_rows, int src_cols,
+                   int outer_row, int outer_col, int outer_rows, int outer_cols,
+                   int inner_rows, int inner_cols);
+
+void NeonPrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                 int layout_cols, int src_rows, int src_cols, int width,
+                 int depth);
+
+void NeonBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                 int n_data, int8_t* quantized_data_ptr,
+                                 float* scaling_factors, int width, int depth,
+                                 int32_t* input_offsets);
+
+void NeonAssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                     const float* batch_scales,
+                                     float* filter_scales,
+                                     const float* bias_ptr, float* output_ptr,
+                                     int output_depth, int batch_size);
+
+template <int Depth, int Width>
+extern void NeonUnpack(float* output_ptr, const int32_t* dst, int batch_size,
+                       int num_units, const float* scaling_factors,
+                       const float* filter_scales, int dst_layout_rows,
+                       int dst_layout_cols);
+
+template <int RowsLeft, int RowsRight, int Cols>
+extern void NeonRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                          int lhs_layout_rows, int lhs_layout_cols,
+                          int rhs_layout_rows, int rhs_layout_cols,
+                          int dst_layout_rows, int dst_layout_cols);
+
+template <int RowsLeft, int RowsRight, int Cols>
+extern void NeonRunKernelNoSDot(const uint8_t* lhs, const int8_t* rhs,
+                                int32_t* dst, int lhs_layout_rows,
+                                int lhs_layout_cols, int rhs_layout_rows,
+                                int rhs_layout_cols, int dst_layout_rows,
+                                int dst_layout_cols);
+template <int RowsLeft, int RowsRight, int Cols>
+extern void NeonRunKernelSDot(const uint8_t* lhs, const int8_t* rhs,
+                              int32_t* dst, int lhs_layout_rows,
+                              int lhs_layout_cols, int rhs_layout_rows,
+                              int rhs_layout_cols, int dst_layout_rows,
+                              int dst_layout_cols);
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // defined(FC_4BIT_NEON)...
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_NEON_FULLY_CONNECTED_IMPL_H_

tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected.cc

@@ -0,0 +1,439 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if defined(FC_4BIT_SSE) && defined(__SSSE3__)
+
+#include <stdint.h>
+#include <stdlib.h>
+
+// NOLINTBEGIN
+#include <tmmintrin.h>
+
+#include <algorithm>
+#include <cstdlib>
+#include <cstring>
+#include <vector>
+
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_common.h"
+#include "tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected_impl.h"
+
+namespace tflite {
+namespace optimized_4bit {
+#define is_aligned(ptr, bytes) ((((size_t)(ptr)) & (bytes - 1)) == 0)
+
+void SsePackInner(const int8_t* src, uint8_t* box, int src_rows, int src_cols,
+                  int outer_row, int outer_col, int outer_rows, int outer_cols,
+                  int inner_rows, int inner_cols) {
+  const int width = inner_rows;
+  const int depth = inner_cols;
+  const int real_depth = depth / 2;
+  const int real_src_cols = src_cols / 2;
+  const int row = outer_row * inner_rows;
+  const int col = outer_col * inner_cols;
+  int src_width = std::min(width, src_rows - row);
+  int src_depth = std::min(depth, src_cols - col);
+  int real_col = col / 2;
+  const int8_t* src_data =
+      src + static_cast<size_t>(row) * real_src_cols + real_col;
+  int real_src_depth = src_depth / 2;
+  const __m128i bitmask_upper = _mm_set1_epi16(255U << 8);
+  const __m128i bitmask_lower = _mm_set1_epi16(255U);
+  const __m128i seven = _mm_set1_epi8(7);
+  for (int m = 0; m < src_width; ++m) {
+    int i = 0;
+    int k = 0;
+    for (; i < (real_src_depth & (~15)); i += 16) {
+      const __m128i values_128i = _mm_loadu_si128((__m128i*)(src_data + i));
+      // sign extend uv1
+      __m128i uv1 = _mm_srai_epi16(values_128i, 4);
+      uv1 = _mm_add_epi8(uv1, seven);
+      uv1 = _mm_and_si128(uv1, bitmask_upper);
+      __m128i uv2 = _mm_slli_epi16(values_128i, 8);
+      uv2 = _mm_srai_epi16(uv2, 12);
+      uv2 = _mm_add_epi8(uv2, seven);
+      uv2 = _mm_and_si128(uv2, bitmask_lower);
+      uv1 = _mm_or_si128(uv1, uv2);
+
+      __m128i lv1 = _mm_slli_epi16(values_128i, 4);
+      lv1 = _mm_srai_epi16(lv1, 4);
+      lv1 = _mm_add_epi8(lv1, seven);
+      lv1 = _mm_and_si128(lv1, bitmask_upper);
+      __m128i lv2 = _mm_slli_epi16(values_128i, 12);
+      lv2 = _mm_srai_epi16(lv2, 12);
+      lv2 = _mm_add_epi8(lv2, seven);
+      lv2 = _mm_and_si128(lv2, bitmask_lower);
+
+      lv1 = _mm_or_si128(lv1, lv2);
+      __m128i u = _mm_or_si128(_mm_slli_epi16(uv1, 4),
+                               _mm_unpackhi_epi64(uv1, _mm_setzero_si128()));
+      __m128i l = _mm_or_si128(_mm_slli_epi16(lv1, 4),
+                               _mm_unpackhi_epi64(lv1, _mm_setzero_si128()));
+      __m128i v = _mm_unpacklo_epi8(l, u);
+      _mm_store_si128((__m128i*)(box + k), v);
+      k += 16;
+    }
+    // Handle remaining values -- if greater than or equal to
+    // 16 values remaining, do the shuffle.
+    if (i < real_src_depth) {
+      int remaining = 8;
+      remaining =
+          remaining < (real_src_depth - i) ? remaining : real_src_depth - i;
+      for (int j = 0; j < remaining; j++) {
+        const int8_t v1 = (int8_t)src_data[i + j];
+        int8_t uv1 = upper(v1);
+        int8_t lv1 = lower(v1);
+        int8_t uv2 = 0;
+        int8_t lv2 = 0;
+        if ((i + j + 8) < real_src_depth) {
+          const int8_t v2 = (int8_t)src_data[i + j + 8];
+          uv2 = upper(v2);
+          lv2 = lower(v2);
+        }
+        box[k] = merge(lv1, lv2);
+        box[k + 1] = merge(uv1, uv2);
+        k += 2;
+      }
+    }
+    box += real_depth;
+    src_data += real_src_cols;
+  }
+}
+
+void SsePrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                int layout_cols, int src_rows, int src_cols, int width,
+                int depth) {
+  size_t size = static_cast<size_t>(layout_rows) * layout_cols / 2;
+  memset(dest, static_cast<uint8_t>(119), sizeof(uint8_t) * size);
+  int outer_cols = layout_cols / depth;
+  int outer_rows = layout_rows / width;
+  int inner_cols = depth;
+  int inner_rows = width;
+  for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
+    for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+      const size_t cluster_index =
+          static_cast<size_t>(outer_row) * outer_cols + outer_col;
+      const int real_depth = inner_cols / 2;
+      uint8_t* box = dest + cluster_index * real_depth * inner_rows;
+      SsePackInner(tensor, box, src_rows, src_cols, outer_row, outer_col,
+                   outer_rows, outer_cols, inner_rows, inner_cols);
+    }
+  }
+}
+
+void SseBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                int n_data, int8_t* quantized_data_ptr,
+                                float* scaling_factors, int width, int depth,
+                                int32_t* input_offsets) {
+  const int rows = n_batch;
+  const int cols = n_data;
+  // depth is always cols
+  const int layout_rows = (rows + (width - 1)) & ~(width - 1);
+  const int layout_cols = (cols + (depth - 1)) & ~(depth - 1);
+  const size_t size = static_cast<size_t>(layout_rows) * layout_cols;
+  int8_t* data = quantized_data_ptr;
+  memset(data, 0, sizeof(int8_t) * size);
+  memset(input_offsets, 0, sizeof(int32_t) * layout_rows);
+  const float* tensor_data = float_data_ptr;
+  // basically, we need to make a new 4D matrix
+  // [rows / width, cols / depth, width, depth] in depth-first
+  const int outer_cols = layout_cols / depth;
+  const int outer_rows = layout_rows / width;
+  float* scaling_factors_ptr = scaling_factors;
+  for (int outer_row = 0; outer_row < outer_rows; outer_row++) {
+    std::vector<float> scale(width);
+    const int row = width * outer_row;
+    scaling_factors_ptr = scaling_factors + row;
+    for (int w = 0; w < width; ++w) {
+      if ((row + w) >= rows) {
+        continue;
+      }
+      const float* start = tensor_data + static_cast<size_t>(row + w) * cols;
+      int c = 0;
+      float scale_denom = 0;
+      for (; c < cols; ++c) {
+        scale_denom = std::max(scale_denom, std::abs(*(start++)));
+      }
+      if (scale_denom == 0) {
+        scale_denom = 127.0;
+      }
+      scale[w] = 127.0 / scale_denom;
+      scaling_factors_ptr[w] = scale_denom / 127.0;
+    }
+    for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+      const int col = depth * outer_col;
+      const int src_width = std::min(width, rows - row);
+      const int src_depth = std::min(depth, cols - col);
+      const size_t cluster_index =
+          static_cast<size_t>(outer_row) * outer_cols + outer_col;
+      int8_t* box = data + cluster_index * depth * width;
+      for (int w = 0; w < src_width; ++w) {
+        const float* float_data =
+            tensor_data + static_cast<size_t>(row + w) * cols + col;
+        int d = 0;
+        for (; d < src_depth; ++d) {
+          int8_t q = static_cast<int8_t>(TfLiteRound(float_data[d] * scale[w]));
+          box[w * depth + d] = q;
+          input_offsets[row + w] += q;
+        }
+      }
+    }
+  }
+  for (int r = 0; r < layout_rows; ++r) {
+    input_offsets[r] = input_offsets[r] * zero_point_4bit;
+  }
+}
+
+void SseAssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                    const float* batch_scales,
+                                    const float* filter_scales,
+                                    const float* bias_ptr, float* output_ptr,
+                                    int output_depth, int batch_size) {
+  if (bias_ptr) {
+    for (int b = 0; b < batch_size; ++b) {
+      const float val = *input_offsets++ * *batch_scales++;
+      const float* filter_scales_ptr = filter_scales;
+      const float* bias_ptr_tmp = bias_ptr;
+      int i = 0;
+      for (; i < output_depth; i++) {
+        *output_ptr++ = (val * *filter_scales_ptr++) + *bias_ptr_tmp++;
+      }
+    }
+    return;
+  }
+  for (int b = 0; b < batch_size; ++b) {
+    const float val = *input_offsets++ * *batch_scales++;
+    const float* filter_scales_ptr = filter_scales;
+    int i = 0;
+    for (; i < output_depth; i++) {
+      *output_ptr++ = (val * *filter_scales_ptr++);
+    }
+  }
+}
+
+template <int Depth, int Width>
+void SseUnpack(float* output_ptr, const int32_t* dst, int batch_size,
+               int num_units, const float* scaling_factors,
+               const float* filter_scales, int dst_layout_rows,
+               int dst_layout_cols) {
+  if (Width == 1) {
+    const int outer_rows = dst_layout_rows / Width;
+    const int outer_cols = dst_layout_cols / Depth;
+    const int32_t* dst_ptr = dst;
+    int unit = 0;
+    for (int outer_col = 0; outer_col < outer_cols;
+         ++outer_col, unit += Depth) {
+      float* tmp_output_ptr = output_ptr + unit;
+      int len = num_units - unit < Depth ? num_units - unit : Depth;
+      int cond = len & ~3;
+      const float* scaling_factors_ptr = scaling_factors;
+      for (int outer_row = 0; outer_row < outer_rows; ++outer_row) {
+        const float scale = *scaling_factors_ptr;
+        const float* filter_scales_ptr = filter_scales + unit;
+        int i = 0;
+        for (; i < cond; i += 4) {
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+        }
+        for (; i < len; ++i) {
+          *(tmp_output_ptr++) += *(dst_ptr++) * scale * (*filter_scales_ptr++);
+        }
+        dst_ptr += (Depth - len);
+        scaling_factors_ptr += Width;
+        tmp_output_ptr += (num_units - len);
+      }
+    }
+    return;
+  }
+  const int outer_rows = dst_layout_rows / Width;
+  const int outer_cols = dst_layout_cols / Depth;
+  for (int outer_col = 0; outer_col < outer_cols; ++outer_col) {
+    const int unit = outer_col * Depth;
+    const int remaining_units = std::min(num_units - unit, Depth);
+    const int depth_offset = Depth - remaining_units;
+    const int width_offset = num_units - remaining_units;
+    int outer_row = 0;
+    for (; outer_row < outer_rows; ++outer_row) {
+      const int batch = outer_row * Width;
+      const int remaining_width = std::min(batch_size - batch, Width);
+      const size_t cluster_index =
+          static_cast<size_t>(outer_col) * outer_rows + outer_row;
+      const int32_t* dst_ptr = dst + cluster_index * Depth * Width;
+      float* tmp_output_ptr =
+          output_ptr + static_cast<size_t>(batch) * num_units + unit;
+      const float* scale = scaling_factors + batch;
+      int w = remaining_width;
+      for (; w > 0; --w, scale++) {
+        int d = remaining_units;
+        const float* filter_scales_ptr = filter_scales + unit;
+        for (; d > 0; --d) {
+          *tmp_output_ptr++ += *dst_ptr++ * (*scale) * (*filter_scales_ptr++);
+        }
+        dst_ptr += depth_offset;
+        tmp_output_ptr += width_offset;
+      }
+    }
+  }
+}
+
+inline __m128i DotProdInt8x4x4(__m128i acc_32x4, __m128i a_8x16,
+                               __m128i b_8x16) {
+  b_8x16 = _mm_sign_epi8(b_8x16, a_8x16);
+  a_8x16 = _mm_abs_epi8(a_8x16);
+  __m128i sumprod_16x8 = _mm_maddubs_epi16(a_8x16, b_8x16);
+  return _mm_add_epi32(acc_32x4,
+                       _mm_madd_epi16(sumprod_16x8, _mm_set1_epi16(1)));
+}
+
+inline __m128i ReduceInt32x4x4(__m128i a, __m128i b, __m128i c, __m128i d) {
+  // Assuming x = [x0, x1, x2, x3]
+  const __m128i a_b_lo_half = _mm_unpacklo_epi32(a, b);  // [a0, b0, a1, b1]
+  const __m128i a_b_hi_half = _mm_unpackhi_epi32(a, b);  // [a2, b2, a3, b3]
+  const __m128i a_plus_b =
+      _mm_add_epi32(a_b_lo_half, a_b_hi_half);  // [a0+a2, b0+b2, a1+a3, b1+b3]
+  const __m128i c_d_lo_half = _mm_unpacklo_epi32(c, d);  // [c0, d0, c1, d1]
+  const __m128i c_d_hi_half = _mm_unpackhi_epi32(c, d);  // [c2, d2, c3, d3]
+  const __m128i c_plus_d =
+      _mm_add_epi32(c_d_lo_half, c_d_hi_half);  // [c0+c2, d0+d2, c1+c3, d1+d3]
+  const __m128i all_evns =
+      _mm_unpacklo_epi64(a_plus_b, c_plus_d);  // [a02, b02, c02, d02]
+  const __m128i all_odds =
+      _mm_unpackhi_epi64(a_plus_b, c_plus_d);  // [a13, b13, c13, d13]
+  return _mm_add_epi32(all_evns, all_odds);    // [a0123, b0123, c0123, d0123]
+}
+
+template <int RowsLeft, int RowsRight, int Cols>
+void SseRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                  int lhs_layout_rows, int lhs_layout_cols, int rhs_layout_rows,
+                  int rhs_layout_cols, int dst_layout_rows,
+                  int dst_layout_cols) {
+  const int start_row = 0;
+  const int start_col = 0;
+  const int end_row = lhs_layout_rows;
+  const int end_col = rhs_layout_rows;
+  const int clamped_end_row = std::min(end_row, dst_layout_cols);
+  const int clamped_end_col = std::min(end_col, dst_layout_rows);
+  int32_t* elementPtr = dst;
+  const int outer_rows = (clamped_end_row + RowsLeft - 1) / RowsLeft;
+  const int outer_cols = (clamped_end_col + RowsRight - 1) / RowsRight;
+  const int depth = std::min(lhs_layout_cols / Cols, rhs_layout_cols / Cols);
+  const __m128i bitmask = _mm_set1_epi8(15);
+  const uintptr_t padding = 15;
+  std::vector<uint8_t> lhs_vec_data(
+      (static_cast<size_t>(RowsLeft) * lhs_layout_cols / 2) + padding);
+  uint8_t* lhs_vec = lhs_vec_data.data();
+  for (int i = start_row; i < outer_rows; ++i) {
+    size_t left_index = static_cast<size_t>(i) * RowsLeft * lhs_layout_cols / 2;
+    const uint8_t* lhs_val_data = lhs + left_index;
+    if (!is_aligned(lhs_val_data, 16)) {
+      size_t size = static_cast<size_t>(RowsLeft) * lhs_layout_cols / 2;
+      uintptr_t aligned =
+          (reinterpret_cast<uintptr_t>(lhs_vec) + padding) & ~(padding);
+      lhs_vec = reinterpret_cast<uint8_t*>(aligned);
+      memcpy(lhs_vec, lhs_val_data, size);
+      lhs_val_data = lhs_vec;
+    }
+    for (int j = start_col; j < outer_cols; ++j) {
+      const uint8_t* lhs_val = lhs_val_data;
+      size_t right_index = static_cast<size_t>(j) * RowsRight * rhs_layout_cols;
+      const int8_t* rhs_val = rhs + right_index;
+      __m128i accum[RowsRight * RowsLeft];
+      for (int m = 0; m < (RowsLeft * RowsRight); ++m) {
+        accum[m] = _mm_set1_epi8(0);
+      }
+      for (int k = 0; k < depth; ++k) {
+        __m128i lhs_row[RowsLeft];
+        for (int m = 0; m < RowsLeft; ++m) {
+          lhs_row[m] = _mm_load_si128((__m128i*)(lhs_val));
+          lhs_val += 16;
+        }
+        __m128i rhs[RowsRight][2];
+        for (int m = 0; m < RowsRight; ++m) {
+          for (int n = 0; n < 2; ++n) {
+            rhs[m][n] = _mm_loadu_si128((__m128i*)(rhs_val));
+            rhs_val += 16;
+          }
+        }
+        __m128i lhs_row_8[RowsLeft][2];
+        for (int m = 0; m < RowsLeft; ++m) {
+          lhs_row_8[m][0] = _mm_srli_epi16(lhs_row[m], 4);
+          lhs_row_8[m][1] = _mm_and_si128(lhs_row[m], bitmask);
+        }
+        for (int m = 0; m < RowsLeft; ++m) {
+          lhs_row_8[m][0] = _mm_and_si128(lhs_row_8[m][0], bitmask);
+        }
+        for (int i = 0; i < 2; ++i) {
+          for (int r = 0; r < RowsRight; ++r) {
+            for (int l = 0; l < RowsLeft; ++l) {
+              accum[r * RowsLeft + l] = DotProdInt8x4x4(
+                  accum[r * RowsLeft + l], lhs_row_8[l][i], rhs[r][i]);
+            }
+          }
+        }
+      }
+      for (int r = 0; r < RowsRight; ++r) {
+        __m128i sum =
+            ReduceInt32x4x4(accum[r * RowsLeft], accum[r * RowsLeft + 1],
+                            accum[r * RowsLeft + 2], accum[r * RowsLeft + 3]);
+        _mm_storeu_si128((__m128i*)elementPtr, sum);
+        elementPtr += 4;
+      }
+    }
+  }
+}
+// NOLINTEND
+
+template void SseUnpack<4, 1>(float* output_ptr, const int32_t* dst,
+                              int batch_size, int num_units,
+                              const float* scaling_factors,
+                              const float* filter_scales, int dst_layout_rows,
+                              int dst_layout_cols);
+
+template void SseUnpack<4, 2>(float* output_ptr, const int32_t* dst,
+                              int batch_size, int num_units,
+                              const float* scaling_factors,
+                              const float* filter_scales, int dst_layout_rows,
+                              int dst_layout_cols);
+
+template void SseUnpack<4, 4>(float* output_ptr, const int32_t* dst,
+                              int batch_size, int num_units,
+                              const float* scaling_factors,
+                              const float* filter_scales, int dst_layout_rows,
+                              int dst_layout_cols);
+
+template void SseRunKernel<4, 1, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                     int32_t* dst, int lhs_layout_rows,
+                                     int lhs_layout_cols, int rhs_layout_rows,
+                                     int rhs_layout_cols, int dst_layout_rows,
+                                     int dst_layout_cols);
+
+template void SseRunKernel<4, 2, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                     int32_t* dst, int lhs_layout_rows,
+                                     int lhs_layout_cols, int rhs_layout_rows,
+                                     int rhs_layout_cols, int dst_layout_rows,
+                                     int dst_layout_cols);
+
+template void SseRunKernel<4, 4, 32>(const uint8_t* lhs, const int8_t* rhs,
+                                     int32_t* dst, int lhs_layout_rows,
+                                     int lhs_layout_cols, int rhs_layout_rows,
+                                     int rhs_layout_cols, int dst_layout_rows,
+                                     int dst_layout_cols);
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // defined(FC_4BIT_SSE) && defined(__SSSE3__)

tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected.h

@@ -0,0 +1,125 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_H_
+#if defined(FC_4BIT_SSE) && defined(__SSSE3__)
+
+#include <stdint.h>
+
+#include "tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected_impl.h"
+
+namespace tflite {
+namespace optimized_4bit {
+
+// Maximum RowsRight compiled RunKernel implementations.
+inline int GetMaxSupportedRows() { return 4; }
+
+// Pack a 4bit inner_rows x inner_cols array from src.
+inline void PackInner(const int8_t* src, uint8_t* box, int src_rows,
+                      int src_cols, int outer_row, int outer_col,
+                      int outer_rows, int outer_cols, int inner_rows,
+                      int inner_cols) {
+  SsePackInner(src, box, src_rows, src_cols, outer_row, outer_col, outer_rows,
+               outer_cols, inner_rows, inner_cols);
+}
+
+// Prepack lhs matrix into dest.
+inline void Prepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                    int layout_cols, int src_rows, int src_cols, int width,
+                    int depth) {
+  SsePrepack(dest, tensor, layout_rows, layout_cols, src_rows, src_cols, width,
+             depth);
+}
+
+// Quantize input floats to 8bit and calculate sum of each column.
+inline void BatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                    int n_data, int8_t* quantized_data_ptr,
+                                    float* scaling_factors, int width,
+                                    int depth, int32_t* input_offsets) {
+  SseBatchQuantizeFloats4Bit(float_data_ptr, n_batch, n_data,
+                             quantized_data_ptr, scaling_factors, width, depth,
+                             input_offsets);
+}
+
+// Write bias + input offset * filter_scale to output_ptr.
+inline void AssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                        const float* batch_scales,
+                                        float* filter_scales,
+                                        const float* bias_ptr,
+                                        float* output_ptr, int output_depth,
+                                        int batch_size) {
+  SseAssignBiasAndComputeOffsets(input_offsets, batch_scales, filter_scales,
+                                 bias_ptr, output_ptr, output_depth,
+                                 batch_size);
+}
+
+// Add accumulated integer sums in dst to float output.
+template <int Depth, int Width>
+void Unpack(float* output_ptr, const int32_t* dst, int batch_size,
+            int num_units, const float* scaling_factors,
+            const float* filter_scales, int dst_layout_rows,
+            int dst_layout_cols) {
+  SseUnpack<Depth, Width>(output_ptr, dst, batch_size, num_units,
+                          scaling_factors, filter_scales, dst_layout_rows,
+                          dst_layout_cols);
+}
+
+// Compute sum of lhs * rhs columnwise.
+template <int RowsLeft, int RowsRight, int Cols>
+void RunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+               int lhs_layout_rows, int lhs_layout_cols, int rhs_layout_rows,
+               int rhs_layout_cols, int dst_layout_rows, int dst_layout_cols) {
+  SseRunKernel<RowsLeft, RowsRight, Cols>(
+      lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols, rhs_layout_rows,
+      rhs_layout_cols, dst_layout_rows, dst_layout_cols);
+}
+
+// Compute sum of lhs * rhs columnwise and write output to output_ptr.
+inline void RunAndUnpack(int rhs_width, const uint8_t* lhs, const int8_t* rhs,
+                         int32_t* dst, int output_depth, int batch_size,
+                         int lhs_layout_rows, int lhs_layout_cols,
+                         int rhs_layout_rows, int rhs_layout_cols,
+                         int dst_layout_rows, int dst_layout_cols,
+                         float* output_ptr, const float* scaling_factors,
+                         const float* filter_scales) {
+  if (rhs_width >= 4) {
+    SseRunKernel<4, 4, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
+                           rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
+                           dst_layout_cols);
+    SseUnpack<4, 4>(output_ptr, dst, batch_size, output_depth, scaling_factors,
+                    filter_scales, dst_layout_rows, dst_layout_cols);
+    return;
+  }
+  if (rhs_width >= 2) {
+    SseRunKernel<4, 2, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
+                           rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
+                           dst_layout_cols);
+    SseUnpack<4, 2>(output_ptr, dst, batch_size, output_depth, scaling_factors,
+                    filter_scales, dst_layout_rows, dst_layout_cols);
+    return;
+  }
+  SseRunKernel<4, 1, 32>(lhs, rhs, dst, lhs_layout_rows, lhs_layout_cols,
+                         rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
+                         dst_layout_cols);
+  SseUnpack<4, 1>(output_ptr, dst, batch_size, output_depth, scaling_factors,
+                  filter_scales, dst_layout_rows, dst_layout_cols);
+}
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // defined(FC_4BIT_SSE) && defined(__SSSE3__)
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_H_

tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected_impl.h

@@ -0,0 +1,64 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_IMPL_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_IMPL_H_
+#if defined(FC_4BIT_SSE) && defined(__SSSE3__)
+
+#include <stdint.h>
+
+#ifndef EIGEN_MAX_ALIGN_BYTES
+#define EIGEN_MAX_ALIGN_BYTES 64
+#endif
+
+namespace tflite {
+namespace optimized_4bit {
+
+void SsePackInner(const int8_t* src, uint8_t* box, int src_rows, int src_cols,
+                  int outer_row, int outer_col, int outer_rows, int outer_cols,
+                  int inner_rows, int inner_cols);
+
+void SsePrepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                int layout_cols, int src_rows, int src_cols, int width,
+                int depth);
+
+void SseBatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                int n_data, int8_t* quantized_data_ptr,
+                                float* scaling_factors, int width, int depth,
+                                int32_t* input_offsets);
+
+void SseAssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                    const float* batch_scales,
+                                    const float* filter_scales,
+                                    const float* bias_ptr, float* output_ptr,
+                                    int output_depth, int batch_size);
+
+template <int Depth, int Width>
+extern void SseUnpack(float* output_ptr, const int32_t* dst, int batch_size,
+                      int num_units, const float* scaling_factors,
+                      const float* filter_scales, int dst_layout_rows,
+                      int dst_layout_cols);
+
+template <int RowsLeft, int RowsRight, int Cols>
+extern void SseRunKernel(const uint8_t* lhs, const int8_t* rhs, int32_t* dst,
+                         int lhs_layout_rows, int lhs_layout_cols,
+                         int rhs_layout_rows, int rhs_layout_cols,
+                         int dst_layout_rows, int dst_layout_cols);
+
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // defined(FC_4BIT_SSE) && defined(__SSSE3__)
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_4BIT_SSE_FULLY_CONNECTED_IMPL_H_

tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h

@@ -0,0 +1,160 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_AVX2_QUANTIZATION_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_AVX2_QUANTIZATION_UTILS_H_
+#ifdef __AVX2__
+
+#include <immintrin.h>
+
+#include <limits>
+
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace tflite {
+namespace avx2_utils {
+
+static inline void mm_storeu_si64(void *dst, __m128i v) {
+#if (defined __clang__) || (defined _MSC_VER)
+  _mm_storeu_si64(dst, v);
+#else
+  // GCC 9 lacks support for _mm_storeu_si64.
+  *static_cast<std::int64_t *>(dst) = _mm_extract_epi64(v, 0);
+#endif
+}
+
+static inline __m256i mm256_blendv_epi32(const __m256i &a, const __m256i &b,
+                                         const __m256i &mask) {
+  __m256 result =
+      _mm256_blendv_ps(_mm256_castsi256_ps(a), _mm256_castsi256_ps(b),
+                       _mm256_castsi256_ps(mask));
+  return _mm256_castps_si256(result);
+}
+
+static inline __m256i rounding_right_shift(const __m256i &value,
+                                           int32_t right_shift) {
+  TFLITE_DCHECK_GT(right_shift, 0);
+  const int32_t one_shift_exp_minus1 = 1 << (right_shift - 1);
+  __m256i nudge = _mm256_set1_epi32(one_shift_exp_minus1);
+  const __m256i r_plus_nudge = _mm256_add_epi32(value, nudge);
+  const __m256i shifted_sum =
+      _mm256_srav_epi32(r_plus_nudge, _mm256_set1_epi32(right_shift));
+
+  // Identify overflow in each lane and create mask.
+  const __m256i mask_num_plus_nudge_overflow = _mm256_cmpgt_epi32(
+      value, _mm256_set1_epi32(0x7fffffff - one_shift_exp_minus1));
+  // Fill results with either (value + nudge) >> exponent or
+  // std::numeric_limits<std::int32_t>::max() in the case of overflow.
+  return mm256_blendv_epi32(
+      shifted_sum, _mm256_set1_epi32(std::numeric_limits<std::int32_t>::max()),
+      mask_num_plus_nudge_overflow);
+}
+
+static inline __m256i rounding_right_shift(const __m256i &value,
+                                           const __m256i right_shift) {
+  const __m256i zeros = _mm256_setzero_si256();
+  const __m256i mask_rightshift_gtz = _mm256_cmpgt_epi32(right_shift, zeros);
+  const __m256i one_shift_exp_minus1 =
+      _mm256_sllv_epi32(_mm256_set1_epi32(1),
+                        _mm256_sub_epi32(right_shift, _mm256_set1_epi32(1)));
+  __m256i nudge =
+      mm256_blendv_epi32(zeros, one_shift_exp_minus1, mask_rightshift_gtz);
+  const __m256i r_plus_nudge = _mm256_add_epi32(value, nudge);
+  const __m256i shifted_sum = _mm256_srav_epi32(r_plus_nudge, right_shift);
+
+  // Identify overflow in each lane and create mask.
+  const __m256i mask_num_plus_nudge_overflow = _mm256_cmpgt_epi32(
+      value, _mm256_sub_epi32(_mm256_set1_epi32(0x7fffffff), nudge));
+  // Fill results with either (value + nudge) >> exponent or
+  // std::numeric_limits<std::int32_t>::max() in the case of overflow.
+  return mm256_blendv_epi32(
+      shifted_sum, _mm256_set1_epi32(std::numeric_limits<std::int32_t>::max()),
+      mask_num_plus_nudge_overflow);
+}
+
+inline void CastInt32ToInt16AndStore(int16 *dst, const __m256i v) {
+  // As _mm256_cvtepi32_epi16 is not supported in AVX2, use the below repack.
+  // Select bytes 0, 1, 4, 5, 8, 9, 12, 13 within each lane, effectively
+  // truncating each 16-bit integer.
+  const __m256i repack_perm = _mm256_set1_epi64x(0x0d0c090805040100);
+  const __m256i shuffled_v = _mm256_shuffle_epi8(v, repack_perm);
+  mm_storeu_si64(dst, _mm256_extracti128_si256(shuffled_v, 0));
+  mm_storeu_si64(dst + 4, _mm256_extracti128_si256(shuffled_v, 1));
+}
+
+inline __m256i MultiplyByQuantizedMultiplier(const __m256i &value,
+                                             const int32_t multiplier,
+                                             const int32_t left_shift) {
+  const __m256i repack_perm = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
+  const __m256i shifted_value =
+      left_shift > 0 ? _mm256_sllv_epi32(value, _mm256_set1_epi32(left_shift))
+                     : value;
+
+  __m256i scaled_v_low = _mm256_mul_epi32(
+      _mm256_cvtepi32_epi64(_mm256_extracti128_si256(shifted_value, 0)),
+      _mm256_set1_epi64x(multiplier));
+  __m256i scaled_v_high = _mm256_mul_epi32(
+      _mm256_cvtepi32_epi64(_mm256_extracti128_si256(shifted_value, 1)),
+      _mm256_set1_epi64x(multiplier));
+
+  scaled_v_low = _mm256_srlv_epi64(scaled_v_low, _mm256_set1_epi64x(31));
+  scaled_v_high = _mm256_srlv_epi64(scaled_v_high, _mm256_set1_epi64x(31));
+  // As _mm256_cvtepi64_epi32 is not supported in AVX2, use the below permute.
+  scaled_v_high = _mm256_slli_epi64(scaled_v_high, 32);
+  __m256i result = _mm256_blend_epi32(scaled_v_low, scaled_v_high, 0xaa);
+  result = _mm256_permutevar8x32_epi32(result, repack_perm);
+  if (left_shift >= 0) {
+    return result;
+  }
+  return rounding_right_shift(result, -left_shift);
+}
+
+inline __m256i MultiplyByQuantizedMultiplier(const __m256i &value,
+                                             const __m256i multiplier,
+                                             const __m256i left_shift) {
+  const __m256i zero_vector = _mm256_setzero_si256();
+  const __m256i positive_left_shift = _mm256_max_epi32(left_shift, zero_vector);
+  const __m256i positive_right_shift =
+      _mm256_max_epi32(_mm256_sub_epi32(zero_vector, left_shift), zero_vector);
+
+  const __m256i repack_perm = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
+  const __m256i shifted_value = _mm256_sllv_epi32(value, positive_left_shift);
+
+  const __m256i multiplier_low =
+      _mm256_cvtepi32_epi64(_mm256_extracti128_si256(multiplier, 0));
+  const __m256i multiplier_high =
+      _mm256_cvtepi32_epi64(_mm256_extracti128_si256(multiplier, 1));
+
+  __m256i scaled_v_low = _mm256_mul_epi32(
+      _mm256_cvtepi32_epi64(_mm256_extracti128_si256(shifted_value, 0)),
+      multiplier_low);
+  __m256i scaled_v_high = _mm256_mul_epi32(
+      _mm256_cvtepi32_epi64(_mm256_extracti128_si256(shifted_value, 1)),
+      multiplier_high);
+
+  scaled_v_low = _mm256_srlv_epi64(scaled_v_low, _mm256_set1_epi64x(31));
+  scaled_v_high = _mm256_srlv_epi64(scaled_v_high, _mm256_set1_epi64x(31));
+  // As _mm256_cvtepi64_epi32 is not supported in AVX2, use the below permute.
+  scaled_v_high = _mm256_slli_epi64(scaled_v_high, 32);
+  __m256i result = _mm256_blend_epi32(scaled_v_low, scaled_v_high, 0xaa);
+  result = _mm256_permutevar8x32_epi32(result, repack_perm);
+
+  return rounding_right_shift(result, positive_right_shift);
+}
+}  // namespace avx2_utils
+}  // namespace tflite
+
+#endif  // __AVX2__
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_AVX2_QUANTIZATION_UTILS_H_

tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils_test.cc

@@ -0,0 +1,138 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h"
+
+#include <gmock/gmock.h>
+#include "tensorflow/lite/kernels/internal/common.h"
+
+#ifdef __AVX2__
+namespace tflite {
+namespace avx2_utils {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+__m256i FillVectorWithInt32(const std::vector<int32_t>& src) {
+  return _mm256_set_epi32(src[7], src[6], src[5], src[4], src[3], src[2],
+                          src[1], src[0]);
+}
+
+void CompareWithReferenceValue(std::vector<int32_t>& reference_values,
+                               const __m256i& result) {
+  // As _mm256_extract_epi32 only supports const int, which should be known
+  // at the comile time, it puts down 8 comparison instead of for-loop.
+  EXPECT_NEAR(reference_values[0], _mm256_extract_epi32(result, 0), 1);
+  EXPECT_NEAR(reference_values[1], _mm256_extract_epi32(result, 1), 1);
+  EXPECT_NEAR(reference_values[2], _mm256_extract_epi32(result, 2), 1);
+  EXPECT_NEAR(reference_values[3], _mm256_extract_epi32(result, 3), 1);
+  EXPECT_NEAR(reference_values[4], _mm256_extract_epi32(result, 4), 1);
+  EXPECT_NEAR(reference_values[5], _mm256_extract_epi32(result, 5), 1);
+  EXPECT_NEAR(reference_values[6], _mm256_extract_epi32(result, 6), 1);
+  EXPECT_NEAR(reference_values[7], _mm256_extract_epi32(result, 7), 1);
+}
+
+TEST(CastInt32ToInt16AndStoreTest, CastInt32ToInt16AndStoreTest) {
+  const std::vector<int16_t> src = {1, 2, 3, 4, 5, 6, 7, 8};
+  int16_t dst[8];
+  const __m256i src_vector = _mm256_set_epi32(src[7], src[6], src[5], src[4],
+                                              src[3], src[2], src[1], src[0]);
+  CastInt32ToInt16AndStore(dst, src_vector);
+  EXPECT_THAT(src, ElementsAreArray(dst));
+}
+
+TEST(MultiplyByQuantizedMultiplierTest, PositiveLeftShiftTest) {
+  std::vector<int32_t> values = {100, 200, 300, 400, 500, 600, 700, 800};
+  const __m256i src_vector = FillVectorWithInt32(values);
+  const int32_t left_shift = 20;
+  const int32_t multiplier = 12345;
+  const __m256i result =
+      MultiplyByQuantizedMultiplier(src_vector, multiplier, left_shift);
+
+  // Get the reference values.
+  for (int i = 0; i < values.size(); i++) {
+    values[i] = tflite::MultiplyByQuantizedMultiplier(values[i], multiplier,
+                                                      left_shift);
+  }
+
+  CompareWithReferenceValue(values, result);
+}
+
+TEST(MultiplyByQuantizedMultiplierTest, NegativeLeftShiftTest) {
+  std::vector<int32_t> values = {1000, 2000, 3000, 4000,
+                                 5000, 6000, 7000, 8000};
+  const __m256i src_vector = FillVectorWithInt32(values);
+  const int32_t left_shift = -3;
+  const int32_t multiplier = 1234567890;
+  const __m256i result =
+      MultiplyByQuantizedMultiplier(src_vector, multiplier, left_shift);
+
+  // Get the reference values.
+  for (int i = 0; i < values.size(); i++) {
+    values[i] = tflite::MultiplyByQuantizedMultiplier(values[i], multiplier,
+                                                      left_shift);
+  }
+
+  CompareWithReferenceValue(values, result);
+}
+
+TEST(MultiplyByQuantizedMultiplierTest, VectorPositiveLeftShiftTest) {
+  std::vector<int32_t> values = {100, 200, 300, 400, 500, 600, 700, 800};
+  const std::vector<int32_t> left_shifts = {20, 19, 18, 17, 16, 15, 14, 13};
+  const std::vector<int32_t> multipliers = {10000, 20000, 30000, 40000,
+                                            50000, 60000, 70000, 80000};
+  const __m256i src_vector = FillVectorWithInt32(values);
+  const __m256i left_shifts_vector = FillVectorWithInt32(left_shifts);
+  const __m256i multipliers_vector = FillVectorWithInt32(multipliers);
+
+  const __m256i result = MultiplyByQuantizedMultiplier(
+      src_vector, multipliers_vector, left_shifts_vector);
+
+  // Get the reference values.
+  for (int i = 0; i < values.size(); i++) {
+    values[i] = tflite::MultiplyByQuantizedMultiplier(values[i], multipliers[i],
+                                                      left_shifts[i]);
+  }
+
+  CompareWithReferenceValue(values, result);
+}
+
+TEST(MultiplyByQuantizedMultiplierTest, VectorNegativeLeftShiftTest) {
+  std::vector<int32_t> values = {1000, 2000, 3000, 4000,
+                                 5000, 6000, 7000, 8000};
+  const std::vector<int32_t> left_shifts = {-3, -4, -5, -6, -7, -8, -9, -10};
+  const std::vector<int32_t> multipliers = {1000000000, 1100000000, 1200000000,
+                                            1300000000, 1400000000, 1500000000,
+                                            1600000000, 1700000000};
+  const __m256i src_vector = FillVectorWithInt32(values);
+  const __m256i left_shifts_vector = FillVectorWithInt32(left_shifts);
+  const __m256i multipliers_vector = FillVectorWithInt32(multipliers);
+
+  const __m256i result = MultiplyByQuantizedMultiplier(
+      src_vector, multipliers_vector, left_shifts_vector);
+
+  // Get the reference values.
+  for (int i = 0; i < values.size(); i++) {
+    values[i] = tflite::MultiplyByQuantizedMultiplier(values[i], multipliers[i],
+                                                      left_shifts[i]);
+  }
+
+  CompareWithReferenceValue(values, result);
+}
+
+}  // namespace
+}  // namespace avx2_utils
+}  // namespace tflite
+
+#endif  //  __AVX2__

tensorflow/lite/kernels/internal/optimized/cpu_check.cc

@@ -0,0 +1,50 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+
+#if defined __linux__ && defined __aarch64__
+#include <sys/auxv.h>
+#endif
+
+namespace tflite {
+
+namespace {
+
+// The implementation of dotprod detection is copied from ruy's internal
+// function DetectDotprod().
+// At the moment it's only implemented on Linux ARM64. Consider syncing again
+// with ruy in the future to share improvements.
+#if defined __linux__ && defined __aarch64__
+bool DetectDotprodByLinuxAuxvMethod() {
+  // This is the value of HWCAP_ASIMDDP in sufficiently recent Linux headers,
+  // however we need to support building against older headers for the time
+  // being.
+  const int kLocalHwcapAsimddp = 1 << 20;
+  return getauxval(AT_HWCAP) & kLocalHwcapAsimddp;
+}
+#endif
+
+}  // namespace
+
+bool DetectArmNeonDotprod() {
+#if defined __linux__ && defined __aarch64__
+  return DetectDotprodByLinuxAuxvMethod();
+#else
+  return false;
+#endif
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/internal/optimized/cpu_check.h

@@ -0,0 +1,40 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_CPU_CHECK_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_CPU_CHECK_H_
+
+// This include is superfluous. However, it's been here for a while, and a
+// number of files have been relying on it to include neon_check.h for them.
+// This should be removed, but with a global run of presubmits to catch
+// any such issues. This requires running more than just TFLite presubmits.
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+
+namespace tflite {
+
+// On A64, returns true if the dotprod extension is present.
+// On other architectures, returns false unconditionally.
+bool DetectArmNeonDotprod();
+
+struct CpuFlags {
+  bool neon_dotprod = false;
+};
+
+inline void GetCpuFlags(CpuFlags* cpu_flags) {
+  cpu_flags->neon_dotprod = DetectArmNeonDotprod();
+}
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_CPU_CHECK_H_

tensorflow/lite/kernels/internal/optimized/depthwiseconv_3x3_filter_common.h

@@ -0,0 +1,590 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_3X3_FILTER_COMMON_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_3X3_FILTER_COMMON_H_
+
+#include <algorithm>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_ops {
+namespace depthwise_conv {
+
+constexpr int kDepthwiseConvScratchWorkspaceSize = 10 * 10 * 64;
+constexpr int kDepthwiseConvAdjustedBiasLimit = 64;
+// In cases such as depth multiplication, we want to be able to load data from
+// the workspace that is beyond the valid range. Macro-block sizes are adjusted
+// to allow for this.
+constexpr int kWorkspaceExtension = 16;
+
+#ifdef USE_NEON
+
+#ifndef __aarch64__
+inline int8x16_t vqtbl4q_s8(int8x16x4_t a, int8x16_t b) {
+  const uint8x16_t mask = vtstq_s8(b, vdupq_n_s8(8));
+
+  // Delete bit 3 from the indices.
+  const int8x16_t high_bits = vshrq_n_s8(b, 4);
+  int8x16_t deleted_bit_3 = b;
+  deleted_bit_3 = vsliq_n_s8(deleted_bit_3, high_bits, 3);
+
+  int8x8x4_t repacked_data;
+
+  // Calculate for lower indices.
+  repacked_data.val[0] = vget_low_s8(a.val[0]);
+  repacked_data.val[1] = vget_low_s8(a.val[1]);
+  repacked_data.val[2] = vget_low_s8(a.val[2]);
+  repacked_data.val[3] = vget_low_s8(a.val[3]);
+  const int8x16_t output_for_lower =
+      vcombine_s8(vtbl4_s8(repacked_data, vget_low_s8(deleted_bit_3)),
+                  vtbl4_s8(repacked_data, vget_high_s8(deleted_bit_3)));
+
+  // Calculate for high indices.
+  repacked_data.val[0] = vget_high_s8(a.val[0]);
+  repacked_data.val[1] = vget_high_s8(a.val[1]);
+  repacked_data.val[2] = vget_high_s8(a.val[2]);
+  repacked_data.val[3] = vget_high_s8(a.val[3]);
+  const int8x16_t output_for_higher =
+      vcombine_s8(vtbl4_s8(repacked_data, vget_low_s8(deleted_bit_3)),
+                  vtbl4_s8(repacked_data, vget_high_s8(deleted_bit_3)));
+
+  // Merge.
+  int8x16_t output = vbslq_s8(mask, output_for_higher, output_for_lower);
+  return output;
+}
+#endif  // !__aarch64__
+
+// Convenience-compatibility functions.
+// Compatibility: Intrinsics reflect a mixture of older and newer ARM
+//     instructions. This actually results in ZIP1 / ZIP2 asm instructions, but
+//     one intrinsic is provided. Also older instructions operated in place,
+//     and it seems more defensive to assume that some versions of intrinsics
+//     might reflect this
+// Convenience: Callers in these kernels want both ZIP1 and ZIP2, and we do not
+//     want the calling code to get cluttered with unpacking int8x16x2_t.
+inline void vzipq_s8_in_place(int8x16_t* a, int8x16_t* b) {
+  int8x16x2_t r8x16;
+  r8x16 = vzipq_s8(*a, *b);
+  *a = r8x16.val[0];
+  *b = r8x16.val[1];
+}
+
+inline void vzipq_s8x2_in_place(int8x16_t* a, int8x16_t* b) {
+  int16x8x2_t r16x8;
+  r16x8 = vzipq_s16(vreinterpretq_s16_s8(*a), vreinterpretq_s16_s8(*b));
+  *a = vreinterpretq_s8_s16(r16x8.val[0]);
+  *b = vreinterpretq_s8_s16(r16x8.val[1]);
+}
+
+// Similar rationale to the zip-in_place functions, but callers only actually
+// need the TRN1 asm instruction result.
+inline void vtrn1_s8x2_in_place(int8x16_t* a, int8x16_t* b) {
+  int16x8x2_t r16x8;
+  r16x8 = vtrnq_s16(vreinterpretq_s16_s8(*a), vreinterpretq_s16_s8(*b));
+  *a = vreinterpretq_s8_s16(r16x8.val[0]);
+}
+
+// Similar rationale to the zip-in_place functions, but callers only actually
+// need the ZIP1 or ZIP2 asm instruction results.
+inline int8x16_t vzip1q_s8(int8x16_t a, int8x16_t b) {
+  return vzipq_s8(a, b).val[0];
+}
+inline int8x16_t vzip2q_s8(int8x16_t a, int8x16_t b) {
+  return vzipq_s8(a, b).val[1];
+}
+
+inline void biregister_rotate_8(int8x16_t* left, int8x16_t* right) {
+  *left = vreinterpretq_s8_u32(vshrq_n_u32(vreinterpretq_u32_s8(*left), 8));
+  *left = vreinterpretq_s8_u32(vsliq_n_u32(vreinterpretq_u32_s8(*left),
+                                           vreinterpretq_u32_s8(*right), 24));
+  *right = vreinterpretq_s8_u32(vshrq_n_u32(vreinterpretq_u32_s8(*right), 8));
+}
+
+#ifndef __aarch64__
+inline int32x4_t vpaddq_s32(int32x4_t a, int32x4_t b) {
+  int32x2_t a0 = vpadd_s32(vget_low_s32(a), vget_high_s32(a));
+  int32x2_t b0 = vpadd_s32(vget_low_s32(b), vget_high_s32(b));
+  return vcombine_s32(a0, b0);
+}
+#endif  // !__aarch64__
+
+#ifdef __ARM_FEATURE_DOTPROD
+// The vdotq_lane_s32 takes int8x8t for the rhs parameter, whereas the actual
+// instruction selects from between 4 32-bit (4x8-bit packed) sub-registers, an
+// unusual interpretation of "lane".
+inline int32x4_t vdotq_four_lane_s32(int32x4_t acc, int8x16_t lhs,
+                                     int8x16_t rhs, const int lane) {
+  switch (lane) {
+    case 0:
+      return vdotq_lane_s32(acc, lhs, vget_low_s8(rhs), 0);
+    case 1:
+      return vdotq_lane_s32(acc, lhs, vget_low_s8(rhs), 1);
+    case 2:
+      return vdotq_lane_s32(acc, lhs, vget_high_s8(rhs), 0);
+    case 3:
+    default:
+      return vdotq_lane_s32(acc, lhs, vget_high_s8(rhs), 1);
+  }
+}
+
+#else
+
+inline int32x4_t vdotq_s32(int32x4_t acc, int8x16_t lhs, int8x16_t rhs) {
+  int32x4_t sum0 = vpaddlq_s16(vmull_s8(vget_low_s8(lhs), vget_low_s8(rhs)));
+  int32x4_t sum1 = vpaddlq_s16(vmull_s8(vget_high_s8(lhs), vget_high_s8(rhs)));
+  int32x4_t sum = vpaddq_s32(sum0, sum1);
+  return vaddq_s32(acc, sum);
+}
+
+inline int32x4_t vdotq_four_lane_s32(int32x4_t acc, int8x16_t lhs,
+                                     int8x16_t rhs, int lane) {
+  int8x8_t lane_rhs;
+  if (lane == 0) {
+    lane_rhs = vreinterpret_s8_s32(
+        vdup_lane_s32(vreinterpret_s32_s8(vget_low_s8(rhs)), 0));
+  } else if (lane == 1) {
+    lane_rhs = vreinterpret_s8_s32(
+        vdup_lane_s32(vreinterpret_s32_s8(vget_low_s8(rhs)), 1));
+  } else if (lane == 2) {
+    lane_rhs = vreinterpret_s8_s32(
+        vdup_lane_s32(vreinterpret_s32_s8(vget_high_s8(rhs)), 0));
+  } else {
+    lane_rhs = vreinterpret_s8_s32(
+        vdup_lane_s32(vreinterpret_s32_s8(vget_high_s8(rhs)), 1));
+  }
+  int32x4_t sum0 = vpaddlq_s16(vmull_s8(vget_low_s8(lhs), lane_rhs));
+  int32x4_t sum1 = vpaddlq_s16(vmull_s8(vget_high_s8(lhs), lane_rhs));
+  int32x4_t sum = vpaddq_s32(sum0, sum1);
+  return vaddq_s32(acc, sum);
+}
+
+#endif  // !__ARM_FEATURE_DOTPROD
+#endif  // ARM NEON
+
+//  This structure is typically used for reducing the magnitude of outputs, and
+//  the historical name reflects that.
+template <DepthwiseConvOutputRounding output_rounding>
+struct DivideByPOT {};
+
+template <>
+struct DivideByPOT<DepthwiseConvOutputRounding::kAwayFromZero> {
+  template <typename IntegerType>
+  static inline IntegerType Run(IntegerType x, int exponent) {
+    return RoundingDivideByPOT(x, exponent);
+  }
+  // Mult versions use the exponents directly, rather than negated.
+  template <typename IntegerType>
+  static inline IntegerType RunMult(IntegerType x, int exponent) {
+    return RoundingDivideByPOT(x, -exponent);
+  }
+};
+
+#ifdef USE_NEON
+template <>
+struct DivideByPOT<DepthwiseConvOutputRounding::kUpward> {
+  template <typename IntegerType>
+  static inline IntegerType Run(IntegerType x, int exponent) {
+    return vqrshlq_s32(x, vdupq_n_s32(static_cast<int32_t>(-exponent)));
+  }
+  template <typename IntegerType>
+  static inline IntegerType RunMult(IntegerType x, IntegerType exponent) {
+    return vqrshlq_s32(x, exponent);
+  }
+  template <typename IntegerType>
+  static inline IntegerType RunMult(IntegerType x, int exponent) {
+    return vqrshlq_s32(x, vdupq_n_s32(static_cast<int32_t>(exponent)));
+  }
+};
+#endif  // ARM NEON
+
+// See CategorizeDotProductKernel for definitive taxonomy.
+enum class DotProduct3x3KernelType {
+  kNone = 0,  // Parameter combination is not supported for dot product kernels.
+  kPlain,
+  kWithDepthMultiplicationStride1,
+  kWithDepthMultiplicationStride2,
+  kStride2,
+};
+
+enum class QuantizationType {
+  kNonPerChannelUint8 = 0,
+  kPerChannelInt8 = 1,
+};
+
+template <QuantizationType quantization_type>
+struct QuantizationTypeImpl {};
+
+template <>
+struct QuantizationTypeImpl<QuantizationType::kNonPerChannelUint8> {
+  typedef uint8_t ExternalType;
+
+  static constexpr int kIntSymmetricZeroPoint = 128;
+  static constexpr uint8_t kUint8SignBit = 0x80;
+};
+
+template <>
+struct QuantizationTypeImpl<QuantizationType::kPerChannelInt8> {
+  typedef int8_t ExternalType;
+
+  static constexpr int kIntSymmetricZeroPoint = 0;
+  static constexpr uint8_t kUint8SignBit = 0x0;
+};
+
+template <
+    QuantizationType quantization_type = QuantizationType::kNonPerChannelUint8>
+inline DotProduct3x3KernelType CategorizeDotProductKernel(
+    const RuntimeShape& input_shape, const RuntimeShape& filter_shape,
+    const RuntimeShape& output_shape, const DepthwiseParams& params,
+    const int32_t* output_shift_ptr = nullptr) {
+  constexpr int kSymmetricZeroPoint =
+      QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
+  const int padding =
+      std::max(params.padding_values.width, params.padding_values.height);
+  const int stride = params.stride_width;
+  const int32_t input_depth = input_shape.Dims(3);
+  const int32_t depth_multiplier = params.depth_multiplier;
+  const int32_t filter_height = filter_shape.Dims(1);
+  const int32_t filter_width = filter_shape.Dims(2);
+
+  bool supported = stride == params.stride_height && stride <= 2 &&
+                   padding <= 1 && filter_width == 3 && filter_height == 3 &&
+                   params.dilation_width_factor == 1 &&
+                   params.dilation_height_factor == 1 &&
+                   (((input_depth % 8) == 0 && depth_multiplier == 1) ||
+                    (input_depth == 1 && depth_multiplier > 1));
+
+  if (!supported) {
+    return DotProduct3x3KernelType::kNone;
+  }
+
+  if (params.weights_offset != -kSymmetricZeroPoint) {
+    return DotProduct3x3KernelType::kNone;
+  }
+
+  if (quantization_type == QuantizationType::kPerChannelInt8) {
+    if (output_shift_ptr == nullptr) {
+      return DotProduct3x3KernelType::kNone;
+    }
+  } else if (params.output_shift > 0) {
+    return DotProduct3x3KernelType::kNone;
+  }
+
+  if (params.depth_multiplier == 1) {
+    if (stride == 1) {
+      return DotProduct3x3KernelType::kPlain;
+    } else if (stride == 2) {
+      return DotProduct3x3KernelType::kStride2;
+    } else {
+      return DotProduct3x3KernelType::kNone;
+    }
+  } else {
+    if (stride == 1) {
+      return DotProduct3x3KernelType::kWithDepthMultiplicationStride1;
+    } else if (stride == 2) {
+      return DotProduct3x3KernelType::kWithDepthMultiplicationStride2;
+    } else {
+      return DotProduct3x3KernelType::kNone;
+    }
+  }
+}
+
+// Encapsulates constant parameters used in DepthwiseConv.
+// 64-bit is used for types that will be added to 64-bit addresses in asm.
+struct DepthwiseConvParams {
+  int64_t input_depth;
+  int64_t input_row_size;
+  int64_t output_depth;
+  int64_t output_row_size;
+  int64_t filter_row_size;
+  int32_t input_offset;
+  int32_t output_offset;
+  int32_t filter_offset;
+  int32_t output_multiplier;
+  int32_t output_activation_min;
+  int32_t output_activation_max;
+  int32_t output_right_shift;
+  int32_t input_width;
+  int32_t input_height;
+  int32_t stride_width;
+  int32_t stride_height;
+  int32_t output_width;
+  int32_t output_height;
+  float float_output_activation_min;
+  float float_output_activation_max;
+};
+
+// Encapsulates constant parameters used in DepthwiseConv using dot-product ops.
+// 64-bit is used for types that will be added to 64-bit addresses in asm.
+//
+// This structure is specifically designed for use in asm.
+struct DepthwiseConvDotProdParams {
+  int64_t input_depth;
+  int64_t output_depth;
+  int32_t stride;
+  int32_t bias_increment;
+  //
+  int32_t input_offset;
+  int32_t output_offset;
+  int32_t output_multiplier;
+  int32_t output_shift;
+  int32_t quantized_activation_min;
+  int32_t quantized_activation_max;
+  //
+  int32_t padding_left;
+  int32_t padding_right;
+  int32_t padding_top;
+  int32_t padding_bottom;
+  //
+  int32_t depth_micro_repeats;
+  //
+  int32_t width_macro_count;
+  int32_t input_width_overall_micro_repeats;
+  int32_t input_width_micro_repeats;
+  int32_t residual_width;
+  int32_t output_width_overall_micro_repeats;
+  int32_t output_width_micro_repeats;
+  int32_t output_residual_width;
+  int32_t workspace_width_micro_repeats;
+  //
+  int32_t height_macro_count;
+  int32_t inbound_block_height;
+  int32_t outbound_block_height;
+  int32_t input_height_stride;
+  int32_t output_height_stride;
+  int32_t workspace_height_stride;
+  //
+  int32_t four_over_stride;
+  //
+  const int32_t* output_multiplier_per_channel;
+  const int32_t* output_shift_per_channel;
+};
+
+template <DepthwiseConvOutputRounding output_rounding, int32_t kDepth,
+          int32_t kStrideWidth, int32_t kStrideHeight>
+struct DepthwiseConvWindow {};
+
+template <DepthwiseConvOutputRounding output_rounding, int32_t kDepth,
+          int32_t kStrideWidth, int32_t kStrideHeight>
+struct DepthwiseConvWindowPerChannel {};
+
+enum class EdgeType { kCorner, kHorizontal, kVertical, kCenter };
+
+template <DepthwiseConvOutputRounding output_rounding, EdgeType kEdgeType,
+          int kPadWidth, int kPadHeight>
+struct DepthwiseConvPartial {};
+
+template <DepthwiseConvOutputRounding output_rounding, EdgeType kEdgeType,
+          int kPadWidth, int kPadHeight>
+struct DepthwiseConvPartialPerChannel {};
+
+// Copies a subset of the input designated by |input_ptr| into |output_ptr|
+// with the specified output dimensions. Supports output depths of 64 only as
+// this is the cache line size.
+template <typename T>
+inline void ShuffleInput(const T* input_ptr, int64_t input_depth,
+                         int32_t input_width, int32_t input_height,
+                         int64_t output_depth, int32_t output_width,
+                         int32_t output_height, T* output_ptr) {
+  const int64_t input_row_size = input_depth * input_width;
+  for (int32_t y = 0; y < output_height; y++) {
+    const T* ptr = input_ptr;
+    for (int32_t x = 0; x < output_width; x++) {
+      memcpy(output_ptr, ptr, output_depth);
+      output_ptr += output_depth;
+      ptr += input_depth;
+    }
+    input_ptr += input_row_size;
+  }
+}
+
+// Calculates the input size depending on stride and output.
+inline int32_t get_shuffle_input_size(int32_t stride, int32_t output) {
+  return stride * (output - 1) + 3;
+}
+
+// Indicates the input and output dimensions used when shuffling input
+// activations.
+struct ShuffleParams {
+  int32_t output_width;
+  int32_t output_height;
+  int32_t input_width;
+  int32_t input_height;
+
+  ShuffleParams() = default;
+  ShuffleParams(int32_t output_width, int32_t output_height,
+                int32_t stride_width, int32_t stride_height)
+      : output_width(output_width),
+        output_height(output_height),
+        input_width(get_shuffle_input_size(stride_width, output_width)),
+        input_height(get_shuffle_input_size(stride_height, output_height)) {}
+};
+
+template <
+    QuantizationType quantization_type = QuantizationType::kNonPerChannelUint8>
+inline bool Fast3x3FilterKernelSupported(
+    const RuntimeShape& input_shape, const RuntimeShape& filter_shape,
+    int32_t stride_width, int32_t stride_height, int32_t dilation_width_factor,
+    int32_t dilation_height_factor, int32_t pad_width, int32_t pad_height,
+    int32_t depth_multiplier, const RuntimeShape& output_shape,
+    int32_t output_shift, const int32_t* output_shift_ptr = nullptr) {
+  const int32_t input_height = input_shape.Dims(1);
+  const int32_t input_width = input_shape.Dims(2);
+  const int32_t input_depth = input_shape.Dims(3);
+  const int32_t filter_height = filter_shape.Dims(1);
+  const int32_t filter_width = filter_shape.Dims(2);
+  const int32_t output_height = output_shape.Dims(1);
+  const int32_t output_width = output_shape.Dims(2);
+
+  bool supported =
+      filter_width == 3 && filter_height == 3 && depth_multiplier == 1 &&
+      (stride_width == 1 || stride_width == 2) &&
+      (stride_height == 1 || stride_height == 2) &&
+      (stride_width == stride_height) && (pad_width == 0 || pad_width == 1) &&
+      (pad_height == 0 || pad_height == 1) && (pad_width == pad_height) &&
+      (input_depth % 8) == 0 && (output_shift <= 0) &&
+      dilation_width_factor == 1 && dilation_height_factor == 1;
+
+  if (!supported) {
+    return false;
+  }
+
+  // Handle case where padding is zero but padding type is not kValid.
+  // This would require special boundary case handling that is not supported.
+
+  const int32_t out_x = output_width - 1;
+  const int32_t out_y = output_height - 1;
+
+  const int32_t in_x_origin = (out_x * stride_width) - pad_width;
+  const int32_t in_y_origin = (out_y * stride_height) - pad_height;
+
+  const int32_t in_x_end = in_x_origin + filter_width;
+  const int32_t in_y_end = in_y_origin + filter_height;
+
+  // Supported only if filter on the right and bottom boundary lies completely
+  // within the input if padding is zero.
+  if (pad_width == 0 && pad_height == 0) {
+    return in_x_end <= input_width && in_y_end <= input_height;
+  }
+
+  // Else if padding is 1, supported if bottom right filter lies +1 past input
+  // width and height.
+  supported = in_x_end <= (input_width + 1) && in_y_end <= (input_height + 1);
+
+  if (!supported) {
+    return false;
+  }
+
+  // Shapes with width 1 and height > 1, and vice versa are not supported yet.
+  if (input_width == 1) {
+    supported = (input_width == input_height);
+  } else if (input_height == 1) {
+    supported = (input_width == input_height);
+  }
+  return supported;
+}
+
+// Permute filter data, and adjust bias data to account for symmetric input
+// offset. Details are provided in the implementation of the
+// kUseCModel3x3DotProduct version.
+//
+// See the comments preceding DepthwiseConvDotProduct3x3() for further notes.
+template <DepthwiseConvImplementation implementation,
+          QuantizationType quantization_type>
+struct ProcessPerDepth {
+  // Routine is contained in a static Run() method. No default template version
+  // is supplied, so that all implementations are deliberate choices of template
+  // specialization.
+  //
+  // Note that the signature of the Run() method will be designed for the asm
+  // implementation rather than conforming to style.
+};
+
+// Copy a macro block of data from the input buffer into the workspace,
+// permuting data within each micro block.
+//
+// (a) Copy a macro block of data, padding as required along the width and
+//     height.
+// (b) Transpose the data within each micro block.
+//
+// See the comments preceding DepthwiseConvDotProduct3x3() for further notes.
+template <DepthwiseConvImplementation implementation,
+          QuantizationType quantization_type,
+          DepthwiseConvDepthMultiplication depth_multiplication,
+          int32_t max_padding>
+struct PackMacroBlock {
+  // Routine is contained in a static Run() method. No default template version
+  // is supplied, so that all implementations are deliberate choices of template
+  // specialization.
+  //
+  // Note that the signature of the Run() method will be designed for the asm
+  // implementation rather than conforming to style.
+};
+
+// Apply filter to macro block of input data and store results. Details are
+// provided in the implementation of the kUseCModel3x3DotProduct version.
+//
+// Parameters for repeats and residual sizes are in terms of outputs.
+//
+// See the comments preceding DepthwiseConvDotProduct3x3() for further notes.
+template <DepthwiseConvImplementation implementation,
+          QuantizationType quantization_type,
+          DepthwiseConvDepthMultiplication depth_multiplication, int32_t stride>
+struct KernelMacroBlock {
+  // Routine is contained in a static Run() method. No default template version
+  // is supplied, so that all implementations are deliberate choices of template
+  // specialization.
+  //
+  // Note that the signature of the Run() method will be designed for the asm
+  // implementation rather than conforming to style.
+};
+
+#if defined(__aarch64__)
+// Experiments suggest that a modest performance improvement is seen, at least
+// on 855 chipset big cores, with cache hints.
+template <typename T>
+inline void PreloadInputBlock(
+    const T* input_block_data,
+    const DepthwiseConvDotProdParams* function_params) {
+  // Preload.
+  const int input_width_micro_repeats =
+      function_params->input_width_micro_repeats;
+  const int block_height = function_params->inbound_block_height;
+  const int residual_width = function_params->residual_width;
+  const int input_height_stride = function_params->input_height_stride;
+  const int input_depth = function_params->input_depth;
+
+  const int total_width = 4 * input_width_micro_repeats + residual_width;
+  const T* row_ptr = input_block_data;
+  for (int k_height = 0; k_height < block_height; ++k_height) {
+    const T* ptr = row_ptr;
+    for (int j = 0; j < total_width; ++j) {
+      // Input data is loaded once.
+      optimized_ops_preload_l1_keep(ptr);
+      ptr += input_depth;
+    }
+    row_ptr += input_height_stride;
+  }
+}
+#endif  // __aarch64__
+
+}  // namespace depthwise_conv
+}  // namespace optimized_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_3X3_FILTER_COMMON_H_

tensorflow/lite/kernels/internal/optimized/depthwiseconv_multithread.h

@@ -0,0 +1,189 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_MULTITHREAD_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_MULTITHREAD_H_
+
+#include <algorithm>
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_float.h"
+#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8.h"
+
+namespace tflite {
+namespace optimized_ops {
+
+// TODO(luwa): add multithread to per-channel depthwise_conv
+// DepthwiseConv can run with multi threads on the dim specified by thread_dim.
+// Each thread processes output elements on dim, thread_dim, in the range of
+// [thread_start, thread_end).
+// For example, assume thread_start = 2, thread_end = 6, and thread_dim = 1, it
+// means that it will calculate DepthwiseConv for output_data[:, 2:5, :, :].
+template <typename T, typename TS>
+struct DepthwiseConvWorkerTask : cpu_backend_threadpool::Task {
+  DepthwiseConvWorkerTask(const DepthwiseParams& params,
+                          const RuntimeShape& input_shape, const T* input_data,
+                          const RuntimeShape& filter_shape,
+                          const T* filter_data, const RuntimeShape& bias_shape,
+                          const TS* bias_data, const RuntimeShape& output_shape,
+                          T* output_data, const CpuFlags& cpu_flags,
+                          int thread_start, int thread_end, int thread_dim)
+      : params_(params),
+        input_shape_(input_shape),
+        input_data_(input_data),
+        filter_shape_(filter_shape),
+        filter_data_(filter_data),
+        bias_shape_(bias_shape),
+        bias_data_(bias_data),
+        output_shape_(output_shape),
+        output_data_(output_data),
+        cpu_flags_(cpu_flags),
+        thread_start_(thread_start),
+        thread_end_(thread_end),
+        thread_dim_(thread_dim) {}
+
+  void Run() override {
+    DepthwiseConvImpl(params_, input_shape_, input_data_, filter_shape_,
+                      filter_data_, bias_shape_, bias_data_, output_shape_,
+                      output_data_, cpu_flags_, thread_start_, thread_end_,
+                      thread_dim_);
+  }
+
+ private:
+  const DepthwiseParams& params_;
+  const RuntimeShape& input_shape_;
+  const T* input_data_;
+  const RuntimeShape& filter_shape_;
+  const T* filter_data_;
+  const RuntimeShape& bias_shape_;
+  const TS* bias_data_;
+  const RuntimeShape& output_shape_;
+  T* output_data_;
+  const CpuFlags& cpu_flags_;
+  int thread_start_;
+  int thread_end_;
+  int thread_dim_;
+};
+
+inline int HowManyConvThreads(const RuntimeShape& output_shape,
+                              const RuntimeShape& filter_shape) {
+  // How many scalar multiplications are needed to make it worth using one
+  // more thread
+  static constexpr int kMinMulPerThread = 1 << 13;  // 8k
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int num_muls = output_shape.FlatSize() * filter_height * filter_width;
+  // Try to avoid real runtime divisions if possible by dividing by a
+  // compile-time constant.
+  int thread_count = std::max(1, num_muls / kMinMulPerThread);
+  return thread_count;
+}
+
+inline bool MultithreadAlongBatches(int thread_count, int batches) {
+  TFLITE_DCHECK_GE(thread_count, 2);
+  // If there are fewer batch entries than the number of threads we want to use,
+  // then better do intra-batch-entry multithreading.
+  if (batches < thread_count) {
+    return false;
+  }
+  // If there are at least 2 batch entries to be handed to each thread, then
+  // it's safe to proceed with batch-wise multithreading: each thread will have
+  // approximately equal number of batch entries to handle, so the load
+  // balancing will be reasonable, and the amount to which the load is not
+  // perfectly balanced will be offset by the inherent advantages of
+  // batch-wise multithreading (each thread is more efficient thanks to working
+  // on larger buffers with less boundary-handling overhead).
+  if (batches >= 2 * thread_count) {
+    return true;
+  }
+  // In the limit case were there are at least 1 but not much more than 1
+  // batch entries per thread, it may be a good idea to do per-batch
+  // multithreading if the number of batch entries is a multiple of the number
+  // of threads, so that each thread will have the same number of batch entries
+  // to process.
+  return ((batches % thread_count) == 0);
+}
+
+template <typename T, typename TS>
+inline void DepthwiseConv(const DepthwiseParams& params,
+                          const RuntimeShape& input_shape, const T* input_data,
+                          const RuntimeShape& filter_shape,
+                          const T* filter_data, const RuntimeShape& bias_shape,
+                          const TS* bias_data, const RuntimeShape& output_shape,
+                          T* output_data,
+                          CpuBackendContext* cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("DepthwiseConv");
+
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+  int thread_count = HowManyConvThreads(output_shape, filter_shape);
+  const int max_threads = cpu_backend_context->max_num_threads();
+  thread_count = std::max(1, std::min(thread_count, max_threads));
+#ifndef TFLITE_WITH_RUY
+  // Cap the number of threads to 2 for float path to avoid regression in
+  // performance (b/132294857).
+  if (std::is_floating_point<T>::value) {
+    thread_count = std::min(thread_count, 2);
+  }
+#endif
+
+  const int output_batches = output_shape.Dims(0);
+  const int output_height = output_shape.Dims(1);
+
+  CpuFlags cpu_flags;
+  GetCpuFlags(&cpu_flags);
+
+  if (thread_count == 1) {
+    DepthwiseConvImpl(params, input_shape, input_data, filter_shape,
+                      filter_data, bias_shape, bias_data, output_shape,
+                      output_data, cpu_flags, /*thread_start=*/0,
+                      /*thread_end=*/output_height, /*thread_dim=*/1);
+    return;
+  }
+
+  int thread_dim, thread_dim_size;
+  if (MultithreadAlongBatches(thread_count, output_batches)) {
+    thread_dim = 0;
+    thread_dim_size = output_batches;
+  } else {
+    thread_dim = 1;
+    thread_dim_size = output_height;
+  }
+
+  std::vector<DepthwiseConvWorkerTask<T, TS>> tasks;
+  // TODO(b/131746020) don't create new heap allocations every time.
+  // At least we make it a single heap allocation by using reserve().
+  tasks.reserve(thread_count);
+  int thread_start = 0;
+  for (int i = 0; i < thread_count; ++i) {
+    int thread_end =
+        thread_start + (thread_dim_size - thread_start) / (thread_count - i);
+    tasks.emplace_back(params, input_shape, input_data, filter_shape,
+                       filter_data, bias_shape, bias_data, output_shape,
+                       output_data, cpu_flags, thread_start, thread_end,
+                       thread_dim);
+    thread_start = thread_end;
+  }
+  cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
+                                  cpu_backend_context);
+}
+
+}  // namespace optimized_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_MULTITHREAD_H_

tensorflow/lite/kernels/internal/optimized/eigen_spatial_convolutions.h

@@ -0,0 +1,45 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_
+
+#define EIGEN_USE_CUSTOM_THREAD_POOL
+#define EIGEN_USE_THREADS
+
+#define Eigen EigenForTFLite
+
+// NOTE: We need to define our own tensor contraction dispatch method before
+// including the unsupported/Eigen/CXX11/Tensor header in order to reduce the
+// total number of kernel instantiations.
+// If you have trouble simply undef out the reducer macro e.g.
+// TFLITE_REDUCE_INSTANTIATIONS, but be aware this will make
+// the binary much bigger!
+#define TFLITE_REDUCE_INSTANTIATIONS
+#if defined(TFLITE_REDUCE_INSTANTIATIONS)
+// Override Eigen tensor contraction dispatch method.
+#define TENSOR_CONTRACTION_DISPATCH(METHOD, ALIGNMENT, ARGS)                  \
+  if (this->m_lhs_inner_dim_contiguous && this->m_rhs_inner_dim_contiguous && \
+      !this->m_rhs_inner_dim_reordered) {                                     \
+    METHOD<true, true, false, ALIGNMENT> ARGS;                                \
+  } else {                                                                    \
+    eigen_assert(false && "Unsupported contraction formats");                 \
+  }
+#endif
+
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h"
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_EIGEN_SPATIAL_CONVOLUTIONS_H_

tensorflow/lite/kernels/internal/optimized/fully_connected_4bit.h

@@ -0,0 +1,152 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_FULLY_CONNECTED_4BIT_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_FULLY_CONNECTED_4BIT_H_
+#include <stdint.h>
+
+#ifndef TFLITE_MMAP_DISABLED
+#include <sys/mman.h>
+#endif
+
+#include <cstdlib>
+#include <memory>
+
+#if defined(FC_4BIT_SSE) && defined(__SSSE3__)
+#include "tensorflow/lite/kernels/internal/optimized/4bit/sse_fully_connected.h"  // IWYU pragma: export
+#elif defined(FC_4BIT_NEON) && (defined(__ARM_NEON__) || defined(__ARM_NEON))
+#include "tensorflow/lite/kernels/internal/optimized/4bit/neon_fully_connected.h"  // IWYU pragma: export
+#else
+#include "tensorflow/lite/kernels/internal/optimized/4bit/fully_connected_reference.h"  // IWYU pragma: export
+#endif
+
+namespace tflite {
+namespace optimized_4bit {
+
+// Define 4-bit filter block size: 4x32 (64 bytes)
+constexpr int FilterWidth = 4;
+constexpr int FilterDepth = 32;
+constexpr int kDefaultAlignmentPadding = 63;
+
+struct Deleter {
+  explicit Deleter(size_t size = 0) : size(size) {}
+  void operator()(uint8_t* memory) {
+    if (!memory) {
+      return;
+    }
+#ifdef TFLITE_MMAP_DISABLED
+    delete[] memory;
+#else
+    munmap(memory, size);
+#endif
+  }
+  size_t size;
+};
+
+struct OpData4Bit {
+  int rows_right = 1;
+  int batch_size = 0;
+  bool needs_prepack = true;
+  uint8_t* prepacked_cache = nullptr;
+  std::unique_ptr<uint8_t[], Deleter> prepacked_cache_buffer;
+  size_t prepacked_cache_buffer_size = 0;
+
+  void AllocatePackedRegion(size_t required_size) {
+#ifdef TFLITE_MMAP_DISABLED
+    uint8_t* region = new uint8_t[required_size];
+    prepacked_cache_buffer =
+        std::unique_ptr<uint8_t[], Deleter>(region, Deleter());
+#else
+    uint8_t* region = reinterpret_cast<uint8_t*>(
+        mmap(nullptr, required_size, PROT_READ | PROT_WRITE,
+             MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
+    prepacked_cache_buffer =
+        std::unique_ptr<uint8_t[], Deleter>(region, Deleter(required_size));
+#ifdef MADV_MERGEABLE
+    madvise(region, required_size, MADV_MERGEABLE);
+#endif
+#endif
+    prepacked_cache = reinterpret_cast<uint8_t*>(
+        (reinterpret_cast<uintptr_t>(prepacked_cache_buffer.get()) +
+         kDefaultAlignmentPadding) &
+        ~kDefaultAlignmentPadding);
+    prepacked_cache_buffer_size = required_size;
+  }
+};
+
+namespace api {
+/* Prepack lhs matrix into dest.
+ * Transform tensor from (src_rows, src_cols) to
+ * (layout_rows / width, layout_cols / depth, width, depth) with possibly
+ * padding, and interleaving values along depth / 2 dimensions.
+ * dest should be aligned and allocated before prepack.
+ */
+inline void Prepack(uint8_t* dest, const int8_t* tensor, int layout_rows,
+                    int layout_cols, int src_rows, int src_cols, int width,
+                    int depth) {
+  optimized_4bit::Prepack(dest, tensor, layout_rows, layout_cols, src_rows,
+                          src_cols, width, depth);
+}
+
+/* Quantize input floats to 8bit and calculate sum of each column.
+ *  Data in float_data_ptr of shape (n_batch x n_data), is quantized and
+ * packed into (n_batch / width, n_data / depth, width, data) into
+ * quantized_data_ptr and input_offsets will contain the product of filter
+ * zero_point and input.
+ */
+inline void BatchQuantizeFloats4Bit(const float* float_data_ptr, int n_batch,
+                                    int n_data, int8_t* quantized_data_ptr,
+                                    float* scaling_factors, int width,
+                                    int depth, int32_t* input_offsets) {
+  optimized_4bit::BatchQuantizeFloats4Bit(float_data_ptr, n_batch, n_data,
+                                          quantized_data_ptr, scaling_factors,
+                                          width, depth, input_offsets);
+}
+
+/* Write bias + input offset * filter_scale to output_ptr.
+ * output_ptr of size (batch_size, output_depth) will have
+ * output_ptr[output_depth * b + o] =
+ *     bias_ptr[o] + input_offsets[b] * batch_scales[b] * filter_scale[o]
+ */
+inline void AssignBiasAndComputeOffsets(const int32_t* input_offsets,
+                                        const float* batch_scales,
+                                        float* filter_scales,
+                                        const float* bias_ptr,
+                                        float* output_ptr, int output_depth,
+                                        int batch_size) {
+  optimized_4bit::AssignBiasAndComputeOffsets(
+      input_offsets, batch_scales, filter_scales, bias_ptr, output_ptr,
+      output_depth, batch_size);
+}
+
+// Compute sum of lhs * rhs columnwise and write output to output_ptr.
+inline void RunAndUnpack(int rhs_width, const uint8_t* lhs, const int8_t* rhs,
+                         int32_t* dst, int output_depth, int batch_size,
+                         int lhs_layout_rows, int lhs_layout_cols,
+                         int rhs_layout_rows, int rhs_layout_cols,
+                         int dst_layout_rows, int dst_layout_cols,
+                         float* output_ptr, const float* scaling_factors,
+                         const float* filter_scales) {
+  optimized_4bit::RunAndUnpack(
+      rhs_width, lhs, rhs, dst, output_depth, batch_size, lhs_layout_rows,
+      lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, dst_layout_rows,
+      dst_layout_cols, output_ptr, scaling_factors, filter_scales);
+}
+
+}  // namespace api
+}  // namespace optimized_4bit
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_FULLY_CONNECTED_4BIT_H_

tensorflow/lite/kernels/internal/optimized/im2col_utils.h

@@ -0,0 +1,511 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_IM2COL_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_IM2COL_UTILS_H_
+
+#include <algorithm>
+#include <cassert>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_ops {
+
+template <typename T>
+inline void ExtractPatchIntoBufferColumn(
+    const RuntimeShape& input_shape, int w, int h, int b, int kheight,
+    int kwidth, int stride_width, int stride_height, int pad_width,
+    int pad_height, int in_width, int in_height, int in_depth,
+    int single_buffer_length, int buffer_id, const T* in_data,
+    T* conv_buffer_data, uint8_t zero_byte) {
+  ruy::profiler::ScopeLabel label("ExtractPatchIntoBufferColumn");
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  // This chunk of code reshapes all the inputs corresponding to
+  // output (b, h, w) to a column vector in conv_buffer(:, buffer_id).
+  const int kwidth_times_indepth = kwidth * in_depth;
+  const int inwidth_times_indepth = in_width * in_depth;
+  const int ih_ungated_start = h * stride_height - pad_height;
+  const int ih_ungated_end = (ih_ungated_start + kheight);
+  const int ih_end = std::min(ih_ungated_end, in_height);
+  const int iw_ungated_start = w * stride_width - pad_width;
+  const int iw_ungated_end = (iw_ungated_start + kwidth);
+  const int iw_end = std::min(iw_ungated_end, in_width);
+  // If the patch is off the edge of the input image, skip writing those rows
+  // and columns from the patch into the output array.
+  const int h_offset = std::max(0, -ih_ungated_start);
+  const int w_offset = std::max(0, -iw_ungated_start);
+  const int ih_start = std::max(0, ih_ungated_start);
+  const int iw_start = std::max(0, iw_ungated_start);
+  const int single_row_num =
+      std::max(0, std::min(kwidth - w_offset, in_width - iw_start)) * in_depth;
+  const int output_row_offset = (buffer_id * single_buffer_length);
+  int out_offset =
+      output_row_offset + (h_offset * kwidth + w_offset) * in_depth;
+  int in_offset = Offset(input_shape, b, ih_start, iw_start, 0);
+
+  // Express all of the calculations as padding around the input patch.
+  const int top_padding = h_offset;
+  const int bottom_padding = (ih_ungated_end - ih_end);
+  const int left_padding = w_offset;
+  const int right_padding = (iw_ungated_end - iw_end);
+  assert(single_row_num ==
+         ((kwidth - (left_padding + right_padding)) * in_depth));
+
+  // Write out zeroes to the elements representing the top rows of the input
+  // patch that are off the edge of the input image.
+  if (top_padding > 0) {
+    const int top_row_elements = (top_padding * kwidth * in_depth);
+    memset(conv_buffer_data + output_row_offset, zero_byte,
+           (top_row_elements * sizeof(T)));
+  }
+
+  // If the patch is on the interior of the input image horizontally, just copy
+  // over the rows sequentially, otherwise add zero padding at the start or end.
+  if ((left_padding == 0) && (right_padding == 0)) {
+    for (int ih = ih_start; ih < ih_end; ++ih) {
+      memcpy(conv_buffer_data + out_offset, in_data + in_offset,
+             single_row_num * sizeof(T));
+      out_offset += kwidth_times_indepth;
+      in_offset += inwidth_times_indepth;
+    }
+  } else {
+    for (int ih = ih_start; ih < ih_end; ++ih) {
+      if (left_padding > 0) {
+        const int left_start = (out_offset - (left_padding * in_depth));
+        memset(conv_buffer_data + left_start, zero_byte,
+               (left_padding * in_depth * sizeof(T)));
+      }
+      memcpy(conv_buffer_data + out_offset, in_data + in_offset,
+             single_row_num * sizeof(T));
+      if (right_padding > 0) {
+        const int right_start = (out_offset + single_row_num);
+        memset(conv_buffer_data + right_start, zero_byte,
+               (right_padding * in_depth * sizeof(T)));
+      }
+      out_offset += kwidth_times_indepth;
+      in_offset += inwidth_times_indepth;
+    }
+  }
+
+  // If the bottom of the patch falls off the input image, pad the values
+  // representing those input rows with zeroes.
+  if (bottom_padding > 0) {
+    const int bottom_row_elements = (bottom_padding * kwidth * in_depth);
+    const int bottom_start =
+        output_row_offset +
+        ((top_padding + (ih_end - ih_start)) * kwidth * in_depth);
+    memset(conv_buffer_data + bottom_start, zero_byte,
+           (bottom_row_elements * sizeof(T)));
+  }
+}
+
+// Supports per-batch zero_byte for per-batch asymmetric quantized inputs.
+template <typename T>
+void DilatedIm2col(const ConvParams& params, const RuntimeShape& input_shape,
+                   const T* input_data, const RuntimeShape& filter_shape,
+                   const RuntimeShape& output_shape, T* im2col_data,
+                   const int32_t* zero_bytes, const int zero_bytes_len) {
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int dilation_width_factor = params.dilation_width_factor;
+  const int dilation_height_factor = params.dilation_height_factor;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+  // For dilated convolution, the input pixels are not contiguous therefore we
+  // can't use the same optimizations as Im2Col(). Though note this code would
+  // work fine for the non-dilated case too (though likely a bit slower).
+  ruy::profiler::ScopeLabel label("DilatedIm2col");
+  TFLITE_DCHECK(dilation_width_factor != 1 || dilation_height_factor != 1);
+  TFLITE_DCHECK(im2col_data);
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int input_depth = MatchingDim(input_shape, 3, filter_shape, 3);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  MatchingDim(output_shape, 3, filter_shape, 0);
+
+  // Construct the MxN sized im2col matrix.
+  // The rows M, are sub-ordered B x H x W
+  const RuntimeShape row_shape({1, batches, output_height, output_width});
+  // The columns, N, are sub-ordered Kh x Kw x Din
+  const RuntimeShape col_shape({1, filter_height, filter_width, input_depth});
+  // Use dimensions M and N to construct dims for indexing directly into im2col
+  const RuntimeShape im2col_shape(
+      {1, 1, row_shape.FlatSize(), col_shape.FlatSize()});
+
+  // Loop through the output rows (B x H x W)
+  for (int batch = 0; batch < batches; ++batch) {
+    const T zero_byte = zero_bytes_len > 1 ? static_cast<T>(zero_bytes[batch])
+                                           : static_cast<T>(zero_bytes[0]);
+    for (int out_y = 0; out_y < output_height; ++out_y) {
+      for (int out_x = 0; out_x < output_width; ++out_x) {
+        // Each im2col row is an output pixel. Arrange the input data in this
+        // row in an order we can conveniently multiply with the filter data.
+        int row_offset = Offset(row_shape, 0, batch, out_y, out_x);
+        const int in_x_origin = (out_x * stride_width) - pad_width;
+        const int in_y_origin = (out_y * stride_height) - pad_height;
+        // Loop through all the pixels of the filter (Kh x Kw)
+        for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
+          const int in_y = in_y_origin + dilation_height_factor * filter_y;
+          if ((in_y >= 0) && (in_y < input_height)) {
+            // Filter row is within the input data.
+            // Loop through all the filter pixels in this row.
+            for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
+              const int in_x = in_x_origin + dilation_width_factor * filter_x;
+              int col_offset = Offset(col_shape, 0, filter_y, filter_x, 0);
+              T* dst = im2col_data +
+                       Offset(im2col_shape, 0, 0, row_offset, col_offset);
+              if ((in_x >= 0) && (in_x < input_width)) {
+                // Filter pixel is within the input, copy the input data.
+                T const* src =
+                    input_data + Offset(input_shape, batch, in_y, in_x, 0);
+                memcpy(dst, src, input_depth * sizeof(T));
+              } else {
+                // Filter pixel is outside the input, zero it out.
+                memset(dst, zero_byte, input_depth * sizeof(T));
+              }
+            }
+          } else {
+            // Filter row is outside the input, zero out the entire filter row.
+            int col_offset = Offset(col_shape, 0, filter_y, 0, 0);
+            T* dst = im2col_data +
+                     Offset(im2col_shape, 0, 0, row_offset, col_offset);
+            memset(dst, zero_byte, filter_width * input_depth * sizeof(T));
+          }
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+void DilatedIm2col(const ConvParams& params, uint8_t zero_byte,
+                   const RuntimeShape& input_shape, const T* input_data,
+                   const RuntimeShape& filter_shape,
+                   const RuntimeShape& output_shape, T* im2col_data) {
+  const int32_t zero_point = static_cast<int32_t>(zero_byte);
+  DilatedIm2col<T>(params, input_shape, input_data, filter_shape, output_shape,
+                   im2col_data, &zero_point, 1);
+}
+
+template <typename T>
+void Im2col(const ConvParams& params, int kheight, int kwidth,
+            uint8_t zero_byte, const RuntimeShape& input_shape,
+            const T* input_data, const RuntimeShape& output_shape,
+            T* output_data) {
+  ruy::profiler::ScopeLabel label("Im2col");
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_depth = input_shape.Dims(3);
+  const int input_width = input_shape.Dims(2);
+  const int input_height = input_shape.Dims(1);
+  const int output_depth = output_shape.Dims(3);
+  const int output_width = output_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+
+  int buffer_id = 0;
+  // Loop over the output nodes.
+  for (int b = 0; b < batches; ++b) {
+    for (int h = 0; h < output_height; ++h) {
+      for (int w = 0; w < output_width; ++w) {
+        ExtractPatchIntoBufferColumn(
+            input_shape, w, h, b, kheight, kwidth, stride_width, stride_height,
+            pad_width, pad_height, input_width, input_height, input_depth,
+            output_depth, buffer_id, input_data, output_data, zero_byte);
+        ++buffer_id;
+      }
+    }
+  }
+}
+
+template <typename T>
+void Im2col(const ConvParams& params, int kheight, int kwidth,
+            const int32_t* input_offsets, const int input_offsets_size,
+            const RuntimeShape& input_shape, const T* input_data,
+            const RuntimeShape& output_shape, T* output_data) {
+  ruy::profiler::ScopeLabel label("Im2col");
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  TFLITE_DCHECK_EQ(batches, input_offsets_size);
+  const int input_depth = input_shape.Dims(3);
+  const int input_width = input_shape.Dims(2);
+  const int input_height = input_shape.Dims(1);
+  const int output_depth = output_shape.Dims(3);
+  const int output_width = output_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+
+  int buffer_id = 0;
+  // Loop over the output nodes.
+  for (int b = 0; b < batches; ++b) {
+    uint8_t zero_byte = static_cast<uint8_t>(input_offsets[b]);
+    for (int h = 0; h < output_height; ++h) {
+      for (int w = 0; w < output_width; ++w) {
+        ExtractPatchIntoBufferColumn(
+            input_shape, w, h, b, kheight, kwidth, stride_width, stride_height,
+            pad_width, pad_height, input_width, input_height, input_depth,
+            output_depth, buffer_id, input_data, output_data, zero_byte);
+        ++buffer_id;
+      }
+    }
+  }
+}
+
+template <typename T>
+inline void ExtractPatchIntoBufferColumn3D(
+    int b, int d, int h, int w,                             // Output indexes.
+    int kdepth, int kheight, int kwidth,                    // Kernel params.
+    int stride_depth, int stride_height, int stride_width,  // Stride params.
+    int pad_depth, int pad_height, int pad_width,           // Padding params.
+    int in_depth, int in_height, int in_width, int in_channel,  // Input shape.
+    int output_row_offset, const T* in_data, T* conv_buffer_data,
+    uint8_t zero_byte) {
+  ruy::profiler::ScopeLabel label("ExtractPatchIntoBufferColumn3D");
+
+  // This chunk of code reshapes all the inputs corresponding to
+  // output (b, d, h, w) to a column vector in conv_buffer(:, buffer_id).
+  const int id_ungated_start = d * stride_depth - pad_depth;
+  const int id_start = std::max(0, id_ungated_start);
+  const int id_ungated_end = (id_ungated_start + kdepth);
+  const int id_end = std::min(id_ungated_end, in_depth);
+
+  const int ih_ungated_start = h * stride_height - pad_height;
+  const int ih_start = std::max(0, ih_ungated_start);
+  const int ih_ungated_end = (ih_ungated_start + kheight);
+  const int ih_end = std::min(ih_ungated_end, in_height);
+
+  const int iw_ungated_start = w * stride_width - pad_width;
+  const int iw_start = std::max(0, iw_ungated_start);
+  const int iw_ungated_end = (iw_ungated_start + kwidth);
+  const int iw_end = std::min(iw_ungated_end, in_width);
+
+  // Calculate the padding sizes.
+  const int d_padding_before = std::max(0, -id_ungated_start);
+  const int d_padding_after = (id_ungated_end - id_end);
+  const int h_padding_before = std::max(0, -ih_ungated_start);
+  const int h_padding_after = (ih_ungated_end - ih_end);
+  const int w_padding_before = std::max(0, -iw_ungated_start);
+  const int w_padding_after = (iw_ungated_end - iw_end);
+
+  // Memset if there are paddings in the depth dimension.
+  const int kd_stride_size = kheight * kwidth * in_channel;
+  const int id_stride_size = in_height * in_width * in_channel;
+
+  if (d_padding_before > 0) {
+    const int d_padding_before_elements = (d_padding_before * kd_stride_size);
+    memset(conv_buffer_data + output_row_offset, zero_byte,
+           (d_padding_before_elements * sizeof(T)));
+  }
+
+  if (d_padding_after > 0) {
+    const int d_padding_after_elements = (d_padding_after * kd_stride_size);
+    const int bottom_start =
+        output_row_offset + (kdepth - d_padding_after) * kd_stride_size;
+    memset(conv_buffer_data + bottom_start, zero_byte,
+           (d_padding_after_elements * sizeof(T)));
+  }
+
+  // If there are paddings in height or width dimension, menset the entire area
+  // to take advantage of sequential memory handling performance.
+  int out_offset = output_row_offset + d_padding_before * kd_stride_size;
+  if (h_padding_before > 0 || h_padding_after > 0 || w_padding_before > 0 ||
+      w_padding_after > 0) {
+    const int middle_elements = (id_end - id_start) * kd_stride_size;
+    memset(conv_buffer_data + out_offset, zero_byte,
+           (middle_elements * sizeof(T)));
+  }
+
+  // Copy the valid data from the input tensor.
+  const int kh_stride_size = kwidth * in_channel;
+  const int ih_stride_size = in_width * in_channel;
+  const int h_padding = h_padding_before + h_padding_after;
+  const int w_padding = w_padding_before + w_padding_after;
+  const int single_row_num = (kwidth - w_padding) * in_channel;
+  out_offset +=
+      h_padding_before * kh_stride_size + w_padding_before * in_channel;
+  const int in_offset_without_d = b * in_depth * id_stride_size +
+                                  ih_start * ih_stride_size +
+                                  iw_start * in_channel;
+  for (int id = id_start; id < id_end; ++id) {
+    int in_offset = in_offset_without_d + id * id_stride_size;
+    for (int ih = ih_start; ih < ih_end; ++ih) {
+      memcpy(conv_buffer_data + out_offset, in_data + in_offset,
+             single_row_num * sizeof(T));
+      out_offset += kh_stride_size;
+      in_offset += ih_stride_size;
+    }
+    out_offset += h_padding * kh_stride_size;
+  }
+}
+
+template <typename T>
+void Im2col3D(const Conv3DParams& params, int kdepth, int kheight, int kwidth,
+              uint8_t zero_byte, const RuntimeShape& input_shape,
+              const T* input_data, const RuntimeShape& im2col_shape,
+              T* im2col_data) {
+  ruy::profiler::ScopeLabel label("Im2col3D");
+  const int stride_depth = params.stride_depth;
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int pad_depth = params.padding_values.depth;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 5);
+  TFLITE_DCHECK_EQ(im2col_shape.DimensionsCount(), 5);
+
+  const int batches = MatchingDim(input_shape, 0, im2col_shape, 0);
+  const int input_depth = input_shape.Dims(1);
+  const int input_height = input_shape.Dims(2);
+  const int input_width = input_shape.Dims(3);
+  const int input_channel = input_shape.Dims(4);
+
+  const int output_depth = im2col_shape.Dims(1);
+  const int output_height = im2col_shape.Dims(2);
+  const int output_width = im2col_shape.Dims(3);
+  const int output_channel = im2col_shape.Dims(4);
+
+  int buffer_id = 0;
+  // Loop over the output nodes.
+  for (int b = 0; b < batches; ++b) {
+    for (int d = 0; d < output_depth; ++d) {
+      for (int h = 0; h < output_height; ++h) {
+        for (int w = 0; w < output_width; ++w) {
+          ExtractPatchIntoBufferColumn3D(
+              b, d, h, w, kdepth, kheight, kwidth, stride_depth, stride_height,
+              stride_width, pad_depth, pad_height, pad_width, input_depth,
+              input_height, input_width, input_channel, buffer_id, input_data,
+              im2col_data, zero_byte);
+          buffer_id += output_channel;
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+inline void DilatedIm2col3D(const Conv3DParams& params, int filter_depth,
+                            int filter_height, int filter_width,
+                            uint8_t zero_byte, const RuntimeShape& input_shape,
+                            const T* input_data,
+                            const RuntimeShape& im2col_shape, T* im2col_data) {
+  ruy::profiler::ScopeLabel label("DilatedIm2col3D");
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 5);
+  TFLITE_DCHECK_EQ(im2col_shape.DimensionsCount(), 5);
+
+  // Only NDHWC format is currently supported.
+  const int batches = MatchingDim(input_shape, 0, im2col_shape, 0);
+  const int input_channels = input_shape.Dims(4);
+  const int input_width = input_shape.Dims(3);
+  const int input_height = input_shape.Dims(2);
+  const int input_depth = input_shape.Dims(1);
+
+  const int output_width = im2col_shape.Dims(3);
+  const int output_height = im2col_shape.Dims(2);
+  const int output_depth = im2col_shape.Dims(1);
+
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  const int pad_depth = params.padding_values.depth;
+
+  // Construct the MxN sized im2col matrix.
+  // The rows M, are sub-ordered B x D x H x W.
+  const RuntimeShape row_shape(
+      {1, batches, output_depth, output_height, output_width});
+  // The columns, N, are sub-ordered Kd x Kh x Kw x Din.
+  const RuntimeShape col_shape(
+      {1, filter_depth, filter_height, filter_width, input_channels});
+  // Use dimensions M and N to construct dims for indexing directly into im2col.
+  const RuntimeShape im2col_reshaped(
+      {1, 1, row_shape.FlatSize(), col_shape.FlatSize()});
+
+  for (int batch = 0; batch < batches; ++batch) {
+    for (int out_d = 0; out_d < output_depth; ++out_d) {
+      const int in_d_origin = (out_d * params.stride_depth) - pad_depth;
+      for (int out_y = 0; out_y < output_height; ++out_y) {
+        const int in_y_origin = (out_y * params.stride_height) - pad_height;
+        for (int out_x = 0; out_x < output_width; ++out_x) {
+          const int in_x_origin = (out_x * params.stride_width) - pad_width;
+          const int row_offset =
+              Offset(row_shape, 0, batch, out_d, out_y, out_x);
+          for (int filter_d = 0; filter_d < filter_depth; ++filter_d) {
+            const int in_d = in_d_origin + params.dilation_depth * filter_d;
+            if ((in_d >= 0) && (in_d < input_depth)) {
+              for (int filter_y = 0; filter_y < filter_height; ++filter_y) {
+                const int in_y =
+                    in_y_origin + params.dilation_height * filter_y;
+                if ((in_y >= 0) && (in_y < input_height)) {
+                  for (int filter_x = 0; filter_x < filter_width; ++filter_x) {
+                    const int in_x =
+                        in_x_origin + params.dilation_width * filter_x;
+                    int col_offset =
+                        Offset(col_shape, 0, filter_d, filter_y, filter_x, 0);
+                    T* dst = im2col_data + Offset(im2col_reshaped, 0, 0,
+                                                  row_offset, col_offset);
+                    if ((in_x >= 0) && (in_x < input_width)) {
+                      // Filter pixel is within the input, copy the input data.
+                      T const* src = input_data + Offset(input_shape, batch,
+                                                         in_d, in_y, in_x, 0);
+                      memcpy(dst, src, input_depth * sizeof(T));
+                    } else {
+                      // Filter pixel is outside the input, zero it out.
+                      memset(dst, zero_byte, input_depth * sizeof(T));
+                    }
+                  }
+                } else {
+                  const int col_offset =
+                      Offset(col_shape, 0, filter_d, filter_y, 0, 0);
+                  T* dst = im2col_data + Offset(im2col_reshaped, 0, 0,
+                                                row_offset, col_offset);
+                  memset(dst, zero_byte,
+                         filter_width * input_depth * sizeof(T));
+                }
+              }
+            } else {
+              const int col_offset = Offset(col_shape, 0, filter_d, 0, 0, 0);
+              T* dst = im2col_data +
+                       Offset(im2col_reshaped, 0, 0, row_offset, col_offset);
+              memset(dst, zero_byte,
+                     filter_height * filter_width * input_depth * sizeof(T));
+            }
+          }
+        }
+      }
+    }
+  }
+}
+
+}  // namespace optimized_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_IM2COL_UTILS_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/README.md

@@ -0,0 +1,8 @@
+This directory contains optimized implementations for int8 fully integer kernels.
+
+Weight filters of convs are expected to be symmetric per-channel quantized in
+the range [-127, 127].
+Inputs/activations are expected to be asymmetric per-layer quantized in the
+range [-128, 127].
+
+THESE ARE EXPERIMENTAL AND PRONE TO CHANGE.

tensorflow/lite/kernels/internal/optimized/integer_ops/add.h

@@ -0,0 +1,513 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_ADD_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_ADD_H_
+
+#include <algorithm>
+
+#include "fixedpoint/fixedpoint.h"
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/add.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+// Element-wise add that can often be used for inner loop of broadcast add as
+// well as the non-broadcast add.
+inline void AddElementwiseInt8(int size, const ArithmeticParams& params,
+                               const int8* input1_data, const int8* input2_data,
+                               int8* output_data) {
+  ruy::profiler::ScopeLabel label("AddElementwiseInt8/8bit");
+  int i = 0;
+  TFLITE_DCHECK_GT(params.input1_offset, -256);
+  TFLITE_DCHECK_GT(params.input2_offset, -256);
+  TFLITE_DCHECK_LT(params.input1_offset, 256);
+  TFLITE_DCHECK_LT(params.input2_offset, 256);
+
+#ifdef USE_NEON
+  const int8x16_t output_activation_min_vector =
+      vdupq_n_s8(params.quantized_activation_min);
+  const int8x16_t output_activation_max_vector =
+      vdupq_n_s8(params.quantized_activation_max);
+
+  const int input1_left_shift = params.left_shift + params.input1_shift;
+  const int input2_left_shift = params.left_shift + params.input2_shift;
+  const int32x4_t input1_left_dup = vdupq_n_s32(input1_left_shift);
+  const int32x4_t input2_left_dup = vdupq_n_s32(input2_left_shift);
+
+  const int16x8_t input1_offset_dup = vdupq_n_s16(params.input1_offset);
+  const int16x8_t input2_offset_dup = vdupq_n_s16(params.input2_offset);
+
+  for (; i <= size - 16; i += 16) {
+    const int8x16_t input1_val_original = vld1q_s8(input1_data + i);
+    const int8x16_t input2_val_original = vld1q_s8(input2_data + i);
+
+    const int16x8_t input1_val_s16_high =
+        vmovl_s8(vget_high_s8(input1_val_original));
+    const int16x8_t input1_val_s16_low =
+        vmovl_s8(vget_low_s8(input1_val_original));
+
+    const int16x8_t input2_val_s16_high =
+        vmovl_s8(vget_high_s8(input2_val_original));
+    const int16x8_t input2_val_s16_low =
+        vmovl_s8(vget_low_s8(input2_val_original));
+    const int16x8_t input1_val_high =
+        vaddq_s16(input1_val_s16_high, input1_offset_dup);
+    const int16x8_t input2_val_high =
+        vaddq_s16(input2_val_s16_high, input2_offset_dup);
+    const int16x8_t input1_val_low =
+        vaddq_s16(input1_val_s16_low, input1_offset_dup);
+    const int16x8_t input2_val_low =
+        vaddq_s16(input2_val_s16_low, input2_offset_dup);
+    const int16x4_t input1_val_high_high = vget_high_s16(input1_val_high);
+    const int16x4_t input1_val_high_low = vget_low_s16(input1_val_high);
+    const int16x4_t input1_val_low_high = vget_high_s16(input1_val_low);
+    const int16x4_t input1_val_low_low = vget_low_s16(input1_val_low);
+    const int16x4_t input2_val_high_high = vget_high_s16(input2_val_high);
+    const int16x4_t input2_val_high_low = vget_low_s16(input2_val_high);
+    const int16x4_t input2_val_low_high = vget_high_s16(input2_val_low);
+    const int16x4_t input2_val_low_low = vget_low_s16(input2_val_low);
+    int32x4_t x111 = vmovl_s16(input1_val_low_low);
+    int32x4_t x112 = vmovl_s16(input1_val_low_high);
+    int32x4_t x121 = vmovl_s16(input1_val_high_low);
+    int32x4_t x122 = vmovl_s16(input1_val_high_high);
+    int32x4_t x211 = vmovl_s16(input2_val_low_low);
+    int32x4_t x212 = vmovl_s16(input2_val_low_high);
+    int32x4_t x221 = vmovl_s16(input2_val_high_low);
+    int32x4_t x222 = vmovl_s16(input2_val_high_high);
+
+    x111 = vshlq_s32(x111, input1_left_dup);
+    x112 = vshlq_s32(x112, input1_left_dup);
+    x121 = vshlq_s32(x121, input1_left_dup);
+    x122 = vshlq_s32(x122, input1_left_dup);
+    x211 = vshlq_s32(x211, input2_left_dup);
+    x212 = vshlq_s32(x212, input2_left_dup);
+    x221 = vshlq_s32(x221, input2_left_dup);
+    x222 = vshlq_s32(x222, input2_left_dup);
+    x111 = vqrdmulhq_n_s32(x111, params.input1_multiplier);
+    x112 = vqrdmulhq_n_s32(x112, params.input1_multiplier);
+    x121 = vqrdmulhq_n_s32(x121, params.input1_multiplier);
+    x122 = vqrdmulhq_n_s32(x122, params.input1_multiplier);
+    x211 = vqrdmulhq_n_s32(x211, params.input2_multiplier);
+    x212 = vqrdmulhq_n_s32(x212, params.input2_multiplier);
+    x221 = vqrdmulhq_n_s32(x221, params.input2_multiplier);
+    x222 = vqrdmulhq_n_s32(x222, params.input2_multiplier);
+    int32x4_t s11 = vaddq_s32(x111, x211);
+    int32x4_t s12 = vaddq_s32(x112, x212);
+    int32x4_t s21 = vaddq_s32(x121, x221);
+    int32x4_t s22 = vaddq_s32(x122, x222);
+    s11 = vqrdmulhq_n_s32(s11, params.output_multiplier);
+    s12 = vqrdmulhq_n_s32(s12, params.output_multiplier);
+    s21 = vqrdmulhq_n_s32(s21, params.output_multiplier);
+    s22 = vqrdmulhq_n_s32(s22, params.output_multiplier);
+    using gemmlowp::RoundingDivideByPOT;
+    s11 = RoundingDivideByPOT(s11, -params.output_shift);
+    s12 = RoundingDivideByPOT(s12, -params.output_shift);
+    s21 = RoundingDivideByPOT(s21, -params.output_shift);
+    s22 = RoundingDivideByPOT(s22, -params.output_shift);
+    const int16x4_t s11_narrowed = vmovn_s32(s11);
+    const int16x4_t s12_narrowed = vmovn_s32(s12);
+    const int16x4_t s21_narrowed = vmovn_s32(s21);
+    const int16x4_t s22_narrowed = vmovn_s32(s22);
+    const int16x8_t s1 = vaddq_s16(vcombine_s16(s11_narrowed, s12_narrowed),
+                                   vdupq_n_s16(params.output_offset));
+    const int16x8_t s2 = vaddq_s16(vcombine_s16(s21_narrowed, s22_narrowed),
+                                   vdupq_n_s16(params.output_offset));
+    const int8x16_t s = vcombine_s8(vqmovn_s16(s1), vqmovn_s16(s2));
+
+    const int8x16_t clamped =
+        vmaxq_s8(output_activation_min_vector,
+                 vminq_s8(output_activation_max_vector, s));
+    vst1q_s8(output_data + i, clamped);
+  }
+#endif  // NEON
+
+  for (; i < size; ++i) {
+    const int32 input1_val = params.input1_offset + input1_data[i];
+    const int32 input2_val = params.input2_offset + input2_data[i];
+    const int32 shifted_input1_val = input1_val * (1 << params.left_shift);
+    const int32 shifted_input2_val = input2_val * (1 << params.left_shift);
+    const int32 scaled_input1_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input1_val, params.input1_multiplier, params.input1_shift);
+    const int32 scaled_input2_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input2_val, params.input2_multiplier, params.input2_shift);
+    const int32 raw_sum = scaled_input1_val + scaled_input2_val;
+    const int32 raw_output =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            raw_sum, params.output_multiplier, params.output_shift) +
+        params.output_offset;
+    const int32 clamped_output =
+        std::min(params.quantized_activation_max,
+                 std::max(params.quantized_activation_min, raw_output));
+    output_data[i] = static_cast<int8>(clamped_output);
+  }
+}
+
+// Element-wise add is used for the non-broadcast add.
+inline void AddElementwiseInt16(int size, const ArithmeticParams& params,
+                                const int16* input1_data,
+                                const int16* input2_data, int16* output_data) {
+  ruy::profiler::ScopeLabel label("AddElementwiseInt16/16bit");
+  int i = 0;
+  TFLITE_DCHECK_GT(params.input1_offset, -32768);
+  TFLITE_DCHECK_GT(params.input2_offset, -32768);
+  TFLITE_DCHECK_LT(params.input1_offset, 32768);
+  TFLITE_DCHECK_LT(params.input2_offset, 32768);
+
+#ifdef __AVX2__
+  const int32_t input1_left_shift = params.left_shift + params.input1_shift;
+  const int32_t input2_left_shift = params.left_shift + params.input2_shift;
+  const __m256i input1_offset = _mm256_set1_epi32(params.input1_offset);
+  const __m256i input2_offset = _mm256_set1_epi32(params.input2_offset);
+  const __m256i output_offset = _mm256_set1_epi32(params.output_offset);
+  const __m256i clamp_max_v =
+      _mm256_set1_epi32(params.quantized_activation_max);
+  const __m256i clamp_min_v =
+      _mm256_set1_epi32(params.quantized_activation_min);
+
+  for (; i <= size - 16; i += 16) {
+    const __m256i input1_val_original =
+        _mm256_loadu_si256(reinterpret_cast<__m256i const*>(input1_data + i));
+    const __m256i input2_val_original =
+        _mm256_loadu_si256(reinterpret_cast<__m256i const*>(input2_data + i));
+
+    __m256i s11 =
+        _mm256_cvtepi16_epi32(_mm256_castsi256_si128(input1_val_original));
+    __m256i s12 =
+        _mm256_cvtepi16_epi32(_mm256_extracti128_si256(input1_val_original, 1));
+    __m256i s21 =
+        _mm256_cvtepi16_epi32(_mm256_castsi256_si128(input2_val_original));
+    __m256i s22 =
+        _mm256_cvtepi16_epi32(_mm256_extracti128_si256(input2_val_original, 1));
+
+    s11 = _mm256_add_epi32(s11, input1_offset);
+    s12 = _mm256_add_epi32(s12, input1_offset);
+    s21 = _mm256_add_epi32(s21, input2_offset);
+    s22 = _mm256_add_epi32(s22, input2_offset);
+
+    s11 = avx2_utils::MultiplyByQuantizedMultiplier(
+        s11, params.input1_multiplier, input1_left_shift);
+    s12 = avx2_utils::MultiplyByQuantizedMultiplier(
+        s12, params.input1_multiplier, input1_left_shift);
+    s21 = avx2_utils::MultiplyByQuantizedMultiplier(
+        s21, params.input2_multiplier, input2_left_shift);
+    s22 = avx2_utils::MultiplyByQuantizedMultiplier(
+        s22, params.input2_multiplier, input2_left_shift);
+
+    __m256i s1 = _mm256_add_epi32(s11, s21);
+    __m256i s2 = _mm256_add_epi32(s12, s22);
+
+    s1 = avx2_utils::MultiplyByQuantizedMultiplier(s1, params.output_multiplier,
+                                                   params.output_shift);
+    s2 = avx2_utils::MultiplyByQuantizedMultiplier(s2, params.output_multiplier,
+                                                   params.output_shift);
+
+    s1 = _mm256_add_epi32(s1, output_offset);
+    s2 = _mm256_add_epi32(s2, output_offset);
+
+    s1 = _mm256_min_epi32(s1, clamp_max_v);
+    s1 = _mm256_max_epi32(s1, clamp_min_v);
+    s2 = _mm256_min_epi32(s2, clamp_max_v);
+    s2 = _mm256_max_epi32(s2, clamp_min_v);
+
+    avx2_utils::CastInt32ToInt16AndStore(output_data + i, s1);
+    avx2_utils::CastInt32ToInt16AndStore(output_data + i + 8, s2);
+  }
+
+#elif defined(USE_NEON)
+  const int32x4_t output_activation_min_vector =
+      vdupq_n_s32(params.quantized_activation_min);
+  const int32x4_t output_activation_max_vector =
+      vdupq_n_s32(params.quantized_activation_max);
+
+  const int input1_left_shift = params.left_shift + params.input1_shift;
+  const int input2_left_shift = params.left_shift + params.input2_shift;
+  const int32x4_t input1_left_dup = vdupq_n_s32(input1_left_shift);
+  const int32x4_t input2_left_dup = vdupq_n_s32(input2_left_shift);
+
+  const int32x4_t input1_offset_dup = vdupq_n_s32(params.input1_offset);
+  const int32x4_t input2_offset_dup = vdupq_n_s32(params.input2_offset);
+  const int32x4_t output_offset_dup = vdupq_n_s32(params.output_offset);
+
+  // Use the size 16 batch as it is effective on pixel 3/4.
+  for (; i <= size - 16; i += 16) {
+    const int16x8_t input11_val_original = vld1q_s16(input1_data + i);
+    const int16x8_t input12_val_original = vld1q_s16(input2_data + i);
+    const int16x8_t input21_val_original = vld1q_s16(input1_data + 8 + i);
+    const int16x8_t input22_val_original = vld1q_s16(input2_data + 8 + i);
+
+    int32x4_t x111 = vmovl_s16(vget_low_s16(input11_val_original));
+    int32x4_t x112 = vmovl_s16(vget_high_s16(input11_val_original));
+    int32x4_t x121 = vmovl_s16(vget_low_s16(input12_val_original));
+    int32x4_t x122 = vmovl_s16(vget_high_s16(input12_val_original));
+    int32x4_t x211 = vmovl_s16(vget_low_s16(input21_val_original));
+    int32x4_t x212 = vmovl_s16(vget_high_s16(input21_val_original));
+    int32x4_t x221 = vmovl_s16(vget_low_s16(input22_val_original));
+    int32x4_t x222 = vmovl_s16(vget_high_s16(input22_val_original));
+
+    x111 = vaddq_s32(x111, input1_offset_dup);
+    x112 = vaddq_s32(x112, input1_offset_dup);
+    x121 = vaddq_s32(x121, input2_offset_dup);
+    x122 = vaddq_s32(x122, input2_offset_dup);
+    x211 = vaddq_s32(x211, input1_offset_dup);
+    x212 = vaddq_s32(x212, input1_offset_dup);
+    x221 = vaddq_s32(x221, input2_offset_dup);
+    x222 = vaddq_s32(x222, input2_offset_dup);
+
+    x111 = vshlq_s32(x111, input1_left_dup);
+    x112 = vshlq_s32(x112, input1_left_dup);
+    x121 = vshlq_s32(x121, input2_left_dup);
+    x122 = vshlq_s32(x122, input2_left_dup);
+    x211 = vshlq_s32(x211, input1_left_dup);
+    x212 = vshlq_s32(x212, input1_left_dup);
+    x221 = vshlq_s32(x221, input2_left_dup);
+    x222 = vshlq_s32(x222, input2_left_dup);
+
+    x111 = vqrdmulhq_n_s32(x111, params.input1_multiplier);
+    x112 = vqrdmulhq_n_s32(x112, params.input1_multiplier);
+    x121 = vqrdmulhq_n_s32(x121, params.input2_multiplier);
+    x122 = vqrdmulhq_n_s32(x122, params.input2_multiplier);
+    x211 = vqrdmulhq_n_s32(x211, params.input1_multiplier);
+    x212 = vqrdmulhq_n_s32(x212, params.input1_multiplier);
+    x221 = vqrdmulhq_n_s32(x221, params.input2_multiplier);
+    x222 = vqrdmulhq_n_s32(x222, params.input2_multiplier);
+
+    int32x4_t s11 = vaddq_s32(x111, x121);
+    int32x4_t s12 = vaddq_s32(x112, x122);
+    int32x4_t s21 = vaddq_s32(x211, x221);
+    int32x4_t s22 = vaddq_s32(x212, x222);
+
+    s11 = vqrdmulhq_n_s32(s11, params.output_multiplier);
+    s12 = vqrdmulhq_n_s32(s12, params.output_multiplier);
+    s21 = vqrdmulhq_n_s32(s21, params.output_multiplier);
+    s22 = vqrdmulhq_n_s32(s22, params.output_multiplier);
+
+    using gemmlowp::RoundingDivideByPOT;
+    s11 = RoundingDivideByPOT(s11, -params.output_shift);
+    s12 = RoundingDivideByPOT(s12, -params.output_shift);
+    s21 = RoundingDivideByPOT(s21, -params.output_shift);
+    s22 = RoundingDivideByPOT(s22, -params.output_shift);
+
+    s11 = vaddq_s32(s11, output_offset_dup);
+    s12 = vaddq_s32(s12, output_offset_dup);
+    s21 = vaddq_s32(s21, output_offset_dup);
+    s22 = vaddq_s32(s22, output_offset_dup);
+
+    s11 = vmaxq_s32(output_activation_min_vector,
+                    vminq_s32(output_activation_max_vector, s11));
+    s12 = vmaxq_s32(output_activation_min_vector,
+                    vminq_s32(output_activation_max_vector, s12));
+    s21 = vmaxq_s32(output_activation_min_vector,
+                    vminq_s32(output_activation_max_vector, s21));
+    s22 = vmaxq_s32(output_activation_min_vector,
+                    vminq_s32(output_activation_max_vector, s22));
+
+    const int16x8_t s1 = vcombine_s16(vqmovn_s32(s11), vqmovn_s32(s12));
+    const int16x8_t s2 = vcombine_s16(vqmovn_s32(s21), vqmovn_s32(s22));
+
+    vst1q_s16(output_data + i, s1);
+    vst1q_s16(output_data + 8 + i, s2);
+  }
+#endif  // NEON
+
+  for (; i < size; ++i) {
+    const int32 input1_val = params.input1_offset + input1_data[i];
+    const int32 input2_val = params.input2_offset + input2_data[i];
+    const int32 shifted_input1_val = input1_val * (1 << params.left_shift);
+    const int32 shifted_input2_val = input2_val * (1 << params.left_shift);
+    const int32 scaled_input1_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input1_val, params.input1_multiplier, params.input1_shift);
+    const int32 scaled_input2_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input2_val, params.input2_multiplier, params.input2_shift);
+    const int32 raw_sum = scaled_input1_val + scaled_input2_val;
+    const int32 raw_output =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            raw_sum, params.output_multiplier, params.output_shift) +
+        params.output_offset;
+    const int32 clamped_output =
+        std::min(params.quantized_activation_max,
+                 std::max(params.quantized_activation_min, raw_output));
+    output_data[i] = static_cast<int16>(clamped_output);
+  }
+}
+
+// Scalar-broadcast add that can be used for inner loop of more general
+// broadcast add, so that, for example, scalar-broadcast with batch will still
+// be fast.
+inline void AddScalarBroadcast(int size, const ArithmeticParams& params,
+                               int8 input1_data, const int8* input2_data,
+                               int8* output_data) {
+  using gemmlowp::RoundingDivideByPOT;
+
+  ruy::profiler::ScopeLabel label("AddScalarBroadcastInt8/8bit");
+  TFLITE_DCHECK_GT(params.input1_offset, -256);
+  TFLITE_DCHECK_GT(params.input2_offset, -256);
+  TFLITE_DCHECK_LT(params.input1_offset, 256);
+  TFLITE_DCHECK_LT(params.input2_offset, 256);
+
+  int i = 0;
+
+#ifdef USE_NEON
+  const int32x4_t left_shift_dup = vdupq_n_s32(params.left_shift);
+  const int8x8_t output_activation_min_vector =
+      vdup_n_s8(params.quantized_activation_min);
+  const int8x8_t output_activation_max_vector =
+      vdup_n_s8(params.quantized_activation_max);
+
+  // Process broadcast scalar.
+  const int8x8_t input1_val_original = vdup_n_s8(input1_data);
+  const int16x8_t input1_val_s16 = vmovl_s8(input1_val_original);
+  const int16x8_t input1_val =
+      vaddq_s16(input1_val_s16, vdupq_n_s16(params.input1_offset));
+  const int16x4_t input1_val_high = vget_high_s16(input1_val);
+  const int16x4_t input1_val_low = vget_low_s16(input1_val);
+  int32x4_t x11 = vmovl_s16(input1_val_low);
+  int32x4_t x12 = vmovl_s16(input1_val_high);
+  x11 = vshlq_s32(x11, left_shift_dup);
+  x12 = vshlq_s32(x12, left_shift_dup);
+  x11 = vqrdmulhq_n_s32(x11, params.input1_multiplier);
+  x12 = vqrdmulhq_n_s32(x12, params.input1_multiplier);
+  const int32x4_t input1_shift_dup = vdupq_n_s32(params.input1_shift);
+  x11 = vshlq_s32(x11, input1_shift_dup);
+  x12 = vshlq_s32(x12, input1_shift_dup);
+
+  for (; i <= size - 8; i += 8) {
+    const int8x8_t input2_val_original = vld1_s8(input2_data + i);
+    const int16x8_t input2_val_s16 = vmovl_s8(input2_val_original);
+    const int16x8_t input2_val =
+        vaddq_s16(input2_val_s16, vdupq_n_s16(params.input2_offset));
+    const int16x4_t input2_val_high = vget_high_s16(input2_val);
+    const int16x4_t input2_val_low = vget_low_s16(input2_val);
+    int32x4_t x21 = vmovl_s16(input2_val_low);
+    int32x4_t x22 = vmovl_s16(input2_val_high);
+    x21 = vshlq_s32(x21, left_shift_dup);
+    x22 = vshlq_s32(x22, left_shift_dup);
+    x21 = vqrdmulhq_n_s32(x21, params.input2_multiplier);
+    x22 = vqrdmulhq_n_s32(x22, params.input2_multiplier);
+    const int32x4_t input2_shift_dup = vdupq_n_s32(params.input2_shift);
+    x21 = vshlq_s32(x21, input2_shift_dup);
+    x22 = vshlq_s32(x22, input2_shift_dup);
+    int32x4_t s1 = vaddq_s32(x11, x21);
+    int32x4_t s2 = vaddq_s32(x12, x22);
+    s1 = vqrdmulhq_n_s32(s1, params.output_multiplier);
+    s2 = vqrdmulhq_n_s32(s2, params.output_multiplier);
+    s1 = RoundingDivideByPOT(s1, -params.output_shift);
+    s2 = RoundingDivideByPOT(s2, -params.output_shift);
+    const int16x4_t s1_narrowed = vmovn_s32(s1);
+    const int16x4_t s2_narrowed = vmovn_s32(s2);
+    const int16x8_t s = vaddq_s16(vcombine_s16(s1_narrowed, s2_narrowed),
+                                  vdupq_n_s16(params.output_offset));
+    const int8x8_t clamped =
+        vmax_s8(output_activation_min_vector,
+                vmin_s8(output_activation_max_vector, vqmovn_s16(s)));
+    vst1_s8(output_data + i, clamped);
+  }
+#endif  // NEON
+
+  if (i < size) {
+    // Process broadcast scalar.
+    const int32 input1_val = params.input1_offset + input1_data;
+    const int32 shifted_input1_val = input1_val * (1 << params.left_shift);
+    const int32 scaled_input1_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input1_val, params.input1_multiplier, params.input1_shift);
+
+    for (; i < size; ++i) {
+      const int32 input2_val = params.input2_offset + input2_data[i];
+      const int32 shifted_input2_val = input2_val * (1 << params.left_shift);
+      const int32 scaled_input2_val =
+          MultiplyByQuantizedMultiplierSmallerThanOneExp(
+              shifted_input2_val, params.input2_multiplier,
+              params.input2_shift);
+      const int32 raw_sum = scaled_input1_val + scaled_input2_val;
+      const int32 raw_output =
+          MultiplyByQuantizedMultiplierSmallerThanOneExp(
+              raw_sum, params.output_multiplier, params.output_shift) +
+          params.output_offset;
+      const int32 clamped_output =
+          std::min(params.quantized_activation_max,
+                   std::max(params.quantized_activation_min, raw_output));
+      output_data[i] = static_cast<int8>(clamped_output);
+    }
+  }
+}
+
+inline void Add(const ArithmeticParams& params,
+                const RuntimeShape& input1_shape, const int8* input1_data,
+                const RuntimeShape& input2_shape, const int8* input2_data,
+                const RuntimeShape& output_shape, int8* output_data) {
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  ruy::profiler::ScopeLabel label("AddInt8/8bit");
+  const int flat_size =
+      MatchingElementsSize(input1_shape, input2_shape, output_shape);
+
+  TFLITE_DCHECK_GT(params.input1_offset, -256);
+  TFLITE_DCHECK_GT(params.input2_offset, -256);
+  TFLITE_DCHECK_LT(params.input1_offset, 256);
+  TFLITE_DCHECK_LT(params.input2_offset, 256);
+  AddElementwiseInt8(flat_size, params, input1_data, input2_data, output_data);
+}
+
+inline void Add(const ArithmeticParams& params,
+                const RuntimeShape& input1_shape, const int16* input1_data,
+                const RuntimeShape& input2_shape, const int16* input2_data,
+                const RuntimeShape& output_shape, int16* output_data) {
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  ruy::profiler::ScopeLabel label("AddInt16/16bit");
+  const int flat_size =
+      MatchingElementsSize(input1_shape, input2_shape, output_shape);
+
+  TFLITE_DCHECK_GT(params.input1_offset, -32768);
+  TFLITE_DCHECK_GT(params.input2_offset, -32768);
+  TFLITE_DCHECK_LT(params.input1_offset, 32768);
+  TFLITE_DCHECK_LT(params.input2_offset, 32768);
+  AddElementwiseInt16(flat_size, params, input1_data, input2_data, output_data);
+}
+
+inline void BroadcastAddDispatch(const ArithmeticParams& params,
+                                 const RuntimeShape& input1_shape,
+                                 const int8* input1_data,
+                                 const RuntimeShape& input2_shape,
+                                 const int8* input2_data,
+                                 const RuntimeShape& output_shape,
+                                 int8* output_data) {
+  if (params.broadcast_category == BroadcastableOpCategory::kGenericBroadcast) {
+    return reference_integer_ops::BroadcastAdd6DSlow(
+        params, input1_shape, input1_data, input2_shape, input2_data,
+        output_shape, output_data);
+  }
+
+  optimized_ops::BinaryBroadcastFiveFold(
+      params, input1_shape, input1_data, input2_shape, input2_data,
+      output_shape, output_data, AddElementwiseInt8, AddScalarBroadcast);
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_ADD_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/conv.h

@@ -0,0 +1,130 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_CONV_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_CONV_H_
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/im2col_utils.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+// Fixed-point per-channel-quantization convolution reference kernel.
+template <typename InputScalar, typename DstScalar>
+inline void ConvPerChannel(
+    const ConvParams& params, const int32* output_multiplier,
+    const int32* output_shift, const RuntimeShape& input_shape,
+    const InputScalar* input_data, const RuntimeShape& filter_shape,
+    const int8* filter_data, const RuntimeShape& bias_shape,
+    const int32* bias_data, const RuntimeShape& output_shape,
+    DstScalar* output_data, const RuntimeShape& im2col_shape,
+    InputScalar* im2col_data, CpuBackendContext* cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("Conv/8bit");
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int dilation_width_factor = params.dilation_width_factor;
+  const int dilation_height_factor = params.dilation_height_factor;
+  const int32 input_offset = params.input_offset;
+  const int32 output_offset = params.output_offset;
+  // Set min and max value of the output.
+  const int32 output_activation_min = params.quantized_activation_min;
+  const int32 output_activation_max = params.quantized_activation_max;
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+  const InputScalar* gemm_input_data = nullptr;
+  const RuntimeShape* gemm_input_shape = nullptr;
+  const int filter_width = filter_shape.Dims(2);
+  const int filter_height = filter_shape.Dims(1);
+  const bool need_dilated_im2col =
+      dilation_width_factor != 1 || dilation_height_factor != 1;
+  const bool need_im2col = stride_width != 1 || stride_height != 1 ||
+                           filter_width != 1 || filter_height != 1;
+  const int8 input_zero_point = -input_offset;
+  const uint8 zero_point_byte =
+      *reinterpret_cast<const uint8*>(&input_zero_point);
+  if (need_dilated_im2col) {
+    TFLITE_DCHECK(im2col_data);
+    optimized_ops::DilatedIm2col(params, zero_point_byte, input_shape,
+                                 input_data, filter_shape, output_shape,
+                                 im2col_data);
+    gemm_input_data = im2col_data;
+    gemm_input_shape = &im2col_shape;
+  } else if (need_im2col) {
+    TFLITE_DCHECK(im2col_data);
+    optimized_ops::Im2col(params, filter_height, filter_width, zero_point_byte,
+                          input_shape, input_data, im2col_shape, im2col_data);
+    gemm_input_data = im2col_data;
+    gemm_input_shape = &im2col_shape;
+  } else {
+    TFLITE_DCHECK(!im2col_data);
+    gemm_input_data = input_data;
+    gemm_input_shape = &input_shape;
+  }
+
+  const int gemm_input_rows = gemm_input_shape->Dims(3);
+  const int gemm_input_cols = FlatSizeSkipDim(*gemm_input_shape, 3);
+  const int filter_rows = filter_shape.Dims(0);
+  const int filter_cols = FlatSizeSkipDim(filter_shape, 0);
+  const int output_rows = output_shape.Dims(3);
+  // See b/79927784.
+  // const int output_cols = FlatSizeSkipDim(output_shape, 3);
+  const int output_cols =
+      output_shape.Dims(0) * output_shape.Dims(1) * output_shape.Dims(2);
+  TFLITE_DCHECK_EQ(output_rows, filter_rows);
+  TFLITE_DCHECK_EQ(output_cols, gemm_input_cols);
+  TFLITE_DCHECK_EQ(filter_cols, gemm_input_rows);
+  TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_rows);
+
+  cpu_backend_gemm::MatrixParams<int8> lhs_params;
+  lhs_params.rows = filter_rows;
+  lhs_params.cols = filter_cols;
+  lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
+  lhs_params.zero_point = 0;  // filter is symmetric-quantized
+  cpu_backend_gemm::MatrixParams<InputScalar> rhs_params;
+  rhs_params.rows = gemm_input_rows;
+  rhs_params.cols = gemm_input_cols;
+  rhs_params.order = cpu_backend_gemm::Order::kColMajor;
+  rhs_params.zero_point = -input_offset;
+  cpu_backend_gemm::MatrixParams<DstScalar> dst_params;
+  dst_params.rows = output_rows;
+  dst_params.cols = output_cols;
+  dst_params.order = cpu_backend_gemm::Order::kColMajor;
+  dst_params.zero_point = output_offset;
+  cpu_backend_gemm::GemmParams<
+      int32, DstScalar,
+      cpu_backend_gemm::QuantizationFlavor::kIntegerWithPerRowMultiplier>
+      gemm_params;
+  gemm_params.bias = bias_data;
+  gemm_params.clamp_min = output_activation_min;
+  gemm_params.clamp_max = output_activation_max;
+  gemm_params.multiplier_fixedpoint_perchannel = output_multiplier;
+  gemm_params.multiplier_exponent_perchannel = output_shift;
+  cpu_backend_gemm::Gemm(lhs_params, filter_data, rhs_params, gemm_input_data,
+                         dst_params, output_data, gemm_params,
+                         cpu_backend_context);
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_CONV_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid.h

@@ -0,0 +1,511 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_HYBRID_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_HYBRID_H_
+
+#include <algorithm>
+#include <memory>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_3x3_filter_common.h"
+#include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv.h"
+#include "tensorflow/lite/kernels/internal/optimized/integer_ops/depthwise_conv_hybrid_3x3_filter.h"
+#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+namespace depthwise_conv {
+
+// Initializes the accumulator buffer with zeros.
+inline void DepthwiseConvInitAccBuffer(int num_output_pixels, int output_depth,
+                                       int32* acc_buffer) {
+  memset(acc_buffer, 0,
+         sizeof(acc_buffer[0]) * output_depth * num_output_pixels);
+}
+
+// Base DWConv Implementation used with both static and dynamic
+// accumulator buffers.
+// Initializes the accumulator buffer with bias values.
+static void DoDepthwiseConvHybridGeneral(
+    const DepthwiseParams& params, const float* input_scales,
+    const RuntimeShape& input_shape, const int8* input_data,
+    const RuntimeShape& filter_shape, const int8* filter_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data,
+    const float* per_channel_scales, const int32_t* input_offsets,
+    int thread_start, int thread_end, int thread_dim, int32* acc_buffer,
+    int32 acc_buffer_size) {
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  const int depth_multiplier = params.depth_multiplier;
+  const float output_activation_min = params.float_activation_min;
+  const float output_activation_max = params.float_activation_max;
+  const int dilation_width_factor = params.dilation_width_factor;
+  const int dilation_height_factor = params.dilation_height_factor;
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int input_depth = input_shape.Dims(3);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int output_rows = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+
+  TFLITE_DCHECK_GE(acc_buffer_size, output_depth);
+  const int kOutputPixelsInAccBuffer = acc_buffer_size / output_depth;
+  const int kAccBufferActualSize = kOutputPixelsInAccBuffer * output_depth;
+  TFLITE_DCHECK_LE(kOutputPixelsInAccBuffer * output_depth,
+                   kAccBufferActualSize);
+  TFLITE_DCHECK_LE(kAccBufferActualSize, acc_buffer_size);
+  TFLITE_DCHECK_GE(kOutputPixelsInAccBuffer, 1);
+  TFLITE_DCHECK(thread_dim == 0 || thread_dim == 1);
+
+  // row_accum_func will point to the core accumulation function to be used
+  // for this DepthwiseConvHybrid op.
+  using row_accum_func_t = decltype(&QuantizedDepthwiseConvAccumRowGeneric);
+  row_accum_func_t row_accum_func = nullptr;
+
+#define TFMINI_USE_DEPTHWISECONV_KERNEL(ALLOW_STRIDED, FIXED_INPUT_DEPTH, \
+                                        FIXED_DEPTH_MULTIPLIER)           \
+  if (!row_accum_func && (stride_width == 1 || ALLOW_STRIDED) &&          \
+      (input_depth == FIXED_INPUT_DEPTH || FIXED_INPUT_DEPTH == 0) &&     \
+      depth_multiplier == FIXED_DEPTH_MULTIPLIER) {                       \
+    row_accum_func =                                                      \
+        QuantizedDepthwiseConvAccumRow<ALLOW_STRIDED, FIXED_INPUT_DEPTH,  \
+                                       FIXED_DEPTH_MULTIPLIER>;           \
+  }
+
+#ifdef USE_NEON
+  // We go over our list of kernels by decreasing order of preference
+  // for the cases where multiple kernels could apply.
+
+  // Start with the fastest kernels: AllowStrided=false, fixed input depth.
+
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 1, 2)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 2)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 2)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 1, 4)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 1)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 4, 4)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 8, 1)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 8)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 2, 1)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(false, 12, 1)
+
+  // Next come the strided kernels: AllowStrided=true, fixed input depth.
+  // They are a bit less efficient, but allow stride!=1.
+
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 2)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 16, 1)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 16)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 20)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 32)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 1, 8)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 8, 1)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 2, 1)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 4, 1)
+
+  // Finally, the kernels allowing a variable input depth,
+  // these are the least efficient but most general kernels.
+
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 1)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 2)
+  TFMINI_USE_DEPTHWISECONV_KERNEL(true, 0, 3)
+#endif  // USE_NEON
+
+  // No matching fast kernel found, use slow fallback.
+  if (!row_accum_func) {
+    row_accum_func = QuantizedDepthwiseConvAccumRowGeneric;
+  }
+
+#undef TFMINI_USE_DEPTHWISECONV_KERNEL
+
+  const int input_height_stride = input_shape.Dims(3) * input_shape.Dims(2);
+  const int input_batch_stride = input_height_stride * input_shape.Dims(1);
+  const int filter_height_stride = filter_shape.Dims(3) * filter_shape.Dims(2);
+
+  // Now that we have determined row_accum_func, we can start work.
+  int batch_start = 0;
+  int batch_end = batches;
+  int row_start = 0;
+  int row_end = output_rows;
+  int output_ptr_offset = 0;
+
+  switch (thread_dim) {
+    case 0:
+      TFLITE_DCHECK_GE(thread_start, 0);
+      TFLITE_DCHECK_LE(thread_end, batches);
+      batch_start = thread_start;
+      batch_end = thread_end;
+      output_ptr_offset = batch_start * FlatSizeSkipDim(output_shape, 0);
+      break;
+    case 1:
+      TFLITE_DCHECK_GE(thread_start, 0);
+      TFLITE_DCHECK_LE(thread_end, output_rows);
+      row_start = thread_start;
+      row_end = thread_end;
+      output_ptr_offset = row_start * output_width * output_depth;
+      break;
+  }
+
+  float* output_ptr = output_data + output_ptr_offset;
+  int batch_step =
+      (output_rows + row_start - row_end) * output_width * output_depth;
+  for (int b = batch_start; b < batch_end; ++b) {
+    float input_scale = input_scales[b];
+    int32_t input_offset = input_offsets[b];
+    for (int out_y = row_start; out_y < row_end; ++out_y) {
+      const int in_y_origin = (out_y * stride_height) - pad_height;
+      const int filter_y_start =
+          std::max(0, (-in_y_origin + dilation_height_factor - 1) /
+                          dilation_height_factor);
+      const int filter_y_end =
+          std::min(filter_height,
+                   (input_height - in_y_origin + dilation_height_factor - 1) /
+                       dilation_height_factor);
+      for (int out_x_buffer_start = 0; out_x_buffer_start < output_width;
+           out_x_buffer_start += kOutputPixelsInAccBuffer) {
+        const int out_x_buffer_end = std::min(
+            output_width, out_x_buffer_start + kOutputPixelsInAccBuffer);
+        // We call a 'pixel' a group of activation that share all but the
+        // 'depth'/'channel' coordinate. num_output_pixels is the number of
+        // output pixels that we will accumulate in this loop iteration.
+        const int num_output_pixels = out_x_buffer_end - out_x_buffer_start;
+        DepthwiseConvInitAccBuffer(num_output_pixels, output_depth, acc_buffer);
+
+        // Accumulation loop. Most of the time should be spent in here.
+        for (int filter_y = filter_y_start; filter_y < filter_y_end;
+             ++filter_y) {
+          const int in_y = in_y_origin + dilation_height_factor * filter_y;
+          row_accum_func(
+              stride_width, dilation_width_factor, input_depth, input_width,
+              input_data + in_y * input_height_stride + b * input_batch_stride,
+              -input_offset, pad_width, depth_multiplier, filter_width,
+              filter_data + filter_y * filter_height_stride, out_x_buffer_start,
+              out_x_buffer_end, output_depth, acc_buffer);
+        }
+        // Finished accumulating int32 values. Just store them as float values
+        gemmlowp::ScopedProfilingLabel label("store");
+        const int num_output_values = output_depth * num_output_pixels;
+        int c = 0;
+        while (c < output_depth) {
+          int target_output_depth = output_depth;
+
+#ifdef USE_NEON
+          const float32x4_t output_activation_min_vec =
+              vdupq_n_f32(output_activation_min);
+          const float32x4_t output_activation_max_vec =
+              vdupq_n_f32(output_activation_max);
+          const float32x4_t input_scale_32x4 = vdupq_n_f32(input_scale);
+          for (; c <= output_depth - 4; c += 4) {
+            if ((c + 4) > output_depth) {
+              break;
+            }
+            const float32x4_t channel_scale_32x4 =
+                vld1q_f32(per_channel_scales + c);
+            const float32x4_t bias_32x4 = vld1q_f32(bias_data + c);
+            for (int n = 0; n < num_output_pixels; ++n) {
+              int loc = n * output_depth + c;
+              int32x4_t acc = vld1q_s32(acc_buffer + loc);
+              float32x4_t float_acc = vcvtq_f32_s32(acc);
+              float_acc = vmulq_f32(float_acc, channel_scale_32x4);
+              float_acc = vmulq_f32(float_acc, input_scale_32x4);
+              float_acc = vaddq_f32(float_acc, bias_32x4);
+              float_acc = vmaxq_f32(float_acc, output_activation_min_vec);
+              float_acc = vminq_f32(float_acc, output_activation_max_vec);
+              vst1q_f32(output_ptr + loc, float_acc);
+            }
+          }
+#endif  // USE_NEON
+
+          for (; c < target_output_depth; c++) {
+            for (int n = 0; n < num_output_pixels; ++n) {
+              int loc = n * output_depth + c;
+              int32 acc = acc_buffer[loc];
+              float float_acc = acc * input_scale * per_channel_scales[c];
+              float_acc += bias_data[c];
+              float_acc = std::max(float_acc, output_activation_min);
+              float_acc = std::min(float_acc, output_activation_max);
+              output_ptr[loc] = float_acc;
+            }
+          }
+        }
+        output_ptr += num_output_values;
+      }
+    }
+    output_ptr += batch_step;
+  }
+}
+
+// Utilize the base implementation of DWConv with a stack allocated accumulator
+// buffer. The static allocation limits the number of depthwise channels that
+// can be processed to kStaticAccBufferMaxSize.
+static void DoDepthwiseConvHybridGeneralStatic(
+    const DepthwiseParams& params, const float* input_scales,
+    const RuntimeShape& input_shape, const int8* input_data,
+    const RuntimeShape& filter_shape, const int8* filter_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data,
+    const float* per_channel_scales, const int32_t* input_offsets,
+    int thread_start, int thread_end, int thread_dim) {
+  static const int kStaticAccBufferMaxSize = 2048;
+  int32 stack_acc_buffer[kStaticAccBufferMaxSize];
+  DoDepthwiseConvHybridGeneral(
+      params, input_scales, input_shape, input_data, filter_shape, filter_data,
+      bias_shape, bias_data, output_shape, output_data, per_channel_scales,
+      input_offsets, thread_start, thread_end, thread_dim, stack_acc_buffer,
+      kStaticAccBufferMaxSize);
+}
+
+// This DWConv function uses static memory for accumulation by default for upto
+// kStaticAccBufferMaxSize channels. Beyound that, a dynamic buffer is used on
+// a per call basis. The function errors out if number of channels is larger
+// than kStaticAccBufferMaxSize and TF_LITE_STATIC_MEMORY is defined.
+inline void DepthwiseConvHybridGeneral(
+    const DepthwiseParams& params, const float* input_scales,
+    const RuntimeShape& input_shape, const int8* input_data,
+    const RuntimeShape& filter_shape, const int8* filter_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data,
+    const float* per_channel_scales, const int32_t* input_offsets,
+    int thread_start, int thread_end, int thread_dim) {
+#ifndef TF_LITE_STATIC_MEMORY
+  static const int kStaticAccBufferMaxSize = 2048;
+  const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3);
+
+  if (kStaticAccBufferMaxSize < output_depth) {
+    std::unique_ptr<int32[]> heap_acc_buffer(new int32[output_depth]);
+    DoDepthwiseConvHybridGeneral(
+        params, input_scales, input_shape, input_data, filter_shape,
+        filter_data, bias_shape, bias_data, output_shape, output_data,
+        per_channel_scales, input_offsets, thread_start, thread_end, thread_dim,
+        heap_acc_buffer.get(), output_depth);
+
+    return;
+  }
+#endif
+
+  DoDepthwiseConvHybridGeneralStatic(
+      params, input_scales, input_shape, input_data, filter_shape, filter_data,
+      bias_shape, bias_data, output_shape, output_data, per_channel_scales,
+      input_offsets, thread_start, thread_end, thread_dim);
+}
+
+}  // namespace depthwise_conv
+
+template <DepthwiseConvOutputRounding kOutputRounding>
+inline void DepthwiseConvHybridWithRounding(
+    const DepthwiseParams& params, const float* input_scales,
+    const RuntimeShape& input_shape, const int8* input_data,
+    const RuntimeShape& filter_shape, const int8* filter_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data,
+    const float* per_channel_scales, const int32_t* input_offsets,
+    int thread_start, int thread_end, int thread_dim) {
+  gemmlowp::ScopedProfilingLabel label("DepthwiseConvHybridInt8/8bit");
+  const int depth_multiplier = params.depth_multiplier;
+  const int dilation_width_factor = params.dilation_width_factor;
+  const int dilation_height_factor = params.dilation_height_factor;
+  TFLITE_DCHECK_GE(dilation_width_factor, 1);
+  TFLITE_DCHECK_GE(dilation_height_factor, 1);
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+  const int output_depth = MatchingDim(filter_shape, 3, output_shape, 3);
+  const int input_depth = input_shape.Dims(3);
+  TFLITE_DCHECK_EQ(output_depth, input_depth * depth_multiplier);
+  TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_depth);
+
+// Enable for arm64 except for the Nvidia Linux 4 Tegra (L4T) running on
+// Jetson TX-2. This compiler does not support the offsetof() macro.
+
+#if defined(__aarch64__) && !defined(GOOGLE_L4T)
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+
+  // Call kernel optimized for depthwise convolutions using 3x3 filters if
+  // parameters are supported.
+  if (optimized_ops::depthwise_conv::Fast3x3FilterKernelSupported<
+      optimized_ops::depthwise_conv::QuantizationType::kNonPerChannelUint8>(
+          input_shape, filter_shape, stride_width, stride_height,
+          dilation_width_factor, dilation_height_factor, pad_width, pad_height,
+          depth_multiplier, output_shape, 0, nullptr)) {
+    gemmlowp::ScopedProfilingLabel specialized_label(
+        "DepthwiseConvHybridInt8/8bit/3x3");
+    optimized_ops::depthwise_conv::DepthwiseConvHybrid3x3FilterPerChannel<
+        DepthwiseConvOutputRounding::kUpward>(
+            params, input_scales, input_shape, input_data,
+            filter_shape, filter_data, bias_shape, bias_data, output_shape,
+            output_data, per_channel_scales, input_offsets,
+            thread_start, thread_end, thread_dim);
+    return;
+  }
+#endif
+
+  gemmlowp::ScopedProfilingLabel specialized_label(
+      "DepthwiseConvHybridInt8/8bit/General");
+  depthwise_conv::DepthwiseConvHybridGeneral(
+      params, input_scales, input_shape, input_data,
+      filter_shape, filter_data, bias_shape, bias_data, output_shape,
+      output_data, per_channel_scales, input_offsets,
+      thread_start, thread_end, thread_dim);
+}
+
+inline void DepthwiseConvHybridImpl(
+    const DepthwiseParams& params, const float* input_scales,
+    const RuntimeShape& input_shape, const int8* input_data,
+    const RuntimeShape& filter_shape, const int8* filter_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data,
+    const float* per_channel_scales, const int32_t* input_offsets,
+    int thread_start, int thread_end, int thread_dim) {
+  return DepthwiseConvHybridWithRounding<
+      DepthwiseConvOutputRounding::kAwayFromZero>(
+          params, input_scales, input_shape, input_data,
+          filter_shape, filter_data, bias_shape, bias_data, output_shape,
+          output_data, per_channel_scales, input_offsets,
+          thread_start, thread_end, thread_dim);
+}
+
+template <typename T, typename TS>
+struct DepthwiseConvHybridWorkerTask : cpu_backend_threadpool::Task {
+  DepthwiseConvHybridWorkerTask(const DepthwiseParams& params,
+                                const float* input_scales,
+                                const RuntimeShape& input_shape,
+                                const T* input_data,
+                                const RuntimeShape& filter_shape,
+                                const T* filter_data,
+                                const RuntimeShape& bias_shape,
+                                const TS* bias_data,
+                                const RuntimeShape& output_shape,
+                                float* output_data,
+                                const float* per_channel_scales,
+                                const int32_t* input_offsets,
+                                int thread_start, int thread_end,
+                                int thread_dim)
+      : params(params),
+        input_scales(input_scales),
+        input_shape(input_shape),
+        input_data(input_data),
+        filter_shape(filter_shape),
+        filter_data(filter_data),
+        bias_shape(bias_shape),
+        bias_data(bias_data),
+        output_shape(output_shape),
+        output_data(output_data),
+        per_channel_scales(per_channel_scales),
+        input_offsets(input_offsets),
+        thread_start(thread_start),
+        thread_end(thread_end),
+        thread_dim(thread_dim) {}
+
+  void Run() override {
+    DepthwiseConvHybridImpl(params, input_scales, input_shape,
+                            input_data, filter_shape, filter_data,
+                            bias_shape, bias_data, output_shape,
+                            output_data, per_channel_scales, input_offsets,
+                            thread_start, thread_end, thread_dim);
+  }
+
+ private:
+  const DepthwiseParams& params;
+  const float* input_scales;
+  const RuntimeShape& input_shape;
+  const T* input_data;
+  const RuntimeShape& filter_shape;
+  const T* filter_data;
+  const RuntimeShape& bias_shape;
+  const TS* bias_data;
+  const RuntimeShape& output_shape;
+  float* output_data;
+  const float* per_channel_scales;
+  const int32_t* input_offsets;
+  int thread_start;
+  int thread_end;
+  int thread_dim;
+};
+
+inline void DepthwiseConvHybridPerChannel(
+    const DepthwiseParams& params, const float* input_scales,
+    const RuntimeShape& input_shape, const int8* input_data,
+    const RuntimeShape& filter_shape, const int8* filter_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data,
+    const float* per_channel_scales, int32_t* input_offsets,
+    CpuBackendContext* cpu_backend_context) {
+  gemmlowp::ScopedProfilingLabel label("DepthwiseConvHybridInt8");
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+  const int output_batches = output_shape.Dims(0);
+  const int output_rows = output_shape.Dims(1);
+  int thread_count_batch = HowManyConvThreads(output_shape, filter_shape, 0);
+  int thread_count_row = HowManyConvThreads(output_shape, filter_shape, 1);
+  int thread_dim, thread_count, thread_dim_size;
+  if (thread_count_batch > thread_count_row) {
+    thread_dim = 0;
+    thread_dim_size = output_batches;
+    thread_count = thread_count_batch;
+  } else {
+    thread_dim = 1;
+    thread_dim_size = output_rows;
+    thread_count = thread_count_row;
+  }
+
+  const int max_threads = cpu_backend_context->max_num_threads();
+  thread_count = std::max(1, std::min(thread_count, max_threads));
+
+  if (thread_count == 1) {
+    DepthwiseConvHybridImpl(params, input_scales, input_shape,
+                            input_data, filter_shape, filter_data, bias_shape,
+                            bias_data, output_shape, output_data,
+                            per_channel_scales, input_offsets,
+                            /*thread_start=*/0, /*thread_end=*/output_rows,
+                            /*thread_dim=*/1);
+  } else {
+    std::vector<DepthwiseConvHybridWorkerTask<int8, float>> tasks;
+    // TODO(b/131746020) don't create new heap allocations every time.
+    // At least we make it a single heap allocation by using reserve().
+    tasks.reserve(thread_count);
+    int thread_start = 0;
+    for (int i = 0; i < thread_count; ++i) {
+      int thread_end =
+          thread_start + (thread_dim_size - thread_start) / (thread_count - i);
+      tasks.emplace_back(params, input_scales, input_shape,
+                         input_data, filter_shape, filter_data, bias_shape,
+                         bias_data, output_shape, output_data,
+                         per_channel_scales, input_offsets, thread_start,
+                         thread_end, thread_dim);
+      thread_start = thread_end;
+    }
+    cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
+                                    cpu_backend_context);
+  }
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_DEPTHWISE_CONV_HYBRID_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/fully_connected.h

@@ -0,0 +1,172 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_FULLY_CONNECTED_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_FULLY_CONNECTED_H_
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/fully_connected.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+template <typename InputScalar, typename DstScalar>
+inline void FullyConnectedPerChannel(
+    const FullyConnectedParams& params, const int32_t* output_multiplier,
+    const int* output_shift, const RuntimeShape& input_shape,
+    const InputScalar* input_data, const RuntimeShape& filter_shape,
+    const int8_t* filter_data, const RuntimeShape& bias_shape,
+    const int32_t* bias_data, const RuntimeShape& output_shape,
+    DstScalar* output_data, CpuBackendContext* cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("FullyConnectedInt8/8bit");
+
+  const int32_t input_offset = params.input_offset;
+  const int32_t output_offset = params.output_offset;
+  const int32_t output_activation_min = params.quantized_activation_min;
+  const int32_t output_activation_max = params.quantized_activation_max;
+  TFLITE_DCHECK_GE(filter_shape.DimensionsCount(), 2);
+  TFLITE_DCHECK_GE(output_shape.DimensionsCount(), 1);
+  // TODO(b/62193649): This really should be:
+  //     const int batches = ArraySize(output_dims, 1);
+  // but the current --variable_batch hack consists in overwriting the 3rd
+  // dimension with the runtime batch size, as we don't keep track for each
+  // array of which dimension is the batch dimension in it.
+  const int output_dim_count = output_shape.DimensionsCount();
+  const int filter_dim_count = filter_shape.DimensionsCount();
+  const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
+  const int filter_rows = filter_shape.Dims(filter_dim_count - 2);
+  const int filter_cols = filter_shape.Dims(filter_dim_count - 1);
+  TFLITE_DCHECK_EQ(filter_shape.FlatSize(), filter_rows * filter_cols);
+  const int output_rows = output_shape.Dims(output_dim_count - 1);
+  TFLITE_DCHECK_EQ(output_rows, filter_rows);
+  if (bias_data) {
+    TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_rows);
+  }
+  const bool use_caching =
+      (cpu_backend_context != nullptr) && cpu_backend_context->use_caching();
+
+  cpu_backend_gemm::MatrixParams<int8_t> lhs_params;
+  lhs_params.rows = filter_rows;
+  lhs_params.cols = filter_cols;
+  lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
+  lhs_params.zero_point = 0;
+  lhs_params.cache_policy =
+      use_caching ? cpu_backend_gemm::DefaultCachePolicy(params.lhs_cacheable)
+                  : cpu_backend_gemm::CachePolicy::kNeverCache;
+  cpu_backend_gemm::MatrixParams<InputScalar> rhs_params;
+  rhs_params.rows = filter_cols;
+  rhs_params.cols = batches;
+  rhs_params.order = cpu_backend_gemm::Order::kColMajor;
+  rhs_params.zero_point = -input_offset;
+  rhs_params.cache_policy =
+      use_caching ? cpu_backend_gemm::DefaultCachePolicy(params.rhs_cacheable)
+                  : cpu_backend_gemm::CachePolicy::kNeverCache;
+  cpu_backend_gemm::MatrixParams<DstScalar> dst_params;
+  dst_params.rows = filter_rows;
+  dst_params.cols = batches;
+  dst_params.order = cpu_backend_gemm::Order::kColMajor;
+  dst_params.zero_point = output_offset;
+  cpu_backend_gemm::GemmParams<
+      int32_t, DstScalar,
+      cpu_backend_gemm::QuantizationFlavor::kIntegerWithPerRowMultiplier>
+      gemm_params;
+  gemm_params.bias = bias_data;
+  gemm_params.clamp_min = output_activation_min;
+  gemm_params.clamp_max = output_activation_max;
+  gemm_params.multiplier_fixedpoint_perchannel = output_multiplier;
+  gemm_params.multiplier_exponent_perchannel = output_shift;
+  cpu_backend_gemm::Gemm(lhs_params, filter_data, rhs_params, input_data,
+                         dst_params, output_data, gemm_params,
+                         cpu_backend_context);
+}
+
+template <typename InputScalar, typename DstScalar>
+inline void FullyConnected(
+    const FullyConnectedParams& params, const RuntimeShape& input_shape,
+    const InputScalar* input_data, const RuntimeShape& filter_shape,
+    const int8_t* filter_data, const RuntimeShape& bias_shape,
+    const int32_t* bias_data, const RuntimeShape& output_shape,
+    DstScalar* output_data, CpuBackendContext* cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("FullyConnectedInt8/8bit");
+
+  const int32_t input_offset = params.input_offset;
+  const int32_t filter_offset = params.weights_offset;
+  const int32_t output_offset = params.output_offset;
+  const int32_t output_multiplier = params.output_multiplier;
+  const int output_shift = params.output_shift;
+  const int32_t output_activation_min = params.quantized_activation_min;
+  const int32_t output_activation_max = params.quantized_activation_max;
+  TFLITE_DCHECK_GE(filter_shape.DimensionsCount(), 2);
+  TFLITE_DCHECK_GE(output_shape.DimensionsCount(), 1);
+  // TODO(b/62193649): This really should be:
+  //     const int batches = ArraySize(output_dims, 1);
+  // but the current --variable_batch hack consists in overwriting the 3rd
+  // dimension with the runtime batch size, as we don't keep track for each
+  // array of which dimension is the batch dimension in it.
+  const int output_dim_count = output_shape.DimensionsCount();
+  const int filter_dim_count = filter_shape.DimensionsCount();
+  const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
+  const int filter_rows = filter_shape.Dims(filter_dim_count - 2);
+  const int filter_cols = filter_shape.Dims(filter_dim_count - 1);
+  TFLITE_DCHECK_EQ(filter_shape.FlatSize(), filter_rows * filter_cols);
+  const int output_rows = output_shape.Dims(output_dim_count - 1);
+  TFLITE_DCHECK_EQ(output_rows, filter_rows);
+  if (bias_data) {
+    TFLITE_DCHECK_EQ(bias_shape.FlatSize(), output_rows);
+  }
+  const bool use_caching =
+      (cpu_backend_context != nullptr) && cpu_backend_context->use_caching();
+
+  cpu_backend_gemm::MatrixParams<int8_t> lhs_params;
+  lhs_params.rows = filter_rows;
+  lhs_params.cols = filter_cols;
+  lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
+  lhs_params.zero_point = -filter_offset;
+  lhs_params.cache_policy =
+      use_caching ? cpu_backend_gemm::DefaultCachePolicy(params.lhs_cacheable)
+                  : cpu_backend_gemm::CachePolicy::kNeverCache;
+  cpu_backend_gemm::MatrixParams<InputScalar> rhs_params;
+  rhs_params.rows = filter_cols;
+  rhs_params.cols = batches;
+  rhs_params.order = cpu_backend_gemm::Order::kColMajor;
+  rhs_params.zero_point = -input_offset;
+  rhs_params.cache_policy =
+      use_caching ? cpu_backend_gemm::DefaultCachePolicy(params.rhs_cacheable)
+                  : cpu_backend_gemm::CachePolicy::kNeverCache;
+  cpu_backend_gemm::MatrixParams<DstScalar> dst_params;
+  dst_params.rows = filter_rows;
+  dst_params.cols = batches;
+  dst_params.order = cpu_backend_gemm::Order::kColMajor;
+  dst_params.zero_point = output_offset;
+  cpu_backend_gemm::GemmParams<int32_t, DstScalar> gemm_params;
+  gemm_params.bias = bias_data;
+  gemm_params.clamp_min = output_activation_min;
+  gemm_params.clamp_max = output_activation_max;
+  gemm_params.multiplier_fixedpoint = output_multiplier;
+  gemm_params.multiplier_exponent = output_shift;
+  cpu_backend_gemm::Gemm(lhs_params, filter_data, rhs_params, input_data,
+                         dst_params, output_data, gemm_params,
+                         cpu_backend_context);
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_FULLY_CONNECTED_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/leaky_relu.h

@@ -0,0 +1,113 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LEAKY_RELU_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LEAKY_RELU_H_
+
+#include <algorithm>
+
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+inline void QuantizeLeakyRelu(const LeakyReluParams& params,
+                              const RuntimeShape& input_shape,
+                              const int16* input_data,
+                              const RuntimeShape& output_shape,
+                              int16* output_data) {
+  const int flat_size = MatchingFlatSize(input_shape, output_shape);
+  const int32_t quantized_min = std::numeric_limits<int16>::min();
+  const int32_t quantized_max = std::numeric_limits<int16>::max();
+  int i = 0;
+
+#ifdef __AVX2__
+  const __m256i input_offset = _mm256_set1_epi32(params.input_offset);
+  const __m256i output_offset = _mm256_set1_epi32(params.output_offset);
+  const __m256i output_muliplier_identity =
+      _mm256_set1_epi32(params.output_multiplier_identity);
+  const __m256i output_shift_identity =
+      _mm256_set1_epi32(params.output_shift_identity);
+  const __m256i output_multiplier_alpha =
+      _mm256_set1_epi32(params.output_multiplier_alpha);
+  const __m256i output_shift_alpha =
+      _mm256_set1_epi32(params.output_shift_alpha);
+  const __m256i clamp_max_v = _mm256_set1_epi32(quantized_max);
+  const __m256i clamp_min_v = _mm256_set1_epi32(quantized_min);
+
+  for (; i <= flat_size - 16; i += 16) {
+    const __m256i input =
+        _mm256_loadu_si256(reinterpret_cast<__m256i const*>(input_data + i));
+    __m256i input_low = _mm256_cvtepi16_epi32(_mm256_castsi256_si128(input));
+    __m256i input_high =
+        _mm256_cvtepi16_epi32(_mm256_extracti128_si256(input, 1));
+    input_low = _mm256_sub_epi32(input_low, input_offset);
+    input_high = _mm256_sub_epi32(input_high, input_offset);
+
+    const __m256i zeros = _mm256_setzero_si256();
+    const __m256i input_low_mask = _mm256_cmpgt_epi32(input_low, zeros);
+    const __m256i input_high_mask = _mm256_cmpgt_epi32(input_high, zeros);
+    const __m256i input_low_output_multiplier = avx2_utils::mm256_blendv_epi32(
+        output_multiplier_alpha, output_muliplier_identity, input_low_mask);
+    const __m256i input_low_output_shift = avx2_utils::mm256_blendv_epi32(
+        output_shift_alpha, output_shift_identity, input_low_mask);
+    const __m256i input_high_output_multiplier = avx2_utils::mm256_blendv_epi32(
+        output_multiplier_alpha, output_muliplier_identity, input_high_mask);
+    const __m256i input_high_output_shift = avx2_utils::mm256_blendv_epi32(
+        output_shift_alpha, output_shift_identity, input_high_mask);
+
+    input_low = avx2_utils::MultiplyByQuantizedMultiplier(
+        input_low, input_low_output_multiplier, input_low_output_shift);
+    input_high = avx2_utils::MultiplyByQuantizedMultiplier(
+        input_high, input_high_output_multiplier, input_high_output_shift);
+
+    input_low = _mm256_add_epi32(input_low, output_offset);
+    input_high = _mm256_add_epi32(input_high, output_offset);
+
+    input_low = _mm256_min_epi32(input_low, clamp_max_v);
+    input_low = _mm256_max_epi32(input_low, clamp_min_v);
+    input_high = _mm256_min_epi32(input_high, clamp_max_v);
+    input_high = _mm256_max_epi32(input_high, clamp_min_v);
+
+    avx2_utils::CastInt32ToInt16AndStore(output_data + i, input_low);
+    avx2_utils::CastInt32ToInt16AndStore(output_data + i + 8, input_high);
+  }
+#endif  // __AVX2__
+
+  for (; i < flat_size; ++i) {
+    const int32_t input_value = input_data[i] - params.input_offset;
+    int32_t unclamped_output;
+    if (input_value >= 0) {
+      unclamped_output = params.output_offset +
+                         MultiplyByQuantizedMultiplier(
+                             input_value, params.output_multiplier_identity,
+                             params.output_shift_identity);
+    } else {
+      unclamped_output = params.output_offset +
+                         MultiplyByQuantizedMultiplier(
+                             input_value, params.output_multiplier_alpha,
+                             params.output_shift_alpha);
+    }
+    const int16 clamped_output =
+        std::min(quantized_max, std::max(quantized_min, unclamped_output));
+    output_data[i] = static_cast<int16>(clamped_output);
+  }
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LEAKY_RELU_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/lut.h

@@ -0,0 +1,72 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LUT_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LUT_H_
+
+#include <cstdint>
+
+#if __aarch64__ && __clang__
+#include <arm_neon.h>
+#endif
+
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+inline void LookupTable(const uint8_t* input_data, int num_elements,
+                        const uint8_t* lut, uint8_t* output_data) {
+  int i = 0;
+#if __aarch64__ && __clang__
+  // This code uses ARM64-only instructions.
+  // TODO(b/143709993): Port to ARMv7
+
+  // Load the tables into registers. (4*4 128-bit registers)
+  uint8x16x4_t table[4];
+  table[0] = vld1q_u8_x4(lut + 16 * 4 * 0);
+  table[1] = vld1q_u8_x4(lut + 16 * 4 * 1);
+  table[2] = vld1q_u8_x4(lut + 16 * 4 * 2);
+  table[3] = vld1q_u8_x4(lut + 16 * 4 * 3);
+
+  // Vectorized loop; process uint8x16_t (16 elements) at a time.
+  constexpr int vectorized_16_loop_step = 16;
+  const int vectorized_16_loop_end =
+      num_elements / vectorized_16_loop_step * vectorized_16_loop_step;
+  for (; i < vectorized_16_loop_end; i += vectorized_16_loop_step) {
+    uint8x16_t input = vld1q_u8(input_data + i);
+    uint8x16_t output = optimized_ops::aarch64_lookup_vector(table, input);
+    vst1q_u8(output_data + i, output);
+  }
+  // Postamble and non-ARM64 code: simple for loop.
+#endif
+  for (; i < num_elements; ++i) {
+    output_data[i] = lut[input_data[i]];
+  }
+}
+
+// LUTPopulate<int8_t> has ordered the LUT so that indexing it with an
+// int8_t is just done by casting it to an uint8_t. We can thus reuse the uint8
+// LookupTable function.
+inline void LookupTable(const int8_t* input_data, int num_elements,
+                        const int8_t* lut, int8_t* output_data) {
+  LookupTable(reinterpret_cast<const uint8_t*>(input_data), num_elements,
+              reinterpret_cast<const uint8_t*>(lut),
+              reinterpret_cast<uint8_t*>(output_data));
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_LUT_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/mean.h

@@ -0,0 +1,251 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MEAN_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MEAN_H_
+
+#include <algorithm>
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+inline void MeanImpl(const tflite::MeanParams& op_params,
+                     const RuntimeShape& input_shape, const int8_t* input_data,
+                     int32 multiplier, int32 shift, int32 bias,
+                     const RuntimeShape& output_shape, int8_t* output_data,
+                     int start_depth, int end_depth) {
+  ruy::profiler::ScopeLabel label("Mean4D/Int8/MeanImpl");
+
+  // Current implementation only supports dimension equals 4 and simultaneous
+  // reduction over width and height.
+  const int output_batch = output_shape.Dims(0);
+  const int output_height = output_shape.Dims(2);
+  const int output_width = output_shape.Dims(2);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+
+  TFLITE_CHECK_EQ(op_params.axis_count, 2);
+  TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
+               (op_params.axis[0] == 2 && op_params.axis[1] == 1));
+  TFLITE_CHECK_EQ(output_height, 1);
+  TFLITE_CHECK_EQ(output_width, 1);
+
+  constexpr static int32_t kMinValue = std::numeric_limits<int8_t>::min();
+  constexpr static int32_t kMaxValue = std::numeric_limits<int8_t>::max();
+
+#ifdef USE_NEON
+  const int32x4_t bias_dup = vdupq_n_s32(bias);
+  const int32x4_t min_dup = vdupq_n_s32(kMinValue);
+  const int32x4_t max_dup = vdupq_n_s32(kMaxValue);
+#endif  // USE_NEON
+  for (int out_b = 0; out_b < output_batch; ++out_b) {
+    int out_d = start_depth;
+#ifdef USE_NEON
+
+    for (; out_d <= end_depth - 16; out_d += 16) {
+      int32x4x4_t temp_sum;
+      temp_sum.val[0] = vdupq_n_s32(0);
+      temp_sum.val[1] = vdupq_n_s32(0);
+      temp_sum.val[2] = vdupq_n_s32(0);
+      temp_sum.val[3] = vdupq_n_s32(0);
+      for (int in_h = 0; in_h < input_height; ++in_h) {
+        for (int in_w = 0; in_w < input_width; ++in_w) {
+          const int8_t* input_data_ptr =
+              input_data + Offset(input_shape, out_b, in_h, in_w, out_d);
+          int8x16_t input_data_val = vld1q_s8(input_data_ptr);
+
+          int16x8_t input_data_low_shift =
+              vmovl_s8(vget_low_s8(input_data_val));
+          int16x8_t input_data_high_shift =
+              vmovl_s8(vget_high_s8(input_data_val));
+
+          int32x4_t input_low_low =
+              vmovl_s16(vget_low_s16(input_data_low_shift));
+          int32x4_t input_high_low =
+              vmovl_s16(vget_high_s16(input_data_low_shift));
+          int32x4_t input_low_high =
+              vmovl_s16(vget_low_s16(input_data_high_shift));
+          int32x4_t input_high_high =
+              vmovl_s16(vget_high_s16(input_data_high_shift));
+
+          temp_sum.val[0] = vaddq_s32(temp_sum.val[0], input_low_low);
+          temp_sum.val[1] = vaddq_s32(temp_sum.val[1], input_high_low);
+          temp_sum.val[2] = vaddq_s32(temp_sum.val[2], input_low_high);
+          temp_sum.val[3] = vaddq_s32(temp_sum.val[3], input_high_high);
+        }
+      }
+
+      temp_sum =
+          MultiplyByQuantizedMultiplier4Rows(temp_sum, multiplier, shift);
+
+      temp_sum.val[0] = vaddq_s32(temp_sum.val[0], bias_dup);
+      temp_sum.val[1] = vaddq_s32(temp_sum.val[1], bias_dup);
+      temp_sum.val[2] = vaddq_s32(temp_sum.val[2], bias_dup);
+      temp_sum.val[3] = vaddq_s32(temp_sum.val[3], bias_dup);
+
+      temp_sum.val[0] = vminq_s32(vmaxq_s32(temp_sum.val[0], min_dup), max_dup);
+      temp_sum.val[1] = vminq_s32(vmaxq_s32(temp_sum.val[1], min_dup), max_dup);
+      temp_sum.val[2] = vminq_s32(vmaxq_s32(temp_sum.val[2], min_dup), max_dup);
+      temp_sum.val[3] = vminq_s32(vmaxq_s32(temp_sum.val[3], min_dup), max_dup);
+
+      int16x4_t narrowed_low_low = vmovn_s32(temp_sum.val[0]);
+      int16x4_t narrowed_high_low = vmovn_s32(temp_sum.val[1]);
+      int16x4_t narrowed_low_high = vmovn_s32(temp_sum.val[2]);
+      int16x4_t narrowed_high_high = vmovn_s32(temp_sum.val[3]);
+
+      int16x8_t combined_low =
+          vcombine_s16(narrowed_low_low, narrowed_high_low);
+      int16x8_t combined_high =
+          vcombine_s16(narrowed_low_high, narrowed_high_high);
+
+      int8x8_t narrowed_low = vmovn_s16(combined_low);
+      int8x8_t narrowed_high = vmovn_s16(combined_high);
+
+      int8x16_t combined_output = vcombine_s8(narrowed_low, narrowed_high);
+
+      int8_t* output_data_ptr =
+          output_data + Offset(output_shape, out_b, 0, 0, out_d);
+      vst1q_s8(output_data_ptr, combined_output);
+    }
+#endif  // USE_NEON
+
+    for (; out_d < end_depth; ++out_d) {
+      int acc = 0;
+      for (int in_h = 0; in_h < input_height; ++in_h) {
+        for (int in_w = 0; in_w < input_width; ++in_w) {
+          acc += input_data[Offset(input_shape, out_b, in_h, in_w, out_d)];
+        }
+      }
+
+      acc = MultiplyByQuantizedMultiplier(acc, multiplier, shift);
+      acc += bias;
+      acc = std::min(std::max(acc, kMinValue), kMaxValue);
+      output_data[Offset(output_shape, out_b, 0, 0, out_d)] =
+          static_cast<int8_t>(acc);
+    }
+  }
+}
+
+struct MeanWorkerTask : cpu_backend_threadpool::Task {
+  MeanWorkerTask(const tflite::MeanParams& op_params,
+                 const RuntimeShape& input_shape, const int8_t* input_data,
+                 int32 multiplier, int32 shift, int32 bias,
+                 const RuntimeShape& output_shape, int8_t* output_data,
+                 int start_height, int end_height)
+      : op_params(op_params),
+        input_shape(input_shape),
+        input_data(input_data),
+        multiplier(multiplier),
+        shift(shift),
+        bias(bias),
+        output_shape(output_shape),
+        output_data(output_data),
+        start_height(start_height),
+        end_height(end_height) {}
+
+  void Run() override {
+    MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
+             output_shape, output_data, start_height, end_height);
+  }
+
+ private:
+  const tflite::MeanParams& op_params;
+  const RuntimeShape& input_shape;
+  const int8_t* input_data;
+  int32 multiplier;
+  int32 shift;
+  int32 bias;
+  const RuntimeShape& output_shape;
+  int8_t* output_data;
+  int start_height;
+  int end_height;
+};
+
+inline void Mean(const tflite::MeanParams& op_params,
+                 const RuntimeShape& unextended_input_shape,
+                 const int8_t* input_data, int32 input_zero_point,
+                 float input_scale, const RuntimeShape& unextended_output_shape,
+                 int8_t* output_data, int32 output_zero_point,
+                 float output_scale, CpuBackendContext* cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("Mean4D/Int8");
+  // Current implementation only supports dimension equals 4 and simultaneous
+  // reduction over width and height.
+  TFLITE_CHECK_EQ(unextended_input_shape.DimensionsCount(), 4);
+  TFLITE_CHECK_LE(unextended_output_shape.DimensionsCount(), 4);
+  const RuntimeShape input_shape =
+      RuntimeShape::ExtendedShape(4, unextended_input_shape);
+  const RuntimeShape output_shape =
+      RuntimeShape::ExtendedShape(4, unextended_output_shape);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  const int output_depth = output_shape.Dims(3);
+
+  TFLITE_CHECK_EQ(op_params.axis_count, 2);
+  TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
+               (op_params.axis[0] == 2 && op_params.axis[1] == 1));
+  TFLITE_CHECK_EQ(output_height, 1);
+  TFLITE_CHECK_EQ(output_width, 1);
+
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const float num_elements_in_axis = input_width * input_height;
+
+  float temp = input_zero_point * input_scale / output_scale;
+  temp = temp > 0 ? temp + 0.5f : temp - 0.5f;
+  int32_t bias = output_zero_point - static_cast<int32_t>(temp);
+  float real_scale = input_scale / (num_elements_in_axis * output_scale);
+
+  int32 multiplier, shift;
+  QuantizeMultiplier(real_scale, &multiplier, &shift);
+
+  constexpr int kMinDepthPerThread = 8;
+  int thread_count = output_depth / kMinDepthPerThread;
+  thread_count = thread_count > 0 ? thread_count : 1;
+  const int capped_thread_count =
+      std::min(thread_count, cpu_backend_context->max_num_threads());
+
+  if (capped_thread_count == 1) {
+    MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
+             output_shape, output_data, 0, output_depth);
+  } else {
+    // Instead parallel for batch, we loop for the output_depth since batch
+    // is typical 1.
+    std::vector<MeanWorkerTask> tasks;
+    // TODO(b/131746020) don't create new heap allocations every time.
+    // At least we make it a single heap allocation by using reserve().
+    tasks.reserve(capped_thread_count);
+    int depth_start = 0;
+    for (int i = 0; i < capped_thread_count; ++i) {
+      // Try to distribute the tasks as even as possible.
+      int depth_end = depth_start +
+                      (output_depth - depth_start) / (capped_thread_count - i);
+      tasks.emplace_back(op_params, input_shape, input_data, multiplier, shift,
+                         bias, output_shape, output_data, depth_start,
+                         depth_end);
+      depth_start = depth_end;
+    }
+    cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
+                                    cpu_backend_context);
+  }
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MEAN_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/mul.h

@@ -0,0 +1,266 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MUL_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MUL_H_
+
+#include <algorithm>
+
+#include "fixedpoint/fixedpoint.h"
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/reference/integer_ops/mul.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+// Element-wise mul that can often be used for inner loop of broadcast Mul as
+// well as the non-broadcast Mul.
+inline void MulElementwise(int size, const ArithmeticParams& params,
+                           const int8* input1_data, const int8* input2_data,
+                           int8* output_data) {
+  ruy::profiler::ScopeLabel label("MulElementwiseInt8/8bit");
+  int i = 0;
+  TFLITE_DCHECK_GT(params.input1_offset, -256);
+  TFLITE_DCHECK_LT(params.input1_offset, 256);
+  TFLITE_DCHECK_GT(params.input2_offset, -256);
+  TFLITE_DCHECK_LT(params.input2_offset, 256);
+  TFLITE_DCHECK_GT(params.output_offset, -256);
+  TFLITE_DCHECK_LT(params.output_offset, 256);
+#ifdef USE_NEON
+  const int16x8_t input1_offset_vector = vdupq_n_s16(params.input1_offset);
+  const int16x8_t input2_offset_vector = vdupq_n_s16(params.input2_offset);
+  const int16x8_t output_offset_vector = vdupq_n_s16(params.output_offset);
+  const auto output_activation_min_vector =
+      vdupq_n_s8(params.quantized_activation_min);
+  const auto output_activation_max_vector =
+      vdupq_n_s8(params.quantized_activation_max);
+  const int left_shift = std::max(0, params.output_shift);
+  const int right_shift = std::max(0, -params.output_shift);
+  const int32x4_t left_shift_vec = vdupq_n_s32(left_shift);
+  for (; i <= size - 16; i += 16) {
+    // We load / store 16 at a time, multiplying as four sets of 4 int32s.
+    const int8x16_t input1_val_original = vld1q_s8(input1_data + i);
+    const int8x16_t input2_val_original = vld1q_s8(input2_data + i);
+
+    const int16x8_t input1_val_s16_high =
+        vmovl_s8(vget_high_s8(input1_val_original));
+    const int16x8_t input1_val_s16_low =
+        vmovl_s8(vget_low_s8(input1_val_original));
+
+    const int16x8_t input2_val_s16_high =
+        vmovl_s8(vget_high_s8(input2_val_original));
+    const int16x8_t input2_val_s16_low =
+        vmovl_s8(vget_low_s8(input2_val_original));
+    const int16x8_t input1_val_high =
+        vaddq_s16(input1_val_s16_high, input1_offset_vector);
+    const int16x8_t input2_val_high =
+        vaddq_s16(input2_val_s16_high, input2_offset_vector);
+    const int16x8_t input1_val_low =
+        vaddq_s16(input1_val_s16_low, input1_offset_vector);
+    const int16x8_t input2_val_low =
+        vaddq_s16(input2_val_s16_low, input2_offset_vector);
+    const int16x4_t input1_val_high_high = vget_high_s16(input1_val_high);
+    const int16x4_t input1_val_high_low = vget_low_s16(input1_val_high);
+    const int16x4_t input1_val_low_high = vget_high_s16(input1_val_low);
+    const int16x4_t input1_val_low_low = vget_low_s16(input1_val_low);
+    const int16x4_t input2_val_high_high = vget_high_s16(input2_val_high);
+    const int16x4_t input2_val_high_low = vget_low_s16(input2_val_high);
+    const int16x4_t input2_val_low_high = vget_high_s16(input2_val_low);
+    const int16x4_t input2_val_low_low = vget_low_s16(input2_val_low);
+
+    auto p1 = vmull_s16(input2_val_high_high, input1_val_high_high);
+    auto p2 = vmull_s16(input2_val_high_low, input1_val_high_low);
+    auto p3 = vmull_s16(input2_val_low_high, input1_val_low_high);
+    auto p4 = vmull_s16(input2_val_low_low, input1_val_low_low);
+
+    p1 = vshlq_s32(p1, left_shift_vec);
+    p2 = vshlq_s32(p2, left_shift_vec);
+    p3 = vshlq_s32(p3, left_shift_vec);
+    p4 = vshlq_s32(p4, left_shift_vec);
+
+    p1 = vqrdmulhq_n_s32(p1, params.output_multiplier);
+    p2 = vqrdmulhq_n_s32(p2, params.output_multiplier);
+    p3 = vqrdmulhq_n_s32(p3, params.output_multiplier);
+    p4 = vqrdmulhq_n_s32(p4, params.output_multiplier);
+    using gemmlowp::RoundingDivideByPOT;
+    p1 = RoundingDivideByPOT(p1, right_shift);
+    p2 = RoundingDivideByPOT(p2, right_shift);
+    p3 = RoundingDivideByPOT(p3, right_shift);
+    p4 = RoundingDivideByPOT(p4, right_shift);
+
+    const auto p1_narrowed = vqmovn_s32(p1);
+    const auto p2_narrowed = vqmovn_s32(p2);
+    const auto p3_narrowed = vqmovn_s32(p3);
+    const auto p4_narrowed = vqmovn_s32(p4);
+
+    const int16x8_t p_part1 =
+        vaddq_s16(vcombine_s16(p2_narrowed, p1_narrowed), output_offset_vector);
+    const int16x8_t p_part2 =
+        vaddq_s16(vcombine_s16(p4_narrowed, p3_narrowed), output_offset_vector);
+    const int8x16_t p = vcombine_s8(vqmovn_s16(p_part2), vqmovn_s16(p_part1));
+
+    const auto clamped = vmaxq_s8(output_activation_min_vector,
+                                  vminq_s8(output_activation_max_vector, p));
+    vst1q_s8(output_data + i, clamped);
+  }
+#endif  // NEON
+
+  for (; i < size; ++i) {
+    const int32 input1_val = params.input1_offset + input1_data[i];
+    const int32 input2_val = params.input2_offset + input2_data[i];
+    const int32 unclamped_result =
+        params.output_offset +
+        MultiplyByQuantizedMultiplier(input1_val * input2_val,
+                                      params.output_multiplier,
+                                      params.output_shift);
+    const int32 clamped_output =
+        std::min(params.quantized_activation_max,
+                 std::max(params.quantized_activation_min, unclamped_result));
+    output_data[i] = static_cast<int8>(clamped_output);
+  }
+}
+
+// Broadcast mul that can often be used for inner loop of broadcast Mul.
+inline void MulSimpleBroadcast(int size, const ArithmeticParams& params,
+                               const int8 broadcast_value,
+                               const int8* input2_data, int8* output_data) {
+  ruy::profiler::ScopeLabel label("BroadMulSimpleBroadcastInt8/8bit");
+  const int16 input1_val = params.input1_offset + broadcast_value;
+
+  int i = 0;
+  TFLITE_DCHECK_GT(params.input1_offset, -256);
+  TFLITE_DCHECK_LT(params.input1_offset, 256);
+  TFLITE_DCHECK_GT(params.input2_offset, -256);
+  TFLITE_DCHECK_LT(params.input2_offset, 256);
+  TFLITE_DCHECK_GT(params.output_offset, -256);
+  TFLITE_DCHECK_LT(params.output_offset, 256);
+#ifdef USE_NEON
+  const auto input2_offset_vector = vdupq_n_s16(params.input2_offset);
+  const auto output_offset_vector = vdupq_n_s16(params.output_offset);
+  const auto output_activation_min_vector =
+      vdupq_n_s8(params.quantized_activation_min);
+  const auto output_activation_max_vector =
+      vdupq_n_s8(params.quantized_activation_max);
+  const int left_shift = std::max(0, params.output_shift);
+  const int right_shift = std::max(0, -params.output_shift);
+  const int32x4_t left_shift_vec = vdupq_n_s32(left_shift);
+  for (; i <= size - 16; i += 16) {
+    // We load / store 16 at a time, multiplying as four sets of 4 int32s.
+    const auto input2_val_original = vld1q_s8(input2_data + i);
+    const auto input2_val_s16_high =
+        vmovl_s8(vget_high_s8(input2_val_original));
+    const auto input2_val_s16_low = vmovl_s8(vget_low_s8(input2_val_original));
+
+    const auto input2_val_high =
+        vaddq_s16(input2_val_s16_high, input2_offset_vector);
+    const auto input2_val_low =
+        vaddq_s16(input2_val_s16_low, input2_offset_vector);
+
+    const auto input2_val_low_low = vget_low_s16(input2_val_low);
+    const auto input2_val_low_high = vget_high_s16(input2_val_low);
+    const auto input2_val_high_low = vget_low_s16(input2_val_high);
+    const auto input2_val_high_high = vget_high_s16(input2_val_high);
+
+    auto p1 = vmull_n_s16(input2_val_high_high, input1_val);
+    auto p2 = vmull_n_s16(input2_val_high_low, input1_val);
+    auto p3 = vmull_n_s16(input2_val_low_high, input1_val);
+    auto p4 = vmull_n_s16(input2_val_low_low, input1_val);
+
+    p1 = vshlq_s32(p1, left_shift_vec);
+    p2 = vshlq_s32(p2, left_shift_vec);
+    p3 = vshlq_s32(p3, left_shift_vec);
+    p4 = vshlq_s32(p4, left_shift_vec);
+
+    p1 = vqrdmulhq_n_s32(p1, params.output_multiplier);
+    p2 = vqrdmulhq_n_s32(p2, params.output_multiplier);
+    p3 = vqrdmulhq_n_s32(p3, params.output_multiplier);
+    p4 = vqrdmulhq_n_s32(p4, params.output_multiplier);
+    using gemmlowp::RoundingDivideByPOT;
+    p1 = RoundingDivideByPOT(p1, right_shift);
+    p2 = RoundingDivideByPOT(p2, right_shift);
+    p3 = RoundingDivideByPOT(p3, right_shift);
+    p4 = RoundingDivideByPOT(p4, right_shift);
+
+    const auto p1_narrowed = vqmovn_s32(p1);
+    const auto p2_narrowed = vqmovn_s32(p2);
+    const auto p3_narrowed = vqmovn_s32(p3);
+    const auto p4_narrowed = vqmovn_s32(p4);
+
+    const int16x8_t p_part1 =
+        vaddq_s16(vcombine_s16(p2_narrowed, p1_narrowed), output_offset_vector);
+    const int16x8_t p_part2 =
+        vaddq_s16(vcombine_s16(p4_narrowed, p3_narrowed), output_offset_vector);
+    const int8x16_t p = vcombine_s8(vqmovn_s16(p_part2), vqmovn_s16(p_part1));
+
+    const auto clamped = vmaxq_s8(output_activation_min_vector,
+                                  vminq_s8(output_activation_max_vector, p));
+    vst1q_s8(output_data + i, clamped);
+  }
+#endif  // NEON
+
+  for (; i < size; ++i) {
+    const int32 input2_val = params.input2_offset + input2_data[i];
+    const int32 unclamped_result =
+        params.output_offset +
+        MultiplyByQuantizedMultiplier(input1_val * input2_val,
+                                      params.output_multiplier,
+                                      params.output_shift);
+    const int32 clamped_output =
+        std::min(params.quantized_activation_max,
+                 std::max(params.quantized_activation_min, unclamped_result));
+    output_data[i] = static_cast<int8>(clamped_output);
+  }
+}
+
+inline void Mul(const ArithmeticParams& params,
+                const RuntimeShape& input1_shape, const int8* input1_data,
+                const RuntimeShape& input2_shape, const int8* input2_data,
+                const RuntimeShape& output_shape, int8* output_data) {
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  ruy::profiler::ScopeLabel label("MulInt8/8bit");
+  const int flat_size =
+      MatchingElementsSize(input1_shape, input2_shape, output_shape);
+
+  MulElementwise(flat_size, params, input1_data, input2_data, output_data);
+}
+
+inline void BroadcastMulDispatch(const ArithmeticParams& params,
+                                 const RuntimeShape& input1_shape,
+                                 const int8* input1_data,
+                                 const RuntimeShape& input2_shape,
+                                 const int8* input2_data,
+                                 const RuntimeShape& output_shape,
+                                 int8* output_data) {
+  if (params.broadcast_category == BroadcastableOpCategory::kGenericBroadcast) {
+    return reference_integer_ops::BroadcastMul6DSlow(
+        params, input1_shape, input1_data, input2_shape, input2_data,
+        output_shape, output_data);
+  }
+
+  optimized_ops::BinaryBroadcastFiveFold(
+      params, input1_shape, input1_data, input2_shape, input2_data,
+      output_shape, output_data, MulElementwise, MulSimpleBroadcast);
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_MUL_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/pooling.h

@@ -0,0 +1,278 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_POOLING_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_POOLING_H_
+
+#include <string.h>
+
+#include <algorithm>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/im2col_utils.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/reference/reference_ops.h"
+#include "tensorflow/lite/kernels/internal/strided_slice_logic.h"
+#include "tensorflow/lite/kernels/internal/tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+inline void MaxPool(const PoolParams& params, const RuntimeShape& input_shape,
+                    const int8_t* input_data, const RuntimeShape& output_shape,
+                    int8_t* output_data) {
+  ruy::profiler::ScopeLabel label("MaxPool/8bit");
+
+  // Here, and in other pooling ops, in order to maintain locality of reference,
+  // to minimize some recalculations, and to load into NEON vector registers, we
+  // use an inner loop down the depth. Since depths can be large and hence we
+  // would need arbitrarily large temporary storage, we divide the work up into
+  // depth tranches just within the batch loop.
+  static constexpr int kPoolingAccTrancheSize = 256;
+
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int depth = MatchingDim(input_shape, 3, output_shape, 3);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  const int stride_height = params.stride_height;
+  const int stride_width = params.stride_width;
+
+  int8_t acc[kPoolingAccTrancheSize];
+  for (int batch = 0; batch < batches; ++batch) {
+    // We proceed through the depth in tranches (see comment above). The
+    // depth_base is the depth at the beginning of the tranche. The
+    // tranche_depth is the depth dimension of the tranche.
+    for (int depth_base = 0; depth_base < depth;
+         depth_base += kPoolingAccTrancheSize) {
+      const int tranche_depth =
+          std::min(depth - depth_base, kPoolingAccTrancheSize);
+      for (int out_y = 0; out_y < output_height; ++out_y) {
+        for (int out_x = 0; out_x < output_width; ++out_x) {
+          const int in_x_origin =
+              (out_x * stride_width) - params.padding_values.width;
+          const int in_y_origin =
+              (out_y * stride_height) - params.padding_values.height;
+          const int filter_x_start = std::max(0, -in_x_origin);
+          const int filter_x_end =
+              std::min(params.filter_width, input_width - in_x_origin);
+          const int filter_y_start = std::max(0, -in_y_origin);
+          const int filter_y_end =
+              std::min(params.filter_height, input_height - in_y_origin);
+          memset(acc, params.quantized_activation_min,
+                 tranche_depth * sizeof(acc[0]));
+          const int8_t* input_ptr =
+              input_data + depth_base +
+              depth * (in_x_origin +
+                       input_width * (in_y_origin + input_height * batch));
+          for (int fy = filter_y_start; fy < filter_y_end; fy++) {
+            const int8_t* input_row_ptr =
+                input_ptr + depth * (fy * input_width + filter_x_start);
+            for (int fx = filter_x_start; fx < filter_x_end; fx++) {
+              const int8_t* input_channel_ptr = input_row_ptr;
+              int channel = 0;
+#ifdef USE_NEON
+              for (; channel <= tranche_depth - 16; channel += 16) {
+                int8x16_t acc_reg = vld1q_s8(acc + channel);
+                int8x16_t input_reg = vld1q_s8(input_channel_ptr);
+                input_channel_ptr += 16;
+                acc_reg = vmaxq_s8(acc_reg, input_reg);
+                vst1q_s8(acc + channel, acc_reg);
+              }
+
+              for (; channel <= tranche_depth - 8; channel += 8) {
+                int8x8_t acc_reg = vld1_s8(acc + channel);
+                int8x8_t input_reg = vld1_s8(input_channel_ptr);
+                input_channel_ptr += 8;
+                acc_reg = vmax_s8(acc_reg, input_reg);
+                vst1_s8(acc + channel, acc_reg);
+              }
+#endif
+              for (; channel < tranche_depth; ++channel) {
+                acc[channel] = std::max(acc[channel], *input_channel_ptr++);
+              }
+              input_row_ptr += depth;
+            }
+          }
+          int8_t* output_ptr = output_data + Offset(output_shape, batch, out_y,
+                                                    out_x, depth_base);
+          int channel = 0;
+#ifdef USE_NEON
+          for (; channel <= tranche_depth - 16; channel += 16) {
+            int8x16_t a = vld1q_s8(acc + channel);
+            a = vminq_s8(a, vdupq_n_s8(params.quantized_activation_max));
+            a = vmaxq_s8(a, vdupq_n_s8(params.quantized_activation_min));
+            vst1q_s8(output_ptr + channel, a);
+          }
+          for (; channel <= tranche_depth - 8; channel += 8) {
+            int8x8_t a = vld1_s8(acc + channel);
+            a = vmin_s8(a, vdup_n_s8(params.quantized_activation_max));
+            a = vmax_s8(a, vdup_n_s8(params.quantized_activation_min));
+            vst1_s8(output_ptr + channel, a);
+          }
+#endif
+          for (; channel < tranche_depth; ++channel) {
+            int8_t a = acc[channel];
+            a = std::max<int8_t>(a, params.quantized_activation_min);
+            a = std::min<int8_t>(a, params.quantized_activation_max);
+            output_ptr[channel] = static_cast<int8_t>(a);
+          }
+        }
+      }
+    }
+  }
+}
+
+inline bool AveragePool(const PoolParams& params,
+                        const RuntimeShape& input_shape,
+                        const int8_t* input_data,
+                        const RuntimeShape& output_shape, int8_t* output_data) {
+  ruy::profiler::ScopeLabel label("AveragePool/8bitWith32bitAccumulator");
+
+  // Here, and in other pooling ops, in order to maintain locality of reference,
+  // to minimize some recalculations, and to load into NEON vector registers, we
+  // use an inner loop down the depth. Since depths can be large and hence we
+  // would need arbitrarily large temporary storage, we divide the work up into
+  // depth tranches just within the batch loop.
+  static constexpr int kPoolingAccTrancheSize = 256;
+
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int depth = MatchingDim(input_shape, 3, output_shape, 3);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  const int stride_height = params.stride_height;
+  const int stride_width = params.stride_width;
+
+  int32_t acc[kPoolingAccTrancheSize];
+  for (int batch = 0; batch < batches; ++batch) {
+    // We proceed through the depth in tranches (see comment above). The
+    // depth_base is the depth at the beginning of the tranche. The
+    // tranche_depth is the depth dimension of the tranche.
+    for (int depth_base = 0; depth_base < depth;
+         depth_base += kPoolingAccTrancheSize) {
+      const int tranche_depth =
+          std::min(depth - depth_base, kPoolingAccTrancheSize);
+      for (int out_y = 0; out_y < output_height; ++out_y) {
+        for (int out_x = 0; out_x < output_width; ++out_x) {
+          const int in_x_origin =
+              (out_x * stride_width) - params.padding_values.width;
+          const int in_y_origin =
+              (out_y * stride_height) - params.padding_values.height;
+          const int filter_x_start = std::max(0, -in_x_origin);
+          const int filter_x_end =
+              std::min(params.filter_width, input_width - in_x_origin);
+          const int filter_y_start = std::max(0, -in_y_origin);
+          const int filter_y_end =
+              std::min(params.filter_height, input_height - in_y_origin);
+          const int filter_count =
+              (filter_x_end - filter_x_start) * (filter_y_end - filter_y_start);
+          if (filter_count == 0) return false;
+          memset(acc, 0, tranche_depth * sizeof(acc[0]));
+          const int8_t* input_ptr =
+              input_data + depth_base +
+              depth * (in_x_origin +
+                       input_width * (in_y_origin + input_height * batch));
+          for (int fy = filter_y_start; fy < filter_y_end; fy++) {
+            const int8_t* input_row_ptr =
+                input_ptr + depth * (fy * input_width + filter_x_start);
+            for (int fx = filter_x_start; fx < filter_x_end; fx++) {
+              const int8_t* input_channel_ptr = input_row_ptr;
+              int channel = 0;
+#ifdef USE_NEON
+              for (; channel <= tranche_depth - 16; channel += 16) {
+                int16x4_t acc_reg[4];
+                int8x16_t input_reg = vld1q_s8(input_channel_ptr);
+                input_channel_ptr += 16;
+                acc_reg[0] = vget_low_s16(vmovl_s8(vget_low_s8(input_reg)));
+                acc_reg[1] = vget_high_s16(vmovl_s8(vget_low_s8(input_reg)));
+                acc_reg[2] = vget_low_s16(vmovl_s8(vget_high_s8(input_reg)));
+                acc_reg[3] = vget_high_s16(vmovl_s8(vget_high_s8(input_reg)));
+                for (int i = 0; i < 4; i++) {
+                  vst1q_s32(
+                      acc + channel + 4 * i,
+                      vaddw_s16(vld1q_s32(acc + channel + 4 * i), acc_reg[i]));
+                }
+              }
+              for (; channel <= tranche_depth - 8; channel += 8) {
+                int16x4_t acc_reg[2];
+                int16x8_t input_reg = vmovl_s8(vld1_s8(input_channel_ptr));
+                input_channel_ptr += 8;
+                acc_reg[0] = vget_low_s16(input_reg);
+                acc_reg[1] = vget_high_s16(input_reg);
+                for (int i = 0; i < 2; i++) {
+                  vst1q_s32(
+                      acc + channel + 4 * i,
+                      vaddw_s16(vld1q_s32(acc + channel + 4 * i), acc_reg[i]));
+                }
+              }
+#endif
+              for (; channel < tranche_depth; ++channel) {
+                acc[channel] += *input_channel_ptr++;
+              }
+              input_row_ptr += depth;
+            }
+          }
+          int8_t* output_ptr = output_data + Offset(output_shape, batch, out_y,
+                                                    out_x, depth_base);
+          int channel = 0;
+#ifdef USE_NEON
+          for (; channel <= tranche_depth - 8; channel += 8) {
+            int16_t buf[8];
+            for (int i = 0; i < 8; i++) {
+              buf[i] =
+                  acc[channel + i] > 0
+                      ? (acc[channel + i] + filter_count / 2) / filter_count
+                      : (acc[channel + i] - filter_count / 2) / filter_count;
+            }
+            int8x8_t buf8 = vqmovn_s16(vld1q_s16(buf));
+            buf8 = vmin_s8(buf8, vdup_n_s8(params.quantized_activation_max));
+            buf8 = vmax_s8(buf8, vdup_n_s8(params.quantized_activation_min));
+            vst1_s8(output_ptr + channel, buf8);
+          }
+#endif
+          for (; channel < tranche_depth; ++channel) {
+            int16_t a = acc[channel] > 0
+                            ? (acc[channel] + filter_count / 2) / filter_count
+                            : (acc[channel] - filter_count / 2) / filter_count;
+            a = std::max<int16_t>(a, params.quantized_activation_min);
+            a = std::min<int16_t>(a, params.quantized_activation_max);
+            output_ptr[channel] = static_cast<int8_t>(a);
+          }
+        }
+      }
+    }
+  }
+  return true;
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_POOLING_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/sub.h

@@ -0,0 +1,237 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_SUB_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_SUB_H_
+
+#include <algorithm>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/optimized/avx2_quantization_utils.h"
+#include "tensorflow/lite/kernels/internal/reference/sub.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+inline void SubElementwiseInt16(int size, const ArithmeticParams& params,
+                                const int16* input1_data,
+                                const int16* input2_data, int16* output_data) {
+  ruy::profiler::ScopeLabel label("SubElementwiseInt16/16bit");
+  int i = 0;
+  TFLITE_DCHECK_GT(params.input1_offset, -32768);
+  TFLITE_DCHECK_GT(params.input2_offset, -32768);
+  TFLITE_DCHECK_LT(params.input1_offset, 32768);
+  TFLITE_DCHECK_LT(params.input2_offset, 32768);
+
+#ifdef __AVX2__
+  const int32_t input1_left_shift = params.left_shift + params.input1_shift;
+  const int32_t input2_left_shift = params.left_shift + params.input2_shift;
+  const __m256i input1_offset = _mm256_set1_epi32(params.input1_offset);
+  const __m256i input2_offset = _mm256_set1_epi32(params.input2_offset);
+  const __m256i output_offset = _mm256_set1_epi32(params.output_offset);
+  const __m256i clamp_max_v =
+      _mm256_set1_epi32(params.quantized_activation_max);
+  const __m256i clamp_min_v =
+      _mm256_set1_epi32(params.quantized_activation_min);
+
+  for (; i <= size - 16; i += 16) {
+    const __m256i input1_val_original =
+        _mm256_loadu_si256(reinterpret_cast<__m256i const*>(input1_data + i));
+    const __m256i input2_val_original =
+        _mm256_loadu_si256(reinterpret_cast<__m256i const*>(input2_data + i));
+
+    __m256i s11 =
+        _mm256_cvtepi16_epi32(_mm256_castsi256_si128(input1_val_original));
+    __m256i s12 =
+        _mm256_cvtepi16_epi32(_mm256_extracti128_si256(input1_val_original, 1));
+    __m256i s21 =
+        _mm256_cvtepi16_epi32(_mm256_castsi256_si128(input2_val_original));
+    __m256i s22 =
+        _mm256_cvtepi16_epi32(_mm256_extracti128_si256(input2_val_original, 1));
+
+    s11 = _mm256_add_epi32(s11, input1_offset);
+    s12 = _mm256_add_epi32(s12, input1_offset);
+    s21 = _mm256_add_epi32(s21, input2_offset);
+    s22 = _mm256_add_epi32(s22, input2_offset);
+
+    s11 = avx2_utils::MultiplyByQuantizedMultiplier(
+        s11, params.input1_multiplier, input1_left_shift);
+    s12 = avx2_utils::MultiplyByQuantizedMultiplier(
+        s12, params.input1_multiplier, input1_left_shift);
+    s21 = avx2_utils::MultiplyByQuantizedMultiplier(
+        s21, params.input2_multiplier, input2_left_shift);
+    s22 = avx2_utils::MultiplyByQuantizedMultiplier(
+        s22, params.input2_multiplier, input2_left_shift);
+
+    __m256i s1 = _mm256_sub_epi32(s11, s21);
+    __m256i s2 = _mm256_sub_epi32(s12, s22);
+
+    s1 = avx2_utils::MultiplyByQuantizedMultiplier(s1, params.output_multiplier,
+                                                   params.output_shift);
+    s2 = avx2_utils::MultiplyByQuantizedMultiplier(s2, params.output_multiplier,
+                                                   params.output_shift);
+
+    s1 = _mm256_add_epi32(s1, output_offset);
+    s2 = _mm256_add_epi32(s2, output_offset);
+
+    s1 = _mm256_min_epi32(s1, clamp_max_v);
+    s1 = _mm256_max_epi32(s1, clamp_min_v);
+    s2 = _mm256_min_epi32(s2, clamp_max_v);
+    s2 = _mm256_max_epi32(s2, clamp_min_v);
+
+    avx2_utils::CastInt32ToInt16AndStore(output_data + i, s1);
+    avx2_utils::CastInt32ToInt16AndStore(output_data + i + 8, s2);
+  }
+#endif  // __AVX2__
+
+  for (; i < size; ++i) {
+    const int32_t input1_val = params.input1_offset + input1_data[i];
+    const int32_t input2_val = params.input2_offset + input2_data[i];
+    const int32_t shifted_input1_val = input1_val * (1 << params.left_shift);
+    const int32_t shifted_input2_val = input2_val * (1 << params.left_shift);
+    const int32_t scaled_input1_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input1_val, params.input1_multiplier, params.input1_shift);
+    const int32_t scaled_input2_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input2_val, params.input2_multiplier, params.input2_shift);
+    const int32_t raw_sum = scaled_input1_val - scaled_input2_val;
+    const int32_t raw_output =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            raw_sum, params.output_multiplier, params.output_shift) +
+        params.output_offset;
+    const int32_t clamped_output =
+        std::min(params.quantized_activation_max,
+                 std::max(params.quantized_activation_min, raw_output));
+    output_data[i] = static_cast<int16>(clamped_output);
+  }
+}
+
+inline void BroadcastSubFiveFold(const ArithmeticParams& unswitched_params,
+                                 const RuntimeShape& input1_shape,
+                                 const int16* unswitched_input1_data,
+                                 const RuntimeShape& input2_shape,
+                                 const int16* unswitched_input2_data,
+                                 const RuntimeShape& output_shape,
+                                 int16* output_data) {
+  ruy::profiler::ScopeLabel label("BroadcastSubFiveFold/16bit");
+
+  ArithmeticParams switched_params = unswitched_params;
+  switched_params.input1_offset = unswitched_params.input2_offset;
+  switched_params.input1_multiplier = unswitched_params.input2_multiplier;
+  switched_params.input1_shift = unswitched_params.input2_shift;
+  switched_params.input2_offset = unswitched_params.input1_offset;
+  switched_params.input2_multiplier = unswitched_params.input1_multiplier;
+  switched_params.input2_shift = unswitched_params.input1_shift;
+
+  const bool use_unswitched =
+      unswitched_params.broadcast_category ==
+      tflite::BroadcastableOpCategory::kFirstInputBroadcastsFast;
+
+  const ArithmeticParams& params =
+      use_unswitched ? unswitched_params : switched_params;
+  const int16_t* input1_data =
+      use_unswitched ? unswitched_input1_data : unswitched_input2_data;
+  const int16_t* input2_data =
+      use_unswitched ? unswitched_input2_data : unswitched_input1_data;
+
+  int16_t* output_data_ptr = output_data;
+  const int16_t* input1_data_ptr = input1_data;
+  const int16_t* input2_data_reset = input2_data;
+  // In the fivefold pattern, y0, y2 and y4 are not broadcast, and so shared
+  // between input shapes. y3 for input 1 is always broadcast, and so the
+  // dimension there is 1, whereas optionally y1 might be broadcast for input 2.
+  // The flatsize for each inputs are as below.
+  // input1.shape.FlatSize = y0 * y1 * y2 * y4,
+  // input2.shape.FlatSize = y0 * y2 * y3 * y4.
+  const int y0 = params.broadcast_shape[0];
+  const int y1 = params.broadcast_shape[1];
+  const int y2 = params.broadcast_shape[2];
+  const int y3 = params.broadcast_shape[3];
+  const int y4 = params.broadcast_shape[4];
+  for (int i0 = 0; i0 < y0; ++i0) {
+    const int16_t* input2_data_ptr = nullptr;
+    for (int i1 = 0; i1 < y1; ++i1) {
+      input2_data_ptr = input2_data_reset;
+      for (int i2 = 0; i2 < y2; ++i2) {
+        for (int i3 = 0; i3 < y3; ++i3) {
+          if (use_unswitched) {
+            SubElementwiseInt16(y4, params, input1_data_ptr, input2_data_ptr,
+                                output_data_ptr);
+          } else {
+            // When input1 and input2 are switched, calculate (input2 - input1)
+            // and use unswitched_params as we switch the switched input here.
+            SubElementwiseInt16(y4, unswitched_params, input2_data_ptr,
+                                input1_data_ptr, output_data_ptr);
+          }
+          input2_data_ptr += y4;
+          output_data_ptr += y4;
+        }
+        // We have broadcast y4 of input1 data y3 times, and now move on.
+        input1_data_ptr += y4;
+      }
+    }
+    // We have broadcast y2*y3*y4 of input2 data y1 times, and now move on.
+    input2_data_reset = input2_data_ptr;
+  }
+}
+
+inline void Sub(const ArithmeticParams& params,
+                const RuntimeShape& input1_shape, const int16* input1_data,
+                const RuntimeShape& input2_shape, const int16* input2_data,
+                const RuntimeShape& output_shape, int16* output_data) {
+  ruy::profiler::ScopeLabel label("SubInt16/16bit");
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  TFLITE_DCHECK_GT(params.input1_offset, -32768);
+  TFLITE_DCHECK_GT(params.input2_offset, -32768);
+  TFLITE_DCHECK_LT(params.input1_offset, 32768);
+  TFLITE_DCHECK_LT(params.input2_offset, 32768);
+
+  const int flat_size =
+      MatchingElementsSize(input1_shape, input2_shape, output_shape);
+  SubElementwiseInt16(flat_size, params, input1_data, input2_data, output_data);
+}
+
+inline void BroadcastSubDispatch(const ArithmeticParams& params,
+                                 const RuntimeShape& input1_shape,
+                                 const int16* input1_data,
+                                 const RuntimeShape& input2_shape,
+                                 const int16* input2_data,
+                                 const RuntimeShape& output_shape,
+                                 int16* output_data) {
+  ruy::profiler::ScopeLabel label("BroadcastSubDispatchInt16/16bit");
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  TFLITE_DCHECK_GT(params.input1_offset, -32768);
+  TFLITE_DCHECK_GT(params.input2_offset, -32768);
+  TFLITE_DCHECK_LT(params.input1_offset, 32768);
+  TFLITE_DCHECK_LT(params.input2_offset, 32768);
+
+  if (params.broadcast_category == BroadcastableOpCategory::kGenericBroadcast) {
+    return reference_ops::BroadcastQuantSubSlow(
+        params, input1_shape, input1_data, input2_shape, input2_data,
+        output_shape, output_data);
+  }
+
+  BroadcastSubFiveFold(params, input1_shape, input1_data, input2_shape,
+                       input2_data, output_shape, output_data);
+}
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_SUB_H_

tensorflow/lite/kernels/internal/optimized/integer_ops/transpose_conv.h

@@ -0,0 +1,118 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_TRANSPOSE_CONV_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_TRANSPOSE_CONV_H_
+
+#include <algorithm>
+
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+
+namespace tflite {
+namespace optimized_integer_ops {
+
+// TransposeConvV2 expect the weights in HWOI order.
+template <typename InputScalar, typename DestinationScalar>
+inline void TransposeConvV2(
+    const ConvParams& params, const int32* output_multiplier,
+    const int32* output_shift, const RuntimeShape& input_shape,
+    const InputScalar* input_data,
+    const RuntimeShape& hwoi_ordered_filter_shape,
+    const int8_t* hwoi_ordered_filter_data, const RuntimeShape& bias_shape,
+    const int32* bias_data, const RuntimeShape& output_shape,
+    DestinationScalar* output_data, const RuntimeShape& col2im_shape,
+    int32_t* col2im_data, int32_t* scratch_data,
+    CpuBackendContext* cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("TransposeConvV2/int8");
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(hwoi_ordered_filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK(col2im_data);
+  TFLITE_DCHECK(hwoi_ordered_filter_data);
+
+  const int batch_size = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_image_size = input_shape.Dims(1) * input_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  const int output_image_size = output_height * output_width;
+  const int input_depth =
+      MatchingDim(input_shape, 3, hwoi_ordered_filter_shape, 3);
+  const int output_depth =
+      MatchingDim(output_shape, 3, hwoi_ordered_filter_shape, 2);
+  const int input_offset = input_image_size * input_depth;
+  const int output_offset = output_image_size * output_depth;
+
+  const int filter_height = hwoi_ordered_filter_shape.Dims(0);
+  const int filter_width = hwoi_ordered_filter_shape.Dims(1);
+  const int padding_top = params.padding_values.height;
+  const int padding_bottom =
+      params.padding_values.height + params.padding_values.height_offset;
+  const int padding_left = params.padding_values.width;
+  const int padding_right =
+      params.padding_values.width + params.padding_values.width_offset;
+  const int stride_height = params.stride_height;
+  const int stride_width = params.stride_width;
+
+  const int32 output_activation_min = params.quantized_activation_min;
+  const int32 output_activation_max = params.quantized_activation_max;
+
+  const int hwoi_ordered_filter_total_size =
+      filter_height * filter_width * output_depth;
+
+  cpu_backend_gemm::MatrixParams<int8_t> lhs_params;
+  lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
+  lhs_params.rows = hwoi_ordered_filter_total_size;
+  lhs_params.cols = input_depth;
+  // Since our weight is symmetric quantized, the zp will always be 0.
+  lhs_params.zero_point = 0;
+
+  int32_t* scratch_data_p = scratch_data;
+  std::fill_n(scratch_data, output_offset * batch_size, static_cast<int32>(0));
+  for (int i = 0; i < batch_size; ++i) {
+    cpu_backend_gemm::MatrixParams<InputScalar> rhs_params;
+    rhs_params.order = cpu_backend_gemm::Order::kColMajor;
+    rhs_params.rows = input_depth;
+    rhs_params.cols = input_image_size;
+    rhs_params.zero_point = -params.input_offset;
+
+    cpu_backend_gemm::MatrixParams<int32_t> dst_params;
+    dst_params.order = cpu_backend_gemm::Order::kColMajor;
+    dst_params.rows = hwoi_ordered_filter_total_size;
+    dst_params.cols = input_image_size;
+
+    cpu_backend_gemm::GemmParams<int32_t, int32_t> gemm_params;
+    cpu_backend_gemm::Gemm(lhs_params, hwoi_ordered_filter_data, rhs_params,
+                           input_data + input_offset * i, dst_params,
+                           col2im_data, gemm_params, cpu_backend_context);
+
+    optimized_ops::Col2im(
+        col2im_data, output_depth, output_height, output_width, filter_height,
+        filter_width, padding_top, padding_left, padding_bottom, padding_right,
+        stride_height, stride_width, scratch_data_p);
+
+    scratch_data_p += output_offset;
+  }
+  scratch_data_p = scratch_data;
+  optimized_ops::BiasAdd(scratch_data_p, bias_data, batch_size, output_height,
+                         output_width, output_depth);
+
+  optimized_ops::Quantize(output_multiplier, output_shift, output_depth,
+                          output_shape.FlatSize(), params.output_offset,
+                          output_activation_min, output_activation_max,
+                          scratch_data, output_data);
+}
+
+}  // namespace optimized_integer_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_INTEGER_OPS_TRANSPOSE_CONV_H_

tensorflow/lite/kernels/internal/optimized/multithreaded_conv.h

@@ -0,0 +1,184 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_MULTITHREADED_CONV_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_MULTITHREADED_CONV_H_
+
+#include <assert.h>
+#include <stdint.h>
+#include <sys/types.h>
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+#include <memory>
+#include <tuple>
+#include <type_traits>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/optimized/eigen_spatial_convolutions.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace multithreaded_ops {
+
+// Shorthands for the types we need when interfacing with the EigenTensor
+// library.
+typedef Eigen::TensorMap<
+    Eigen::Tensor<float, 2, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned>
+    EigenMatrix;
+typedef Eigen::TensorMap<
+    const Eigen::Tensor<float, 2, Eigen::RowMajor, Eigen::DenseIndex>,
+    Eigen::Aligned>
+    ConstEigenMatrix;
+
+typedef Eigen::TensorMap<
+    Eigen::Tensor<float, 4, Eigen::RowMajor, Eigen::DenseIndex>, Eigen::Aligned>
+    EigenTensor;
+typedef Eigen::TensorMap<
+    const Eigen::Tensor<float, 4, Eigen::RowMajor, Eigen::DenseIndex>,
+    Eigen::Aligned>
+    ConstEigenTensor;
+
+// Utility functions we need for the EigenTensor API.
+template <typename Device, typename T>
+struct MatMulConvFunctor {
+  // Computes on device "d": out = in0 * in1, where * is matrix
+  // multiplication.
+  void operator()(
+      const Device& d, EigenMatrix out, ConstEigenMatrix in0,
+      ConstEigenMatrix in1,
+      const Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1>& dim_pair) {
+    out.device(d) = in0.contract(in1, dim_pair);
+  }
+};
+
+template <class T>
+class EigenTensorConvFunctor {
+ private:
+  Eigen::PaddingType RuntimePadding2EigenPadding(PaddingType padding) {
+    switch (padding) {
+      case PaddingType::kValid:
+        return Eigen::PADDING_VALID;
+      case PaddingType::kSame:
+        return Eigen::PADDING_SAME;
+      case PaddingType::kNone:
+        assert(false);  // should never get here.
+        return Eigen::PADDING_VALID;
+    }
+    return Eigen::PADDING_SAME;  // Prevent compiler warning about missing
+                                 // return
+  }
+
+ public:
+  void operator()(const Eigen::ThreadPoolDevice& device, const T* input_data,
+                  int input_batches, int input_height, int input_width,
+                  int input_depth, const T* filter_data, int filter_height,
+                  int filter_width, int filter_count, int stride_rows,
+                  int stride_cols, int pad_width, int pad_height,
+                  PaddingType padding, T* output_data, int output_height,
+                  int output_width) {
+    const bool is_1x1_kernel = (filter_height == 1 && filter_width == 1 &&
+                                stride_rows == 1 && stride_cols == 1);
+    if (is_1x1_kernel) {
+      // For 1x1 kernel, the 2D convolution is reduced to matrix
+      // multiplication.
+      const int conv_width = output_height * output_width;
+      Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> dim_pair;
+      dim_pair[0] = Eigen::IndexPair<Eigen::DenseIndex>(1, 0);
+      EigenMatrix output(output_data, input_batches * conv_width, filter_count);
+      ConstEigenMatrix input(input_data, input_batches * conv_width,
+                             input_depth);
+      ConstEigenMatrix filter(filter_data, input_depth, filter_count);
+      MatMulConvFunctor<Eigen::ThreadPoolDevice, T>()(device, output, input,
+                                                      filter, dim_pair);
+    } else if (filter_height == input_height && filter_width == input_width &&
+               pad_width == 0 && pad_height == 0) {
+      // If the input data and filter have the same height/width,
+      // the 2D convolution is reduced to matrix multiplication.
+      const int k =  // Length of reduction dimension.
+          filter_width * filter_height * input_depth;
+      Eigen::array<Eigen::IndexPair<Eigen::DenseIndex>, 1> dim_pair;
+      dim_pair[0] = Eigen::IndexPair<Eigen::DenseIndex>(1, 0);
+      EigenMatrix output(output_data, input_batches, filter_count);
+      ConstEigenMatrix input(input_data, input_batches, k);
+      ConstEigenMatrix filter(filter_data, k, filter_count);
+      MatMulConvFunctor<Eigen::ThreadPoolDevice, T>()(device, output, input,
+                                                      filter, dim_pair);
+    } else {
+      EigenTensor output(output_data, input_batches, output_height,
+                         output_width, filter_count);
+      ConstEigenTensor input(input_data, input_batches, input_height,
+                             input_width, input_depth);
+      ConstEigenTensor filter(filter_data, filter_height, filter_width,
+                              input_depth, filter_count);
+      output.device(device) =
+          Eigen::SpatialConvolution(input, filter, stride_cols, stride_rows,
+                                    RuntimePadding2EigenPadding(padding));
+    }
+  }
+};
+
+inline void Conv(const Eigen::ThreadPoolDevice& device,
+                 const ConvParams& params, const RuntimeShape& input_shape,
+                 const float* input_data, const RuntimeShape& filter_shape,
+                 const float* filter_data, const RuntimeShape& bias_shape,
+                 const float* bias_data, const RuntimeShape& output_shape,
+                 float* output_data, const RuntimeShape& im2col_shape,
+                 float* im2col_data) {
+  // Nest profiling under "Conv", to aggregate with other kernels.
+  ruy::profiler::ScopeLabel label("Conv");
+  ruy::profiler::ScopeLabel inner_label("Multithreaded EigenTensor");
+
+  // im2col data should not be generated for the multi-thread supporting case.
+  TFLITE_DCHECK(!im2col_data);
+  (void)im2col_shape;
+  const int stride_width = params.stride_width;
+  const int stride_height = params.stride_height;
+  const PaddingType padding = params.padding_type;
+  const int pad_width = params.padding_values.width;
+  const int pad_height = params.padding_values.height;
+  const float output_activation_min = params.float_activation_min;
+  const float output_activation_max = params.float_activation_max;
+  TFLITE_DCHECK_EQ(input_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(filter_shape.DimensionsCount(), 4);
+  TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 4);
+
+  const int batches = MatchingDim(input_shape, 0, output_shape, 0);
+  const int input_depth = MatchingDim(input_shape, 3, filter_shape, 3);
+  const int output_depth = MatchingDim(filter_shape, 0, output_shape, 3);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const int filter_height = filter_shape.Dims(1);
+  const int filter_width = filter_shape.Dims(2);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  EigenTensorConvFunctor<float> conv_functor;
+  conv_functor(device, input_data, batches, input_height, input_width,
+               input_depth, filter_data, filter_height, filter_width,
+               output_depth, stride_height, stride_width, pad_height, pad_width,
+               padding, output_data, output_height, output_width);
+
+  optimized_ops::AddBiasAndEvalActivationFunction(
+      output_activation_min, output_activation_max, bias_shape, bias_data,
+      output_shape, output_data);
+}
+
+}  // namespace multithreaded_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_MULTITHREADED_CONV_H_

tensorflow/lite/kernels/internal/optimized/neon_check.h

@@ -0,0 +1,40 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_
+
+#if defined(__ARM_NEON__) || defined(__ARM_NEON)
+#define USE_NEON
+#include <arm_neon.h>  // IWYU pragma: export
+#endif
+
+#if defined __GNUC__ && defined __SSE4_1__ && !defined TF_LITE_DISABLE_X86_NEON
+#define USE_NEON
+#include "NEON_2_SSE.h"  // IWYU pragma: export
+#endif
+
+// NEON_OR_PORTABLE(SomeFunc, args) calls NeonSomeFunc(args) if USE_NEON is
+// defined, PortableSomeFunc(args) otherwise.
+#ifdef USE_NEON
+// Always use Neon code
+#define NEON_OR_PORTABLE(funcname, ...) Neon##funcname(__VA_ARGS__)
+
+#else
+// No NEON available: Use Portable code
+#define NEON_OR_PORTABLE(funcname, ...) Portable##funcname(__VA_ARGS__)
+
+#endif  // defined(USE_NEON)
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_

tensorflow/lite/kernels/internal/optimized/neon_tensor_utils.h

@@ -0,0 +1,334 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_tensor_utils_impl.h"
+#include "tensorflow/lite/kernels/internal/reference/portable_tensor_utils_impl.h"
+
+namespace tflite {
+namespace tensor_utils {
+
+void MatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
+                                         int m_cols, const float* vector,
+                                         int n_batch, float* result) {
+  NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                   vector, n_batch, result);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(const int8_t* __restrict__ matrix,
+                                         const int m_rows, const int m_cols,
+                                         const int8_t* __restrict__ vectors,
+                                         const float* scaling_factors,
+                                         int n_batch,
+                                         float* __restrict__ result) {
+  NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                   vectors, scaling_factors, n_batch, result);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(const int8_t* __restrict__ matrix,
+                                         const int m_rows, const int m_cols,
+                                         const int8_t* __restrict__ vectors,
+                                         const float* scaling_factors,
+                                         int n_batch, int32_t* scratch,
+                                         float* __restrict__ result,
+                                         CpuBackendContext* context) {
+  NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                   vectors, scaling_factors, n_batch, scratch, result, context);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
+    bool* compute_row_sums, CpuBackendContext* context) {
+  NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                   vectors, scaling_factors, n_batch, result, per_channel_scale,
+                   input_offset, scratch, row_sums, compute_row_sums, context);
+}
+
+void SparseMatrixBatchVectorMultiplyAccumulate1x4(
+    const float* __restrict__ matrix, const int32_t* __restrict__ segments,
+    const int32_t* __restrict__ indices, int m_rows, int m_cols,
+    const float* __restrict__ vector, int n_batch, float* __restrict__ result) {
+  NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate1x4, matrix,
+                   segments, indices, m_rows, m_cols, vector, n_batch, result);
+}
+
+void SparseMatrixBatchVectorMultiplyAccumulate(
+    const float* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    int m_rows, int m_cols, const float* __restrict__ vector, int n_batch,
+    float* __restrict__ result) {
+  NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate, matrix, ledger,
+                   m_rows, m_cols, vector, n_batch, result);
+}
+
+void SparseMatrixBatchVectorMultiplyAccumulate1x16(
+    const int8_t* __restrict__ matrix, const int32_t* __restrict__ segments,
+    const int32_t* __restrict__ indices, int m_rows, int m_cols,
+    const int8_t* __restrict__ vector, const int32_t* __restrict__ bias_vector,
+    int n_batch, const int32_t input_offset, const int32_t output_multiplier,
+    const int32_t output_shift, const int32_t* per_channel_scale,
+    const int32_t* per_channel_shift, const int32_t output_offset,
+    const int32_t output_activation_min, const int32_t output_activation_max,
+    int8_t* __restrict__ result) {
+  NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate1x16, matrix,
+                   segments, indices, m_rows, m_cols, vector, bias_vector,
+                   n_batch, input_offset, output_multiplier, output_shift,
+                   per_channel_scale, per_channel_shift, output_offset,
+                   output_activation_min, output_activation_max, result);
+}
+
+void SparseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const uint8_t* ledger, const int m_rows,
+    const int m_cols, const int8_t* __restrict__ vectors,
+    const float* scaling_factors, int n_batch, float* __restrict__ result,
+    const float* per_channel_scale) {
+  NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate, matrix, ledger,
+                   m_rows, m_cols, vectors, scaling_factors, n_batch, result,
+                   per_channel_scale);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* input, const int32_t* bias,
+    const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
+    int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
+    int32_t* scratch, int16_t* output, CpuBackendContext* context) {
+  NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, input, bias,
+                   input_to_gate_weights, multiplier, shift, n_batch, n_input,
+                   n_output, output_zp, scratch, output, context);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* input, const int32_t* bias,
+    const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
+    int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
+    int32_t* scratch, int8_t* output, CpuBackendContext* context) {
+  NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, input, bias,
+                   input_to_gate_weights, multiplier, shift, n_batch, n_input,
+                   n_output, output_zp, scratch, output, context);
+}
+
+void MatrixBatchVectorMultiply(const int8_t* input, int32_t input_zeropoint,
+                               const int8_t* input_to_gate_weights,
+                               int32_t input_to_gate_effective_scale_a,
+                               int32_t input_to_gate_effective_scale_b,
+                               int32_t n_batch, int32_t n_input, int32_t n_cell,
+                               int8_t* gate_output, int8_t gate_output_zp) {
+  PortableMatrixBatchVectorMultiply(
+      input, input_zeropoint, input_to_gate_weights,
+      input_to_gate_effective_scale_a, input_to_gate_effective_scale_b, n_batch,
+      n_input, n_cell, gate_output, gate_output_zp);
+}
+
+void MatrixBatchVectorMultiply(const int16_t* hidden,
+                               const int8_t* hidden_to_output_weights,
+                               int32_t proj_effective_scale_a,
+                               int32_t proj_effective_scale_b,
+                               const int32_t* gate_bias, int32_t n_batch,
+                               int32_t n_hidden, int32_t n_output,
+                               int32_t output_zp, int8_t* proj_output) {
+  PortableMatrixBatchVectorMultiply(hidden, hidden_to_output_weights,
+                                    proj_effective_scale_a,
+                                    proj_effective_scale_b, gate_bias, n_batch,
+                                    n_hidden, n_output, output_zp, proj_output);
+}
+
+void MatrixScalarMultiplyAccumulate(const int8_t* matrix, int32_t scalar,
+                                    int32_t n_row, int32_t n_col,
+                                    int32_t* output) {
+  NEON_OR_PORTABLE(MatrixScalarMultiplyAccumulate, matrix, scalar, n_row, n_col,
+                   output);
+}
+
+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) {
+  NEON_OR_PORTABLE(ApplyLayerNorm, input, layer_norm_weights, bias,
+                   layer_norm_scale_a, layer_norm_scale_b, variance_limit,
+                   n_batch, n_input, output);
+}
+
+void ApplyLayerNormFloat(const int16_t* input,
+                         const int16_t* layer_norm_weights,
+                         int32_t layer_norm_scale_a, int32_t layer_norm_scale_b,
+                         const int32_t* bias, int n_batch, int n_input,
+                         int16_t* output) {
+  PortableApplyLayerNormFloat(input, layer_norm_weights, layer_norm_scale_a,
+                              layer_norm_scale_b, bias, n_batch, n_input,
+                              output);
+}
+
+void ApplySigmoid(const int16_t* input, int32_t n_batch, int32_t n_input,
+                  int16_t* output) {
+  NEON_OR_PORTABLE(ApplySigmoid, input, n_batch, n_input, output);
+}
+
+void ApplySigmoidFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
+                       int16_t* output) {
+  PortableApplySigmoidFloat(input, n_batch, n_input, output);
+}
+
+void ApplyTanh(int32_t integer_bits, const int16_t* input, int32_t n_batch,
+               int32_t n_input, int16_t* output) {
+  NEON_OR_PORTABLE(ApplyTanh, integer_bits, input, n_batch, n_input, output);
+}
+
+void ApplyTanhFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
+                    int32_t integer_bits, int16_t* output) {
+  PortableApplyTanhFloat(input, n_batch, n_input, integer_bits, output);
+}
+
+void CwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch,
+              int n_input, int shift, int16_t* output) {
+  NEON_OR_PORTABLE(CwiseMul, input_1, input_2, n_batch, n_input, shift, output);
+}
+
+void CwiseMul(const int16_t* input_1, const int16_t* input_2,
+              int32_t multiplier, int shift, int n_batch, int n_input,
+              int32_t output_zp, int8_t* output) {
+  NEON_OR_PORTABLE(CwiseMul, input_1, input_2, multiplier, shift, n_batch,
+                   n_input, output_zp, output);
+}
+
+void CwiseAdd(const int16_t* input_1, const int16_t* input_2, int n_batch,
+              int n_input, int16_t* output) {
+  NEON_OR_PORTABLE(CwiseAdd, input_1, input_2, n_batch, n_input, output);
+}
+
+void CwiseClipping(float* vector, const int v_size,
+                   const float clipping_value) {
+  NEON_OR_PORTABLE(CwiseClipping, vector, v_size, clipping_value);
+}
+void CwiseClipping(int16_t* vector, const int v_size,
+                   const int16_t clipping_value) {
+  NEON_OR_PORTABLE(CwiseClipping, vector, v_size, clipping_value);
+}
+void CwiseClipping(int8_t* vector, const int v_size,
+                   const int8_t clipping_value) {
+  NEON_OR_PORTABLE(CwiseClipping, vector, v_size, clipping_value);
+}
+
+void BatchVectorBatchVectorDotProduct(const int16_t* vector1,
+                                      const int16_t* vector2, int v_size,
+                                      int n_batch, int32_t* result) {
+  PortableBatchVectorBatchVectorDotProduct(vector1, vector2, v_size, n_batch,
+                                           result);
+}
+
+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) {
+  NEON_OR_PORTABLE(VectorBatchVectorCwiseProductAccumulate, vector, v_size,
+                   batch_vector, n_batch, multiplier, shift, result);
+}
+
+float VectorVectorDotProduct(const float* vector1, const float* vector2,
+                             int v_size) {
+  return NEON_OR_PORTABLE(VectorVectorDotProduct, vector1, vector2, v_size);
+}
+
+void Sub1Vector(const float* vector, int v_size, float* result) {
+  NEON_OR_PORTABLE(Sub1Vector, vector, v_size, result);
+}
+
+void Sub1Vector(const int16_t* vector, int v_size, int16_t* result) {
+  NEON_OR_PORTABLE(Sub1Vector, vector, v_size, result);
+}
+
+// Check if all entries of a vector are zero for float.
+bool IsZeroVector(const float* vector, int v_size) {
+  return NEON_OR_PORTABLE(IsZeroVector, vector, v_size);
+}
+
+// Check if all entries of a vector are zero for int8.
+bool IsZeroVector(const int8_t* vector, int v_size) {
+  return NEON_OR_PORTABLE(IsZeroVector, vector, v_size);
+}
+
+void VectorScalarMultiply(const int8_t* vector, int v_size, float scale,
+                          float* result) {
+  NEON_OR_PORTABLE(VectorScalarMultiply, vector, v_size, scale, result);
+}
+
+void SymmetricQuantizeFloats(const float* values, const int size,
+                             int8_t* quantized_values, float* min_value,
+                             float* max_value, float* scaling_factor) {
+  NEON_OR_PORTABLE(SymmetricQuantizeFloats, values, size, quantized_values,
+                   min_value, max_value, scaling_factor);
+}
+
+void SymmetricQuantizeFloats(const float* values, const int size,
+                             int8_t* quantized_values, float min_value,
+                             float max_value, float* scaling_factor) {
+  NEON_OR_PORTABLE(SymmetricQuantizeFloats, values, size, quantized_values,
+                   min_value, max_value, scaling_factor);
+}
+
+void AsymmetricQuantizeFloats(const float* values, const int size,
+                              int8_t* quantized_values, float* scaling_factor,
+                              int32_t* offset) {
+  NEON_OR_PORTABLE(AsymmetricQuantizeFloats, values, size, quantized_values,
+                   scaling_factor, offset);
+}
+
+void ReductionSumVector(const float* input_vector, float* output_vector,
+                        int output_size, int reduction_size) {
+  NEON_OR_PORTABLE(ReductionSumVector, input_vector, output_vector, output_size,
+                   reduction_size);
+}
+
+void ReductionSumVector(const int32_t* input_vector, int32_t* output_vector,
+                        int output_size, int reduction_size) {
+  PortableReductionSumVector(input_vector, output_vector, output_size,
+                             reduction_size);
+}
+
+void ReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
+                        int output_size, int reduction_size) {
+  NEON_OR_PORTABLE(ReductionSumVector, input_vector, output_vector, output_size,
+                   reduction_size);
+}
+
+void MeanStddevNormalization(const float* __restrict__ input_vector,
+                             float* __restrict__ output_vector, int v_size,
+                             int n_batch) {
+  NEON_OR_PORTABLE(MeanStddevNormalization, input_vector, output_vector, v_size,
+                   n_batch);
+}
+
+void TwoGateSaturatingAdd(const int8_t* input, int8_t input_zp,
+                          const int8_t* recurrent, int8_t recurrent_zp,
+                          int32_t input_effective_scale_a,
+                          int32_t input_effective_scale_b,
+                          int32_t recurrent_effective_scale_a,
+                          int32_t recurrent_effective_scale_b, int32_t n_batch,
+                          int32_t n_cell, int16_t* output) {
+  PortableTwoGateSaturatingAdd(
+      input, input_zp, recurrent, recurrent_zp, input_effective_scale_a,
+      input_effective_scale_b, recurrent_effective_scale_a,
+      recurrent_effective_scale_b, n_batch, n_cell, output);
+}
+
+}  // namespace tensor_utils
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_H_

tensorflow/lite/kernels/internal/optimized/neon_tensor_utils_impl.h

@@ -0,0 +1,198 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_IMPL_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_IMPL_H_
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
+
+#if defined(_MSC_VER)
+#define __restrict__ __restrict
+#endif
+
+namespace tflite {
+namespace tensor_utils {
+
+#ifdef USE_NEON
+
+// Multiply a matrix by a batch vector, and store results in a batch-size
+// vector.
+void NeonMatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
+                                             int m_cols, const float* vector,
+                                             int n_batch, float* result);
+
+// Matrix multiplication for quantized values using symmetric quantization.
+void NeonMatrixBatchVectorMultiplyAccumulate(const int8_t* __restrict__ matrix,
+                                             const int m_rows, const int m_cols,
+                                             const int8_t* __restrict__ vectors,
+                                             const float* scaling_factors,
+                                             int n_batch,
+                                             float* __restrict__ result);
+
+// Same as above but with a scratch buffer and CpuBackendContext for the
+// int8 x int8 -> int32 accumulation computation
+void NeonMatrixBatchVectorMultiplyAccumulate(const int8_t* __restrict__ matrix,
+                                             const int m_rows, const int m_cols,
+                                             const int8_t* __restrict__ vectors,
+                                             const float* scaling_factors,
+                                             int n_batch, int32_t* scratch,
+                                             float* __restrict__ result,
+                                             CpuBackendContext* context);
+
+// Matrix multiplication for quantized values using asymmetric quantization.
+void NeonMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
+    bool* compute_row_sums, CpuBackendContext* context);
+
+void NeonApplyLayerNorm(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);
+
+void NeonApplySigmoid(const int16_t* input, int32_t n_batch, int32_t n_input,
+                      int16_t* output);
+
+void NeonApplyTanh(int32_t integer_bits, const int16_t* input, int32_t n_batch,
+                   int32_t n_input, int16_t* output);
+
+void NeonCwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch,
+                  int n_input, int shift, int16_t* output);
+
+void NeonCwiseMul(const int16_t* input_1, const int16_t* input_2,
+                  int32_t multiplier, int shift, int n_batch, int n_input,
+                  int32_t output_zp, int8_t* output);
+
+void NeonCwiseAdd(const int16_t* input_1, const int16_t* input_2, int n_batch,
+                  int n_input, int16_t* output);
+
+void NeonCwiseClipping(float* vector, const int v_size,
+                       const float clipping_value);
+void NeonCwiseClipping(int16_t* vector, const int v_size,
+                       const int16_t clipping_value);
+void NeonCwiseClipping(int8_t* vector, const int v_size,
+                       const int8_t clipping_value);
+
+void NeonMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* input, const int32_t* bias,
+    const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
+    int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
+    int32_t* scratch, int8_t* output, CpuBackendContext* context);
+
+void NeonMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* input, const int32_t* bias,
+    const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
+    int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
+    int32_t* scratch, int16_t* output, CpuBackendContext* context);
+
+void NeonMatrixScalarMultiplyAccumulate(const int8_t* matrix, int32_t scalar,
+                                        int32_t n_row, int32_t n_col,
+                                        int32_t* output);
+
+void NeonSparseMatrixBatchVectorMultiplyAccumulate1x4(
+    const float* __restrict__ matrix, const int32_t* __restrict__ segments,
+    const int32_t* __restrict__ indices, int m_rows, int m_cols,
+    const float* __restrict__ vector, int n_batch, float* __restrict__ result);
+
+// Multiply a matrix by a batch vector, and store results in a batch-size
+// vector. Sparse version.
+void NeonSparseMatrixBatchVectorMultiplyAccumulate(
+    const float* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    int m_rows, int m_cols, const float* __restrict__ vector, int n_batch,
+    float* __restrict__ result);
+
+// Multiplies a symmetric quantized matrix by a quantized batch vector. The
+// matrix is stored in sparse format.
+void NeonSparseMatrixBatchVectorMultiplyAccumulate1x16(
+    const int8_t* __restrict__ matrix, const int32_t* __restrict__ segments,
+    const int32_t* __restrict__ indices, int m_rows, int m_cols,
+    const int8_t* __restrict__ vector, const int32_t* __restrict__ bias_vector,
+    int n_batch, const int32_t input_offset, const int32_t output_multiplier,
+    int32_t output_shift, const int32_t* per_channel_scale,
+    const int32_t* per_channel_shift, int32_t output_offset,
+    const int32_t output_activation_min, const int32_t output_activation_max,
+    int8_t* __restrict__ result);
+
+// Matrix multiplication for quantized values using symmetric quantization.
+// Sparse version.
+void NeonSparseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const uint8_t* ledger, const int m_rows,
+    const int m_cols, const int8_t* __restrict__ vectors,
+    const float* scaling_factors, int n_batch, float* __restrict__ result,
+    const float* per_channel_scale);
+
+// Dot product of two vectors.
+float NeonVectorVectorDotProduct(const float* vector1, const float* vector2,
+                                 int v_size);
+
+// Compute "1.0f - elements of vector" (used in CIFG).
+void NeonSub1Vector(const float* vector, int v_size, float* result);
+
+void NeonSub1Vector(const int16_t* vector, int v_size, int16_t* result);
+
+// Multiply all elements of vector with a scalar.
+void NeonVectorScalarMultiply(const int8_t* vector, int v_size, float scale,
+                              float* result);
+
+// Check if all entries of a vector are zero.
+bool NeonIsZeroVector(const float* vector, int v_size);
+
+// Check if all entries of a vector are zero.
+bool NeonIsZeroVector(const int8_t* vector, int v_size);
+
+// Symmetric quantizer.
+void NeonSymmetricQuantizeFloats(const float* values, const int size,
+                                 int8_t* quantized_values, float* min,
+                                 float* max, float* scaling_factor);
+
+// Symmetric quantizer.
+void NeonSymmetricQuantizeFloats(const float* values, const int size,
+                                 int8_t* quantized_values, float min, float max,
+                                 float* scaling_factor);
+
+// Asymmetric quantizer.
+void NeonAsymmetricQuantizeFloats(const float* values, const int size,
+                                  int8_t* quantized_values,
+                                  float* scaling_factor, int32_t* offset);
+
+// Reduce-sum on a float input vector:
+// input_vector: float pointer to input vector.
+// output_vector: float pointer to vector.
+// output_size: output vector size.
+// reduction_size: number of consecutive elements from input vector which are
+// added to get one element of output.
+void NeonReductionSumVector(const float* input_vector, float* output_vector,
+                            int output_size, int reduction_size);
+
+void NeonReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
+                            int output_size, int reduction_size);
+
+void NeonVectorBatchVectorCwiseProductAccumulate(
+    const int16_t* vector, int v_size, const int16_t* batch_vector, int n_batch,
+    int32_t multiplier, int shift, int16_t* result);
+
+// Layer norm for each batch.
+void NeonMeanStddevNormalization(const float* __restrict__ input_vector,
+                                 float* __restrict__ output_vector, int v_size,
+                                 int n_batch);
+
+#endif  // USE_NEON
+
+}  // namespace tensor_utils
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_TENSOR_UTILS_IMPL_H_

tensorflow/lite/kernels/internal/optimized/optimized_4bit_test.cc

@@ -0,0 +1,597 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <cstdlib>
+#include <functional>
+#include <memory>
+#include <random>
+#include <tuple>
+#include <type_traits>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/optimized/fully_connected_4bit.h"
+
+namespace tflite {
+namespace {
+
+std::mt19937 random_engine(2023);
+std::uniform_real_distribution<float> real_dist(0.f, 1.f);
+std::uniform_int_distribution<int32_t> int_dist(-7, 7);
+
+struct TestPack {
+  TestPack(std::vector<int8_t> src_data, int src_rows, int src_cols, int width,
+           int depth)
+      : src_data(src_data),
+        src_rows(src_rows),
+        src_cols(src_cols),
+        width(width),
+        depth(depth),
+        rows((src_rows + (width - 1)) & ~(width - 1)),
+        cols((src_cols + (depth - 1)) & ~(depth - 1)),
+        // Must be vector-aligned.
+        packed_data_buffer(
+            [=]() -> uint8_t* {
+              void* ptr;
+              if (posix_memalign(&ptr, 64, rows * cols + padding)) {
+                abort();
+              }
+              return static_cast<uint8_t*>(ptr);
+            }(),
+            [](uint8_t* ptr) { free(ptr); }) {}
+
+  void Prepack() {
+    packed_data = packed_data_buffer.get();
+    optimized_4bit::Prepack(packed_data, src_data.data(), rows, cols, src_rows,
+                            src_cols, width, depth);
+  }
+
+  std::vector<uint8_t> AsVector() {
+    int size = rows * cols / 2;
+    std::vector<uint8_t> values(size);
+    for (int i = 0; i < size; i++) {
+      values[i] = packed_data[i];
+    }
+    packed_data = nullptr;
+    return values;
+  }
+
+  std::vector<int8_t> src_data;
+  uint8_t* packed_data;
+  int src_rows;
+  int src_cols;
+  int width;
+  int depth;
+  int rows;
+  int cols;
+  int padding = optimized_4bit::kDefaultAlignmentPadding;
+  std::unique_ptr<uint8_t, std::function<void(uint8_t*)>> packed_data_buffer;
+};
+
+class RunPackTests
+    : public ::testing::TestWithParam<::testing::tuple<int, int>> {};
+
+TEST_P(RunPackTests, RunPackTests) {
+  auto params = GetParam();
+  int src_rows = std::get<0>(params);
+  int src_cols = std::get<1>(params);
+  std::vector<int8_t> test_data;
+  test_data.reserve(src_rows * src_cols / 2);
+  for (int i = 0; i < src_rows; ++i) {
+    int stride = optimized_4bit::FilterDepth / 4;
+    int strides = src_cols / stride / 2;
+    int v = -7;
+    int l = 0;
+    for (int j = 0; j < strides; j++) {
+      for (int k = 0; k < stride; k++) {  // 8
+        int lower = static_cast<uint8_t>(v) & UINT8_C(15);
+        int upper = static_cast<uint8_t>(v) << 4;
+        test_data.push_back(upper | lower);
+        l++;
+      }
+      v++;
+    }
+    while (l < (src_cols / 2)) {
+      int lower = static_cast<uint8_t>(v) & UINT8_C(15);
+      test_data.push_back(lower << 4 | lower);
+      l++;
+    }
+  }
+  TestPack test(test_data, src_rows, src_cols, optimized_4bit::FilterWidth,
+                optimized_4bit::FilterDepth);
+  test.Prepack();
+  std::vector<uint8_t> result = test.AsVector();
+  int outer_rows = test.rows / optimized_4bit::FilterWidth;
+  int outer_cols = test.cols / optimized_4bit::FilterDepth;
+  int k = 0;
+  for (int i = 0; i < outer_rows; ++i) {
+    int v = -7;
+    for (int j = 0; j < outer_cols; ++j) {
+      for (int w = 0; w < optimized_4bit::FilterWidth; w++) {
+        int c = 0;
+        for (; c < optimized_4bit::FilterDepth / 2; c++) {
+          uint8_t res = result[k++];
+          uint8_t res0 = res >> 4;
+          uint8_t res1 = res & UINT8_C(15);
+          int res00 = res0 - 7;
+          int res11 = res1 - 7;
+          if ((i * optimized_4bit::FilterWidth + w) < src_rows) {
+            if ((j * optimized_4bit::FilterDepth / 2 + c) < src_cols / 2) {
+              EXPECT_EQ(res00, v % 8);
+            }
+            if ((j * optimized_4bit::FilterDepth / 2 + c + 16) < src_cols / 2) {
+              EXPECT_EQ(res11, v + 1 % 8);
+            }
+          }
+        }
+      }
+      v += 2;
+    }
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(RunPackTests, RunPackTests,
+                         ::testing::ValuesIn({
+                             std::make_tuple(4, 32),
+                             std::make_tuple(4, 46),
+                             std::make_tuple(4, 56),
+                             std::make_tuple(5, 64),
+                             std::make_tuple(5, 72),
+                             std::make_tuple(5, 80),
+                             std::make_tuple(5, 84),
+                         }));
+
+struct TestQuantize {
+  TestQuantize(std::vector<float> src_data, int src_rows, int src_cols,
+               int width, int depth)
+      : src_data(src_data),
+        src_rows(src_rows),
+        src_cols(src_cols),
+        width(width),
+        depth(depth) {
+    rows = (src_rows + (width - 1)) & ~(width - 1);
+    cols = (src_cols + (depth - 1)) & ~(depth - 1);
+    scaling_factors.assign(rows, 1.0);
+    input_offsets.assign(rows, 0);
+    output_data.assign(rows * cols, 0);
+  }
+
+  void BatchQuantizeFloats4Bit() {
+    optimized_4bit::BatchQuantizeFloats4Bit(
+        src_data.data(), src_rows, src_cols, output_data.data(),
+        scaling_factors.data(), width, depth, input_offsets.data());
+  }
+
+  std::vector<float> src_data;
+  int src_rows;
+  int src_cols;
+  int rows;
+  int cols;
+  int width;
+  int depth;
+  std::vector<int8_t> output_data;
+  std::vector<float> scaling_factors;
+  std::vector<int32_t> input_offsets;
+};
+
+class RunQuantizeInputTests
+    : public ::testing::TestWithParam<::testing::tuple<int, int, int>> {};
+
+TEST_P(RunQuantizeInputTests, RunQuantizeInputsTests) {
+  auto params = GetParam();
+  int width = std::get<0>(params);
+  int src_rows = std::get<1>(params);
+  int src_cols = std::get<2>(params);
+  std::vector<float> test_data;
+  test_data.reserve(src_rows * src_cols);
+
+  float v = -127.0;
+  for (int i = 0; i < src_rows; ++i) {
+    for (int j = 0; j < src_cols; ++j) {
+      test_data.push_back(v / (i + 1));
+      v = -v;
+    }
+  }
+  TestQuantize test(test_data, src_rows, src_cols, width,
+                    optimized_4bit::FilterDepth);
+  test.BatchQuantizeFloats4Bit();
+  int8_t* result = test.output_data.data();
+  int k = 0;
+  int outer_rows = test.rows / width;
+  int outer_cols = test.cols / optimized_4bit::FilterDepth;
+  for (int i = 0; i < outer_rows; ++i) {
+    for (int j = 0; j < outer_cols; ++j) {
+      for (int w = 0; w < width; w++) {
+        int c = 0;
+        v = -127;
+        for (; c < optimized_4bit::FilterDepth; c++) {
+          int8_t res = result[k++];
+          int res0 = static_cast<int>(res);
+          if ((i * width + w) < src_rows) {
+            if ((j * optimized_4bit::FilterDepth + c) < src_cols) {
+              EXPECT_EQ(res0, v);
+            }
+          }
+          v = -v;
+        }
+      }
+      v += 2;
+    }
+  }
+  for (int i = 0; i < test.rows; i++) {
+    if (i >= src_rows) {
+      continue;
+    }
+    EXPECT_EQ(test.input_offsets[i], 0);
+    EXPECT_NEAR(test.scaling_factors[i], 1.0 / (1 + i), 1e-3);
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(RunQuantizeInputTests, RunQuantizeInputTests,
+                         ::testing::ValuesIn({
+                             std::make_tuple(1, 1, 32),
+                             std::make_tuple(1, 3, 46),
+                             std::make_tuple(1, 9, 64),
+                             std::make_tuple(1, 25, 72),
+                             std::make_tuple(2, 2, 32),
+                             std::make_tuple(2, 3, 46),
+                             std::make_tuple(2, 9, 64),
+                             std::make_tuple(2, 25, 72),
+                             std::make_tuple(4, 4, 32),
+                             std::make_tuple(4, 5, 46),
+                             std::make_tuple(4, 9, 64),
+                             std::make_tuple(4, 25, 72),
+                         }));
+
+struct TestAssignBiasAndComputeOffset {
+  TestAssignBiasAndComputeOffset(std::vector<float> output_data,
+                                 std::vector<int32_t> input_offsets,
+                                 std::vector<float> input_scales,
+                                 std::vector<float> filter_scales,
+                                 std::vector<float> bias, int output_rows,
+                                 int output_cols, bool use_bias)
+      : output_data(output_data),
+        input_offsets(input_offsets),
+        input_scales(input_scales),
+        filter_scales(filter_scales),
+        bias(bias),
+        output_rows(output_rows),
+        output_cols(output_cols),
+        use_bias(use_bias) {}
+
+  void AssignBiasAndComputeOffsets() {
+    optimized_4bit::AssignBiasAndComputeOffsets(
+        input_offsets.data(), input_scales.data(), filter_scales.data(),
+        use_bias ? bias.data() : nullptr, output_data.data(), output_cols,
+        output_rows);
+  }
+  std::vector<float> output_data;
+  std::vector<int32_t> input_offsets;
+  std::vector<float> input_scales;
+  std::vector<float> filter_scales;
+  std::vector<float> bias;
+  int output_rows;
+  int output_cols;
+  bool use_bias;
+};
+
+class RunAssignBiasAndOffsetsTests
+    : public ::testing::TestWithParam<::testing::tuple<int, int, bool>> {};
+
+TEST_P(RunAssignBiasAndOffsetsTests, RunAssignBiasAndOffsetssTests) {
+  auto params = GetParam();
+  int output_rows = std::get<0>(params);
+  int output_cols = std::get<1>(params);
+  bool use_bias = std::get<2>(params);
+  std::vector<float> test_data(output_rows * output_cols, 0);
+  std::vector<float> test_input_scales(output_rows);
+  std::vector<int32_t> test_input_offsets(output_rows);
+  std::vector<float> test_filter_scales(output_cols);
+  std::vector<float> test_bias(output_cols);
+
+  for (int i = 0; i < output_rows; ++i) {
+    test_input_scales[i] = real_dist(random_engine);
+    test_input_offsets[i] = int_dist(random_engine);
+  }
+  for (int i = 0; i < output_cols; ++i) {
+    test_filter_scales[i] = real_dist(random_engine);
+    test_bias[i] = real_dist(random_engine);
+  }
+
+  TestAssignBiasAndComputeOffset test(
+      test_data, test_input_offsets, test_input_scales, test_filter_scales,
+      test_bias, output_rows, output_cols, use_bias);
+  test.AssignBiasAndComputeOffsets();
+  float* result = test.output_data.data();
+  for (int i = 0; i < output_rows; ++i) {
+    for (int j = 0; j < output_cols; ++j) {
+      float val = result[i * output_cols + j];
+      float expected = use_bias ? test_bias[j] : 0;
+      expected +=
+          test_input_offsets[i] * test_input_scales[i] * test_filter_scales[j];
+      EXPECT_NEAR(val, expected, 1e-3);
+    }
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(RunAssignBiasAndOffsetsTests,
+                         RunAssignBiasAndOffsetsTests,
+                         ::testing::ValuesIn({
+                             std::make_tuple(1, 8, true),
+                             std::make_tuple(4, 17, false),
+                             std::make_tuple(4, 17, true),
+                             std::make_tuple(11, 33, false),
+                             std::make_tuple(11, 33, true),
+                         }));
+
+struct TestUnpack {
+  TestUnpack(std::vector<int32_t> src_data, std::vector<float> input_scales,
+             std::vector<float> filter_scales, int src_rows, int src_cols,
+             int output_rows, int output_cols)
+      : src_data(src_data),
+        input_scales(input_scales),
+        filter_scales(filter_scales),
+        src_rows(src_rows),
+        src_cols(src_cols),
+        output_rows(output_rows),
+        output_cols(output_cols) {
+    output_data.assign(output_rows * output_cols, 0.0);
+  }
+
+  template <int Depth, int Width>
+  void Unpack() {
+    optimized_4bit::Unpack<Depth, Width>(
+        output_data.data(), src_data.data(), output_rows, output_cols,
+        input_scales.data(), filter_scales.data(), src_rows, src_cols);
+  }
+
+  std::vector<int32_t> src_data;
+  std::vector<float> input_scales;
+  std::vector<float> filter_scales;
+  std::vector<float> output_data;
+  int src_rows;
+  int src_cols;
+  int output_rows;
+  int output_cols;
+};
+
+class RunUnpackTests
+    : public ::testing::TestWithParam<::testing::tuple<int, int, int>> {};
+
+TEST_P(RunUnpackTests, RunUnpackTests) {
+  auto params = GetParam();
+  int src_rows = std::get<0>(params);
+  int src_cols = std::get<1>(params);
+  // In this case, we only unpack 1 rhs row,
+  // so the batch_size and accumulator rows must match.
+  int output_rows = src_rows;
+  int output_cols = std::get<2>(params);
+  std::vector<float> test_input_scales(src_rows);
+  std::vector<float> test_filter_scales(src_cols);
+  std::vector<int32_t> test_data;
+  test_data.reserve(src_rows * src_cols);
+  int outer_cols = src_cols / optimized_4bit::FilterWidth;
+  int outer_rows = src_rows;
+  for (int j = 0; j < outer_cols; ++j) {
+    for (int i = 0; i < outer_rows; ++i) {
+      for (int k = 0; k < optimized_4bit::FilterWidth; ++k) {
+        test_data.push_back(i);
+      }
+    }
+  }
+  for (int i = 0; i < src_rows; ++i) {
+    test_input_scales[i] = real_dist(random_engine);
+  }
+  for (int i = 0; i < src_cols; ++i) {
+    test_filter_scales[i] = real_dist(random_engine);
+  }
+  TestUnpack test(test_data, test_input_scales, test_filter_scales, src_rows,
+                  src_cols, output_rows, output_cols);
+  test.Unpack<4, 1>();
+  std::vector<float> result = test.output_data;
+  for (int i = 0; i < output_rows; ++i) {
+    for (int j = 0; j < output_cols; ++j) {
+      float res = result[i * output_cols + j];
+      EXPECT_EQ(res, i * test_input_scales[i] * test_filter_scales[j]);
+    }
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(RunUnpackTests, RunUnpackTests,
+                         ::testing::ValuesIn({
+                             std::make_tuple(1, 8, 5),
+                             std::make_tuple(3, 4, 4),
+                         }));
+
+class RunKernelTests
+    : public ::testing::TestWithParam<::testing::tuple<int, int, int, int>> {};
+
+TEST_P(RunKernelTests, RunKernelTests) {
+  auto params = GetParam();
+  int rhs_width = std::get<0>(params);
+  int lhs_layout_rows = std::get<1>(params);
+  int rhs_layout_rows = std::get<2>(params);
+  int lhs_layout_cols = std::get<3>(params);
+  int rhs_layout_cols = lhs_layout_cols;
+
+  std::vector<uint8_t> test_lhs(lhs_layout_rows * lhs_layout_cols / 2, 0.0);
+  std::vector<int8_t> test_rhs(rhs_layout_rows * rhs_layout_cols, 0.0);
+  std::vector<int32_t> test_accum(lhs_layout_rows * rhs_layout_rows, 0.0);
+
+  int lhs_outer_rows = lhs_layout_rows / optimized_4bit::FilterWidth;
+  int lhs_outer_cols = lhs_layout_cols / optimized_4bit::FilterDepth;
+
+  // pack lhs
+  for (int i = 0; i < lhs_outer_rows; ++i) {
+    for (int j = 0; j < lhs_outer_cols; ++j) {
+      for (int k = 0; k < optimized_4bit::FilterWidth; ++k) {
+        for (int l = 0; l < optimized_4bit::FilterDepth / 2; ++l) {
+          uint8_t u = static_cast<uint8_t>(int_dist(random_engine) + 7);
+          uint8_t v = static_cast<uint8_t>(int_dist(random_engine) + 7);
+          int lower = static_cast<uint8_t>(v) & UINT8_C(15);
+          int upper = static_cast<uint8_t>(u) << 4;
+          int cluster_index = (i * lhs_outer_cols + j) *
+                              optimized_4bit::FilterDepth *
+                              optimized_4bit::FilterWidth / 2;
+          int index = cluster_index + k * (optimized_4bit::FilterDepth / 2);
+          test_lhs[index + l] = (upper | lower);
+        }
+      }
+    }
+  }
+  int rhs_outer_rows = rhs_layout_rows / rhs_width;
+  int rhs_outer_cols = rhs_layout_cols / optimized_4bit::FilterDepth;
+
+  for (int i = 0; i < rhs_outer_rows; ++i) {
+    for (int j = 0; j < rhs_outer_cols; ++j) {
+      for (int k = 0; k < rhs_width; ++k) {
+        for (int l = 0; l < optimized_4bit::FilterDepth; ++l) {
+          int8_t u = static_cast<int8_t>(int_dist(random_engine));
+          int cluster_index = (i * rhs_outer_cols + j) *
+                              optimized_4bit::FilterDepth * rhs_width;
+          int index = cluster_index + k * optimized_4bit::FilterDepth;
+          test_rhs[index + l] = u;
+        }
+      }
+    }
+  }
+
+  int index = 0;
+  std::vector<int32_t> expected_accum(lhs_layout_rows * rhs_layout_rows, 0.0);
+  int outer_cols = rhs_outer_cols;
+  for (int i = 0; i < lhs_outer_rows; ++i) {
+    for (int j = 0; j < rhs_outer_rows; ++j) {
+      for (int k = 0; k < rhs_width; ++k) {
+        for (int l = 0; l < optimized_4bit::FilterWidth; ++l) {
+          for (int m = 0; m < outer_cols; ++m) {
+            for (int n = 0; n < optimized_4bit::FilterDepth; ++n) {
+              int right_index = ((j * outer_cols + m) * rhs_width + k) *
+                                optimized_4bit::FilterDepth;
+              int8_t rhs = test_rhs[right_index + n];
+              int left_index =
+                  ((i * outer_cols + m) * optimized_4bit::FilterWidth + l) *
+                  (optimized_4bit::FilterDepth / 2);
+              uint8_t lhs = 0;
+              if (n < optimized_4bit::FilterDepth / 2) {
+                int a = n % 16;
+                lhs = static_cast<uint8_t>(test_lhs[left_index + a] >> 4);
+              } else {
+                int a = n % 16;
+                lhs = static_cast<uint8_t>(test_lhs[left_index + a] &
+                                           UINT8_C(15));
+              }
+              int accum_index = ((i * rhs_outer_rows + j) * rhs_width + k) *
+                                optimized_4bit::FilterWidth;
+              expected_accum[accum_index + l] += rhs * lhs;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  index = 0;
+  switch (rhs_width) {
+#if defined(FC_4BIT_NEON) && defined(__aarch64__)
+    case 4:
+      optimized_4bit::RunKernel<optimized_4bit::FilterWidth, 4,
+                                optimized_4bit::FilterDepth>(
+          test_lhs.data(), test_rhs.data(), test_accum.data(), lhs_layout_rows,
+          lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, rhs_layout_rows,
+          lhs_layout_rows);
+      break;
+    case 2:
+      optimized_4bit::RunKernel<optimized_4bit::FilterWidth, 2,
+                                optimized_4bit::FilterDepth>(
+          test_lhs.data(), test_rhs.data(), test_accum.data(), lhs_layout_rows,
+          lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, rhs_layout_rows,
+          lhs_layout_rows);
+      break;
+#endif
+    case 1:
+      [[fallthrough]];
+    default:
+      optimized_4bit::RunKernel<optimized_4bit::FilterWidth, 1,
+                                optimized_4bit::FilterDepth>(
+          test_lhs.data(), test_rhs.data(), test_accum.data(), lhs_layout_rows,
+          lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, rhs_layout_rows,
+          lhs_layout_rows);
+      break;
+  }
+
+  for (int i = 0; i < (rhs_layout_rows * lhs_layout_rows); ++i) {
+    int32_t expected_val = expected_accum[i];
+    int32_t val = test_accum[i];
+    EXPECT_EQ(val, expected_val);
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    RunKernelTests, RunKernelTests, ::testing::ValuesIn({
+      std::make_tuple(1, 4, 1, 32), std::make_tuple(1, 8, 1, 32),
+          std::make_tuple(1, 16, 1, 32), std::make_tuple(1, 4, 1, 64),
+          std::make_tuple(1, 8, 1, 64), std::make_tuple(1, 16, 1, 64),
+          std::make_tuple(1, 4, 5, 64), std::make_tuple(1, 8, 9, 64),
+          std::make_tuple(1, 16, 17, 64),
+#if defined(FC_4BIT_NEON) && defined(__aarch64__)
+          std::make_tuple(2, 8, 2, 32), std::make_tuple(2, 16, 2, 32),
+          std::make_tuple(2, 4, 4, 64), std::make_tuple(2, 8, 4, 64),
+          std::make_tuple(2, 16, 4, 64), std::make_tuple(2, 4, 4, 64),
+          std::make_tuple(2, 8, 8, 64), std::make_tuple(2, 16, 16, 64),
+          std::make_tuple(4, 4, 4, 32), std::make_tuple(4, 8, 4, 32),
+          std::make_tuple(4, 16, 4, 32), std::make_tuple(4, 4, 8, 64),
+          std::make_tuple(4, 8, 8, 64), std::make_tuple(4, 16, 8, 64),
+          std::make_tuple(4, 4, 8, 64), std::make_tuple(4, 8, 12, 64),
+          std::make_tuple(4, 16, 32, 64),
+#endif
+    }));
+
+template <typename T>
+class OverflowKernelTests : public ::testing::Test {};
+
+using RhsWidths = ::testing::Types<std::integral_constant<int, 1>
+#if defined(FC_4BIT_NEON) && defined(__aarch64__)
+                                   ,
+                                   std::integral_constant<int, 2>,
+                                   std::integral_constant<int, 4>
+#endif
+                                   >;
+TYPED_TEST_SUITE(OverflowKernelTests, RhsWidths);
+
+TYPED_TEST(OverflowKernelTests, OverflowTest) {
+  constexpr int rhs_width = TypeParam::value;
+  int lhs_layout_rows = 4;
+  int rhs_layout_rows = rhs_width;
+  int lhs_layout_cols = 1024;
+  int rhs_layout_cols = 1024;
+
+  // 0xff = two packed int4s with value of 15.
+  std::vector<uint8_t> test_lhs(lhs_layout_rows * lhs_layout_cols / 2, 0xff);
+  std::vector<int8_t> test_rhs(rhs_layout_rows * rhs_layout_cols, 127);
+  std::vector<int32_t> test_accum(lhs_layout_rows * rhs_layout_rows);
+
+  optimized_4bit::RunKernel<optimized_4bit::FilterWidth, rhs_width,
+                            optimized_4bit::FilterDepth>(
+      test_lhs.data(), test_rhs.data(), test_accum.data(), lhs_layout_rows,
+      lhs_layout_cols, rhs_layout_rows, rhs_layout_cols, rhs_layout_rows,
+      lhs_layout_rows);
+
+  for (int i = 0; i < (rhs_layout_rows * lhs_layout_rows); ++i) {
+    EXPECT_EQ(test_accum[i], 1950720);
+  }
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/optimized/optimized_ops_utils.h

@@ -0,0 +1,91 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_ops {
+// Make a local VectorMap typedef allowing to map a float array
+// as a Eigen vector expression. The std::conditional here is to
+// construct the suitable Eigen type for the constness of the
+// data. Indeed, for const data, we need to produce
+//    Eigen::Map<const Eigen::Matrix<float, ...>>
+// and not the more straightforward
+//    Eigen::Map<Eigen::Matrix<const float, ...>>
+template <typename Scalar>
+using VectorMap = typename std::conditional<
+    std::is_const<Scalar>::value,
+    Eigen::Map<const Eigen::Matrix<typename std::remove_const<Scalar>::type,
+                                   Eigen::Dynamic, 1>>,
+    Eigen::Map<Eigen::Matrix<Scalar, Eigen::Dynamic, 1>>>::type;
+
+template <typename Scalar>
+VectorMap<Scalar> MapAsVector(Scalar* data, const RuntimeShape& shape) {
+  const int size = shape.FlatSize();
+  return VectorMap<Scalar>(data, size, 1);
+}
+
+// Make a local VectorMap typedef allowing to map a float array
+// as a Eigen matrix expression. The same explanation as for VectorMap
+// above also applies here.
+template <typename Scalar>
+using MatrixMap = typename std::conditional<
+    std::is_const<Scalar>::value,
+    Eigen::Map<const Eigen::Matrix<typename std::remove_const<Scalar>::type,
+                                   Eigen::Dynamic, Eigen::Dynamic>>,
+    Eigen::Map<Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>>>::type;
+
+template <typename Scalar>
+MatrixMap<Scalar> MapAsMatrixWithLastDimAsRows(Scalar* data,
+                                               const RuntimeShape& shape) {
+  const int dims_count = shape.DimensionsCount();
+  const int rows = shape.Dims(dims_count - 1);
+  const int cols = FlatSizeSkipDim(shape, dims_count - 1);
+  return MatrixMap<Scalar>(data, rows, cols);
+}
+
+template <typename Scalar>
+MatrixMap<Scalar> MapAsMatrixWithFirstDimAsCols(Scalar* data,
+                                                const RuntimeShape& shape) {
+  const int cols = shape.Dims(0);
+  const int rows = FlatSizeSkipDim(shape, 0);
+  return MatrixMap<Scalar>(data, rows, cols);
+}
+
+template <typename Scalar>
+using ArrayMap = typename std::conditional<
+    std::is_const<Scalar>::value,
+    Eigen::Map<const Eigen::Array<typename std::remove_const<Scalar>::type,
+                                  Eigen::Dynamic, Eigen::Dynamic>>,
+    Eigen::Map<Eigen::Array<Scalar, Eigen::Dynamic, Eigen::Dynamic>>>::type;
+
+template <typename Scalar>
+ArrayMap<Scalar> MapAsArrayWithLastDimAsRows(Scalar* data,
+                                             const RuntimeShape& shape) {
+  const int dims_count = shape.DimensionsCount();
+  const int rows = shape.Dims(dims_count - 1);
+  const int cols = FlatSizeSkipDim(shape, dims_count - 1);
+  return ArrayMap<Scalar>(data, rows, cols);
+}
+
+}  // namespace optimized_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_OPTIMIZED_OPS_UTILS_H_

tensorflow/lite/kernels/internal/optimized/reduce.h

@@ -0,0 +1,808 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_
+
+#include <stdint.h>
+
+#include <algorithm>
+#include <limits>
+#include <vector>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops_utils.h"
+#include "tensorflow/lite/kernels/internal/optimized/reduce_utils.h"
+#include "tensorflow/lite/kernels/internal/reduce_common.h"
+#include "tensorflow/lite/kernels/internal/reference/reduce.h"
+#include "tensorflow/lite/kernels/internal/runtime_shape.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/kernel_util.h"
+
+namespace tflite {
+namespace optimized_ops {
+
+inline void MeanImpl(const tflite::MeanParams& op_params,
+                     const RuntimeShape& input_shape, const uint8_t* input_data,
+                     int32 multiplier, int32 shift, int32 bias,
+                     const RuntimeShape& output_shape, uint8_t* output_data,
+                     int start_depth, int end_depth) {
+  ruy::profiler::ScopeLabel label("Mean4D/Uint8/MeanImpl");
+
+  // Current implementation only supports dimension equals 4 and simultaneous
+  // reduction over width and height.
+  const int output_batch = output_shape.Dims(0);
+  const int output_height = output_shape.Dims(2);
+  const int output_width = output_shape.Dims(2);
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+
+  TFLITE_CHECK_EQ(op_params.axis_count, 2);
+  TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
+               (op_params.axis[0] == 2 && op_params.axis[1] == 1));
+  TFLITE_CHECK_EQ(output_height, 1);
+  TFLITE_CHECK_EQ(output_width, 1);
+
+  constexpr int32_t kMinValue = std::numeric_limits<uint8_t>::min();
+  constexpr int32_t kMaxValue = std::numeric_limits<uint8_t>::max();
+
+#ifdef USE_NEON
+  const int32x4_t bias_dup = vdupq_n_s32(bias);
+  const int32x4_t min_dup = vdupq_n_s32(kMinValue);
+  const int32x4_t max_dup = vdupq_n_s32(kMaxValue);
+#endif  // USE_NEON
+
+  for (int out_b = 0; out_b < output_batch; ++out_b) {
+    int out_d = start_depth;
+#ifdef USE_NEON
+
+    for (; out_d <= end_depth - 16; out_d += 16) {
+      int32x4x4_t temp_sum;
+      temp_sum.val[0] = vdupq_n_s32(0);
+      temp_sum.val[1] = vdupq_n_s32(0);
+      temp_sum.val[2] = vdupq_n_s32(0);
+      temp_sum.val[3] = vdupq_n_s32(0);
+      for (int in_h = 0; in_h < input_height; ++in_h) {
+        for (int in_w = 0; in_w < input_width; ++in_w) {
+          const uint8_t* input_data_ptr =
+              input_data + Offset(input_shape, out_b, in_h, in_w, out_d);
+          uint8x16_t input_data_val = vld1q_u8(input_data_ptr);
+
+          int16x8_t input_data_low_shift =
+              vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(input_data_val)));
+          int16x8_t input_data_high_shift =
+              vreinterpretq_s16_u16(vmovl_u8(vget_high_u8(input_data_val)));
+
+          int32x4_t input_low_low =
+              vmovl_s16(vget_low_s16(input_data_low_shift));
+          int32x4_t input_high_low =
+              vmovl_s16(vget_high_s16(input_data_low_shift));
+          int32x4_t input_low_high =
+              vmovl_s16(vget_low_s16(input_data_high_shift));
+          int32x4_t input_high_high =
+              vmovl_s16(vget_high_s16(input_data_high_shift));
+
+          temp_sum.val[0] = vaddq_s32(temp_sum.val[0], input_low_low);
+          temp_sum.val[1] = vaddq_s32(temp_sum.val[1], input_high_low);
+          temp_sum.val[2] = vaddq_s32(temp_sum.val[2], input_low_high);
+          temp_sum.val[3] = vaddq_s32(temp_sum.val[3], input_high_high);
+        }
+      }
+
+      temp_sum =
+          MultiplyByQuantizedMultiplier4Rows(temp_sum, multiplier, shift);
+
+      temp_sum.val[0] = vaddq_s32(temp_sum.val[0], bias_dup);
+      temp_sum.val[1] = vaddq_s32(temp_sum.val[1], bias_dup);
+      temp_sum.val[2] = vaddq_s32(temp_sum.val[2], bias_dup);
+      temp_sum.val[3] = vaddq_s32(temp_sum.val[3], bias_dup);
+
+      temp_sum.val[0] = vminq_s32(vmaxq_s32(temp_sum.val[0], min_dup), max_dup);
+      temp_sum.val[1] = vminq_s32(vmaxq_s32(temp_sum.val[1], min_dup), max_dup);
+      temp_sum.val[2] = vminq_s32(vmaxq_s32(temp_sum.val[2], min_dup), max_dup);
+      temp_sum.val[3] = vminq_s32(vmaxq_s32(temp_sum.val[3], min_dup), max_dup);
+
+      uint16x4_t narrowed_low_low =
+          vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[0]));
+      uint16x4_t narrowed_high_low =
+          vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[1]));
+      uint16x4_t narrowed_low_high =
+          vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[2]));
+      uint16x4_t narrowed_high_high =
+          vmovn_u32(vreinterpretq_u32_s32(temp_sum.val[3]));
+
+      uint16x8_t combined_low =
+          vcombine_u16(narrowed_low_low, narrowed_high_low);
+      uint16x8_t combined_high =
+          vcombine_u16(narrowed_low_high, narrowed_high_high);
+
+      uint8x8_t narrowed_low = vmovn_u16(combined_low);
+      uint8x8_t narrowed_high = vmovn_u16(combined_high);
+
+      uint8x16_t combined_output = vcombine_u8(narrowed_low, narrowed_high);
+
+      uint8_t* output_data_ptr =
+          output_data + Offset(output_shape, out_b, 0, 0, out_d);
+      vst1q_u8(output_data_ptr, combined_output);
+    }
+#endif  // USE_NEON
+
+    for (; out_d < end_depth; ++out_d) {
+      int acc = 0;
+      for (int in_h = 0; in_h < input_height; ++in_h) {
+        for (int in_w = 0; in_w < input_width; ++in_w) {
+          acc += input_data[Offset(input_shape, out_b, in_h, in_w, out_d)];
+        }
+      }
+
+      acc = MultiplyByQuantizedMultiplier(acc, multiplier, shift);
+      acc += bias;
+      acc = std::min(std::max(acc, kMinValue), kMaxValue);
+      output_data[Offset(output_shape, out_b, 0, 0, out_d)] =
+          static_cast<uint8_t>(acc);
+    }
+  }
+}
+
+struct MeanWorkerTask : cpu_backend_threadpool::Task {
+  MeanWorkerTask(const tflite::MeanParams& op_params,
+                 const RuntimeShape& input_shape, const uint8_t* input_data,
+                 int32 multiplier, int32 shift, int32 bias,
+                 const RuntimeShape& output_shape, uint8_t* output_data,
+                 int start_height, int end_height)
+      : op_params(op_params),
+        input_shape(input_shape),
+        input_data(input_data),
+        multiplier(multiplier),
+        shift(shift),
+        bias(bias),
+        output_shape(output_shape),
+        output_data(output_data),
+        start_height(start_height),
+        end_height(end_height) {}
+
+  void Run() override {
+    MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
+             output_shape, output_data, start_height, end_height);
+  }
+
+ private:
+  const tflite::MeanParams& op_params;
+  const RuntimeShape& input_shape;
+  const uint8_t* input_data;
+  int32 multiplier;
+  int32 shift;
+  int32 bias;
+  const RuntimeShape& output_shape;
+  uint8_t* output_data;
+  int start_height;
+  int end_height;
+};
+
+inline void Mean(const tflite::MeanParams& op_params,
+                 const RuntimeShape& unextended_input_shape,
+                 const uint8_t* input_data, int32 input_zero_point,
+                 float input_scale, const RuntimeShape& unextended_output_shape,
+                 uint8_t* output_data, int32 output_zero_point,
+                 float output_scale, CpuBackendContext* cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("Mean4D/Uint8");
+  // Current implementation only supports dimension equals 4 and simultaneous
+  // reduction over width and height.
+  TFLITE_CHECK_EQ(unextended_input_shape.DimensionsCount(), 4);
+  TFLITE_CHECK_LE(unextended_output_shape.DimensionsCount(), 4);
+  const RuntimeShape input_shape =
+      RuntimeShape::ExtendedShape(4, unextended_input_shape);
+  const RuntimeShape output_shape =
+      RuntimeShape::ExtendedShape(4, unextended_output_shape);
+  const int output_height = output_shape.Dims(1);
+  const int output_width = output_shape.Dims(2);
+  const int output_depth = output_shape.Dims(3);
+
+  TFLITE_CHECK_EQ(op_params.axis_count, 2);
+  TFLITE_CHECK((op_params.axis[0] == 1 && op_params.axis[1] == 2) ||
+               (op_params.axis[0] == 2 && op_params.axis[1] == 1));
+  TFLITE_CHECK_EQ(output_height, 1);
+  TFLITE_CHECK_EQ(output_width, 1);
+
+  const int input_height = input_shape.Dims(1);
+  const int input_width = input_shape.Dims(2);
+  const float num_elements_in_axis = input_width * input_height;
+
+  float temp = input_zero_point * input_scale / output_scale;
+  temp = temp > 0 ? temp + 0.5f : temp - 0.5f;
+  int32_t bias = output_zero_point - static_cast<int32_t>(temp);
+  float real_scale = input_scale / (num_elements_in_axis * output_scale);
+
+  int32 multiplier, shift;
+  QuantizeMultiplier(real_scale, &multiplier, &shift);
+
+  constexpr int kMinDepthPerThread = 8;
+  int thread_count = output_depth / kMinDepthPerThread;
+  thread_count = thread_count > 0 ? thread_count : 1;
+  const int capped_thread_count =
+      std::min(thread_count, cpu_backend_context->max_num_threads());
+
+  if (capped_thread_count == 1) {
+    MeanImpl(op_params, input_shape, input_data, multiplier, shift, bias,
+             output_shape, output_data, 0, output_depth);
+  } else {
+    // Instead parallel for batch, we loop for the output_depth since batch
+    // is typical 1.
+    std::vector<MeanWorkerTask> tasks;
+    // TODO(b/131746020) don't create new heap allocations every time.
+    // At least we make it a single heap allocation by using reserve().
+    tasks.reserve(capped_thread_count);
+    int depth_start = 0;
+    for (int i = 0; i < capped_thread_count; ++i) {
+      // Try to distribute the tasks as even as possible.
+      int depth_end = depth_start +
+                      (output_depth - depth_start) / (capped_thread_count - i);
+      tasks.emplace_back(op_params, input_shape, input_data, multiplier, shift,
+                         bias, output_shape, output_data, depth_start,
+                         depth_end);
+      depth_start = depth_end;
+    }
+    cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
+                                    cpu_backend_context);
+  }
+}
+
+template <typename T>
+struct SumOp {
+  inline T operator()(const T& a) const { return a; }
+  inline T operator()(const T& a, const T& b) const { return a + b; }
+  static constexpr T kNeutralElement = T(0);
+};
+
+template <typename T, typename U>
+struct CastSumOp {
+  inline U operator()(const T& a) const { return static_cast<U>(a); }
+  inline U operator()(const U& a, const T& b) const {
+    return a + static_cast<U>(b);
+  }
+  static constexpr U kNeutralElement = U(0);
+};
+
+template <typename T>
+struct ProdOp {
+  inline T operator()(const T& a) const { return a; }
+  inline T operator()(const T& a, const T& b) const { return a * b; }
+  static constexpr T kNeutralElement = T(1);
+};
+
+template <typename T>
+struct MaxOp {
+  inline T operator()(const T& a) const { return a; }
+  inline T operator()(const T& a, const T& b) const { return (a > b) ? a : b; }
+  static constexpr T kNeutralElement = std::numeric_limits<T>::lowest();
+};
+
+template <typename T>
+struct MinOp {
+  inline T operator()(const T& a) const { return a; }
+  inline T operator()(const T& a, const T& b) const { return (a < b) ? a : b; }
+  static constexpr T kNeutralElement = std::numeric_limits<T>::max();
+};
+
+struct AndOp {
+  inline bool operator()(bool a) const { return a; }
+  inline bool operator()(bool a, bool b) const { return a && b; }
+  static constexpr bool kNeutralElement = true;
+};
+
+struct OrOp {
+  inline bool operator()(bool a) const { return a; }
+  inline bool operator()(bool a, bool b) const { return a || b; }
+  static constexpr bool kNeutralElement = false;
+};
+
+// When the number of axis is zero, the reduction is simply a copy.
+template <typename T>
+void ReduceIsCopy(const T* input_data, const int* input_dims,
+                  const int input_num_dims, T* output_data) {
+  int num_elems = NumElements(input_dims, input_num_dims);
+  memcpy(output_data, input_data, num_elems * sizeof(T));
+}
+
+// Reduces the input over either odd or even dimensions using Op.
+// One recursive call for each dimension is made.
+// 'depth' is the depth of recursion.
+// 'parity' indicates whether odd or even dimensions are being reduced.
+// ReducerFirst is applied to the first element to be written to each output
+// position.
+// ReducerNext is applied to each subsequent element to be written to each
+// output position.
+template <typename T, typename U, typename ReducerFirst, typename ReducerNext>
+inline std::pair<const T*, U*> ReduceImpl(const T* input_data,
+                                          const int* input_dims, U* output_data,
+                                          int depth, int parity, bool next,
+                                          const ReducerFirst& reducer_first,
+                                          const ReducerNext& reducer_next) {
+  // The output pointer is incremented conditionally depending on whether the
+  // odd or even dimension is being reduced.
+  // The input pointer is always incremented as each input is read once.
+  if (depth > 0) {
+    U* future_output = output_data;
+    bool update_output = (depth % 2) == parity;
+    for (int i = 0; i < input_dims[0]; ++i) {
+      if (i > 0 && !update_output) {
+        next = true;
+      }
+      std::tie(input_data, future_output) =
+          ReduceImpl(input_data, &input_dims[1], output_data, depth - 1, parity,
+                     next, reducer_first, reducer_next);
+      if (update_output) {
+        output_data = future_output;
+      }
+    }
+    output_data = future_output;
+  } else {
+    // Reduce the final dimension.
+    if (parity) {
+      // Reduce the even dimension. The entire dimension is reduced into one
+      // value.
+      U res = next ? reducer_next(*output_data, *input_data++)
+                   : reducer_first(*input_data++);
+      for (int i = 1; i < input_dims[0]; ++i) {
+        res = reducer_next(res, *input_data++);
+      }
+      *output_data++ = res;
+    } else {
+      // Reduce the odd dimension. Each input is accumulated into a separate
+      // output.
+      if (!next) {
+        for (int i = 0; i < input_dims[0]; ++i) {
+          U res = reducer_first(*input_data++);
+          *output_data++ = res;
+        }
+      } else {
+        for (int i = 0; i < input_dims[0]; ++i) {
+          U res = *output_data;
+          res = reducer_next(res, *input_data++);
+          *output_data++ = res;
+        }
+      }
+    }
+  }
+  return {input_data, output_data};
+}
+
+// A generic reduce method that can be used for reduce_sum, reduce_mean, etc.
+// This method iterates through input data and reduce elements along the
+// dimensions given in axis. ReducerFirst is used the first time each output
+// element is written and ReducerNext is used for all subsequent writes.
+template <typename In, typename Out, typename ReducerFirst,
+          typename ReducerNext>
+inline bool Reduce(const In* input_data, const int* input_dims,
+                   const int input_num_dims, const int* axis,
+                   const int num_axis, Out* output_data,
+                   const ReducerFirst& reducer_first,
+                   const ReducerNext& reducer_next) {
+  const int parity = (axis[num_axis - 1] == input_num_dims - 1) ? 1 : 0;
+  ReduceImpl(input_data, input_dims, output_data, input_num_dims - 1, parity,
+             /*next=*/false, reducer_first, reducer_next);
+  return true;
+}
+
+// Computes the mean or sum of elements across dimensions given in axis.
+// It does so in two stages, first calculates the sum of elements along the axis
+// then divides it by the number of element in axis for quantized values.
+template <typename T, typename U>
+bool QuantizedMeanOrSum(const T* input_data, int32_t input_zero_point,
+                        float input_scale, const int* input_dims,
+                        const int input_num_dims, T* output_data,
+                        int32_t output_zero_point, float output_scale,
+                        const int* output_dims, const int output_num_dims,
+                        const int* axis, const int num_axis_dimensions,
+                        bool keep_dims, int* normalized_dims,
+                        int* resolved_axis, U* temp_sum, bool compute_sum) {
+  const int32_t kMinValue = std::numeric_limits<T>::min();
+  const int32_t kMaxValue = std::numeric_limits<T>::max();
+  ruy::profiler::ScopeLabel label(compute_sum ? "QuantizedSum"
+                                              : "QuantizedMean");
+  // Reset output data.
+  size_t num_outputs = 1;
+  for (int idx = 0; idx < output_num_dims; ++idx) {
+    size_t current = static_cast<size_t>(output_dims[idx]);
+    // Overflow prevention.
+    if (num_outputs > std::numeric_limits<size_t>::max() / current) {
+      return false;
+    }
+    num_outputs *= current;
+  }
+
+  // Return early when input shape has zero dim. This is done after initializing
+  // data for output tensor because there are cases that the input tensor is
+  // empty but output tensor is not. In that case, output tensor should be
+  // filled with init_value.
+  for (int i = 0; i < input_num_dims; ++i) {
+    if (input_dims[i] == 0) return true;
+  }
+
+  // Resolve axis.
+  int num_resolved_axis = 0;
+  int normalized_num_dims = 0;
+  if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
+                                 resolved_axis, num_resolved_axis, input_dims,
+                                 normalized_dims, normalized_num_dims)) {
+    return false;
+  }
+
+  if (num_resolved_axis == 0) {
+    int count = NumElements(input_dims, input_num_dims);
+    for (int i = 0; i < count; ++i) {
+      temp_sum[i] = U(input_data[i]);
+    }
+  } else {
+    if (!Reduce<T, U, CastSumOp<T, U>, CastSumOp<T, U>>(
+            input_data, normalized_dims, normalized_num_dims, resolved_axis,
+            num_resolved_axis, temp_sum, CastSumOp<T, U>(),
+            CastSumOp<T, U>())) {
+      return false;
+    }
+  }
+
+  // Calculate mean by dividing output_data by num of aggregated element.
+  size_t num_elements_in_axis = 1;
+  for (int idx = 0; idx < num_resolved_axis; ++idx) {
+    size_t current = static_cast<size_t>(normalized_dims[resolved_axis[idx]]);
+    // Overflow prevention.
+    if (current > (std::numeric_limits<size_t>::max() / num_elements_in_axis)) {
+      return false;
+    }
+    num_elements_in_axis *= current;
+  }
+
+  if (num_elements_in_axis > 0) {
+    const float scale = input_scale / output_scale;
+    if (compute_sum) {
+      const float bias = -input_zero_point * scale * num_elements_in_axis;
+      for (size_t idx = 0; idx < num_outputs; ++idx) {
+        U value = static_cast<U>(TfLiteRound(temp_sum[idx] * scale + bias)) +
+                  output_zero_point;
+        value = std::min(std::max(value, kMinValue), kMaxValue);
+        output_data[idx] = static_cast<T>(value);
+      }
+    } else {
+      const float bias = -input_zero_point * scale;
+      for (size_t idx = 0; idx < num_outputs; ++idx) {
+        float float_mean = static_cast<float>(temp_sum[idx]) /
+                           static_cast<float>(num_elements_in_axis);
+        float result = TfLiteMin(
+            TfLiteRound(float_mean * scale + bias) + output_zero_point,
+            static_cast<float>(std::numeric_limits<T>::max()));
+        result = TfLiteMax(result,
+                           static_cast<float>(std::numeric_limits<T>::min()));
+        output_data[idx] = static_cast<T>(result);
+      }
+    }
+  }
+  return true;
+}
+
+using ops::builtin::reduce::ReduceType;
+
+template <typename T>
+inline bool ReduceDispatcher(const T* input_data, const int* input_dims,
+                             const int input_num_dims, const int* output_dims,
+                             int output_num_dims, T* output_data,
+                             const int* axis, const int64_t num_axis_dimensions,
+                             ReduceType reduce_type) {
+  T init_value;
+  switch (reduce_type) {
+    case ReduceType::kProd:
+      init_value = ProdOp<T>::kNeutralElement;
+      break;
+    case ReduceType::kSum:
+      init_value = SumOp<T>::kNeutralElement;
+      break;
+    case ReduceType::kMin:
+      init_value = MinOp<T>::kNeutralElement;
+      break;
+    case ReduceType::kMax:
+      init_value = MaxOp<T>::kNeutralElement;
+      break;
+    default:
+      return false;
+  }
+  // Return early when input shape has zero dim. This is done after initializing
+  // data for output tensor because there are cases that the input tensor is
+  // empty but output tensor is not. In that case, output tensor should be
+  // filled with Op::kNeutralElement.
+  for (int i = 0; i < input_num_dims; ++i) {
+    if (input_dims[i] == 0) {
+      return reference_ops::InitTensorDataForReduce(
+          output_dims, output_num_dims, init_value, output_data);
+    }
+  }
+
+  switch (reduce_type) {
+    case ReduceType::kProd:
+      return Reduce<T, T, ProdOp<T>, ProdOp<T>>(
+          input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
+          output_data, ProdOp<T>(), ProdOp<T>());
+    case ReduceType::kSum:
+      return Reduce<T, T, SumOp<T>, SumOp<T>>(
+          input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
+          output_data, SumOp<T>(), SumOp<T>());
+    case ReduceType::kMin:
+      return Reduce<T, T, MinOp<T>, MinOp<T>>(
+          input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
+          output_data, MinOp<T>(), MinOp<T>());
+    case ReduceType::kMax:
+      return Reduce<T, T, MaxOp<T>, MaxOp<T>>(
+          input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
+          output_data, MaxOp<T>(), MaxOp<T>());
+    default:
+      return false;
+  }
+}
+
+template <>
+inline bool ReduceDispatcher<bool>(const bool* input_data,
+                                   const int* input_dims,
+                                   const int input_num_dims,
+                                   const int* output_dims, int output_num_dims,
+                                   bool* output_data, const int* axis,
+                                   const int64_t num_axis_dimensions,
+                                   ReduceType reduce_type) {
+  bool init_value;
+  switch (reduce_type) {
+    case ReduceType::kAny:
+      init_value = OrOp::kNeutralElement;
+      break;
+    case ReduceType::kAll:
+      init_value = AndOp::kNeutralElement;
+      break;
+    default:
+      return false;
+  }
+  // Return early when input shape has zero dim. This is done after initializing
+  // data for output tensor because there are cases that the input tensor is
+  // empty but output tensor is not. In that case, output tensor should be
+  // filled with Op::kNeutralElement.
+  for (int i = 0; i < input_num_dims; ++i) {
+    if (input_dims[i] == 0) {
+      return reference_ops::InitTensorDataForReduce(
+          output_dims, output_num_dims, init_value, output_data);
+    }
+  }
+  switch (reduce_type) {
+    case ReduceType::kAll:
+      return Reduce<bool, bool, AndOp, AndOp>(
+          input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
+          output_data, AndOp(), AndOp());
+    case ReduceType::kAny:
+      return Reduce<bool, bool, OrOp, OrOp>(
+          input_data, input_dims, input_num_dims, axis, num_axis_dimensions,
+          output_data, OrOp(), OrOp());
+    default:
+      return false;
+  }
+}
+
+// Calculate the reduced product by rescaling each multiplication step to
+// avoid an overflow.
+template <typename T>
+struct ReducerFirst {
+  explicit ReducerFirst(int input_zero_point_arg)
+      : input_zero_point(input_zero_point_arg) {}
+  int32_t operator()(T in) const { return in - input_zero_point; }
+  int input_zero_point;
+};
+
+template <typename T>
+struct ReducerNext {
+  ReducerNext(int32_t input_zero_point_arg, int32_t scaling_multiplier_arg,
+              int32_t scaling_shift_arg)
+      : input_zero_point(input_zero_point_arg),
+        scaling_multiplier(scaling_multiplier_arg),
+        scaling_shift(scaling_shift_arg) {}
+  int32_t operator()(int32_t current, T in) const {
+    const int64_t result =
+        static_cast<int64_t>(current) * (in - input_zero_point);
+    return MultiplyByQuantizedMultiplier(result, scaling_multiplier,
+                                         scaling_shift);
+  }
+  int32_t input_zero_point, scaling_multiplier, scaling_shift;
+};
+
+template <typename T>
+inline bool QuantizedReduceProd(
+    const T* input_data, int32_t input_zero_point,
+    const RuntimeShape& input_shape, T* output_data, int32_t output_zero_point,
+    const RuntimeShape& output_shape, const int* axis,
+    const int64_t num_axis_dimensions, int* resolved_axis, int* normalized_dims,
+    int32_t* temp_prod, int32_t scaling_multiplier, int scaling_shift) {
+  const int32_t kMinValue = std::numeric_limits<T>::min();
+  const int32_t kMaxValue = std::numeric_limits<T>::max();
+
+  // Resolve axis.
+  int num_resolved_axis = 0;
+  int normalized_num_dims = 0;
+  if (!reduce_utils::ResolveAxis(input_shape.DimensionsCount(), axis,
+                                 num_axis_dimensions, resolved_axis,
+                                 num_resolved_axis, input_shape.DimsData(),
+                                 normalized_dims, normalized_num_dims)) {
+    return false;
+  }
+
+  if (!Reduce<T, int32_t, ReducerFirst<T>, ReducerNext<T>>(
+          input_data, normalized_dims, normalized_num_dims, resolved_axis,
+          num_resolved_axis, temp_prod, ReducerFirst<T>(input_zero_point),
+          ReducerNext<T>(input_zero_point, scaling_multiplier,
+                         scaling_shift))) {
+    return false;
+  }
+
+  for (int i = 0; i < output_shape.FlatSize(); i++) {
+    int32_t result =
+        MultiplyByQuantizedMultiplier(static_cast<int64_t>(temp_prod[i]),
+                                      scaling_multiplier, scaling_shift) +
+        output_zero_point;
+    result = std::min(std::max(result, kMinValue), kMaxValue);
+    output_data[i] = static_cast<T>(result);
+  }
+
+  return true;
+}
+
+template <typename T>
+inline void Mean(const tflite::MeanParams& op_params,
+                 const RuntimeShape& input_shape, const T* input_data,
+                 const RuntimeShape& output_shape, T* output_data) {
+  return reference_ops::Mean(op_params, input_shape, input_data, output_shape,
+                             output_data);
+}
+
+// Computes the mean of elements across dimensions given in axis.
+// It does so in two stages, first calculates the sum of elements along the axis
+// then divides it by the number of element in axis.
+template <typename T, typename U>
+inline bool MeanGeneral(const T* input_data, const int* input_dims,
+                        const int input_num_dims, T* output_data,
+                        const int* output_dims, const int output_num_dims,
+                        const int* axis, const int num_axis_dimensions,
+                        bool keep_dims, int* normalized_dims,
+                        int* resolved_axis, U* temp_sum) {
+  ruy::profiler::ScopeLabel label("Mean");
+  // Resolve axis.
+  int num_resolved_axis = 0;
+  int normalized_num_dims = 0;
+  if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
+                                 resolved_axis, num_resolved_axis, input_dims,
+                                 normalized_dims, normalized_num_dims)) {
+    return false;
+  }
+  if (num_resolved_axis == 0) {
+    optimized_ops::ReduceIsCopy(input_data, input_dims, input_num_dims,
+                                output_data);
+    return true;
+  }
+  // Reset output data.
+  size_t num_outputs = 1;
+  for (int idx = 0; idx < output_num_dims; ++idx) {
+    size_t current = static_cast<size_t>(output_dims[idx]);
+    // Overflow prevention.
+    if (num_outputs > std::numeric_limits<size_t>::max() / current) {
+      return false;
+    }
+    num_outputs *= current;
+  }
+
+  if (!Reduce<T, U, CastSumOp<T, U>, CastSumOp<T, U>>(
+          input_data, normalized_dims, normalized_num_dims, resolved_axis,
+          num_resolved_axis, temp_sum, CastSumOp<T, U>(), CastSumOp<T, U>())) {
+    return false;
+  }
+
+  // Calculate mean by dividing output_data by num of aggregated element.
+  size_t num_elements_in_axis = 1;
+  for (int idx = 0; idx < num_resolved_axis; ++idx) {
+    size_t current = static_cast<size_t>(normalized_dims[resolved_axis[idx]]);
+    // Overflow prevention.
+    if (current > (std::numeric_limits<size_t>::max() / num_elements_in_axis)) {
+      return false;
+    }
+    num_elements_in_axis *= current;
+  }
+
+  if (num_elements_in_axis > 0) {
+    for (size_t idx = 0; idx < num_outputs; ++idx) {
+      output_data[idx] =
+          static_cast<T>(temp_sum[idx] / static_cast<U>(num_elements_in_axis));
+    }
+  }
+  return true;
+}
+
+template <typename T, typename U>
+inline bool Mean(const T* input_data, const int* input_dims,
+                 const int input_num_dims, T* output_data,
+                 const int* output_dims, const int output_num_dims,
+                 const int* axis, const int num_axis_dimensions, bool keep_dims,
+                 int* normalized_dims, int* resolved_axis, U* temp_sum) {
+  return MeanGeneral(input_data, input_dims, input_num_dims, output_data,
+                     output_dims, output_num_dims, axis, num_axis_dimensions,
+                     false, normalized_dims, resolved_axis, temp_sum);
+}
+
+// Use Eigen when Mean is calculated over the last dimension only of a float
+// tensor.
+template <>
+inline bool Mean<float, float>(const float* input_data, const int* input_dims,
+                               const int input_num_dims, float* output_data,
+                               const int* output_dims,
+                               const int output_num_dims, const int* axis,
+                               const int num_axis_dimensions, bool keep_dims,
+                               int* normalized_dims, int* resolved_axis,
+                               float* temp_sum) {
+  // Handle reduce_mean for the last dimensions.
+  int num_resolved_axis = 0;
+  int normalized_num_dims = 0;
+  if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
+                                 resolved_axis, num_resolved_axis, input_dims,
+                                 normalized_dims, normalized_num_dims)) {
+    return false;
+  }
+  if (normalized_num_dims > 1 && num_resolved_axis == 1 &&
+      resolved_axis[0] == (normalized_num_dims - 1)) {
+    ruy::profiler::ScopeLabel label("MeanLastDim/Float");
+    int output_size = normalized_dims[0];
+    const int last_input_dim = normalized_dims[1];
+
+    // TODO(b/152563685): Consider use eigen to cover more general cases.
+    const MatrixMap<const float> in_mat(input_data, last_input_dim,
+                                        output_size);
+    VectorMap<float> out(output_data, output_size, 1);
+    out = (in_mat.array().colwise().sum()) / static_cast<float>(last_input_dim);
+    return true;
+  }
+
+  return MeanGeneral(input_data, input_dims, input_num_dims, output_data,
+                     output_dims, output_num_dims, axis, num_axis_dimensions,
+                     false, normalized_dims, resolved_axis, temp_sum);
+}
+
+// Computes the generic value (i.e., sum/max/min/prod) of elements across
+// dimensions given in axis. It needs to pass in init_value and reducer.
+template <typename T>
+inline bool ReduceGeneric(const T* input_data, const int* input_dims,
+                          const int input_num_dims, T* output_data,
+                          const int* output_dims, const int output_num_dims,
+                          const int* axis, const int64_t num_axis_dimensions,
+                          int* resolved_axis, int* normalized_dims,
+                          ReduceType reduce_type) {
+  int num_resolved_axis = 0;
+  int normalized_num_dims = 0;
+  if (!reduce_utils::ResolveAxis(input_num_dims, axis, num_axis_dimensions,
+                                 resolved_axis, num_resolved_axis, input_dims,
+                                 normalized_dims, normalized_num_dims)) {
+    return false;
+  }
+  if (num_resolved_axis == 0) {
+    optimized_ops::ReduceIsCopy(input_data, input_dims, input_num_dims,
+                                output_data);
+    return true;
+  }
+  return ReduceDispatcher(input_data, normalized_dims, normalized_num_dims,
+                          output_dims, output_num_dims, output_data,
+                          resolved_axis, num_resolved_axis, reduce_type);
+}
+
+}  // namespace optimized_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_H_

tensorflow/lite/kernels/internal/optimized/reduce_utils.h

@@ -0,0 +1,138 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_UTILS_H_
+
+#include <stdint.h>
+
+#include <algorithm>
+#include <cstring>
+
+namespace tflite {
+namespace reduce_utils {
+
+inline void RemoveSize1Dims(int* shape_out, int& out_num_dims, int* axis_out,
+                            int& out_num_axis) {
+  for (int64_t i = 0; i < out_num_dims;) {
+    if (shape_out[i] == 1) {
+      for (int64_t j = i + 1; j < out_num_dims; ++j) {
+        shape_out[j - 1] = shape_out[j];
+      }
+      for (int64_t j = 0; j < out_num_axis; ++j) {
+        if (axis_out[j] == i) {
+          for (int64_t k = j + 1; k < out_num_axis; ++k) {
+            axis_out[k - 1] = axis_out[k];
+          }
+          out_num_axis -= 1;
+          break;
+        }
+      }
+      for (int64_t j = 0; j < out_num_axis; ++j) {
+        if (axis_out[j] > i) {
+          axis_out[j] -= 1;
+        }
+      }
+      --out_num_dims;
+    } else {
+      ++i;
+    }
+  }
+}
+
+// This method parses the input 'axis' to remove duplicates, handle negative
+// values and remove redundant dimensions. It returns a valid 'axis_out' and
+// 'shape_out' contains the flattened input shape. 'out_num_dims' contains the
+// reduced number of dimensions.
+inline bool ResolveAxis(const int num_dims, const int* axis,
+                        const int64_t num_axis, int* axis_out,
+                        int& out_num_axis, const int* shape_in, int* shape_out,
+                        int& out_num_dims) {
+  // Short-circuit axis resolution for scalars; the axis will go unused.
+  if (num_dims == 0) {
+    out_num_axis = 0;
+    out_num_dims = 0;
+    return true;
+  }
+  out_num_axis = 0;
+  out_num_dims = num_dims;
+  // o(n^2) is fine since out_num_axis should be really small, mostly <= 4
+  for (int64_t idx = 0; idx < num_axis; ++idx) {
+    // Handle negative index. A positive index 'p_idx' can be represented as a
+    // negative index 'n_idx' as: n_idx = p_idx-num_dims
+    // eg: For num_dims=3, [0, 1, 2] is the same as [-3, -2, -1]  */
+    int current = axis[idx] < 0 ? (axis[idx] + num_dims) : axis[idx];
+    if (current < 0 || current >= num_dims) {
+      return false;
+    }
+    bool is_dup = false;
+    for (int j = 0; j < out_num_axis; ++j) {
+      if (axis_out[j] == current) {
+        is_dup = true;
+        break;
+      }
+    }
+    if (!is_dup) {
+      axis_out[out_num_axis] = current;
+      out_num_axis += 1;
+    }
+  }
+  // If two or more adjacent dimensions are either reduced
+  // over or not, then the second and subsequent dimensions may be flattened.
+  memcpy(shape_out, shape_in, num_dims * sizeof(int));
+  std::sort(&axis_out[0], &axis_out[out_num_axis]);
+
+  RemoveSize1Dims(shape_out, out_num_dims, axis_out, out_num_axis);
+  if (out_num_axis > 0) {
+    int64_t j = out_num_axis - 1;
+    // true if the previous index is present in axis_out.
+    bool previous_here = (axis_out[j] == out_num_dims - 1);
+    if (previous_here) {
+      j -= 1;
+    }
+
+    for (int64_t i = out_num_dims - 2; i >= 0; --i) {
+      // true if the current index is present in axis_out.
+      bool current_here = j >= 0 ? (axis_out[j] == i) : false;
+      if (current_here == previous_here) {
+        shape_out[i] *= shape_out[i + 1];
+        for (int64_t k = i + 1; k + 1 < out_num_dims; ++k) {
+          shape_out[k] = shape_out[k + 1];
+        }
+        // All axis bigger than this need to be reduced by 1.
+        for (int64_t k = 0; k < out_num_axis; ++k) {
+          if (axis_out[k] > i) {
+            axis_out[k] -= 1;
+          }
+        }
+        if (current_here) {
+          for (int64_t k = j + 1; k + 1 < out_num_axis; ++k) {
+            axis_out[k] = axis_out[k + 1];
+          }
+          out_num_axis -= 1;
+        }
+        out_num_dims -= 1;
+      }
+      if (current_here) {
+        j -= 1;
+      }
+      previous_here = current_here;
+    }
+  }
+  return true;
+}
+}  // namespace reduce_utils
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_REDUCE_UTILS_H_

tensorflow/lite/kernels/internal/optimized/reduce_utils_test.cc

@@ -0,0 +1,137 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/optimized/reduce_utils.h"
+
+#include <gmock/gmock.h>
+
+namespace tflite {
+namespace reduce_utils {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+void TestFunction(const std::vector<int>& axis_in,
+                  const std::vector<int>& shape_in,
+                  const std::vector<int>& expected_axis_out,
+                  const std::vector<int>& expected_shape_out) {
+  int num_dims = shape_in.size();
+  int expected_out_num_dims = expected_shape_out.size();
+  int actual_out_num_dims;
+  int expected_out_num_axis = expected_axis_out.size();
+  int actual_out_num_axis;
+  std::vector<int> actual_shape_out(num_dims);
+  std::vector<int> actual_axis_out(num_dims);
+  ResolveAxis(shape_in.size(), axis_in.data(), axis_in.size(),
+              actual_axis_out.data(), actual_out_num_axis, shape_in.data(),
+              actual_shape_out.data(), actual_out_num_dims);
+  EXPECT_EQ(expected_out_num_dims, actual_out_num_dims);
+  EXPECT_EQ(expected_out_num_axis, actual_out_num_axis);
+  EXPECT_THAT(expected_shape_out,
+              ElementsAreArray(actual_shape_out.data(), expected_out_num_dims));
+  EXPECT_THAT(expected_axis_out,
+              ElementsAreArray(actual_axis_out.data(), expected_out_num_axis));
+}
+
+TEST(ResolveAxisTest, Flatten_0_1_2) {
+  const std::vector<int> axis_in = {0, 1, 2};
+  const std::vector<int> shape_in = {2, 3, 4, 5};
+  const std::vector<int> expected_shape_out{24, 5};
+  const std::vector<int> expected_axis_out{0};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, Flatten_0_1_2_3) {
+  const std::vector<int> axis_in = {3, 2};
+  const std::vector<int> shape_in = {2, 3, 4, 5};
+  const std::vector<int> expected_shape_out{6, 20};
+  const std::vector<int> expected_axis_out{1};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, ZeroDims) {
+  const std::vector<int> axis_in = {};
+  const std::vector<int> shape_in = {};
+  const std::vector<int> expected_shape_out{};
+  const std::vector<int> expected_axis_out{};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, DoNothing) {
+  const std::vector<int> axis_in = {0};
+  const std::vector<int> shape_in = {4, 5};
+  const std::vector<int> expected_shape_out{4, 5};
+  const std::vector<int> expected_axis_out{0};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, NegativeAxis) {
+  const std::vector<int> axis_in = {-2};
+  const std::vector<int> shape_in = {4, 3};
+  const std::vector<int> expected_shape_out{4, 3};
+  const std::vector<int> expected_axis_out{0};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, NegativeAxisFold) {
+  const std::vector<int> axis_in = {-1};
+  const std::vector<int> shape_in = {4, 3, 5};
+  const std::vector<int> expected_shape_out{12, 5};
+  const std::vector<int> expected_axis_out{1};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, DuplicateAxis) {
+  const std::vector<int> axis_in = {2, 1, 2, 1, 2, 1};
+  const std::vector<int> shape_in = {4, 3, 2};
+  const std::vector<int> expected_shape_out{4, 6};
+  const std::vector<int> expected_axis_out{1};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, DuplicateNegativeAxis) {
+  const std::vector<int> axis_in = {2, -1, -2, -1, 2, 1};
+  const std::vector<int> shape_in = {4, 3, 2};
+  const std::vector<int> expected_shape_out{4, 6};
+  const std::vector<int> expected_axis_out{1};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, RemoveSize1Dim) {
+  const std::vector<int> axis_in = {0};
+  const std::vector<int> shape_in = {1, 4, 3, 1};
+  const std::vector<int> expected_shape_out{4, 3};
+  const std::vector<int> expected_axis_out{};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, OneSize1DimToScalar) {
+  const std::vector<int> axis_in = {0};
+  const std::vector<int> shape_in = {1};
+  const std::vector<int> expected_shape_out{};
+  const std::vector<int> expected_axis_out{};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+TEST(ResolveAxisTest, InterleavedSize1Dim) {
+  const std::vector<int> axis_in = {1, 3};
+  const std::vector<int> shape_in = {1, 2, 1, 4, 1, 7};
+  const std::vector<int> expected_shape_out{8, 7};
+  const std::vector<int> expected_axis_out{0};
+  TestFunction(axis_in, shape_in, expected_axis_out, expected_shape_out);
+}
+
+}  // namespace
+}  // namespace reduce_utils
+}  // namespace tflite

tensorflow/lite/kernels/internal/optimized/sparse_ops/fully_connected.h

@@ -0,0 +1,268 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SPARSE_OPS_FULLY_CONNECTED_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SPARSE_OPS_FULLY_CONNECTED_H_
+
+#include <algorithm>
+#include <cstdint>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/cpu_backend_threadpool.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+#include "tensorflow/lite/kernels/internal/optimized/optimized_ops.h"
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+#include "tensorflow/lite/kernels/internal/tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+namespace optimized_ops {
+
+inline void FullyConnectedSparseWeight(
+    const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
+    const RuntimeShape& input_shape, const float* input_data,
+    const RuntimeShape& weights_shape, const float* weights_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data) {
+  ruy::profiler::ScopeLabel label("FullyConnected");
+  ruy::profiler::ScopeLabel inner_label("Random Sparse");
+  const float output_activation_min = params.float_activation_min;
+  const float output_activation_max = params.float_activation_max;
+
+  const int output_elements = output_shape.FlatSize();
+  const int output_dims_count = output_shape.DimensionsCount();
+  const int weights_dims_count = weights_shape.DimensionsCount();
+  const int batches = FlatSizeSkipDim(output_shape, output_dims_count - 1);
+  const int output_depth = MatchingDim(weights_shape, weights_dims_count - 2,
+                                       output_shape, output_dims_count - 1);
+  const int accum_depth = weights_shape.Dims(weights_dims_count - 1);
+  const int w0_size = sparsity.dim_metadata[0].dense_size;
+  const int* w1_segments = sparsity.dim_metadata[1].array_segments->data;
+  const int* w1_indices = sparsity.dim_metadata[1].array_indices->data;
+
+  for (int i = 0; i < output_elements; ++i) {
+    output_data[i] = 0.f;
+  }
+
+  for (int b = 0; b < batches; ++b) {
+    for (int idx_0 = 0; idx_0 < w0_size; ++idx_0) {
+      for (int pw1 = w1_segments[idx_0]; pw1 < w1_segments[idx_0 + 1]; ++pw1) {
+        int idx_1 = w1_indices[pw1];
+        output_data[b * output_depth + idx_0] +=
+            weights_data[pw1] * input_data[b * accum_depth + idx_1];
+      }
+    }
+  }
+
+  for (int b = 0; b < batches; ++b) {
+    for (int i = 0; i < output_depth; ++i) {
+      float total = output_data[b * output_depth + i];
+      const float bias_value = bias_data ? bias_data[i] : 0;
+      output_data[b * output_depth + i] = ActivationFunctionWithMinMax(
+          total + bias_value, output_activation_min, output_activation_max);
+    }
+  }
+}
+
+inline void FullyConnectedSparseWeight1x16Impl(
+    const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
+    const RuntimeShape& input_shape, const int8_t* input_data,
+    const RuntimeShape& weights_shape, const int8_t* weights_data,
+    const int32_t* per_channel_scale, const int32_t* per_channel_shift,
+    const RuntimeShape& bias_shape, const int32_t* bias_data,
+    const RuntimeShape& output_shape, int8_t* output_data, int thread_start,
+    int thread_end, const CpuBackendContext& cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("FullyConnected");
+  ruy::profiler::ScopeLabel inner_label("1x16 Block Sparse");
+
+  const int input_dims_count = input_shape.DimensionsCount();
+  const int output_dims_count = output_shape.DimensionsCount();
+  const int weights_dims_count = weights_shape.DimensionsCount();
+  const int batches = thread_end - thread_start;
+  const int input_depth = MatchingDim(weights_shape, weights_dims_count - 1,
+                                      input_shape, input_dims_count - 1);
+  const int output_depth = MatchingDim(weights_shape, weights_dims_count - 2,
+                                       output_shape, output_dims_count - 1);
+  const int32_t input_offset = params.input_offset;
+  const int32_t output_offset = params.output_offset;
+  const int32_t output_multiplier = params.output_multiplier;
+  const int32_t output_shift = params.output_shift;
+  const int32_t output_activation_min = params.quantized_activation_min;
+  const int32_t output_activation_max = params.quantized_activation_max;
+
+  const int* w1_segments = sparsity.dim_metadata[1].array_segments->data;
+  const int* w1_indices = sparsity.dim_metadata[1].array_indices->data;
+
+  tensor_utils::SparseMatrixBatchVectorMultiplyAccumulate1x16(
+      weights_data, w1_segments, w1_indices, weights_shape.Dims(0),
+      weights_shape.Dims(1), input_data + thread_start * input_depth, bias_data,
+      batches, input_offset, output_multiplier, output_shift, per_channel_scale,
+      per_channel_shift, output_offset, output_activation_min,
+      output_activation_max, output_data + thread_start * output_depth);
+}
+
+inline void FullyConnectedSparseWeight1x4Impl(
+    const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
+    const RuntimeShape& input_shape, const float* input_data,
+    const RuntimeShape& weights_shape, const float* weights_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data, int thread_start,
+    int thread_end, const CpuBackendContext& cpu_backend_context) {
+  ruy::profiler::ScopeLabel label("FullyConnected");
+  ruy::profiler::ScopeLabel inner_label("1x4 Block Sparse");
+  const float output_activation_min = params.float_activation_min;
+  const float output_activation_max = params.float_activation_max;
+
+  const int input_dims_count = input_shape.DimensionsCount();
+  const int output_dims_count = output_shape.DimensionsCount();
+  const int weights_dims_count = weights_shape.DimensionsCount();
+  const int batches = thread_end - thread_start;
+  const int input_depth = MatchingDim(weights_shape, weights_dims_count - 1,
+                                      input_shape, input_dims_count - 1);
+  const int output_depth = MatchingDim(weights_shape, weights_dims_count - 2,
+                                       output_shape, output_dims_count - 1);
+  const int* w1_segments = sparsity.dim_metadata[1].array_segments->data;
+  const int* w1_indices = sparsity.dim_metadata[1].array_indices->data;
+
+  tensor_utils::SparseMatrixBatchVectorMultiplyAccumulate1x4(
+      weights_data, w1_segments, w1_indices, weights_shape.Dims(0),
+      weights_shape.Dims(1), input_data + thread_start * input_depth, batches,
+      output_data + thread_start * output_depth);
+
+  ruy::profiler::ScopeLabel activation_label("activation function");
+  for (int b = thread_start; b < thread_end; ++b) {
+    for (int i = 0; i < output_depth; ++i) {
+      float total = output_data[b * output_depth + i];
+      const float bias_value = bias_data ? bias_data[i] : 0;
+      output_data[b * output_depth + i] = ActivationFunctionWithMinMax(
+          total + bias_value, output_activation_min, output_activation_max);
+    }
+  }
+}
+
+struct FullyConnectedSparseWeight1x4Task : cpu_backend_threadpool::Task {
+  FullyConnectedSparseWeight1x4Task(
+      const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
+      const RuntimeShape& input_shape, const float* input_data,
+      const RuntimeShape& weights_shape, const float* weights_data,
+      const RuntimeShape& bias_shape, const float* bias_data,
+      const RuntimeShape& output_shape, float* output_data, int thread_start,
+      int thread_end, const CpuBackendContext& cpu_backend_context_x)
+      : sparsity(sparsity),
+        params(params),
+        input_shape(input_shape),
+        input_data(input_data),
+        weights_shape(weights_shape),
+        weights_data(weights_data),
+        bias_shape(bias_shape),
+        bias_data(bias_data),
+        output_shape(output_shape),
+        output_data(output_data),
+        thread_start(thread_start),
+        thread_end(thread_end),
+        cpu_backend_context(cpu_backend_context_x) {}
+
+  void Run() override {
+    FullyConnectedSparseWeight1x4Impl(
+        sparsity, params, input_shape, input_data, weights_shape, weights_data,
+        bias_shape, bias_data, output_shape, output_data, thread_start,
+        thread_end, cpu_backend_context);
+  }
+
+ private:
+  const TfLiteSparsity& sparsity;
+  const FullyConnectedParams& params;
+  const RuntimeShape& input_shape;
+  const float* input_data;
+  const RuntimeShape& weights_shape;
+  const float* weights_data;
+  const RuntimeShape& bias_shape;
+  const float* bias_data;
+  const RuntimeShape& output_shape;
+  float* output_data;
+  int thread_start;
+  int thread_end;
+  const CpuBackendContext& cpu_backend_context;
+};
+
+inline void FullyConnectedSparseWeight1x16(
+    const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
+    const RuntimeShape& input_shape, const int8_t* input_data,
+    const RuntimeShape& weights_shape, const int8_t* weights_data,
+    const int32_t* per_channel_scale, const int32_t* per_channel_shift,
+    const RuntimeShape& bias_shape, const int32_t* bias_data,
+    const RuntimeShape& output_shape, int8_t* output_data,
+    CpuBackendContext* cpu_backend_context) {
+  const int output_elements = output_shape.FlatSize();
+  memset(output_data, 0, output_elements * sizeof(int8_t));
+
+  const int batches =
+      FlatSizeSkipDim(output_shape, output_shape.DimensionsCount() - 1);
+
+  // TODO(b/220851507): Add multi-thread support for quantized sparse kernel.
+  return FullyConnectedSparseWeight1x16Impl(
+      sparsity, params, input_shape, input_data, weights_shape, weights_data,
+      per_channel_scale, per_channel_shift, bias_shape, bias_data, output_shape,
+      output_data, 0, batches, *cpu_backend_context);
+}
+
+// The multi-threaded kernel slices the workload along the batch dimension. If
+// there's not enough batches of data, the number of threads used is equal to
+// the batch size. We can improve this later with slicing along the row
+// dimension of the weight.
+inline void FullyConnectedSparseWeight1x4(
+    const TfLiteSparsity& sparsity, const FullyConnectedParams& params,
+    const RuntimeShape& input_shape, const float* input_data,
+    const RuntimeShape& weights_shape, const float* weights_data,
+    const RuntimeShape& bias_shape, const float* bias_data,
+    const RuntimeShape& output_shape, float* output_data,
+    CpuBackendContext* cpu_backend_context) {
+  const int output_elements = output_shape.FlatSize();
+  memset(output_data, 0, output_elements * sizeof(float));
+
+  const int max_threads = cpu_backend_context->max_num_threads();
+  const int batches =
+      FlatSizeSkipDim(output_shape, output_shape.DimensionsCount() - 1);
+  const int thread_count = std::max(1, std::min(batches, max_threads));
+  if (thread_count == 1) {
+    return FullyConnectedSparseWeight1x4Impl(
+        sparsity, params, input_shape, input_data, weights_shape, weights_data,
+        bias_shape, bias_data, output_shape, output_data, 0, batches,
+        *cpu_backend_context);
+  }
+  std::vector<FullyConnectedSparseWeight1x4Task> tasks;
+  tasks.reserve(thread_count);
+  int thread_start = 0;
+  for (int i = 0; i < thread_count; ++i) {
+    // This makes sure the workload is relatively balanced when batches is not a
+    // multiple of thread_count. The first mod(batches, thread_count) tasks need
+    // to process one more batch than the rest.
+    int thread_end = thread_start + batches / thread_count;
+    if (i < batches % thread_count) thread_end++;
+
+    tasks.emplace_back(sparsity, params, input_shape, input_data, weights_shape,
+                       weights_data, bias_shape, bias_data, output_shape,
+                       output_data, thread_start, thread_end,
+                       *cpu_backend_context);
+    thread_start = thread_end;
+  }
+  cpu_backend_threadpool::Execute(tasks.size(), tasks.data(),
+                                  cpu_backend_context);
+}
+
+}  // namespace optimized_ops
+}  // namespace tflite
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SPARSE_OPS_FULLY_CONNECTED_H_

tensorflow/lite/kernels/internal/optimized/sse_check.h

@@ -0,0 +1,28 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_CHECK_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_CHECK_H_
+
+#if defined(__SSSE3__)
+// SSSE 3 available: Use the SSE code.
+#define SSE_OR_PORTABLE(funcname, ...) Sse##funcname(__VA_ARGS__)
+
+#else
+
+// No SSSE 3 available: Use Portable code
+#define SSE_OR_PORTABLE(funcname, ...) Portable##funcname(__VA_ARGS__)
+#endif
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_CHECK_H_

tensorflow/lite/kernels/internal/optimized/sse_tensor_utils.cc

@@ -0,0 +1,676 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/optimized/sse_tensor_utils_impl.h"
+
+#ifdef __SSSE3__
+
+#include <emmintrin.h>  // SSE2
+#include <tmmintrin.h>  // SSSE3
+#ifdef __SSE4_1__
+#include <smmintrin.h>  // SSE4.1
+#endif
+#ifdef __AVX2__
+#include <immintrin.h>
+
+#include "absl/base/prefetch.h"
+#endif
+
+#include <cstdint>
+
+#include "ruy/profiler/instrumentation.h"  // from @ruy
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm.h"
+#include "tensorflow/lite/kernels/cpu_backend_gemm_params.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+
+namespace tflite {
+namespace tensor_utils {
+namespace {
+
+#if defined(__SSE2__)
+// Note: this part is copied from XNNPACK/src/xnnpack/intrinsics-polyfill.h
+// w.r.t the defition of '_mm_loadu_si32' intrinsic.
+// GCC any, Clang pre-8, Android NDK Clang pre-8.0.7, Apple Clang pre-11, and
+// ICC pre-16
+#if (defined(__GNUC__) && !defined(__clang__) &&                             \
+     !defined(__INTEL_COMPILER)) ||                                          \
+    (defined(__clang__) && !defined(__apple_build_version__) &&              \
+     (__clang_major__ < 8)) ||                                               \
+    (defined(__clang__) && defined(__ANDROID__) && (__clang_major__ == 8) && \
+     (__clang_minor__ == 0) && (__clang_patchlevel__ < 7)) ||                \
+    (defined(__clang__) && defined(__apple_build_version__) &&               \
+     (__apple_build_version__ < 11000000)) ||                                \
+    (defined(__INTEL_COMPILER) && (__INTEL_COMPILER < 1600))
+
+static inline __m128i _mm_loadu_si32(const void* address) {
+  return _mm_cvtsi32_si128(*((const int*)address));
+}
+#endif  // GCC any, Clang pre-8, Android NDK Clang pre-8.0.7, Apple Clang pre-11
+        // and ICC pre-16
+#endif  // __SSE2__
+
+// Dot product of four int8 vectors of 4 elements packed into a XMM register.
+// Result is four int32 scalars packed into a XMM register.
+// int8x4x4 · int8x4x4 => int32x4
+static inline __m128i DotProdInt8x4x4(__m128i a_8x16, __m128i b_8x16) {
+  // Transfer sign from 'a' to 'b', as _mm_maddubs_epi16 treats 'a' unsigned.
+  b_8x16 = _mm_sign_epi8(b_8x16, a_8x16);
+  a_8x16 = _mm_abs_epi8(a_8x16);
+  // sumprod[i] = a[2*i]*b[2*i] + a[2*i+1]*b[2*i+1] (i = 0..7)
+  __m128i sumprod_16x8 = _mm_maddubs_epi16(a_8x16, b_8x16);
+  // sumprod[i] = sumprod[2*i]*1 + sumprod[2*i+1]*1 (i = 0..3)
+  return _mm_madd_epi16(sumprod_16x8, _mm_set1_epi16(1));
+}
+
+// Horizontally add 4 int32 values stored in a single XMM register to int32_t.
+static inline int32_t ReduceInt32x4(__m128i acc) {
+  // Shuffle to contain high half of acc (both in high and low halfs).
+  __m128i shuffle = _mm_unpackhi_epi64(acc, acc);
+  // Add shuffle and acc; low half is sums of twos (high half is ignored).
+  acc = _mm_add_epi32(acc, shuffle);
+  // Shuffle the two elements in low half (ignore high half).
+  shuffle = _mm_shuffle_epi32(acc, _MM_SHUFFLE(2, 3, 0, 1));
+  // Add shuffle and acc; lowest element is sum of all 4 input.
+  acc = _mm_add_epi32(acc, shuffle);
+  // Return lowest element as int32_t.
+  return _mm_cvtsi128_si32(acc);
+}
+
+#ifdef __AVX2__
+// Horizontally add 4 float values stored in a single XMM register to float.
+static inline float ReduceFloat32x4(__m128 acc) {
+  __m128 shuffle = _mm_movehdup_ps(acc);
+  acc = _mm_add_ps(acc, shuffle);
+  shuffle = _mm_movehl_ps(shuffle, acc);
+  acc = _mm_add_ss(acc, shuffle);
+  return _mm_cvtss_f32(acc);
+}
+
+// Horizontally add 8 float values stored in a single XMM register to float.
+static inline float ReduceFloat32x8(__m256 acc) {
+  __m128 low = _mm256_extractf128_ps(acc, 0);
+  __m128 high = _mm256_extractf128_ps(acc, 1);
+  return ReduceFloat32x4(_mm_add_ps(low, high));
+}
+
+// Dot product of four int8 vectors of 4 elements packed into a YMM register.
+// Result is eight int32 scalars packed into a YMM register.
+// int8x4x8 · int8x4x8 => int32x8
+static inline __m256i DotProdInt8x4x8(__m256i a_16x16, __m256i b_16x16) {
+  // Transfer sign from 'a' to 'b', as _mm256_maddubs_epi16 treats 'a' unsigned.
+  b_16x16 = _mm256_sign_epi8(b_16x16, a_16x16);
+  a_16x16 = _mm256_abs_epi8(a_16x16);
+  // sumprod[i] = a[2*i]*b[2*i] + a[2*i+1]*b[2*i+1] (i = 0..15)
+  __m256i sumprod_16x16 = _mm256_maddubs_epi16(a_16x16, b_16x16);
+  // sumprod[i] = sumprod[2*i]*1 + sumprod[2*i+1]*1 (i = 0..7)
+  return _mm256_madd_epi16(sumprod_16x16, _mm256_set1_epi16(1));
+}
+#endif  // __AVX2__
+
+// Horizontally add each of 4 XMM registers with 4 int32 values, pack result
+// into a single XMM register. Similar to ReduceInt32x4, but with 4x inputs.
+static inline __m128i ReduceInt32x4x4(__m128i a, __m128i b, __m128i c,
+                                      __m128i d) {
+  // Assuming x = [x0, x1, x2, x3]
+  const __m128i a_b_lo_half = _mm_unpacklo_epi32(a, b);  // [a0, b0, a1, b1]
+  const __m128i a_b_hi_half = _mm_unpackhi_epi32(a, b);  // [a2, b2, a3, b3]
+  const __m128i a_plus_b =
+      _mm_add_epi32(a_b_lo_half, a_b_hi_half);  // [a0+a2, b0+b2, a1+a3, b1+b3]
+  const __m128i c_d_lo_half = _mm_unpacklo_epi32(c, d);  // [c0, d0, c1, d1]
+  const __m128i c_d_hi_half = _mm_unpackhi_epi32(c, d);  // [c2, d2, c3, d3]
+  const __m128i c_plus_d =
+      _mm_add_epi32(c_d_lo_half, c_d_hi_half);  // [c0+c2, d0+d2, c1+c3, d1+d3]
+  const __m128i all_evns =
+      _mm_unpacklo_epi64(a_plus_b, c_plus_d);  // [a02, b02, c02, d02]
+  const __m128i all_odds =
+      _mm_unpackhi_epi64(a_plus_b, c_plus_d);  // [a13, b13, c13, d13]
+  return _mm_add_epi32(all_evns, all_odds);    // [a0123, b0123, c0123, d0123]
+}
+
+// Returns the ith element of a XMM register holding float numbers.
+template <int i>
+float GetFloatVectorElement(__m128 v) {
+  static_assert(i >= 0 && i < 4, "The index must be 0 <= i < 4.");
+  // Note, _mm_extract_ps returns int, so we can't use it here.
+  // These lines will be optimized to extractps anyway.
+  v = _mm_shuffle_ps(v, v, _MM_SHUFFLE(i, i, i, i));
+  return _mm_cvtss_f32(v);
+}
+
+}  // namespace
+
+#ifdef __AVX2__
+constexpr int kFloatValuesPerAvx2Vector = 8;
+template <int PerVectorSize>
+inline int RoundDownVectors(int size) {
+  return size & ~(PerVectorSize - 1);
+}
+
+void Avx2MatrixBatchVectorMultiplyAccumulateImpl(
+    const float* __restrict__ matrix, int m_rows, int m_cols,
+    const float* __restrict__ vector, int n_batch, float* __restrict__ result) {
+  // If v_size is not divisible by the vector size, then we need to process the
+  // final few elements sequentially. postamble_start shows the start index
+  // where this should happen.
+  const int postamble_start =
+      RoundDownVectors<kFloatValuesPerAvx2Vector>(m_cols);
+
+  for (int b = 0; b < n_batch; ++b) {
+    float* result_in_batch = result + b * m_rows;
+    const float* vector_in_batch = vector + b * m_cols;
+    const float* matrix_row = matrix;
+
+    // Main matrix by vector multiplication loop
+    for (int r = 0; r < m_rows; ++r) {
+      __m256 acc_32x8 = _mm256_setzero_ps();
+      int c = 0;
+      for (; c < postamble_start; c += kFloatValuesPerAvx2Vector) {
+        // Load 8 float values from vector and matrix row.
+        __m256 vector_f32x8 = _mm256_loadu_ps(vector_in_batch + c);
+        __m256 matrix_f32x8 = _mm256_loadu_ps(matrix_row + c);
+
+        // Multiply the vector and matrix row and add to accumulator.
+        __m256 res = _mm256_mul_ps(vector_f32x8, matrix_f32x8);
+        acc_32x8 = _mm256_add_ps(acc_32x8, res);
+      }
+      // Add the 8 intermediate sum values to get the final dot-prod value for
+      // this column.
+      float sum = ReduceFloat32x8(acc_32x8);
+      for (; (c < m_cols); c++) {
+        sum += matrix_row[c] * vector_in_batch[c];
+      }
+      *result_in_batch += sum;
+      ++result_in_batch;
+      matrix_row += m_cols;
+    }
+  }
+}
+
+void Avx2MatrixBatchVectorMultiplyAccumulateImpl(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, const int32_t* row_sums) {
+  for (std::intptr_t batch = 0; batch < n_batch; ++batch) {
+    const float batch_scaling_factor = scaling_factors[batch];
+    const int32_t batch_offset = input_offset ? input_offset[batch] : 0;
+    // Compute dot-product for every column.
+    for (std::intptr_t row = 0; row < m_rows; ++row) {
+      // Get the address of the first element of the row.
+      const int8_t* __restrict__ row_ptr = matrix + row * m_cols;
+      const float row_scale =
+          per_channel_scale ? per_channel_scale[row] * batch_scaling_factor
+                            : batch_scaling_factor;
+      const int32_t row_offset =
+          row_sums && batch_offset ? batch_offset * row_sums[row] : 0;
+      // Initialize the dot product sum for the row to 0.
+      __m256i dotprod_32x8 = _mm256_setzero_si256();
+      std::intptr_t col = 0;
+      constexpr int prefetch_distance = 704;
+      // For every block of 32x 8-bit inputs.
+      while (col < (m_cols & ~31)) {
+        absl::PrefetchToLocalCache(vectors + col + prefetch_distance);
+        absl::PrefetchToLocalCache(row_ptr + col + prefetch_distance);
+        const __m256i vec_16x16 =
+            _mm256_loadu_si256(reinterpret_cast<const __m256i*>(vectors + col));
+        const __m256i row_16x16 =
+            _mm256_loadu_si256(reinterpret_cast<const __m256i*>(row_ptr + col));
+        // dotprod += vec · row
+        dotprod_32x8 = _mm256_add_epi32(dotprod_32x8,
+                                        DotProdInt8x4x8(vec_16x16, row_16x16));
+        col += 32;
+      }
+      // Sum lower and upper halves of 32x8 vector into 32x4 vector
+      __m128i low = _mm256_extracti128_si256(dotprod_32x8, 0);
+      __m128i high = _mm256_extracti128_si256(dotprod_32x8, 1);
+      __m128i dotprod_32x4 = _mm_add_epi32(low, high);
+      // Postamble for 16x 8-bit inputs.
+      if (col < (m_cols & ~15)) {
+        const __m128i vec_16x8 =
+            _mm_loadu_si128(reinterpret_cast<const __m128i*>(vectors + col));
+        const __m128i row_16x8 =
+            _mm_loadu_si128(reinterpret_cast<const __m128i*>(row_ptr + col));
+        // dotprod += vec · row
+        dotprod_32x4 =
+            _mm_add_epi32(dotprod_32x4, DotProdInt8x4x4(vec_16x8, row_16x8));
+        col += 16;
+      }
+      // Postamble for 8x 8-bit inputs.
+      if (col < (m_cols & ~7)) {
+        const __m128i vec_16x8 = _mm_cvtepi8_epi16(
+            _mm_loadl_epi64(reinterpret_cast<const __m128i*>(vectors + col)));
+        const __m128i row_16x8 = _mm_cvtepi8_epi16(
+            _mm_loadl_epi64(reinterpret_cast<const __m128i*>(row_ptr + col)));
+        // dotprod += vec · row
+        dotprod_32x4 =
+            _mm_add_epi32(dotprod_32x4, _mm_madd_epi16(vec_16x8, row_16x8));
+        col += 8;
+      }
+      // Postamble for 4x 8-bit inputs.
+      if (col < (m_cols & ~3)) {
+        const __m128i vec_32x4 = _mm_cvtepi8_epi32(
+            _mm_loadu_si32(reinterpret_cast<const __m128i*>(vectors + col)));
+        const __m128i row_32x4 = _mm_cvtepi8_epi32(
+            _mm_loadu_si32(reinterpret_cast<const __m128i*>(row_ptr + col)));
+        // dotprod += vec · row
+        dotprod_32x4 =
+            _mm_add_epi32(dotprod_32x4, _mm_mullo_epi32(vec_32x4, row_32x4));
+        col += 4;
+      }
+
+      // Horizontally add the 4 intermediate sum values to get the final
+      // dot-prod value for this row.
+      int32_t sum = ReduceInt32x4(dotprod_32x4);
+
+#pragma clang loop unroll(disable) vectorize(disable)
+      // Postamble loop for <4x remaining 8-bit inputs.
+      for (; col < m_cols; ++col) {
+        sum += row_ptr[col] * vectors[col];
+      }  // for col
+      if (row_offset) {
+        sum -= row_offset;
+      }
+      *result += sum * row_scale;
+      ++result;
+    }  // for row
+
+    vectors += m_cols;
+  }  // for batch
+}
+
+#endif  // __AVX2__
+
+void SseMatrixBatchVectorMultiplyAccumulateImpl(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, const int32_t* row_sums) {
+#ifdef __AVX2__
+  Avx2MatrixBatchVectorMultiplyAccumulateImpl(
+      matrix, m_rows, m_cols, vectors, scaling_factors, n_batch, result,
+      per_channel_scale, input_offset, row_sums);
+  return;
+#else
+  for (std::intptr_t batch = 0; batch < n_batch; ++batch) {
+    const float batch_scaling_factor = scaling_factors[batch];
+    const int32_t batch_offset = input_offset ? input_offset[batch] : 0;
+    // Compute dot-product for every column.
+    for (std::intptr_t row = 0; row < m_rows; ++row) {
+      // Get the address of the first element of the row.
+      const int8_t* __restrict__ row_ptr = matrix + row * m_cols;
+      const float row_scale =
+          per_channel_scale ? per_channel_scale[row] * batch_scaling_factor
+                            : batch_scaling_factor;
+      const int32_t row_offset =
+          row_sums && batch_offset ? batch_offset * row_sums[row] : 0;
+      // Initialize the dot product sum for the row to 0.
+      __m128i dotprod_32x4 = _mm_setzero_si128();
+      std::intptr_t col = 0;
+      // For every block of 16x 8-bit inputs.
+      while (col < (m_cols & ~15)) {
+        const __m128i vec_8x16 =
+            _mm_loadu_si128(reinterpret_cast<const __m128i*>(vectors + col));
+        const __m128i row_8x16 =
+            _mm_loadu_si128(reinterpret_cast<const __m128i*>(row_ptr + col));
+        // dotprod += vec · row
+        dotprod_32x4 =
+            _mm_add_epi32(dotprod_32x4, DotProdInt8x4x4(vec_8x16, row_8x16));
+        col += 16;
+      }
+#ifdef __SSE4_1__
+      // Postamble for 8x 8-bit inputs.
+      if (col < (m_cols & ~7)) {
+        const __m128i vec_16x8 = _mm_cvtepi8_epi16(
+            _mm_loadl_epi64(reinterpret_cast<const __m128i*>(vectors + col)));
+        const __m128i row_16x8 = _mm_cvtepi8_epi16(
+            _mm_loadl_epi64(reinterpret_cast<const __m128i*>(row_ptr + col)));
+        // dotprod += vec · row
+        dotprod_32x4 =
+            _mm_add_epi32(dotprod_32x4, _mm_madd_epi16(vec_16x8, row_16x8));
+        col += 8;
+      }
+      // Postamble for 4x 8-bit inputs.
+      if (col < (m_cols & ~3)) {
+        const __m128i vec_32x4 = _mm_cvtepi8_epi32(
+            _mm_loadu_si32(reinterpret_cast<const __m128i*>(vectors + col)));
+        const __m128i row_32x4 = _mm_cvtepi8_epi32(
+            _mm_loadu_si32(reinterpret_cast<const __m128i*>(row_ptr + col)));
+        // dotprod += vec · row
+        dotprod_32x4 =
+            _mm_add_epi32(dotprod_32x4, _mm_mullo_epi32(vec_32x4, row_32x4));
+        col += 4;
+      }
+#endif
+
+      // Horizontally add the 4 intermediate sum values to get the final
+      // dot-prod value for this row.
+      int32_t sum = ReduceInt32x4(dotprod_32x4);
+
+#if defined(__SSE4_1__) && defined(__clang__)
+      // SSE 4.1: Don't try to unroll and vectorize this, already done above.
+#pragma clang loop unroll(disable) vectorize(disable)
+#endif
+      // Postamble loop for <4x (<16x without SSE 4.1) remaining 8-bit inputs.
+      for (; col < m_cols; ++col) {
+        sum += row_ptr[col] * vectors[col];
+      }  // for col
+      if (row_offset) {
+        sum -= row_offset;
+      }
+      *result += sum * row_scale;
+      ++result;
+    }  // for row
+
+    vectors += m_cols;
+  }  // for batch
+#endif  // ifdef __AVX2__
+}
+
+void SseCpuBackendGemm(const int8_t* input, const int32_t* bias,
+                       const int8_t* input_to_gate_weights, int32_t n_batch,
+                       int32_t n_input, int32_t n_output, int32_t output_zp,
+                       int32_t* scratch, CpuBackendContext* context) {
+  using ::tflite::cpu_backend_gemm::Gemm;
+  using ::tflite::cpu_backend_gemm::GemmParams;
+  using ::tflite::cpu_backend_gemm::MatrixParams;
+
+  MatrixParams<int8_t> lhs_params;
+  lhs_params.order = cpu_backend_gemm::Order::kRowMajor;
+  lhs_params.rows = n_output;
+  lhs_params.cols = n_input;
+  lhs_params.cache_policy = cpu_backend_gemm::CachePolicy::kCacheIfLargeSpeedup;
+
+  MatrixParams<int8_t> rhs_params;
+  rhs_params.order = cpu_backend_gemm::Order::kColMajor;
+  rhs_params.rows = n_input;
+  rhs_params.cols = n_batch;
+
+  MatrixParams<int32_t> dst_params;
+  dst_params.order = cpu_backend_gemm::Order::kColMajor;
+  dst_params.rows = n_output;
+  dst_params.cols = n_batch;
+
+  GemmParams<int32, int32> gemm_params;
+  if (bias) {
+    gemm_params.bias = bias;
+  }
+  cpu_backend_gemm::Gemm(lhs_params, input_to_gate_weights, rhs_params, input,
+                         dst_params, scratch, gemm_params, context);
+}
+
+void SseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result) {
+  SseMatrixBatchVectorMultiplyAccumulateImpl(
+      matrix, m_rows, m_cols, vectors, scaling_factors, n_batch, result,
+      /*per_channel_scale=*/nullptr, /*input_offset=*/nullptr,
+      /*row_sums=*/nullptr);
+}
+
+void SseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch, int32_t* scratch,
+    float* __restrict__ result, CpuBackendContext* context) {
+  // TODO(b/183178387): Use a proper query to detect AVX/optimized paths.
+  if (m_rows % 4 == 0 && !context->PreferGemmlowpOnX86()) {
+    const int32_t* bias = static_cast<const int32_t*>(nullptr);
+    SseCpuBackendGemm(vectors, bias, matrix, n_batch, m_cols, m_rows,
+                      /*output_zp=*/0, scratch, context);
+
+    {
+      ruy::profiler::ScopeLabel label("HybridMultiplyScalingFactor");
+      // Multiply by float scaling factors and write to result
+      const int total_size = n_batch * m_rows;
+      int i = 0;
+      for (; i <= total_size - 8; i += 8, result += 8) {
+        const float batch_scaling_factor0 = scaling_factors[i / m_rows];
+        const float batch_scaling_factor1 = scaling_factors[(i + 4) / m_rows];
+        const __m128 scaling_factor0 = _mm_set1_ps(batch_scaling_factor0);
+        const __m128 scaling_factor1 = _mm_set1_ps(batch_scaling_factor1);
+        const __m128i scratch_val0 =
+            _mm_loadu_si128(reinterpret_cast<const __m128i*>(scratch + i));
+        const __m128i scratch_val1 =
+            _mm_loadu_si128(reinterpret_cast<const __m128i*>(scratch + i + 4));
+        const __m128 float_val0 = _mm_cvtepi32_ps(scratch_val0);
+        const __m128 float_val1 = _mm_cvtepi32_ps(scratch_val1);
+        const __m128 prod0 = _mm_mul_ps(float_val0, scaling_factor0);
+        const __m128 result0 = _mm_add_ps(_mm_load1_ps(result), prod0);
+        const __m128 prod1 = _mm_mul_ps(float_val1, scaling_factor1);
+        const __m128 result1 = _mm_add_ps(_mm_load1_ps(result + 4), prod1);
+        _mm_store_ps(result, result0);
+        _mm_store_ps(result + 4, result1);
+      }
+      scratch += i;
+      for (; i < total_size; i++) {
+        const float batch_scaling_factor = scaling_factors[i / m_rows];
+        int32_t x = *(scratch++);
+        *result += x * batch_scaling_factor;
+        ++result;
+      }
+    }
+    return;
+  }
+
+  SseMatrixBatchVectorMultiplyAccumulateImpl(
+      matrix, m_rows, m_cols, vectors, scaling_factors, n_batch, result,
+      /*per_channel_scale=*/nullptr, /*input_offset=*/nullptr,
+      /*row_sums=*/nullptr);
+}
+
+void SseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
+    bool* compute_row_sums, CpuBackendContext* context) {
+  if ((input_offset != nullptr) && (!compute_row_sums || *compute_row_sums)) {
+    SseReductionSumVector(matrix, row_sums, m_rows, m_cols);
+    if (compute_row_sums) {
+      *compute_row_sums = false;
+    }
+  }
+  SseMatrixBatchVectorMultiplyAccumulateImpl(
+      matrix, m_rows, m_cols, vectors, scaling_factors, n_batch, result,
+      per_channel_scale, input_offset, row_sums);
+}
+
+namespace {
+
+// Implements sparse-matrix - vector multiply-accumulate.
+inline void SseSparseMatrixVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    const int m_rows, const int m_cols, const int8_t* __restrict__ vector,
+    const float batch_scaling_factor, float* __restrict__ result,
+    const float* per_channel_scale) {
+  static const std::intptr_t kBlockSize = 16;
+  TFLITE_DCHECK_EQ(m_cols % kBlockSize, 0);
+  const uint8_t* __restrict__ ledger_ptr = ledger;
+  for (std::intptr_t row = 0; row < m_rows; ++row) {
+    // Initialize the dot product sum for the row to 0.
+    __m128i dotprod_32x4 = _mm_setzero_si128();
+    std::intptr_t num_nonzero_blocks = *ledger_ptr++;
+    for (std::intptr_t i = 0; i < num_nonzero_blocks; i++) {
+      const std::intptr_t col_index = *ledger_ptr++ * kBlockSize;
+      const __m128i vec_8x16 =
+          _mm_loadu_si128(reinterpret_cast<const __m128i*>(vector + col_index));
+      const __m128i row_8x16 =
+          _mm_loadu_si128(reinterpret_cast<const __m128i*>(matrix));
+      // dotprod += vec · row
+      dotprod_32x4 =
+          _mm_add_epi32(dotprod_32x4, DotProdInt8x4x4(vec_8x16, row_8x16));
+      matrix += kBlockSize;
+    }  // for col
+    // Horizontally add the 4 intermediate sum values to get the final
+    // dot-prod value for this row.
+    int32_t dotprod = ReduceInt32x4(dotprod_32x4);
+    const float total_scaling_factor =
+        per_channel_scale ? per_channel_scale[row] * batch_scaling_factor
+                          : batch_scaling_factor;
+    result[row] += dotprod * total_scaling_factor;
+  }  // for row
+}
+
+// Implements sparse-matrix - batch-of-4-vectors multiply-accumulate.
+// The stride between vectors and results must be equal to m_cols.
+// Parameter 'batch' is the index of the first batch, must be a multiple of 4.
+inline void SseSparseMatrix4VectorsMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    const int m_rows, const int m_cols,
+    const int8_t* __restrict__ const vectors,
+    const __m128 batch_scaling_factors_fx4, float* __restrict__ const results,
+    const float* per_channel_scale) {
+  static const std::intptr_t kBlockSize = 16;
+  TFLITE_DCHECK_EQ(m_cols % kBlockSize, 0);
+
+  const int8_t* __restrict__ vector0 = vectors + 0 * m_cols;
+  const int8_t* __restrict__ vector1 = vectors + 1 * m_cols;
+  const int8_t* __restrict__ vector2 = vectors + 2 * m_cols;
+  const int8_t* __restrict__ vector3 = vectors + 3 * m_cols;
+  float* __restrict__ result0 = results + 0 * m_rows;
+  float* __restrict__ result1 = results + 1 * m_rows;
+  float* __restrict__ result2 = results + 2 * m_rows;
+  float* __restrict__ result3 = results + 3 * m_rows;
+
+  for (std::intptr_t row = 0; row < m_rows; ++row) {
+    // Initialize the dot product sum for the row to 0.
+    __m128i dp0_32x4 = _mm_setzero_si128();
+    __m128i dp1_32x4 = _mm_setzero_si128();
+    __m128i dp2_32x4 = _mm_setzero_si128();
+    __m128i dp3_32x4 = _mm_setzero_si128();
+
+    std::intptr_t num_nonzero_blocks = *ledger++;
+    for (std::intptr_t i = 0; i < num_nonzero_blocks; i++) {
+      const std::intptr_t col_index = *ledger++ * kBlockSize;
+      // vecN are for different batches
+      const __m128i vec0_8x16 = _mm_loadu_si128(
+          reinterpret_cast<const __m128i*>(vector0 + col_index));
+      const __m128i vec1_8x16 = _mm_loadu_si128(
+          reinterpret_cast<const __m128i*>(vector1 + col_index));
+      const __m128i vec2_8x16 = _mm_loadu_si128(
+          reinterpret_cast<const __m128i*>(vector2 + col_index));
+      const __m128i vec3_8x16 = _mm_loadu_si128(
+          reinterpret_cast<const __m128i*>(vector3 + col_index));
+      const __m128i row_8x16 =
+          _mm_loadu_si128(reinterpret_cast<const __m128i*>(matrix));
+      // dp += vec · row
+      // dpN are for different batches
+      dp0_32x4 = _mm_add_epi32(dp0_32x4, DotProdInt8x4x4(vec0_8x16, row_8x16));
+      dp1_32x4 = _mm_add_epi32(dp1_32x4, DotProdInt8x4x4(vec1_8x16, row_8x16));
+      dp2_32x4 = _mm_add_epi32(dp2_32x4, DotProdInt8x4x4(vec2_8x16, row_8x16));
+      dp3_32x4 = _mm_add_epi32(dp3_32x4, DotProdInt8x4x4(vec3_8x16, row_8x16));
+      matrix += kBlockSize;
+    }  // for col
+
+    // Horizontally add the 4 intermediate values.
+    const __m128i dp_32x4 =
+        ReduceInt32x4x4(dp0_32x4, dp1_32x4, dp2_32x4, dp3_32x4);
+    // Convert to float
+    const __m128 dp_fx4 = _mm_cvtepi32_ps(dp_32x4);
+    // Load the results (This is an Accumulate function..)
+    __m128 result_fx4 =
+        _mm_set_ps(result3[row], result2[row], result1[row], result0[row]);
+
+    const __m128 total_scaling_factors_fx4 =
+        per_channel_scale ? _mm_mul_ps(batch_scaling_factors_fx4,
+                                       _mm_set1_ps(per_channel_scale[row]))
+                          : batch_scaling_factors_fx4;
+    // result += dp .* scaling
+    result_fx4 =
+        _mm_add_ps(result_fx4, _mm_mul_ps(dp_fx4, total_scaling_factors_fx4));
+    // Save the results
+    result0[row] = GetFloatVectorElement<0>(result_fx4);
+    result1[row] = GetFloatVectorElement<1>(result_fx4);
+    result2[row] = GetFloatVectorElement<2>(result_fx4);
+    result3[row] = GetFloatVectorElement<3>(result_fx4);
+  }  // for row
+}
+
+}  // namespace
+
+void SseSparseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    const int m_rows, const int m_cols, const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ results, const float* per_channel_scale) {
+  int batch = 0;
+  const int kBatchSize4 = 4;
+  const int n_batch_rounddown_to_batchsize_4 = n_batch & ~(kBatchSize4 - 1);
+  while (batch < n_batch_rounddown_to_batchsize_4) {
+    const __m128 scaling_factors_fx4 = _mm_loadu_ps(scaling_factors + batch);
+    SseSparseMatrix4VectorsMultiplyAccumulate(matrix, ledger, m_rows, m_cols,
+                                              vectors, scaling_factors_fx4,
+                                              results, per_channel_scale);
+    batch += kBatchSize4;
+    vectors += kBatchSize4 * m_cols;
+    results += kBatchSize4 * m_rows;
+  }  // for batch
+  while (batch < n_batch) {
+    SseSparseMatrixVectorMultiplyAccumulate(matrix, ledger, m_rows, m_cols,
+                                            vectors, scaling_factors[batch],
+                                            results, per_channel_scale);
+    ++batch;
+    vectors += m_cols;
+    results += m_rows;
+  }  // for batch
+}
+
+void SseReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
+                           const int output_size, const int reduction_size) {
+  static constexpr std::intptr_t kBlockSize = 16;
+  for (std::intptr_t row = 0; row < output_size; ++row) {
+    const int8_t* __restrict__ row_ptr = input_vector + row * reduction_size;
+    __m128i row_sum_16x8 = _mm_setzero_si128();
+    std::intptr_t col = 0;
+    for (; col < (reduction_size & ~(kBlockSize - 1)); col += kBlockSize) {
+      const __m128i row_8x16 =
+          _mm_loadu_si128(reinterpret_cast<const __m128i*>(row_ptr + col));
+      const __m128i row_16x8 = _mm_maddubs_epi16(_mm_set1_epi8(1), row_8x16);
+      row_sum_16x8 = _mm_add_epi16(row_sum_16x8, row_16x8);
+    }  // for col
+#ifdef __SSE4_1__
+    // Postamble for 8x 8-bit inputs.
+    if (col < (reduction_size & ~7)) {
+      // _mm_loadu_si64 not supported in gcc versions < 9, breaks kokoro build.
+      const __m128i row_16x8 = _mm_cvtepi8_epi16(
+          _mm_loadl_epi64(reinterpret_cast<const __m128i*>(row_ptr + col)));
+      // dotprod += vec · row
+      row_sum_16x8 = _mm_add_epi16(row_sum_16x8, row_16x8);
+      col += 8;
+    }
+#endif
+    const __m128i row_sum_32x4 =
+        _mm_madd_epi16(row_sum_16x8, _mm_set1_epi16(1));
+    int32_t row_sum = ReduceInt32x4(row_sum_32x4);
+#if defined(__SSE4_1__) && defined(__clang__)
+    // SSE 4.1: Don't try to unroll and vectorize this, already done above.
+#pragma clang loop unroll(disable) vectorize(disable)
+#endif
+    for (; col < reduction_size; col++) {
+      row_sum += row_ptr[col];
+    }
+    output_vector[row] = row_sum;
+  }
+}
+
+}  // namespace tensor_utils
+}  // namespace tflite
+
+#endif  // __SSSE3__

tensorflow/lite/kernels/internal/optimized/sse_tensor_utils.h

@@ -0,0 +1,347 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_H_
+
+// Note: This file is a copy-paste version of neon_tensor_utils.h, only
+// difference is in MatrixBatchVectorMultiplyAccumulate and
+// SparseMatrixBatchVectorMultiplyAccumulate (other functions do not have SSE
+// implementation yet).
+
+// Note: Most of the functions below use NEON_OR_PORTABLE, through the Intel
+// NEON_2_SSE translator library. If a native SSE version of a function is
+// implemented, replace the appropriate one to SSE_OR_PORTABLE.
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/neon_tensor_utils_impl.h"
+#include "tensorflow/lite/kernels/internal/optimized/sse_check.h"
+#include "tensorflow/lite/kernels/internal/optimized/sse_tensor_utils_impl.h"
+#include "tensorflow/lite/kernels/internal/reference/portable_tensor_utils_impl.h"
+
+namespace tflite {
+namespace tensor_utils {
+
+void MatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
+                                         int m_cols, const float* vector,
+                                         int n_batch, float* result) {
+#if defined(__AVX2__)
+  Avx2MatrixBatchVectorMultiplyAccumulateImpl(matrix, m_rows, m_cols, vector,
+                                              n_batch, result);
+#else
+  NEON_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                   vector, n_batch, result);
+#endif
+}
+
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result) {
+  SSE_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                  vectors, scaling_factors, n_batch, result);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
+    bool* compute_row_sums, CpuBackendContext* context) {
+  SSE_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                  vectors, scaling_factors, n_batch, result, per_channel_scale,
+                  input_offset, scratch, row_sums, compute_row_sums, context);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    int32_t* __restrict__ scratch, float* __restrict__ result,
+    CpuBackendContext* __restrict__ context) {
+  SSE_OR_PORTABLE(MatrixBatchVectorMultiplyAccumulate, matrix, m_rows, m_cols,
+                  vectors, scaling_factors, n_batch, scratch, result, context);
+}
+
+void SparseMatrixBatchVectorMultiplyAccumulate1x4(
+    const float* __restrict__ matrix, const int32_t* __restrict__ segments,
+    const int32_t* __restrict__ indices, int m_rows, int m_cols,
+    const float* __restrict__ vector, int n_batch, float* __restrict__ result) {
+  NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate1x4, matrix,
+                   segments, indices, m_rows, m_cols, vector, n_batch, result);
+}
+
+void SparseMatrixBatchVectorMultiplyAccumulate1x16(
+    const int8_t* __restrict__ matrix, const int32_t* __restrict__ segments,
+    const int32_t* __restrict__ indices, int m_rows, int m_cols,
+    const int8_t* __restrict__ vector, const int32_t* __restrict__ bias_vector,
+    int n_batch, const int32_t input_offset, const int32_t output_multiplier,
+    const int32_t output_shift, const int32_t* per_channel_scale,
+    const int32_t* per_channel_shift, const int32_t output_offset,
+    const int32_t output_activation_min, const int32_t output_activation_max,
+    int8_t* __restrict__ result) {
+  NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate1x16, matrix,
+                   segments, indices, m_rows, m_cols, vector, bias_vector,
+                   n_batch, input_offset, output_multiplier, output_shift,
+                   per_channel_scale, per_channel_shift, output_offset,
+                   output_activation_min, output_activation_max, result);
+}
+
+void SparseMatrixBatchVectorMultiplyAccumulate(
+    const float* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    int m_rows, int m_cols, const float* __restrict__ vector, int n_batch,
+    float* __restrict__ result) {
+  NEON_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate, matrix, ledger,
+                   m_rows, m_cols, vector, n_batch, result);
+}
+
+void SparseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    const int m_rows, const int m_cols, const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale) {
+  SSE_OR_PORTABLE(SparseMatrixBatchVectorMultiplyAccumulate, matrix, ledger,
+                  m_rows, m_cols, vectors, scaling_factors, n_batch, result,
+                  per_channel_scale);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* input, const int32_t* input_zeropoint_times_weights,
+    const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
+    int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
+    int32_t* scratch, int16_t* output, CpuBackendContext* context) {
+  PortableMatrixBatchVectorMultiplyAccumulate(
+      input, input_zeropoint_times_weights, input_to_gate_weights, multiplier,
+      shift, n_batch, n_input, n_output, output_zp, scratch, output, context);
+}
+
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* input, const int32_t* input_zeropoint_times_weights,
+    const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
+    int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
+    int32_t* scratch, int8_t* output, CpuBackendContext* context) {
+  PortableMatrixBatchVectorMultiplyAccumulate(
+      input, input_zeropoint_times_weights, input_to_gate_weights, multiplier,
+      shift, n_batch, n_input, n_output, output_zp, scratch, output, context);
+}
+
+void MatrixBatchVectorMultiply(const int8_t* input, int32_t input_zeropoint,
+                               const int8_t* input_to_gate_weights,
+                               int32_t input_to_gate_effective_scale_a,
+                               int32_t input_to_gate_effective_scale_b,
+                               int32_t n_batch, int32_t n_input, int32_t n_cell,
+                               int8_t* gate_output, int8_t gate_output_zp) {
+  PortableMatrixBatchVectorMultiply(
+      input, input_zeropoint, input_to_gate_weights,
+      input_to_gate_effective_scale_a, input_to_gate_effective_scale_b, n_batch,
+      n_input, n_cell, gate_output, gate_output_zp);
+}
+
+void MatrixBatchVectorMultiply(const int16_t* hidden,
+                               const int8_t* hidden_to_output_weights,
+                               int32_t proj_effective_scale_a,
+                               int32_t proj_effective_scale_b,
+                               const int32_t* gate_bias, int32_t n_batch,
+                               int32_t n_hidden, int32_t n_output,
+                               int32_t output_zp, int8_t* proj_output) {
+  PortableMatrixBatchVectorMultiply(hidden, hidden_to_output_weights,
+                                    proj_effective_scale_a,
+                                    proj_effective_scale_b, gate_bias, n_batch,
+                                    n_hidden, n_output, output_zp, proj_output);
+}
+
+void MatrixScalarMultiplyAccumulate(const int8_t* matrix, int32_t scalar,
+                                    int32_t n_row, int32_t n_col,
+                                    int32_t* output) {
+  PortableMatrixScalarMultiplyAccumulate(matrix, scalar, n_row, n_col, output);
+}
+
+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) {
+  PortableApplyLayerNorm(input, layer_norm_weights, bias, layer_norm_scale_a,
+                         layer_norm_scale_b, variance_limit, n_batch, n_input,
+                         output);
+}
+
+void ApplyLayerNormFloat(const int16_t* input,
+                         const int16_t* layer_norm_weights,
+                         int32_t layer_norm_scale_a, int32_t layer_norm_scale_b,
+                         const int32_t* bias, int n_batch, int n_input,
+                         int16_t* output) {
+  PortableApplyLayerNormFloat(input, layer_norm_weights, layer_norm_scale_a,
+                              layer_norm_scale_b, bias, n_batch, n_input,
+                              output);
+}
+
+void ApplySigmoid(const int16_t* input, int32_t n_batch, int32_t n_input,
+                  int16_t* output) {
+  PortableApplySigmoid(input, n_batch, n_input, output);
+}
+
+void ApplySigmoidFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
+                       int16_t* output) {
+  PortableApplySigmoidFloat(input, n_batch, n_input, output);
+}
+
+void ApplyTanh(int32_t intger_bits, const int16_t* input, int32_t n_batch,
+               int32_t n_input, int16_t* output) {
+  PortableApplyTanh(intger_bits, input, n_batch, n_input, output);
+}
+
+void ApplyTanhFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
+                    int32_t integer_bits, int16_t* output) {
+  PortableApplyTanhFloat(input, n_batch, n_input, integer_bits, output);
+}
+
+void CwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch,
+              int n_input, int shift, int16_t* output) {
+  PortableCwiseMul(input_1, input_2, n_batch, n_input, shift, output);
+}
+
+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) {
+  PortableCwiseMul(input_1, input_2, multiplier, shift, n_batch, n_input,
+                   output_zp, output);
+}
+
+void CwiseAdd(const int16_t* input_1, const int16_t* input_2, int n_batch,
+              int n_input, int16_t* output) {
+  PortableCwiseAdd(input_1, input_2, n_batch, n_input, output);
+}
+
+void CwiseClipping(float* vector, const int v_size,
+                   const float clipping_value) {
+  PortableCwiseClipping(vector, v_size, clipping_value);
+}
+
+void CwiseClipping(int16_t* vector, const int v_size,
+                   const int16_t clipping_value) {
+  PortableCwiseClipping(vector, v_size, clipping_value);
+}
+
+void CwiseClipping(int8_t* vector, const int v_size,
+                   const int8_t clipping_value) {
+  PortableCwiseClipping(vector, v_size, clipping_value);
+}
+
+void BatchVectorBatchVectorDotProduct(const int16_t* vector1,
+                                      const int16_t* vector2, int v_size,
+                                      int n_batch, int32_t* result) {
+  PortableBatchVectorBatchVectorDotProduct(vector1, vector2, v_size, n_batch,
+                                           result);
+}
+
+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) {
+  NEON_OR_PORTABLE(VectorBatchVectorCwiseProductAccumulate, vector, v_size,
+                   batch_vector, n_batch, multiplier, shift, result);
+}
+
+float VectorVectorDotProduct(const float* vector1, const float* vector2,
+                             int v_size) {
+  return NEON_OR_PORTABLE(VectorVectorDotProduct, vector1, vector2, v_size);
+}
+
+void Sub1Vector(const float* vector, int v_size, float* result) {
+  NEON_OR_PORTABLE(Sub1Vector, vector, v_size, result);
+}
+
+void Sub1Vector(const int16_t* vector, int v_size, int16_t* result) {
+  PortableSub1Vector(vector, v_size, result);
+}
+
+// Check if all entries of a vector are zero for float.
+bool IsZeroVector(const float* vector, int v_size) {
+  return NEON_OR_PORTABLE(IsZeroVector, vector, v_size);
+}
+
+// Check if all entries of a vector are zero for int8.
+bool IsZeroVector(const int8_t* vector, int v_size) {
+  return PortableIsZeroVector(vector, v_size);
+}
+
+void VectorScalarMultiply(const int8_t* vector, int v_size, float scale,
+                          float* result) {
+  NEON_OR_PORTABLE(VectorScalarMultiply, vector, v_size, scale, result);
+}
+
+void SymmetricQuantizeFloats(const float* values, const int size,
+                             int8_t* quantized_values, float* min_value,
+                             float* max_value, float* scaling_factor) {
+  NEON_OR_PORTABLE(SymmetricQuantizeFloats, values, size, quantized_values,
+                   min_value, max_value, scaling_factor);
+}
+
+void SymmetricQuantizeFloats(const float* values, const int size,
+                             int8_t* quantized_values, float min_value,
+                             float max_value, float* scaling_factor) {
+  NEON_OR_PORTABLE(SymmetricQuantizeFloats, values, size, quantized_values,
+                   min_value, max_value, scaling_factor);
+}
+
+void AsymmetricQuantizeFloats(const float* values, const int size,
+                              int8_t* quantized_values, float* scaling_factor,
+                              int32_t* offset) {
+  NEON_OR_PORTABLE(AsymmetricQuantizeFloats, values, size, quantized_values,
+                   scaling_factor, offset);
+}
+
+void ReductionSumVector(const float* input_vector, float* output_vector,
+                        int output_size, int reduction_size) {
+  NEON_OR_PORTABLE(ReductionSumVector, input_vector, output_vector, output_size,
+                   reduction_size);
+}
+
+void ReductionSumVector(const int32_t* input_vector, int32_t* output_vector,
+                        int output_size, int reduction_size) {
+  PortableReductionSumVector(input_vector, output_vector, output_size,
+                             reduction_size);
+}
+
+void ReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
+                        int output_size, int reduction_size) {
+  SSE_OR_PORTABLE(ReductionSumVector, input_vector, output_vector, output_size,
+                  reduction_size);
+}
+
+void MeanStddevNormalization(const float* __restrict__ input_vector,
+                             float* __restrict__ output_vector, int v_size,
+                             int n_batch) {
+  PortableMeanStddevNormalization(input_vector, output_vector, v_size, n_batch);
+}
+
+void TwoGateSaturatingAdd(const int8_t* input, int8_t input_zp,
+                          const int8_t* recurrent, int8_t recurrent_zp,
+                          int32_t input_effective_scale_a,
+                          int32_t input_effective_scale_b,
+                          int32_t recurrent_effective_scale_a,
+                          int32_t recurrent_effective_scale_b, int32_t n_batch,
+                          int32_t n_cell, int16_t* output) {
+  PortableTwoGateSaturatingAdd(
+      input, input_zp, recurrent, recurrent_zp, input_effective_scale_a,
+      input_effective_scale_b, recurrent_effective_scale_a,
+      recurrent_effective_scale_b, n_batch, n_cell, output);
+}
+
+}  // namespace tensor_utils
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_H_

tensorflow/lite/kernels/internal/optimized/sse_tensor_utils_impl.h

@@ -0,0 +1,87 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_IMPL_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_IMPL_H_
+
+#include <cstdint>
+
+#include "tensorflow/lite/kernels/cpu_backend_context.h"
+
+#if defined(_MSC_VER)
+#define __restrict__ __restrict
+#endif
+
+namespace tflite {
+namespace tensor_utils {
+
+#if defined(__AVX2__)
+// Matrix multiplication for float values.
+void Avx2MatrixBatchVectorMultiplyAccumulateImpl(
+    const float* __restrict__ matrix, int m_rows, int m_cols,
+    const float* __restrict__ vector, int n_batch, float* __restrict__ result);
+
+// Matrix multiplication for quantized values using asymmetric quantization.
+void Avx2MatrixBatchVectorMultiplyAccumulateImpl(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
+    bool* compute_row_sums, CpuBackendContext* context);
+#endif  // defined(__AVX2__)
+
+#ifdef __SSSE3__
+
+// Matrix multiplication for quantized values using symmetric quantization.
+void SseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result);
+
+// Matrix multiplication for quantized values using symmetric quantization
+// with additional scratch memory for GEMM operation prior to scaling.
+void SseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch, int32_t* scratch,
+    float* __restrict__ result, CpuBackendContext* context);
+
+// Matrix multiplication for quantized values using asymmetric quantization.
+void SseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
+    bool* compute_row_sums, CpuBackendContext* context);
+
+// Matrix multiplication for quantized values using symmetric quantization.
+// Sparse version.
+void SseSparseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    const int m_rows, const int m_cols, const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale);
+
+void SseReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
+                           const int output_size, const int reduction_size);
+
+#endif  // __SSSE3__
+
+}  // namespace tensor_utils
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_SSE_TENSOR_UTILS_IMPL_H_

tensorflow/lite/kernels/internal/per_channel_dequantize_test.cc

@@ -0,0 +1,155 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+#include "tensorflow/lite/kernels/internal/reference/dequantize.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/test_util.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+TEST(PerChannelDequantize, TestInt8ToFloat_2D) {
+  const std::vector<float> scales = {0.5, 0.25};
+  const std::vector<int> zero_points = {-1, -1};
+  const int quantized_dimension = 0;
+
+  const RuntimeShape shape({2, 5});
+
+  const std::vector<int8_t> input = {-128, -127, -126, -125, -124,
+                                     123,  124,  125,  126,  127};
+  std::vector<float> output(10, -1);
+
+  PerChannelDequantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelDequantize(op_params, shape, input.data(), shape,
+                                      output.data());
+  EXPECT_THAT(output,
+              ElementsAreArray(ArrayFloatNear({-63.5, -63, -62.5, -62, -61.5,
+                                               31, 31.25, 31.5, 31.75, 32})));
+}
+
+TEST(PerChannelDequantize, TestInt8ToFloat_3D) {
+  const std::vector<float> scales = {0.5, 0.25, 0.5, 0.25, 1.0};
+  const std::vector<int> zero_points = {-1, 1, -1, 1, 0};
+  const int quantized_dimension = 2;
+
+  const RuntimeShape shape({1, 2, 5});
+
+  const std::vector<int8_t> input = {-128, -127, -126, -125, -124,
+                                     123,  124,  125,  126,  127};
+  std::vector<float> output(10, -1);
+
+  PerChannelDequantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelDequantize(op_params, shape, input.data(), shape,
+                                      output.data());
+  EXPECT_THAT(output,
+              ElementsAreArray(ArrayFloatNear({-63.5, -32, -62.5, -31.5, -124,
+                                               62, 30.75, 63, 31.25, 127})));
+}
+
+TEST(PerChannelDequantize, TestInt8ToFloat_4DDim0) {
+  const std::vector<float> scales = {0.5, 0.25};
+  const std::vector<int> zero_points = {-1, 1};
+  const int quantized_dimension = 0;
+
+  RuntimeShape shape({2, 2, 5, 1});
+
+  const std::vector<int8_t> input = {-128, -127, -126, -125, -124, 123,  124,
+                                     125,  126,  127,  -128, -127, -126, -125,
+                                     -124, 123,  124,  125,  126,  127};
+  std::vector<float> output(20, -1);
+
+  PerChannelDequantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelDequantize(op_params, shape, input.data(), shape,
+                                      output.data());
+  EXPECT_THAT(output, ElementsAreArray(ArrayFloatNear(
+                          {-63.5,  -63,  -62.5, -62,    -61.5, 62,     62.5,
+                           63,     63.5, 64,    -32.25, -32,   -31.75, -31.5,
+                           -31.25, 30.5, 30.75, 31,     31.25, 31.5})));
+}
+
+TEST(PerChannelDequantize, TestInt8ToFloat_4DDim3) {
+  const std::vector<float> scales = {0.5, 0.25, 0.5, 0.25, 1.0};
+  const std::vector<int> zero_points = {-1, 1, -1, 1, 0};
+  const int quantized_dimension = 3;
+
+  RuntimeShape shape({1, 2, 2, 5});
+
+  const std::vector<int8_t> input = {-128, -127, -126, -125, -124, 123,  124,
+                                     125,  126,  127,  -128, -127, -126, -125,
+                                     -124, 123,  124,  125,  126,  127};
+  std::vector<float> output(20, -1);
+
+  PerChannelDequantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelDequantize(op_params, shape, input.data(), shape,
+                                      output.data());
+  EXPECT_THAT(output, ElementsAreArray(ArrayFloatNear(
+                          {-63.5, -32,   -62.5, -31.5, -124,  62,    30.75,
+                           63,    31.25, 127,   -63.5, -32,   -62.5, -31.5,
+                           -124,  62,    30.75, 63,    31.25, 127})));
+}
+
+TEST(PerChannelDequantize, TestInt4ToFloat_2D) {
+  const std::vector<float> scales = {0.5, 0.25};
+  const std::vector<int> zero_points = {-1, -1};
+  const int quantized_dimension = 0;
+
+  const RuntimeShape unpacked_shape({2, 4});
+
+  const std::vector<int8_t> packed_int4_input = {-1, 0, 65, -127};
+  std::vector<float> output(8, -1);
+  const size_t bytes_unpacked = packed_int4_input.size() * 2;
+  auto unpacked_input_data = std::make_unique<int8_t[]>(bytes_unpacked);
+  tflite::tensor_utils::UnpackPackedIntToInt8(packed_int4_input.data(),
+                                              bytes_unpacked, /*bit_width=*/4,
+                                              unpacked_input_data.get());
+  EXPECT_THAT(std::vector<int8_t>(unpacked_input_data.get(),
+                                  unpacked_input_data.get() + bytes_unpacked),
+              ElementsAreArray(ArrayFloatNear({-1, -1, 0, 0, 1, 4, 1, -8})));
+
+  PerChannelDequantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelDequantize(op_params, unpacked_shape,
+                                      unpacked_input_data.get(), unpacked_shape,
+                                      output.data());
+  // This comes from (UNPACKED - zero_point) * scale.
+  EXPECT_THAT(output, ElementsAreArray(ArrayFloatNear(
+                          {0, 0, 0.5, 0.5, 0.5, 1.25, 0.5, -1.75})));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/per_channel_quantize_test.cc

@@ -0,0 +1,118 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/reference/quantize.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+#include "tensorflow/lite/kernels/test_util.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+
+TEST(PerChannelQuantize, TestFloatToInt8_2D) {
+  const std::vector<float> scales = {0.5, 0.25};
+  const std::vector<int> zero_points = {-1, -1};
+  const int quantized_dimension = 0;
+
+  const RuntimeShape shape({2, 5});
+
+  const std::vector<float> input = {-63.5, -63,   -62.5, -62,   -61.5,
+                                    31,    31.25, 31.5,  31.75, 32};
+  std::vector<int8_t> output(10, -1);
+
+  PerChannelQuantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelQuantize(op_params, shape, input.data(), shape,
+                                    output.data());
+  EXPECT_THAT(output, ElementsAreArray({-128, -127, -126, -125, -124, 123, 124,
+                                        125, 126, 127}));
+}
+
+TEST(PerChannelQuantize, TestFloatToInt8_3D) {
+  const std::vector<float> scales = {0.5, 0.25, 0.5, 0.25, 1.0};
+  const std::vector<int> zero_points = {-1, 1, -1, 1, 0};
+  const int quantized_dimension = 2;
+
+  const RuntimeShape shape({1, 2, 5});
+
+  const std::vector<float> input = {-63.5, -32,   -62.5, -31.5, -124,
+                                    62,    30.75, 63,    31.25, 127};
+  std::vector<int8_t> output(10, -1);
+
+  PerChannelQuantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelQuantize(op_params, shape, input.data(), shape,
+                                    output.data());
+  EXPECT_THAT(output, ElementsAreArray({-128, -127, -126, -125, -124, 123, 124,
+                                        125, 126, 127}));
+}
+
+TEST(PerChannelQuantize, TestFloatToInt8_4DDim0) {
+  const std::vector<float> scales = {0.5, 0.25};
+  const std::vector<int> zero_points = {-1, 1};
+  const int quantized_dimension = 0;
+
+  RuntimeShape shape({2, 2, 5, 1});
+
+  const std::vector<float> input = {
+      -63.5,  -63, -62.5,  -62,   -61.5,  62,   62.5,  63, 63.5,  64,
+      -32.25, -32, -31.75, -31.5, -31.25, 30.5, 30.75, 31, 31.25, 31.5};
+  std::vector<int8_t> output(20, -1);
+
+  PerChannelQuantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelQuantize(op_params, shape, input.data(), shape,
+                                    output.data());
+  EXPECT_THAT(output,
+              ElementsAreArray({-128, -127, -126, -125, -124, 123,  124,
+                                125,  126,  127,  -128, -127, -126, -125,
+                                -124, 123,  124,  125,  126,  127}));
+}
+
+TEST(PerChannelQuantize, TestFloatToInt8_4DDim3) {
+  const std::vector<float> scales = {0.5, 0.25, 0.5, 0.25, 1.0};
+  const std::vector<int> zero_points = {-1, 1, -1, 1, 0};
+  const int quantized_dimension = 3;
+
+  RuntimeShape shape({1, 2, 2, 5});
+
+  const std::vector<float> input = {
+      -63.5, -32, -62.5, -31.5, -124, 62, 30.75, 63, 31.25, 127,
+      -63.5, -32, -62.5, -31.5, -124, 62, 30.75, 63, 31.25, 127};
+  std::vector<int8_t> output(20, -1);
+
+  PerChannelQuantizationParams op_params;
+  op_params.zero_point = zero_points.data();
+  op_params.scale = scales.data();
+  op_params.quantized_dimension = quantized_dimension;
+  reference_ops::PerChannelQuantize(op_params, shape, input.data(), shape,
+                                    output.data());
+  EXPECT_THAT(output,
+              ElementsAreArray({-128, -127, -126, -125, -124, 123,  124,
+                                125,  126,  127,  -128, -127, -126, -125,
+                                -124, 123,  124,  125,  126,  127}));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/portable_tensor.h

@@ -0,0 +1,141 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_PORTABLE_TENSOR_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_PORTABLE_TENSOR_H_
+
+#include <cstddef>
+#include <vector>
+
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+
+// A list of tensors in a format that can be used by kernels like split and
+// concatenation.
+template <typename T>
+class VectorOfTensors {
+ public:
+  // Build with the tensors in 'tensor_list'.
+  VectorOfTensors(const TfLiteContext& context,
+                  const TfLiteIntArray& tensor_list) {
+    int num_tensors = tensor_list.size;
+
+    all_data_.reserve(num_tensors);
+    all_shape_.reserve(num_tensors);
+    all_shape_ptr_.reserve(num_tensors);
+
+    for (int i = 0; i < num_tensors; ++i) {
+      TfLiteTensor* t = &context.tensors[tensor_list.data[i]];
+      all_data_.push_back(GetTensorData<T>(t));
+      all_shape_.push_back(GetTensorShape(t));
+    }
+
+    // Taking the pointer from inside a std::vector is only OK if the vector is
+    // never modified, so we populate all_shape in the previous loop and then we
+    // are free to grab iterators here.
+    for (int i = 0; i < num_tensors; ++i) {
+      all_shape_ptr_.push_back(&all_shape_[i]);
+    }
+  }
+
+  explicit VectorOfTensors(const std::vector<TfLiteTensor*>& tensors) {
+    int num_tensors = tensors.size();
+
+    all_data_.reserve(num_tensors);
+    all_shape_.reserve(num_tensors);
+    all_shape_ptr_.reserve(num_tensors);
+
+    for (auto* t : tensors) {
+      all_data_.push_back(GetTensorData<T>(t));
+      all_shape_.push_back(GetTensorShape(t));
+    }
+
+    // Taking the pointer from inside a std::vector is only OK if the vector is
+    // never modified, so we populate all_shape in the previous loop and then we
+    // are free to grab iterators here.
+    for (int i = 0; i < num_tensors; ++i) {
+      all_shape_ptr_.push_back(&all_shape_[i]);
+    }
+  }
+  // Return a pointer to the data pointers of all tensors in the list. For
+  // example:
+  //   float* const* f = v.data();
+  //   f[0][1] is the second element of the first tensor.
+  T* const* data() const { return all_data_.data(); }
+
+  // Return a pointer the shape pointers of all tensors in the list. For
+  // example:
+  //   const RuntimeShape* const* d = v.dims();
+  //   dims[1] are the dimensions of the second tensor in the list.
+  const RuntimeShape* const* shapes() const { return all_shape_ptr_.data(); }
+
+  size_t size() const { return all_data_.size(); }
+
+ private:
+  std::vector<T*> all_data_;
+  std::vector<RuntimeShape> all_shape_;
+  std::vector<RuntimeShape*> all_shape_ptr_;
+};
+
+// A list of quantized tensors in a format that can be used by kernels like
+// split and concatenation.
+class VectorOfQuantizedTensors : public VectorOfTensors<uint8_t> {
+ public:
+  // Build with the tensors in 'tensor_list'.
+  VectorOfQuantizedTensors(const TfLiteContext& context,
+                           const TfLiteIntArray& tensor_list)
+      : VectorOfTensors<uint8_t>(context, tensor_list) {
+    for (int i = 0; i < tensor_list.size; ++i) {
+      TfLiteTensor* t = &context.tensors[tensor_list.data[i]];
+      zero_point_.push_back(t->params.zero_point);
+      scale_.push_back(t->params.scale);
+    }
+  }
+
+  const float* scale() const { return scale_.data(); }
+  const int32_t* zero_point() const { return zero_point_.data(); }
+
+ private:
+  std::vector<int32_t> zero_point_;
+  std::vector<float> scale_;
+};
+
+// Writes randomly accessed values from `input` sequentially into `output`.
+template <typename T>
+class SequentialTensorWriter {
+ public:
+  SequentialTensorWriter(const TfLiteTensor* input, TfLiteTensor* output) {
+    input_data_ = GetTensorData<T>(input);
+    output_ptr_ = GetTensorData<T>(output);
+  }
+  SequentialTensorWriter(const T* input_data, T* output_data)
+      : input_data_(input_data), output_ptr_(output_data) {}
+
+  void Write(int position) { *output_ptr_++ = input_data_[position]; }
+  void WriteN(int position, int len) {
+    memcpy(output_ptr_, &input_data_[position], sizeof(T) * len);
+    output_ptr_ += len;
+  }
+
+ private:
+  const T* input_data_;
+  T* output_ptr_;
+};
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_PORTABLE_TENSOR_H_

tensorflow/lite/kernels/internal/portable_tensor_utils.cc

@@ -0,0 +1,206 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_
+
+#include "tensorflow/lite/kernels/internal/portable_tensor_utils.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cmath>
+#include <cstdint>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+
+#if defined(_MSC_VER)
+#define __restrict__ __restrict
+#endif
+
+namespace tflite {
+
+// Not all backends support CpuBackendContext usage, so forward declare to avoid
+// pulling in its implementation. Use of CpuBackendContext in method
+// implementations is purely optional.
+class CpuBackendContext;
+
+namespace tensor_utils {
+
+// Apply Rectified Linear to elements of a vector.
+void ApplyReluToVector(const float* __restrict__ vector, int v_size,
+                       float* __restrict__ result) {
+  for (int v = 0; v < v_size; v++) {
+    result[v] = std::max(0.0f, vector[v]);
+  }
+}
+
+// Apply Rectified Linear 1 (cap to [-1;1]) to elements of a vector
+void ApplyRelu1ToVector(const float* __restrict__ vector, int v_size,
+                        float* __restrict__ result) {
+  for (int v = 0; v < v_size; v++) {
+    result[v] = std::max(-1.0f, std::min(vector[v], 1.0f));
+  }
+}
+
+// Apply Rectified Linear 6 (cap to [0;6]) to elements of a vector
+void ApplyRelu6ToVector(const float* __restrict__ vector, int v_size,
+                        float* __restrict__ result) {
+  for (int v = 0; v < v_size; v++) {
+    result[v] = std::max(0.0f, std::min(vector[v], 6.0f));
+  }
+}
+
+// Apply signbit to elements of a vector
+void ApplySignbitToVector(const float* __restrict__ vector, int v_size,
+                          float* __restrict__ result) {
+  for (int v = 0; v < v_size; v++) {
+    result[v] = std::signbit(vector[v]);
+  }
+}
+
+void UnpackDenseInt4IntoInt8(const int8_t* src_buffer, int num_elements,
+                             int8_t* dst_buffer) {
+  // num_elements means the number of elements regardless of packed or unpacked.
+  // For example, 3 elements means both
+  //   1) Packed: 3 int4's = 12 bit -> 16 bits (padded) = 2 bytes.
+  //      stored in src_buffer[0] and src_buffer[1] (i = 0..1)
+  //   2) Unpacked: 3 int8's = 3 bytes.
+  //.     stored in dst_buffer[0], dst_buffer[1] and dst_buffer[2] (j = 0..2)
+  for (int i = 0; i < num_elements / 2; i++) {
+    int8_t byte = src_buffer[i];
+    // Shift left first so that sign is properly extended when shifted right
+    int8_t lower = static_cast<int8_t>(byte << 4) >> 4;
+    int8_t higher = byte >> 4;
+    dst_buffer[2 * i] = lower;
+    dst_buffer[2 * i + 1] = higher;
+  }
+
+  // If the buffer size is odd, extract the final lower nibble.
+  if (num_elements % 2 != 0) {
+    dst_buffer[num_elements - 1] =
+        static_cast<int8_t>(src_buffer[num_elements / 2] << 4) >> 4;
+  }
+}
+
+void UnpackPackedIntToInt8(const int8_t* src_buffer, int num_elements,
+                           int bit_width, int8_t* dst_buffer,
+                           bool unpack_unsigned) {
+  assert(bit_width == 2 || bit_width == 4);
+  if (bit_width == 4) {
+    // num_elements means the number of elements regardless of packed or
+    // unpacked. For example, 3 elements means both
+    //   1) Packed: 3 int4's = 12 bit -> 16 bits (padded) = 2 bytes.
+    //      stored in src_buffer[0] and src_buffer[1] (i = 0..1)
+    //   2) Unpacked: 3 int8's = 3 bytes.
+    //.     stored in dst_buffer[0], dst_buffer[1] and dst_buffer[2] (j = 0..2)
+    for (int i = 0; i < num_elements / 2; i++) {
+      int8_t byte = src_buffer[i];
+      int8_t lower, higher;
+      if (unpack_unsigned) {
+        lower = byte & 0x0F;
+        higher = (byte >> 4) & 0x0F;
+      } else {
+        // Shift left first so that sign is properly extended when shifted right
+        lower = static_cast<int8_t>(byte << 4) >> 4;
+        higher = byte >> 4;
+      }
+      dst_buffer[2 * i] = lower;
+      dst_buffer[2 * i + 1] = higher;
+    }
+
+    // If the buffer size is odd, extract the final lower nibble.
+    if (num_elements % 2 != 0) {
+      int8_t byte = src_buffer[num_elements / 2];
+      dst_buffer[num_elements - 1] =
+          unpack_unsigned ? (byte & 0x0F) : static_cast<int8_t>(byte << 4) >> 4;
+    }
+  } else if (bit_width == 2) {
+    for (int i = 0; i < num_elements / 4; i++) {
+      int8_t byte = src_buffer[i];
+      if (unpack_unsigned) {
+        dst_buffer[4 * i] = byte & 0x03;
+        dst_buffer[4 * i + 1] = (byte >> 2) & 0x03;
+        dst_buffer[4 * i + 2] = (byte >> 4) & 0x03;
+        dst_buffer[4 * i + 3] = (byte >> 6) & 0x03;
+      } else {
+        // Shift left first so that sign is properly extended when shifted right
+        int8_t val1 = static_cast<int8_t>(byte << 6) >> 6;
+        int8_t val2 = static_cast<int8_t>((byte << 4) & 0xFF) >> 6;
+        int8_t val3 = static_cast<int8_t>((byte << 2) & 0xFF) >> 6;
+        int8_t val4 = byte >> 6;
+        dst_buffer[4 * i] = val1;
+        dst_buffer[4 * i + 1] = val2;
+        dst_buffer[4 * i + 2] = val3;
+        dst_buffer[4 * i + 3] = val4;
+      }
+    }
+
+    // Handle the remaining elements.
+    int remaining_elements = num_elements % 4;
+    if (remaining_elements > 0) {
+      int8_t byte = src_buffer[num_elements / 4];
+      for (int i = 0; i < remaining_elements; i++) {
+        if (unpack_unsigned) {
+          dst_buffer[num_elements - remaining_elements + i] =
+              (byte >> (2 * i)) & 0x03;
+        } else {
+          dst_buffer[num_elements - remaining_elements + i] =
+              static_cast<int8_t>((byte << (6 - 2 * i)) & 0xFF) >> 6;
+        }
+      }
+    }
+  }
+}
+
+void PackInt8IntoDenseInt(const int8_t* src_buffer, int num_elements,
+                          int bit_width, int8_t* dst_buffer) {
+  assert(bit_width == 2 || bit_width == 4);
+  if (bit_width == 4) {
+    // num_elements means the number of elements regardless of packed or
+    // unpacked. For example, 3 elements means both
+    //   1) Unpacked: 3 int8's = 3 bytes.
+    //      stored in src_buffer[0], src_buffer[1] and src_buffer[2] (j = 0..2)
+    //   2) Packed: 3 int4's = 12 bit -> 16 bits (padded) = 2 bytes.
+    //      stored in dst_buffer[0] and dst_buffer[1] (i = 0..1)
+    for (int i = 0; i < num_elements / 2; ++i) {
+      dst_buffer[i] = (src_buffer[2 * i] & 0x0F) | (src_buffer[2 * i + 1] << 4);
+    }
+    // If the buffer size is odd, pack the final nibble.
+    if (num_elements % 2 != 0) {
+      dst_buffer[num_elements / 2] = src_buffer[num_elements - 1] & 0x0F;
+    }
+  } else if (bit_width == 2) {
+    for (int i = 0; i < num_elements / 4; ++i) {
+      dst_buffer[i] = (src_buffer[4 * i] & 0x03) |
+                      ((src_buffer[4 * i + 1] & 0x03) << 2) |
+                      ((src_buffer[4 * i + 2] & 0x03) << 4) |
+                      ((src_buffer[4 * i + 3] & 0x03) << 6);
+    }
+    // Handle the remaining elements.
+    int remaining_elements = num_elements % 4;
+    if (remaining_elements > 0) {
+      int8_t packed_val = 0;
+      for (int i = 0; i < remaining_elements; ++i) {
+        packed_val |= (src_buffer[num_elements - remaining_elements + i] & 0x03)
+                      << (i * 2);
+      }
+      dst_buffer[num_elements / 4] = packed_val;
+    }
+  }
+}
+
+}  // namespace tensor_utils
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_TENSOR_UTILS_H_

tensorflow/lite/kernels/internal/portable_tensor_utils.h

@@ -0,0 +1,661 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_PORTABLE_TENSOR_UTILS_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_PORTABLE_TENSOR_UTILS_H_
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+
+#include "tensorflow/lite/core/c/builtin_op_data.h"
+#include "tensorflow/lite/core/c/common.h"
+
+#if defined(_MSC_VER)
+#define __restrict__ __restrict
+#endif
+
+namespace tflite {
+
+// Not all backends support CpuBackendContext usage, so forward declare to avoid
+// pulling in its implementation. Use of CpuBackendContext in method
+// implementations is purely optional.
+class CpuBackendContext;
+
+namespace tensor_utils {
+
+// Multiplies a matrix with a scalar and reduce the result on each row to a
+// scalar.
+// Parameters:
+//     - matrix: matrix of size n_row * n_col
+//     - scalar: the scalar that is multiplied to each element in the matrix
+//     - n_row:  the row count of the matrix
+//     - n_col:  the column count of the matrix
+//     - output: the 32bit output
+// Note: We do not need saturation because the int8 * int8 is safe from overflow
+// in (2^31-1) / (2^14) = 131072, which is bigger than the n_row. Non-zero
+// initial output value is not exceptionally large.
+void MatrixScalarMultiplyAccumulate(const int8_t* matrix, int32_t scalar,
+                                    int32_t n_row, int32_t n_col,
+                                    int32_t* output);
+
+// Add another vector for each batch in the batch vector.
+template <typename T>
+void VectorBatchVectorAdd(const T* vector, int v_size, int n_batch,
+                          T* batch_vector) {
+  for (int b = 0; b < n_batch; b++) {
+    for (int i = 0; i < v_size; ++i) {
+      batch_vector[i] += vector[i];
+    }
+    batch_vector += v_size;
+  }
+}
+
+// Cwise product of two vectors.
+template <typename T>
+inline void VectorVectorCwiseProduct(const T* vector1, const T* vector2,
+                                     int v_size, T* result) {
+  for (int v = 0; v < v_size; v++) {
+    *result++ = *vector1++ * *vector2++;
+  }
+}
+
+// Cwise product of a vector and a batch-vector.
+template <typename T>
+inline void VectorBatchVectorCwiseProduct(const T* vector, int v_size,
+                                          const T* batch_vector, int n_batch,
+                                          T* result) {
+  for (int b = 0; b < n_batch; b++) {
+    VectorVectorCwiseProduct(vector, batch_vector, v_size, result);
+    // Update the pointers.
+    result += v_size;
+    batch_vector += v_size;
+  }
+}
+
+// Cwise product and accumulate of two vectors. Since it's a MAC operation, the
+// assumption here is that result array is initialized to valid values.
+template <typename T>
+inline void VectorVectorCwiseProductAccumulate(const T* __restrict__ vector1,
+                                               const T* __restrict__ vector2,
+                                               int v_size,
+                                               T* __restrict__ result) {
+  for (int v = 0; v < v_size; v++) {
+    *result++ += *vector1++ * *vector2++;
+  }
+}
+
+// Cwise product and accumulate of a vector and a batch-vector. Since it's a MAC
+// operation, the assumption here is that result array is initialized to valid
+// values.
+template <typename T>
+inline void VectorBatchVectorCwiseProductAccumulate(const T* vector, int v_size,
+                                                    const T* batch_vector,
+                                                    int n_batch, T* result) {
+  for (int b = 0; b < n_batch; b++) {
+    VectorVectorCwiseProductAccumulate(vector, batch_vector, v_size, result);
+    // Update the pointers.
+    result += v_size;
+    batch_vector += v_size;
+  }
+}
+
+// Batch vector initialization with another vector.
+template <typename T>
+void VectorBatchVectorAssign(const T* vector, int v_size, int n_batch,
+                             T* batch_vector) {
+  for (int b = 0; b < n_batch; b++) {
+    std::copy_n(vector, v_size, batch_vector + b * v_size);
+  }
+}
+
+// Checks if all entries of vector are zero for float.
+bool IsZeroVector(const float* vector, int v_size);
+
+// Checks if all entries of vector are zero for int8.
+bool IsZeroVector(const int8_t* vector, int v_size);
+
+// Quantizes a buffer of floating point values using a symmetric quantization
+// (i.e. linear quantization without an offset) to 8-bit signed integers.
+// It also outputs the range (min, max) of the floating point buffer, and the
+// scaling factor used to quantize the values.
+void SymmetricQuantizeFloats(const float* values, const int size,
+                             int8_t* quantized_values, float* min_value,
+                             float* max_value, float* scaling_factor);
+
+// Quantizes a buffer of floating point values using a symmetric quantization
+// (i.e. linear quantization without an offset) to 8-bit signed integers.
+// It uses the range (min, max) provided to the function to calculate the
+// appropriate scaling factor to quantize the values.
+void SymmetricQuantizeFloats(const float* values, const int size,
+                             int8_t* quantized_values, float min_value,
+                             float max_value, float* scaling_factor);
+
+void AsymmetricQuantizeFloats(const float* values, const int size,
+                              int8_t* quantized_values, float* scaling_factor,
+                              int32_t* offset);
+
+// Helper function to quantize floats.
+// float_data_ptr     input float vectors
+// n_batch            number of input vectors
+// n_data             size of a single input vector
+// quantized_data_ptr (out) vector with quantized data
+// scaling_factors    (out) scaling factors (one per vector)
+// zero_points        (out) zero points (one per vector)
+// do_asymmetric      controls if the quantization should be asymmetric.
+inline void BatchQuantizeFloats(const float* float_data_ptr, int n_batch,
+                                int n_data, int8_t* quantized_data_ptr,
+                                float* scaling_factors, int32_t* zero_points,
+                                bool do_asymmetric) {
+  for (int b = 0; b < n_batch; ++b) {
+    const int offset = b * n_data;
+    if (do_asymmetric) {
+      tensor_utils::AsymmetricQuantizeFloats(
+          float_data_ptr + offset, n_data, quantized_data_ptr + offset,
+          &scaling_factors[b], &zero_points[b]);
+    } else {
+      float unused_min, unused_max;
+      tensor_utils::SymmetricQuantizeFloats(
+          float_data_ptr + offset, n_data, quantized_data_ptr + offset,
+          &unused_min, &unused_max, &scaling_factors[b]);
+      if (zero_points) zero_points[b] = 0;
+    }
+  }
+}
+
+// Multiplies a matrix by a "batched" vector (i.e. a matrix with a batch
+// dimension composed by input vectors independent from each other). The result
+// of the multiplication is accumulated to the passed result buffer.
+// More specifically, for a matrix M of shape [n, i] and a batched-vector
+// of shape [i, batch] it will first compute the product of shape [n, batch].
+// This product will be accumulated to the result buffer.
+void MatrixBatchVectorMultiplyAccumulate(const float* matrix, int m_rows,
+                                         int m_cols, const float* vector,
+                                         int n_batch, float* result);
+
+// Same as the function above, but the matrix is a sparse tensor with block
+// pattern 1x4.
+// This function assumes that m_cols is a multiple of the block size (4 in this
+// case) so that there's no incomplete block.
+void SparseMatrixBatchVectorMultiplyAccumulate1x4(
+    const float* __restrict__ matrix, const int32_t* __restrict__ segments,
+    const int32_t* __restrict__ indices, int m_rows, int m_cols,
+    const float* __restrict__ vector, int n_batch, float* __restrict__ result);
+
+// Same as the function above, but the matrix is stored in block compressed
+// sparse row format with block pattern 1x16 which consists of two arrays:
+//   1. A matrix array stores non-zero blocks of the matrix in row major.
+//   2. A ledger array stores nrows groups, one group per row. Each group starts
+//      with an integer representing the number of non-zero blocks for the
+//      corresponding row and follows with column indexes of the first element
+//      of each non-zero block.
+// This function assumes that
+//   1. m_cols is a multiple of 16 so that all blocks are full blocks.
+//   2. m_cols < 254 * 16 so that block index can be represented by uint8.
+void SparseMatrixBatchVectorMultiplyAccumulate(
+    const float* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    int m_rows, int m_cols, const float* __restrict__ vector, int n_batch,
+    float* __restrict__ result);
+
+// Same as the function above, but for values quantized using symmetric
+// quantization (e.g. by calling SymmetricQuantizeFloats).
+// The passed scaling factors is a buffer of the quantization scaling factors
+// that will be used to dequentize the products into the final result buffer.
+// These scaling factors are the multiplication of the matrix scaling factor
+// by the vector's scaling factor, one per batch (i.e. this allows quantizing
+// each batch in the batch-vector matrix independently).
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result);
+
+// Same as the function above except that vector values
+// are quantized with asymmetric quantization per-batch and the matrix
+// is quantized per row.
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* __restrict__ per_channel_scale,
+    const int32_t* __restrict__ input_offset);
+
+// Same as the function above, but the matrix is a sparse tensor with block
+// pattern 1x16.
+// This function assumes that m_cols is a multiple of the block size (16 in this
+// case) so that there's no incomplete block. Also, it assumes all offsets of
+// input, output and filter are zero.
+void SparseMatrixBatchVectorMultiplyAccumulate1x16(
+    const int8_t* __restrict__ matrix, const int32_t* __restrict__ segments,
+    const int32_t* __restrict__ indices, int m_rows, int m_cols,
+    const int8_t* __restrict__ vector, const int32_t* __restrict__ bias_vector,
+    int n_batch, const int32_t input_offset, const int32_t output_multiplier,
+    int32_t output_shift, const int32_t* per_channel_scale,
+    const int32_t* per_channel_shift, int32_t output_offset,
+    const int32_t output_activation_min, const int32_t output_activation_max,
+    int8_t* __restrict__ result);
+
+// Same as the function above, but the matrix is stored in block compressed
+// sparse row format with block pattern 1x16 which consists of two arrays:
+//   1. A matrix array stores non-zero blocks of the matrix in row major.
+//   2. A ledger array stores nrows groups, one group per row. Each group starts
+//      with an integer representing the number of non-zero blocks for the
+//      corresponding row followed by column index of the first element of
+//      each non-zero block.
+// This function assumes that
+//   1. m_cols is a multiple of 16 so that all blocks are full blocks.
+//   2. m_cols < 254 * 16 so that block index can be represented by uint8.
+void SparseMatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const uint8_t* __restrict__ ledger,
+    const int m_rows, const int m_cols, const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale = nullptr);
+
+// Same as the above 8, 8, 8 integer matmul except for the presence of zero
+// point and non-accumulative.
+// TODO(b/148688698): remove this function by folding zero point calculation in
+// prepare() function.
+void MatrixBatchVectorMultiply(const int8_t* input, int32_t input_zeropoint,
+                               const int8_t* input_to_gate_weights,
+                               int32_t input_to_gate_effective_scale_a,
+                               int32_t input_to_gate_effective_scale_b,
+                               int32_t n_batch, int32_t n_input, int32_t n_cell,
+                               int8_t* gate_output, int8_t gate_output_zp);
+
+// Same as above but has 16 bit and 8 bit input and 8 bit output.
+// Used in projection when hidden is 16bit.
+void MatrixBatchVectorMultiply(const int16_t* hidden,
+                               const int8_t* hidden_to_output_weights,
+                               int32_t proj_effective_scale_a,
+                               int32_t proj_effective_scale_b,
+                               const int32_t* gate_bias, int32_t n_batch,
+                               int32_t n_hidden, int32_t n_output,
+                               int32_t output_zp, int8_t* proj_output);
+
+// Apply Layer Normalization (https://arxiv.org/abs/1607.06450) to a Quantized
+// vector.
+// Parameters:
+//     - input: batch vector of size n_batch * n_input; 16 bit.
+//     - layer_norm_weights:  the quantized layer normalization weights.
+//     - bias: the bias for the layer normalization.
+//     - layer_norm_scale_a: multiplier for scale factor.
+//     - layer_norm_scale_b: shift for scale factor.
+//     - variance_limit: the guard to make sure the inverse does not overflow.
+//     - n_batch: the number of batches.
+//     - n_input: the size for input and output.
+//     - output:  the 16 bit output
+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);
+
+// Same as above but the internal calculation is done in float.
+void ApplyLayerNormFloat(const int16_t* input,
+                         const int16_t* layer_norm_weights,
+                         int32_t layer_norm_scale_a, int32_t layer_norm_scale_b,
+                         const int32_t* bias, int n_batch, int n_input,
+                         int16_t* output);
+
+// Apply Sigmoid to a quantized vector.
+// Parameters:
+//     - input: batch vector of size n_batch * n_input; 16 bit.
+//     - n_batch: the number of batches.
+//     - n_input: the size for input and output.
+//     - output:  the 16 bit output
+// The input is in Q3.12 format and the output is in Q0.15 format.
+void ApplySigmoid(const int16_t* input, int32_t n_batch, int32_t n_input,
+                  int16_t* output);
+
+// Same as above but the internal calculation is float.
+void ApplySigmoidFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
+                       int16_t* output);
+
+// Apply Tanh to a quantized vector.
+// Parameters:
+//     - integer_bits: the integer bits of the input.
+//                     Currently supports 0, 1, 2, 3, 4, 5, 6.
+//     - input: batch vector of size n_batch * n_input; 16 bit.
+//     - n_batch: the number of batches.
+//     - n_input: the size for input and output.
+//     - output:  the 16 bit output
+// The input is in Qm.15-m format and the output is in Q0.15 format.
+void ApplyTanh(int32_t intger_bits, const int16_t* input, int32_t n_batch,
+               int32_t n_input, int16_t* output);
+
+// Apply Tanh to a quantized vector. The internal calculation is in float.
+//    - Input has 2^(integer_bits) as scale.
+//    - Output has Q0.15 as scale.
+void ApplyTanhFloat(const int16_t* input, int32_t n_batch, int32_t n_input,
+                    int32_t integer_bits, int16_t* output);
+
+// Element-wise multiplication of two quantized vectors.
+// Parameters:
+//     - input_1: batch vector of size n_batch * n_input; 16 bit.
+//     - input_2: batch vector of size n_batch * n_input; 16 bit.
+//     - n_batch: the number of batches.
+//     - n_input: the size for input and output.
+//     - shift:   the shift needed to produce the output.
+//     - output:  the 16 bit output of size n_batch * n_input.
+// Output does not need to be initialized.
+void CwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch,
+              int n_input, int shift, int16_t* output);
+
+// Element-wise multiplication of two quantized vectors.
+// Parameters:
+//     - input_1: batch vector of size n_batch * n_input; 16 bit.
+//     - input_2: batch vector of size n_batch * n_input; 16 bit.
+//     - n_batch: the number of batches.
+//     - n_input: the size for input and output.
+//     - shift:   the shift needed to produce the output.
+//     - output:  the 8 bit output of size n_batch * n_input.
+// Output does not need to be initialized.
+void CwiseMul(const int16_t* input_1, const int16_t* input_2, int n_batch,
+              int n_input, int shift, int8_t* output);
+
+// Element-wise multiplication of two quantized vectors with rescaling.
+// Parameters:
+//     - input_1:    batch vector of size n_batch * n_input; 16 bit.
+//     - input_2:    batch vector of size n_batch * n_input; 16 bit.
+//     - multiplier: the multiplier part of scale.
+//     - shift:      the shift part of scale.
+//     - n_batch:    the number of batches.
+//     - n_input:    the size for input and output.
+//     - output:     the 8 bit output of size n_batch * n_input.
+//     - output_zp:  the zero point of output.
+// Output does not need to be initialized.
+// Multiplier ("m") and shift ("s") are connected to scale ("s") with s = m *
+// 2^(s - 31).
+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);
+
+// Element-wise saturating addition of two quantized vectors without rescaling.
+// Parameters:
+//     - input_1:    batch vector of size n_batch * n_input; 16 bit.
+//     - input_2:    batch vector of size n_batch * n_input; 16 bit.
+//     - n_batch:    the number of batches.
+//     - n_input:    the size for input and output.
+//     - output:     the 8 bit output of size n_batch * n_input.
+// Output does not need to be initialized.
+void CwiseAdd(const int16_t* input_1, const int16_t* input_2, int n_batch,
+              int n_input, int16_t* output);
+
+// Element-wise in-place clipping of a vector. Overloaded for float, int16_t,
+// int8_t. Parameters:
+//     - vector:         vector of size v_size.
+//     - v_size:         the size of the vector.
+//     - clipping_value: the value used for clipping.
+void CwiseClipping(float* vector, const int v_size, const float clipping_value);
+void CwiseClipping(int16_t* vector, const int v_size,
+                   const int16_t clipping_value);
+void CwiseClipping(int8_t* vector, const int v_size,
+                   const int8_t clipping_value);
+
+// Dot product of two vectors.
+float VectorVectorDotProduct(const float* vector1, const float* vector2,
+                             int v_size);
+
+// Dot product of two batch vectors of size n_batch * v_size:
+// vector1 = [x_1_1, x_1_2, ..., x_1_vsize,
+//            x_2_1, x_2_2, ..., x_2_vsize,
+//            ...
+//            x_nbatch_1,..., x_nbatch_vsize]
+// vector2 = [y_1_1, y_1_2, ..., y_1_vsize,
+//            y_2_1, y_2_2, ..., y_2_vsize,
+//            ...
+//            y_nbatch_1,..., y_nbatch_vsize]
+// Then result will be a vector of n_batch size starting from 'result':
+// [x_1_1 * y_1_1 + x_1_2 * y_1_2 + ... + x_1_vsize * y_1_vsize,
+//  x_2_1 * y_2_1 + x_2_2 * y_2_2 + ... + x_2_vsize * y_2_vsize,
+//  ...
+//  x_nbatch_1 * y_nbatch_1 + ... + x_nbatch_vsize * y_nbatch_vsize]
+template <typename T>
+inline void BatchVectorBatchVectorDotProduct(const T* vector1, const T* vector2,
+                                             int v_size, int n_batch,
+                                             T* result) {
+  for (int b = 0; b < n_batch; b++) {
+    result[b] = VectorVectorDotProduct(vector1, vector2, v_size);
+    vector1 += v_size;
+    vector2 += v_size;
+  }
+}
+
+// Same as above but input is 16bit and output is 32bit.
+void BatchVectorBatchVectorDotProduct(const int16_t* vector1,
+                                      const int16_t* vector2, int v_size,
+                                      int n_batch, int32_t* result);
+
+// Same as above, but inputs are 16bit integer and output is 16bit integer.
+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);
+
+// Compute "1.0f - elements of vector" (used in CIFG).
+void Sub1Vector(const float* vector, int v_size, float* result);
+
+// Compute "1.0f - elements of vector" (used in CIFG) for int16 input.
+// "vector" has range [0, 32767] because it is the output of sigmoid function.
+void Sub1Vector(const int16_t* vector, int v_size, int16_t* result);
+
+// Reduce-sum on a float input vector:
+// input_vector: float pointer to input vector.
+// output_vector: float pointer to vector.
+// output_size: output vector size.
+// reduction_size: number of consecutive elements from input vector which are
+// added to get one element of output.
+void ReductionSumVector(const float* input_vector, float* output_vector,
+                        int output_size, int reduction_size);
+
+// Same as above but input/output is 32 bit integer.
+void ReductionSumVector(const int32_t* input_vector, int32_t* output_vector,
+                        int output_size, int reduction_size);
+
+// Same as above but input is 8 bit integer.
+void ReductionSumVector(const int8_t* input_vector, int32_t* output_vector,
+                        int output_size, int reduction_size);
+
+// Multiply all elements of vector with a scalar.
+void VectorScalarMultiply(const int8_t* vector, int v_size, float scale,
+                          float* result);
+
+// Layer norm for each batch.
+void MeanStddevNormalization(const float* input_vector, float* output_vector,
+                             int v_size, int n_batch);
+
+// Saturate Add with rescale on both inputs.
+void TwoGateSaturatingAdd(const int8_t* input, int8_t input_zp,
+                          const int8_t* recurrent, int8_t recurrent_zp,
+                          int32_t input_effective_scale_a,
+                          int32_t input_effective_scale_b,
+                          int32_t recurrent_effective_scale_a,
+                          int32_t recurrent_effective_scale_b, int32_t n_batch,
+                          int32_t n_cell, int16_t* output);
+
+// Same as the function above, but provide a scratch buffer for the
+// int8 x int8 -> int32 and a CpuBackendContext for the accumulator
+// computation.
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors,
+    const float* __restrict__ scaling_factors, int n_batch,
+    int32_t* __restrict__ scratch, float* __restrict__ result,
+    CpuBackendContext* __restrict__ context);
+
+// Same as the function above except that can make use of cached row sums.
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float* scaling_factors,
+    int n_batch, float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
+    bool* compute_row_sums, CpuBackendContext* context);
+
+// Same as the function above, but provides separate scaling factor for the
+// matrix and the vectors. The scaling factors are multiplied in the
+// scaling_factor_scratch buffer.
+inline void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* __restrict__ matrix, const int m_rows, const int m_cols,
+    const int8_t* __restrict__ vectors, const float matrix_scaling_factor,
+    const float* vector_scaling_factors, int n_batch,
+    float* __restrict__ result, const float* per_channel_scale,
+    const int32_t* input_offset, int32_t* scratch, int32_t* row_sums,
+    bool* compute_row_sums, float* scaling_factor_scratch,
+    CpuBackendContext* context) {
+  for (int b = 0; b < n_batch; ++b) {
+    scaling_factor_scratch[b] =
+        vector_scaling_factors[b] * matrix_scaling_factor;
+  }
+  MatrixBatchVectorMultiplyAccumulate(matrix, m_rows, m_cols, vectors,
+                                      scaling_factor_scratch, n_batch, result,
+                                      per_channel_scale, input_offset, scratch,
+                                      row_sums, compute_row_sums, context);
+}
+
+// Multiplies a matrix by a "batched" vector (i.e. a matrix with a batch
+// dimension composed by input vectors independent from each other). The result
+// of the multiplication is accumulated to the passed result buffer.
+// More specifically, for a matrix M of shape [n, i] and a batched-vector
+// of shape [i, batch] it will first compute the product of shape [n, batch].
+// This product will be accumulated to the result buffer,
+// Parameters:
+//     - input: batch vector of size n_batch * n_input
+//     - bias:  vector of size b_input
+//     - input_to_gate_weights: matrix of size n_input * n_output
+//     - multiplier: scalar
+//     - shift: scalar
+//     - n_batch: the batch size
+//     - n_input: the input size
+//     - n_output: the output size
+//     - output_zp: the zero point of the output.
+//     - scratch: batch vector of size n_batch * n_output
+//     - output: the 16 bit output
+// Notes:
+//     - this is used for gate matmul: for non-cifg it is for input, forget,
+//       cell, output gates; for cifg, it is for forget, cell, output gates.
+//     - multiplier and shift combined gives the scale.
+//     - assumes input zero point is 0.
+//     - scratch is created for optimization purpose only.
+// TODO(b/152066492): this can be removed if some future optimization
+// work makes it unnecessary.
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* input, const int32_t* bias,
+    const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
+    int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
+    int32_t* scratch, int16_t* output, CpuBackendContext* context);
+
+// Multiplies a matrix by a "batched" vector (i.e. a matrix with a batch
+// dimension composed by input vectors independent from each other). The result
+// of the multiplication is accumulated to the passed result buffer.
+// More specifically, for a matrix M of shape [n, i] and a batched-vector
+// of shape [i, batch] it will first compute the product of shape [n, batch].
+// This product will be accumulated to the result buffer,
+// Parameters:
+//     - input: batch vector of size n_batch * n_input
+//     - bias:  vector of size b_input
+//     - input_to_gate_weights: matrix of size n_input * n_output
+//     - multiplier: scalar
+//     - shift: scalar
+//     - n_batch: the batch size
+//     - n_input: the input size
+//     - n_output: the output size
+//     - output_zp: the zero point of the output.
+//     - scratch: batch vector of size n_batch * n_output
+//     - output: the 8 bit output
+// Notes:
+//     - this is used for projection matmul.
+//     - multiplier and shift combined gives the scale.
+//     - assumes input zero point is 0.
+//     - scratch is created for optimization purpose only.
+// TODO(b/152066492): this can be removed if some future optimization
+// work makes it unnecessary.
+void MatrixBatchVectorMultiplyAccumulate(
+    const int8_t* input, const int32_t* bias,
+    const int8_t* input_to_gate_weights, int32_t multiplier, int32_t shift,
+    int32_t n_batch, int32_t n_input, int32_t n_output, int32_t output_zp,
+    int32_t* scratch, int8_t* output, CpuBackendContext* context);
+
+// Apply Rectified Linear to elements of a vector.
+void ApplyReluToVector(const float* __restrict__ vector, int v_size,
+                       float* __restrict__ result);
+
+// Apply Rectified Linear 1 (cap to [-1;1]) to elements of a vector
+void ApplyRelu1ToVector(const float* __restrict__ vector, int v_size,
+                        float* __restrict__ result);
+
+// Apply Rectified Linear 6 (cap to [0;6]) to elements of a vector
+void ApplyRelu6ToVector(const float* __restrict__ vector, int v_size,
+                        float* __restrict__ result);
+
+// Apply signbit to elements of a vector
+void ApplySignbitToVector(const float* __restrict__ vector, int v_size,
+                          float* __restrict__ result);
+
+// Unpack or inflate `src_buffer` by taking each element and splitting it as
+// two elements into `dst_buffer`.
+// Parameters:
+//   src_buffer   : Densely packed buffer containing int4 values
+//   num_elements : Number of elements stored in the buffer. Note that this can
+//                  be smaller than the size of `src_buffer` by 1 if it's odd,
+//                  in which case the last nibble in `src_buffer` is ignored.
+//                  This should be equal to the size of `dst_buffer`.
+//   dst_buffer   : Buffer to unpack into. Should be allocated by the caller.
+//                  Size should be at least `num_elements`.
+// Notes:
+//   For example, given `src_buffer = {0x12, 0x34};`, calling this function
+//   will return `dst_buffer = {0x02, 0x01, 0x04, 0x03}`.
+void UnpackDenseInt4IntoInt8(const int8_t* src_buffer, int num_elements,
+                             int8_t* dst_buffer);
+
+// Unpack or inflate `src_buffer` by taking each byte and splitting it into
+// multiple elements into `dst_buffer`. Supports 2-bit and 4-bit packed integers
+// Parameters:
+//   src_buffer   : Densely packed buffer containing int2 or int4 values.
+//   num_elements : Number of unpacked elements to be read from the buffer.
+//                  This should be equal to the size of `dst_buffer`.
+//   bit_width    : The bit width of the packed elements (either 2 or 4).
+//   dst_buffer   : Buffer to unpack into. Should be allocated by the caller.
+//                  Size should be at least `num_elements`.
+// Notes:
+//   For 4-bit unpacking: e.g., `src_buffer = {0x12, 0x34};` (num_elements = 4)
+//   will return `dst_buffer = {0x02, 0x01, 0x04, 0x03}`.
+//   For 2-bit unpacking: e.g., `src_buffer = {0x12};` (num_elements = 4)
+//   will return `dst_buffer = {0x02, 0x00, 0x01, 0x00}` (sign extended).
+void UnpackPackedIntToInt8(const int8_t* src_buffer, int num_elements,
+                           int bit_width, int8_t* dst_buffer,
+                           bool unpack_unsigned = false);
+
+// Pack `src_buffer` into a densely packed buffer of int2 or int4 values.
+// Parameters:
+//   src_buffer   : Buffer containing int2 or int4 values stored in int8
+//                  memory.
+//   num_elements : Number of elements stored in the buffer. Note that this can
+//                  be smaller than the size of `src_buffer` by 1 if it's odd,
+//                  in which case the last nibble in `src_buffer` is ignored.
+//                  This should be equal to the size of `dst_buffer`.
+//   bit_width    : The bit width of the packed elements (either 2 or 4).
+//   dst_buffer   : Buffer to pack into. Should be allocated by the caller.
+//                  Size should be at least `num_elements`.
+// Notes:
+//   For 4-bit packing: e.g., given `src_buffer = {0x02, 0x01, 0x04, 0x03}`,
+//   calling this function will return `dst_buffer = {0x12, 0x34}`.
+//   For 2-bit packing: e.g., given `src_buffer = {0x00, 0x01, 0x00, 0x02}`,
+//   calling this function will return `dst_buffer = {0x84}`.
+void PackInt8IntoDenseInt(const int8_t* src_buffer, int num_elements,
+                          int bit_width, int8_t* dst_buffer);
+}  // namespace tensor_utils
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_PORTABLE_TENSOR_UTILS_H_

tensorflow/lite/kernels/internal/quantization_util.cc

@@ -0,0 +1,416 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/cppmath.h"
+
+namespace tflite {
+
+namespace {
+// These constants are used to manipulate the binary representation of doubles.
+// Double-precision binary64 floating point format is:
+// Bit |  63  |  62-52   |   51-0   |
+//     | Sign | Exponent | Fraction |
+// To avoid 64-bit integers as much as possible, I break this into high and
+// low 32-bit chunks. High is:
+// Bit |  31  |  30-20   |      19-0     |
+//     | Sign | Exponent | High Fraction |
+// Low is:
+// Bit |     31-0     |
+//     | Low Fraction |
+// We then access the components through logical bit-wise operations to
+// extract the parts needed, with the positions and masks derived from the
+// layout shown above.
+constexpr uint64_t kSignMask = 0x8000000000000000LL;
+constexpr uint64_t kExponentMask = 0x7ff0000000000000LL;
+constexpr int32_t kExponentShift = 52;
+constexpr int32_t kExponentBias = 1023;
+constexpr uint32_t kExponentIsBadNum = 0x7ff;
+constexpr uint64_t kFractionMask = 0x000fffffffc00000LL;
+constexpr uint32_t kFractionShift = 22;
+constexpr uint32_t kFractionRoundingMask = 0x003fffff;
+constexpr uint32_t kFractionRoundingThreshold = 0x00200000;
+}  // namespace
+
+void QuantizeMultiplier(double double_multiplier, int32_t* quantized_multiplier,
+                        int* shift) {
+#if TFLITE_SINGLE_ROUNDING
+  // Single-rounding MultiplyByQuantizedMultiplier only supports positive
+  // multipliers.
+  // TFLITE_DCHECK(double_multiplier >= 0);
+#endif
+  if (double_multiplier == 0.) {
+    *quantized_multiplier = 0;
+    *shift = 0;
+    return;
+  }
+#ifdef TFLITE_EMULATE_FLOAT
+  // If we're trying to avoid the use of floating-point instructions (for
+  // example on microcontrollers) then use an alternative implementation
+  // that only requires integer and bitwise operations. To enable this, you
+  // need to set the define during the build process for your platform.
+  int64_t q_fixed = IntegerFrExp(double_multiplier, shift);
+#else   // TFLITE_EMULATE_FLOAT
+  const double q = std::frexp(double_multiplier, shift);
+  auto q_fixed = static_cast<int64_t>(TfLiteRound(q * (1LL << 31)));
+#endif  // TFLITE_EMULATE_FLOAT
+  TFLITE_CHECK(q_fixed <= (1LL << 31));
+  if (q_fixed == (1LL << 31)) {
+    q_fixed /= 2;
+    ++*shift;
+  }
+  TFLITE_CHECK_LE(q_fixed, std::numeric_limits<int32_t>::max());
+  // A shift amount smaller than -31 would cause all bits to be shifted out
+  // and thus all results would be zero. We implement that instead with
+  // q_fixed==0, so as to avoid hitting issues with right-shift
+  // operations with shift amounts greater than 31. Note that this happens
+  // roughly when abs(double_multiplier) < 2^-31 and the present handling means
+  // that we're effectively flushing tiny double_multiplier's to zero.
+  // We could conceivably handle values in the range (roughly) [32, 63]
+  // as 'denormals' i.e. (shift==0, q_fixed < 2^30). In that point of view
+  // the present handling is just doing 'flush denormals to zero'. We could
+  // reconsider and actually generate nonzero denormals if a need arises.
+  if (*shift < -31) {
+    *shift = 0;
+    q_fixed = 0;
+  }
+#if TFLITE_SINGLE_ROUNDING
+  // Single-rounding MultiplyByQuantizedMultiplier doesn't support a shift > 30,
+  // saturate it.
+  if (*shift > 30) {
+    *shift = 30;
+    q_fixed = (1LL << 31) - 1;
+  }
+#endif
+  *quantized_multiplier = static_cast<int32_t>(q_fixed);
+}
+
+void QuantizeMultiplierGreaterThanOne(double double_multiplier,
+                                      int32_t* quantized_multiplier,
+                                      int* left_shift) {
+  TFLITE_CHECK_GT(double_multiplier, 1.);
+  QuantizeMultiplier(double_multiplier, quantized_multiplier, left_shift);
+  TFLITE_CHECK_GE(*left_shift, 0);
+}
+
+void QuantizeMultiplierSmallerThanOneExp(double double_multiplier,
+                                         int32_t* quantized_multiplier,
+                                         int* left_shift) {
+  TFLITE_CHECK_LT(double_multiplier, 1.);
+  TFLITE_CHECK_GT(double_multiplier, 0.);
+  int shift;
+  QuantizeMultiplier(double_multiplier, quantized_multiplier, &shift);
+  TFLITE_CHECK_LE(shift, 0);
+  *left_shift = shift;
+}
+
+int64_t IntegerFrExp(double input, int* shift) {
+  // Make sure our assumptions about the double layout hold.
+  TFLITE_CHECK_EQ(8, sizeof(double));
+
+  // We want to access the bits of the input double value directly, which is
+  // tricky to do safely, so use a union to handle the casting.
+  union {
+    double double_value;
+    uint64_t double_as_uint;
+  } cast_union;
+  cast_union.double_value = input;
+  const uint64_t u = cast_union.double_as_uint;
+
+  // If the bitfield is all zeros apart from the sign bit, this is a normalized
+  // zero value, so return standard values for this special case.
+  if ((u & ~kSignMask) == 0) {
+    *shift = 0;
+    return 0;
+  }
+
+  // Deal with NaNs and Infs, which are always indicated with a fixed pattern in
+  // the exponent, and distinguished by whether the fractions are zero or
+  // non-zero.
+  const uint32_t exponent_part = ((u & kExponentMask) >> kExponentShift);
+  if (exponent_part == kExponentIsBadNum) {
+    *shift = std::numeric_limits<int>::max();
+    if (u & kFractionMask) {
+      // NaN, so just return zero (with the exponent set to INT_MAX).
+      return 0;
+    } else {
+      // Infinity, so return +/- INT_MAX.
+      if (u & kSignMask) {
+        return std::numeric_limits<int64_t>::min();
+      } else {
+        return std::numeric_limits<int64_t>::max();
+      }
+    }
+  }
+
+  // The shift is fairly easy to extract from the high bits of the double value,
+  // just by masking it out and applying a bias. The std::frexp() implementation
+  // always returns values between 0.5 and 1.0 though, whereas the exponent
+  // assumes 1.0 to 2.0 is the standard range, so I add on one to match that
+  // interface.
+  *shift = (exponent_part - kExponentBias) + 1;
+
+  // There's an implicit high bit in the double format definition, so make sure
+  // we include that at the top, and then reconstruct the rest of the fractional
+  // value from the remaining fragments.
+  int64_t fraction = 0x40000000 + ((u & kFractionMask) >> kFractionShift);
+
+  // We're cutting off some bits at the bottom, so to exactly match the standard
+  // frexp implementation here we'll apply rounding by adding one to the least
+  // significant bit of the result if the discarded portion is over half of the
+  // maximum.
+  if ((u & kFractionRoundingMask) > kFractionRoundingThreshold) {
+    fraction += 1;
+  }
+  // Negate the fraction if the sign bit was set.
+  if (u & kSignMask) {
+    fraction *= -1;
+  }
+
+  return fraction;
+}
+
+double DoubleFromFractionAndShift(int64_t fraction, int shift) {
+  union {
+    double double_value;
+    uint64_t double_as_uint;
+  } result;
+
+  // Detect NaNs and infinities.
+  if (shift == std::numeric_limits<int>::max()) {
+    if (fraction == 0) {
+      return std::numeric_limits<double>::quiet_NaN();
+    } else if (fraction > 0) {
+      return std::numeric_limits<double>::infinity();
+    } else {
+      return -std::numeric_limits<double>::infinity();
+    }
+  }
+
+  // Return a normalized zero for a zero fraction.
+  if (fraction == 0) {
+    result.double_as_uint = 0;
+    return result.double_value;
+  }
+
+  bool is_negative = (fraction < 0);
+  int64_t encoded_fraction = is_negative ? -fraction : fraction;
+  int64_t encoded_shift = (shift - 1);
+  while (encoded_fraction < 0x40000000) {
+    encoded_fraction *= 2;
+    encoded_shift -= 1;
+  }
+  while (encoded_fraction > 0x80000000) {
+    encoded_fraction /= 2;
+    encoded_shift += 1;
+  }
+  encoded_fraction -= 0x40000000;
+  if (encoded_shift < -1022) {
+    encoded_shift = -1023;
+  } else if (encoded_shift > 1022) {
+    encoded_shift = 1023;
+  }
+  encoded_shift += kExponentBias;
+  uint64_t encoded_sign = is_negative ? kSignMask : 0;
+  result.double_as_uint = encoded_sign | (encoded_shift << kExponentShift) |
+                          (encoded_fraction << kFractionShift);
+  return result.double_value;
+}
+
+double IntegerDoubleMultiply(double a, double b) {
+  int a_shift;
+  const int64_t a_fraction = IntegerFrExp(a, &a_shift);
+  int b_shift;
+  const int64_t b_fraction = IntegerFrExp(b, &b_shift);
+  // Detect NaNs and infinities.
+  if (a_shift == std::numeric_limits<int>::max() ||
+      (b_shift == std::numeric_limits<int>::max())) {
+    return std::numeric_limits<double>::quiet_NaN();
+  }
+  const int result_shift = a_shift + b_shift + 1;
+  const int64_t result_fraction = (a_fraction * b_fraction) >> 32;
+  return DoubleFromFractionAndShift(result_fraction, result_shift);
+}
+
+int IntegerDoubleCompare(double a, double b) {
+  int a_shift;
+  const int64_t a_fraction = IntegerFrExp(a, &a_shift);
+  int b_shift;
+  const int64_t b_fraction = IntegerFrExp(b, &b_shift);
+
+  // Detect NaNs and infinities.
+  if (a_shift == std::numeric_limits<int>::max() ||
+      (b_shift == std::numeric_limits<int>::max())) {
+    return 1;
+  }
+
+  if ((a_fraction == 0) && (b_fraction < 0)) {
+    return 1;
+  } else if ((a_fraction < 0) && (b_fraction == 0)) {
+    return -1;
+  } else if (a_shift < b_shift) {
+    return -1;
+  } else if (a_shift > b_shift) {
+    return 1;
+  } else if (a_fraction < b_fraction) {
+    return -1;
+  } else if (a_fraction > b_fraction) {
+    return 1;
+  } else {
+    return 0;
+  }
+}
+
+void PreprocessSoftmaxScaling(double beta, double input_scale,
+                              int input_integer_bits,
+                              int32_t* quantized_multiplier, int* left_shift) {
+  // If the overall multiplier (input and beta) is large, then exp() of an
+  // input difference of 1 scaled by this will be large.  In other words, we
+  // can cap the multiplier and know that, when it is used, the output will be
+  // (round to) zero wherever the input is not at the maximum value.
+
+  // If the overall scale is less than one, and input_integer_bits=0, then the
+  // result is double equivalent of Q0.31 (actually with more precision). Thus
+  // this generates a Q(input_integer_bits).(31-input_integer_bits)
+  // representation.
+#if TFLITE_SINGLE_ROUNDING
+  const double max_real_multiplier = (1LL << 30) - 1.0;
+#else
+  const double max_real_multiplier = (1LL << 31) - 1.0;
+#endif
+
+#ifdef TFLITE_EMULATE_FLOAT
+  const double input_beta = IntegerDoubleMultiply(beta, input_scale);
+  int shift;
+  int64_t fraction = IntegerFrExp(input_beta, &shift);
+  shift += (31 - input_integer_bits);
+  double input_beta_real_multiplier =
+      DoubleFromFractionAndShift(fraction, shift);
+  if (IntegerDoubleCompare(input_beta_real_multiplier, max_real_multiplier) >
+      0) {
+    input_beta_real_multiplier = max_real_multiplier;
+  }
+#else   // TFLITE_EMULATE_FLOAT
+  const double input_beta_real_multiplier =
+      std::min<double>(beta * input_scale * (1 << (31 - input_integer_bits)),
+                       max_real_multiplier);
+#endif  // TFLITE_EMULATE_FLOAT
+
+  QuantizeMultiplierGreaterThanOne(input_beta_real_multiplier,
+                                   quantized_multiplier, left_shift);
+}
+
+void PreprocessLogSoftmaxScalingExp(double beta, double input_scale,
+                                    int input_integer_bits,
+                                    int32_t* quantized_multiplier,
+                                    int* left_shift,
+                                    int32_t* reverse_scaling_divisor,
+                                    int* reverse_scaling_left_shift) {
+  PreprocessSoftmaxScaling(beta, input_scale, input_integer_bits,
+                           quantized_multiplier, left_shift);
+
+  // Also calculate what amounts to the inverse scaling factor for the input.
+  const double real_reverse_scaling_divisor =
+      (1 << (31 - *left_shift)) / static_cast<double>(*quantized_multiplier);
+  tflite::QuantizeMultiplierSmallerThanOneExp(real_reverse_scaling_divisor,
+                                              reverse_scaling_divisor,
+                                              reverse_scaling_left_shift);
+}
+
+int CalculateInputRadius(int input_integer_bits, int input_left_shift,
+                         int total_signed_bits) {
+#ifdef TFLITE_EMULATE_FLOAT
+  int64_t result = (1 << input_integer_bits) - 1;
+  result <<= (total_signed_bits - input_integer_bits);
+  result >>= input_left_shift;
+  return result;
+#else   // TFLITE_EMULATE_FLOAT
+  const double max_input_rescaled =
+      1.0 * ((1 << input_integer_bits) - 1) *
+      (1LL << (total_signed_bits - input_integer_bits)) /
+      (1LL << input_left_shift);
+  // Tighten bound using floor.  Suppose that we could use the exact value.
+  // After scaling the difference, the result would be at the maximum.  Thus we
+  // must ensure that our value has lower magnitude.
+  return static_cast<int>(std::floor(max_input_rescaled));
+#endif  // TFLITE_EMULATE_FLOAT
+}
+
+void NudgeQuantizationRange(const float min, const float max,
+                            const int quant_min, const int quant_max,
+                            float* nudged_min, float* nudged_max,
+                            float* nudged_scale) {
+  // This code originates from tensorflow/core/kernels/fake_quant_ops_functor.h.
+  const float quant_min_float = static_cast<float>(quant_min);
+  const float quant_max_float = static_cast<float>(quant_max);
+  *nudged_scale = (max - min) / (quant_max_float - quant_min_float);
+  const float zero_point_from_min = quant_min_float - min / *nudged_scale;
+  uint16_t nudged_zero_point;
+  if (zero_point_from_min < quant_min_float) {
+    nudged_zero_point = static_cast<uint16_t>(quant_min);
+  } else if (zero_point_from_min > quant_max_float) {
+    nudged_zero_point = static_cast<uint16_t>(quant_max);
+  } else {
+    nudged_zero_point = static_cast<uint16_t>(TfLiteRound(zero_point_from_min));
+  }
+  *nudged_min = (quant_min_float - nudged_zero_point) * (*nudged_scale);
+  *nudged_max = (quant_max_float - nudged_zero_point) * (*nudged_scale);
+}
+
+void FakeQuantizeArray(const float nudged_scale, const float nudged_min,
+                       const float nudged_max, const float* input_data,
+                       float* output_data, const float size) {
+  // This code originates from tensorflow/core/kernels/fake_quant_ops_functor.h.
+  const float inv_nudged_scale = 1.0f / nudged_scale;
+
+  for (int i = 0; i < size; i++) {
+    const float src_val = input_data[i];
+    const float clamped = std::min(nudged_max, std::max(nudged_min, src_val));
+    const float clamped_shifted = clamped - nudged_min;
+    const float dst_val =
+        TfLiteRound(clamped_shifted * inv_nudged_scale) * nudged_scale +
+        nudged_min;
+    output_data[i] = dst_val;
+  }
+}
+
+bool CheckedLog2(const float x, int* log2_result) {
+  // Using TfLiteRound instead of std::round and std::log instead of
+  // std::log2 to work around these functions being missing in a toolchain
+  // used in some TensorFlow tests as of May 2018.
+  const float x_log2 = std::log(x) * (1.0f / std::log(2.0f));
+  const float x_log2_rounded = TfLiteRound(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-3f;
+}
+
+void QuantizeMultiplierArray(const double* effective_scales, size_t size,
+                             int32_t* effective_scale_significand,
+                             int* effective_shift) {
+  for (size_t i = 0; i < size; ++i) {
+    QuantizeMultiplier(effective_scales[i], &effective_scale_significand[i],
+                       &effective_shift[i]);
+  }
+}
+
+}  // namespace tflite

tensorflow/lite/kernels/internal/quantization_util.h

@@ -0,0 +1,293 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_
+
+#include <cmath>
+#include <cstdint>
+#include <limits>
+
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/types.h"
+
+namespace tflite {
+
+// Given the min and max values of a float array, return
+// reasonable quantization parameters to use for this array.
+template <typename T>
+QuantizationParams ChooseQuantizationParams(double rmin, double rmax,
+                                            bool narrow_range) {
+  const T qmin = std::numeric_limits<T>::min() + (narrow_range ? 1 : 0);
+  const T qmax = std::numeric_limits<T>::max();
+  const double qmin_double = qmin;
+  const double qmax_double = qmax;
+  // 0 should always be a representable value. Let's assume that the initial
+  // min,max range contains 0.
+  TFLITE_CHECK_LE(rmin, 0.);
+  TFLITE_CHECK_GE(rmax, 0.);
+  if (rmin == rmax) {
+    // Special case where the min,max range is a point. Should be {0}.
+    TFLITE_CHECK_EQ(rmin, 0.);
+    TFLITE_CHECK_EQ(rmax, 0.);
+    QuantizationParams quantization_params;
+    quantization_params.zero_point = 0;
+    quantization_params.scale = 0.;
+    return quantization_params;
+  }
+
+  // General case.
+  //
+  // First determine the scale.
+  const double scale = (rmax - rmin) / (qmax_double - qmin_double);
+
+  // Zero-point computation.
+  // First the initial floating-point computation. The zero-point can be
+  // determined from solving an affine equation for any known pair
+  // (real value, corresponding quantized value).
+  // We know two such pairs: (rmin, qmin) and (rmax, qmax).
+  // The arithmetic error on the zero point computed from either pair
+  // will be roughly machine_epsilon * (sum of absolute values of terms)
+  // so we want to use the variant that adds the smaller terms.
+  const double zero_point_from_min = qmin_double - rmin / scale;
+  const double zero_point_from_max = qmax_double - rmax / scale;
+  const double zero_point_from_min_error =
+      std::abs(qmin_double) + std::abs(rmin / scale);
+  const double zero_point_from_max_error =
+      std::abs(qmax_double) + std::abs(rmax / scale);
+
+  const double zero_point_double =
+      zero_point_from_min_error < zero_point_from_max_error
+          ? zero_point_from_min
+          : zero_point_from_max;
+
+  // Now we need to nudge the zero point to be an integer
+  // (our zero points are integer, and this is motivated by the requirement
+  // to be able to represent the real value "0" exactly as a quantized value,
+  // which is required in multiple places, for example in Im2col with SAME
+  // padding).
+  T nudged_zero_point = 0;
+  if (zero_point_double < qmin_double) {
+    nudged_zero_point = qmin;
+  } else if (zero_point_double > qmax_double) {
+    nudged_zero_point = qmax;
+  } else {
+    nudged_zero_point = static_cast<T>(round(zero_point_double));
+  }
+  // The zero point should always be in the range of quantized value,
+  // [qmin, qmax].
+  TFLITE_CHECK_GE(nudged_zero_point, qmin);
+  TFLITE_CHECK_LE(nudged_zero_point, qmax);
+
+  // Finally, store the result nudged quantization params.
+  QuantizationParams quantization_params;
+  quantization_params.zero_point = nudged_zero_point;
+  quantization_params.scale = scale;
+  return quantization_params;
+}
+
+template <typename T>
+QuantizationParams ChooseQuantizationParams(double rmin, double rmax) {
+  return ChooseQuantizationParams<T>(rmin, rmax, false);
+}
+
+// LINT.IfChange
+// Converts a floating-point number to an integer. For all inputs x where
+// static_cast<IntOut>(x) is legal according to the C++ standard, the result
+// is identical to that cast (i.e. the result is x with its fractional part
+// truncated whenever that is representable as IntOut).
+//
+// static_cast would cause undefined behavior for the following cases, which
+// have well-defined behavior for this function:
+//
+//  1. If x is NaN, the result is zero.
+//
+//  2. If the truncated form of x is above the representable range of IntOut,
+//     the result is std::numeric_limits<IntOut>::max().
+//
+//  3. If the truncated form of x is below the representable range of IntOut,
+//     the result is std::numeric_limits<IntOut>::min().
+//
+// Note that cases #2 and #3 cover infinities as well as finite numbers.
+//
+// The range of FloatIn must include the range of IntOut, otherwise
+// the results are undefined.
+// TODO(sfeuz): Replace by absl::SafeCast once available.
+template <class IntOut, class FloatIn>
+IntOut SafeCast(FloatIn x) {
+  static_assert(!std::numeric_limits<FloatIn>::is_integer,
+                "FloatIn is integer");
+  static_assert(std::numeric_limits<IntOut>::is_integer,
+                "IntOut is not integer");
+  static_assert(std::numeric_limits<IntOut>::radix == 2, "IntOut is base 2");
+
+  // Special case NaN, for which the logic below doesn't work.
+  if (std::isnan(x)) {
+    return 0;
+  }
+
+  // Negative values all clip to zero for unsigned results.
+  if (!std::numeric_limits<IntOut>::is_signed && x < 0) {
+    return 0;
+  }
+
+  // Handle infinities.
+  if (std::isinf(x)) {
+    return x < 0 ? std::numeric_limits<IntOut>::min()
+                 : std::numeric_limits<IntOut>::max();
+  }
+
+  // Set exp such that x == f * 2^exp for some f with |f| in [0.5, 1.0),
+  // unless x is zero in which case exp == 0. Note that this implies that the
+  // magnitude of x is strictly less than 2^exp.
+  int exp = 0;
+  std::frexp(x, &exp);
+
+  // Let N be the number of non-sign bits in the representation of IntOut. If
+  // the magnitude of x is strictly less than 2^N, the truncated version of x
+  // is representable as IntOut. The only representable integer for which this
+  // is not the case is kMin for signed types (i.e. -2^N), but that is covered
+  // by the fall-through below.
+  if (exp <= std::numeric_limits<IntOut>::digits) {
+    return x;
+  }
+
+  // Handle numbers with magnitude >= 2^N.
+  return x < 0 ? std::numeric_limits<IntOut>::min()
+               : std::numeric_limits<IntOut>::max();
+}
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/kernels/internal/quantization_util.h)
+
+// Decompose a double multiplier into a Q0.31 int32 representation of its
+// significand, and shift representation of NEGATIVE its exponent ---
+// this is intended as a RIGHT-shift.
+//
+// Restricted to the case where the multiplier < 1 (and non-negative).
+void QuantizeMultiplierSmallerThanOneExp(double double_multiplier,
+                                         int32_t* quantized_multiplier,
+                                         int* left_shift);
+
+// Decompose a double multiplier into a Q0.31 int32 representation of its
+// significand, and shift representation of its exponent.
+//
+// Restricted to the case where the multiplier > 1.
+void QuantizeMultiplierGreaterThanOne(double double_multiplier,
+                                      int32_t* quantized_multiplier,
+                                      int* left_shift);
+
+// Decompose a double multiplier into a Q0.31 int32 representation of its
+// significand, and shift representation of its exponent.
+//
+// Handles an arbitrary positive multiplier. The 'shift' output-value is
+// basically the 'floating-point exponent' of the multiplier:
+// Negative for a right-shift (when the multiplier is <1), positive for a
+// left-shift (when the multiplier is >1)
+void QuantizeMultiplier(double double_multiplier, int32_t* quantized_multiplier,
+                        int* shift);
+
+// Splits a double input value into a returned fraction, and a shift value from
+// the exponent, using only bitwise and integer operations to support
+// microcontrollers and other environments without floating-point support.
+//
+// This is designed to be a replacement for how std::frexp() is used within the
+// QuantizeMultiplier() function, and so has a different signature than the
+// standard version, returning a 64-bit integer rather than a double. This
+// result has a maximum value of 1<<31, with the fraction expressed as a
+// proportion of that maximum.
+//
+// std::frexp() returns NaNs and infinities unmodified, but since we're
+// returning integers that can't represent those values, instead we return
+// a shift of std::numeric_limits<int>::max() for all bad numbers, with an int64
+// result of 0 for NaNs, std:numeric_limits<int64_t>::max() for +INFINITY, and
+// std::numeric_limits<int64_t>::min() for -INFINITY. Denormalized inputs will
+// result in return values that end up truncating some bits at the end,
+// reflecting the loss of precision inherent in denormalization.
+int64_t IntegerFrExp(double input, int* shift);
+
+// Converts an integer fraction in the format produced by IntegerFrExp (where
+// 0x40000000 is 1.0) and an exponent shift (between -1022 and +1022) into an
+// IEEE binary64 double format result. The implementation uses only integer and
+// bitwise operators, so no floating point hardware support or emulation is
+// needed. This is here so quantized operations can run non-time-critical
+// preparation calculations on microcontrollers and other platforms without
+// float support.
+double DoubleFromFractionAndShift(int64_t fraction, int shift);
+
+// Performs a multiplication of two numbers in double format, using only integer
+// and bitwise instructions. This is aimed at supporting housekeeping functions
+// for quantized operations on microcontrollers without floating-point hardware.
+double IntegerDoubleMultiply(double a, double b);
+
+// Returns -1 if a is less than b, 0 if a and b are equal, and +1 if a is
+// greater than b. It is implemented using only integer and logical instructions
+// so that it can be easily run on microcontrollers for quantized operations.
+int IntegerDoubleCompare(double a, double b);
+
+// This first creates a multiplier in a double equivalent of
+// Q(input_integer_bits).(31-input_integer_bits) representation, with extra
+// precision in the double's fractional bits.  It then splits the result into
+// significand and exponent.
+void PreprocessSoftmaxScaling(double beta, double input_scale,
+                              int input_integer_bits,
+                              int32_t* quantized_multiplier, int* left_shift);
+// Like PreprocessSoftmaxScaling, but inverse scaling factors also calculated.
+void PreprocessLogSoftmaxScalingExp(double beta, double input_scale,
+                                    int input_integer_bits,
+                                    int32_t* quantized_multiplier,
+                                    int* left_shift,
+                                    int32_t* reverse_scaling_divisor,
+                                    int* reverse_scaling_left_shift);
+// Calculate the largest input that will result in a within-bounds intermediate
+// result within MultiplyByQuantizedMultiplierGreaterThanOne.  In other words,
+// it must not overflow before we reduce the value by multiplication by the
+// input multiplier.  The negative radius is used as the minimum difference in
+// Softmax.
+int CalculateInputRadius(int input_integer_bits, int input_left_shift,
+                         int total_signed_bits = 31);
+
+// Nudges a min/max quantization range to ensure zero is zero.
+// Gymnastics with nudged zero point is to ensure that real zero maps to
+// an integer, which is required for e.g. zero-padding in convolutional layers.
+// Outputs nudged_min, nudged_max, nudged_scale.
+void NudgeQuantizationRange(const float min, const float max,
+                            const int quant_min, const int quant_max,
+                            float* nudged_min, float* nudged_max,
+                            float* nudged_scale);
+
+// Fake quantizes (quantizes and dequantizes) input_data using the scale,
+// nudged_min, and nudged_max from NudgeQuantizationRange. This matches the code
+// in TensorFlow's FakeQuantizeWithMinMaxVarsFunctor.
+void FakeQuantizeArray(const float nudged_scale, const float nudged_min,
+                       const float nudged_max, const float* input_data,
+                       float* output_data, const float size);
+
+// If x is approximately a power of two (with any positive or negative
+// exponent), stores that exponent (i.e. log2(x)) in *log2_result, otherwise
+// returns false.
+bool CheckedLog2(const float x, int* log2_result);
+
+// Decomposes an array of double multipliers into a Q0.31 int32 representation
+// of its significand, and shift representation of its exponent.
+//
+// Handles an arbitrary multiplier. The 'shift' output-value is
+// basically the 'floating-point exponent' of the multiplier:
+// Negative for a right-shift (when the multiplier is <1), positive for a
+// left-shift (when the multiplier is >1)
+void QuantizeMultiplierArray(const double* effective_scales, size_t size,
+                             int32_t* effective_scale_significand,
+                             int* effective_shift);
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_

tensorflow/lite/kernels/internal/quantization_util_test.cc

@@ -0,0 +1,579 @@
+/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/quantization_util.h"
+
+#include <limits>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/lite/kernels/internal/common.h"
+
+namespace tflite {
+namespace {
+
+using ::testing::ElementsAreArray;
+using ::testing::Pair;
+
+template <class FloatIn, class IntOut>
+void RunSafeCastTests() {
+  const IntOut imax = std::numeric_limits<IntOut>::max();
+  EXPECT_GT(imax, 0);
+  const IntOut imin = std::numeric_limits<IntOut>::min();
+  const bool s = std::numeric_limits<IntOut>::is_signed;
+  if (s) {
+    EXPECT_LT(imin, 0);
+  } else {
+    EXPECT_EQ(0, imin);
+  }
+
+  // Some basic tests.
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(0.0)), 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-0.0)), 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(0.99)), 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(1.0)), 1);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(1.01)), 1);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(1.99)), 1);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(2.0)), 2);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(2.01)), 2);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-0.99)), 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-1.0)), s ? -1 : 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-1.01)), s ? -1 : 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-1.99)), s ? -1 : 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-2.0)), s ? -2 : 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-2.01)), s ? -2 : 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(117.9)), 117);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(118.0)), 118);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(118.1)), 118);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-117.9)), s ? -117 : 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-118.0)), s ? -118 : 0);
+  EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(-118.1)), s ? -118 : 0);
+
+  // Some edge cases.
+  EXPECT_EQ(SafeCast<IntOut>(std::numeric_limits<FloatIn>::max()), imax);
+  EXPECT_EQ(SafeCast<IntOut>(std::numeric_limits<FloatIn>::lowest()), imin);
+  EXPECT_EQ(SafeCast<IntOut>(std::numeric_limits<FloatIn>::infinity()), imax);
+  EXPECT_EQ(SafeCast<IntOut>(-std::numeric_limits<FloatIn>::infinity()), imin);
+  EXPECT_EQ(SafeCast<IntOut>(std::numeric_limits<FloatIn>::quiet_NaN()), 0);
+
+  // Some larger numbers.
+  if (sizeof(IntOut) >= 4 && sizeof(FloatIn) > 4) {
+    EXPECT_EQ(SafeCast<IntOut>(static_cast<FloatIn>(0x76543210)), 0x76543210);
+  }
+
+  if (sizeof(FloatIn) > sizeof(IntOut)) {
+    // Check values near imax.
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) + 0.1)),
+              imax);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) + 0.99)),
+              imax);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) + 1.0)),
+              imax);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) + 1.99)),
+              imax);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) + 2.0)),
+              imax);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) - 0.1)),
+              imax - 1);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) - 0.99)),
+              imax - 1);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) - 1.0)),
+              imax - 1);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) - 1.01)),
+              imax - 2);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) - 1.99)),
+              imax - 2);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) - 2.0)),
+              imax - 2);
+    EXPECT_EQ(SafeCast<IntOut>(
+                  static_cast<FloatIn>(static_cast<FloatIn>(imax) - 2.01)),
+              imax - 3);
+  }
+
+  // Check values considerably larger in magnitude than imin and imax
+  EXPECT_EQ(
+      SafeCast<IntOut>(static_cast<FloatIn>(static_cast<FloatIn>(imax) * 2)),
+      imax);
+  EXPECT_EQ(
+      SafeCast<IntOut>(static_cast<FloatIn>(static_cast<FloatIn>(imax) * 20)),
+      imax);
+  EXPECT_EQ(
+      SafeCast<IntOut>(static_cast<FloatIn>(static_cast<FloatIn>(imax) * 100)),
+      imax);
+  EXPECT_EQ(
+      SafeCast<IntOut>(static_cast<FloatIn>(static_cast<FloatIn>(imin) * 2)),
+      imin);
+  EXPECT_EQ(
+      SafeCast<IntOut>(static_cast<FloatIn>(static_cast<FloatIn>(imin) * 20)),
+      imin);
+  EXPECT_EQ(
+      SafeCast<IntOut>(static_cast<FloatIn>(static_cast<FloatIn>(imin) * 100)),
+      imin);
+}
+
+TEST(QuantizationUtilTest, SafeCast) {
+  RunSafeCastTests<float, int8_t>();
+  RunSafeCastTests<double, int8_t>();
+  RunSafeCastTests<float, int16_t>();
+  RunSafeCastTests<double, int16_t>();
+  RunSafeCastTests<float, int32_t>();
+  RunSafeCastTests<double, int32_t>();
+  RunSafeCastTests<float, int64_t>();
+  RunSafeCastTests<double, int64_t>();
+  RunSafeCastTests<float, uint8_t>();
+  RunSafeCastTests<double, uint8_t>();
+  RunSafeCastTests<float, uint16_t>();
+  RunSafeCastTests<double, uint16_t>();
+  RunSafeCastTests<float, uint32_t>();
+  RunSafeCastTests<double, uint32_t>();
+  RunSafeCastTests<float, uint64_t>();
+  RunSafeCastTests<double, uint64_t>();
+}
+
+// Example taken from http://www.tensorflow.org/performance/quantization
+//
+//  Quantized | Float
+//  --------- | -----
+//  0         | -10.0
+//  255       | 30.0
+//  128       | 10.0
+TEST(QuantizationUtilTest, ChooseQuantizationParams) {
+  QuantizationParams qp = ChooseQuantizationParams<uint8_t>(-10.0, 30.0);
+  EXPECT_NEAR(qp.scale, 0.156863, 1e-5);
+  EXPECT_EQ(qp.zero_point, 64);
+}
+
+TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroPointOnMinBoundary) {
+  QuantizationParams qp = ChooseQuantizationParams<uint8_t>(0.0, 30.0);
+  EXPECT_NEAR(qp.scale, 0.117647, 1e-5);
+  EXPECT_EQ(qp.zero_point, 0);
+}
+
+#if GTEST_HAS_DEATH_TEST
+TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroNotInRange) {
+  // Assumption is that zero is within the range.
+  EXPECT_DEATH(ChooseQuantizationParams<uint8_t>(10.0, 30.0), "");
+}
+
+TEST(QuantizationUtilTest, ChooseQuantizationParamsEmptyRangePositive) {
+  // Assumption is that zero is within the range.
+  EXPECT_DEATH(ChooseQuantizationParams<uint8_t>(30.0, 30.0), "");
+}
+#endif  // GTEST_HAS_DEATH_TEST
+
+TEST(QuantizationUtilTest, ChooseQuantizationParamsEmptyRangeZero) {
+  QuantizationParams qp = ChooseQuantizationParams<uint8_t>(0.0, 0.0);
+  EXPECT_NEAR(qp.scale, 0.0, 1e-5);
+  EXPECT_EQ(qp.zero_point, 0);
+}
+
+TEST(QuantizationUtilTest, ChooseQuantizationParamsZeroPointOnMaxBoundary) {
+  QuantizationParams qp = ChooseQuantizationParams<uint8_t>(-10.0, 0.0);
+  EXPECT_NEAR(qp.scale, 0.039216, 1e-5);
+  EXPECT_EQ(qp.zero_point, 255);
+}
+
+TEST(QuantizationUtilTest, IntegerFrExp) {
+  int shift;
+  int64_t result = IntegerFrExp(0.0, &shift);
+  EXPECT_EQ(0, result);
+  EXPECT_EQ(0, shift);
+
+  result = IntegerFrExp(1.0, &shift);
+  EXPECT_NEAR(0x40000000, result, 1);
+  EXPECT_EQ(1, shift);
+
+  result = IntegerFrExp(0.25, &shift);
+  EXPECT_NEAR(0x40000000, result, 1);
+  EXPECT_EQ(-1, shift);
+
+  result = IntegerFrExp(-1.0, &shift);
+  EXPECT_NEAR(-(1 << 30), result, 1);
+  EXPECT_EQ(1, shift);
+
+  result = IntegerFrExp(123.45, &shift);
+  EXPECT_NEAR(2071147315, result, 1);
+  EXPECT_EQ(7, shift);
+
+  result = IntegerFrExp(NAN, &shift);
+  EXPECT_NEAR(0, result, 1);
+  EXPECT_EQ(0x7fffffff, shift);
+
+  result = IntegerFrExp(INFINITY, &shift);
+  EXPECT_NEAR(std::numeric_limits<int64_t>::max(), result, 1);
+  EXPECT_EQ(0x7fffffff, shift);
+
+  result = IntegerFrExp(-INFINITY, &shift);
+  EXPECT_NEAR(std::numeric_limits<int64_t>::min(), result, 1);
+  EXPECT_EQ(0x7fffffff, shift);
+}
+
+TEST(QuantizationUtilTest, IntegerFrExpVersusDouble) {
+  int shift;
+  int32_t result = IntegerFrExp(0.0, &shift);
+  EXPECT_EQ(result, 0);
+  EXPECT_EQ(shift, 0);
+
+  int double_shift;
+  double double_result = std::frexp(0.0, &double_shift);
+  EXPECT_EQ(double_result, 0);
+  EXPECT_EQ(double_shift, 0);
+
+  result = IntegerFrExp(1.0, &shift);
+  EXPECT_NEAR(result, 0x40000000, 1);
+  EXPECT_EQ(shift, 1);
+  double_result = std::frexp(1.0, &double_shift);
+  EXPECT_NEAR(double_result, 0.5, 1e-5);
+  EXPECT_EQ(double_shift, 1);
+
+  result = IntegerFrExp(0.25, &shift);
+  EXPECT_NEAR(result, 0x40000000, 1);
+  EXPECT_EQ(shift, -1);
+  double_result = std::frexp(0.25, &double_shift);
+  EXPECT_NEAR(double_result, 0.5, 1e-5);
+  EXPECT_EQ(double_shift, -1);
+
+  result = IntegerFrExp(-1.0, &shift);
+  EXPECT_NEAR(result, -(1 << 30), 1);
+  EXPECT_EQ(shift, 1);
+  double_result = std::frexp(-1.0, &double_shift);
+  EXPECT_NEAR(double_result, -0.5, 1e-5);
+  EXPECT_EQ(double_shift, 1);
+
+  result = IntegerFrExp(123.45, &shift);
+  EXPECT_NEAR(result, (0.964453 * (1LL << 31)), 1000);
+  EXPECT_EQ(shift, 7);
+  double_result = std::frexp(123.45, &double_shift);
+  EXPECT_NEAR(double_result, 0.964453, 1e-5);
+  EXPECT_EQ(double_shift, 7);
+}
+
+TEST(QuantizationUtilTest, DoubleFromFractionAndShift) {
+  double result = DoubleFromFractionAndShift(0, 0);
+  EXPECT_EQ(0, result);
+
+  result = DoubleFromFractionAndShift(0x40000000, 1);
+  EXPECT_NEAR(1.0, result, 1e-5);
+
+  result = DoubleFromFractionAndShift(0x40000000, 2);
+  EXPECT_NEAR(2.0, result, 1e-5);
+
+  int shift;
+  int64_t fraction = IntegerFrExp(3.0, &shift);
+  result = DoubleFromFractionAndShift(fraction, shift);
+  EXPECT_NEAR(3.0, result, 1e-5);
+
+  fraction = IntegerFrExp(123.45, &shift);
+  result = DoubleFromFractionAndShift(fraction, shift);
+  EXPECT_NEAR(123.45, result, 1e-5);
+
+  fraction = IntegerFrExp(-23.232323, &shift);
+  result = DoubleFromFractionAndShift(fraction, shift);
+  EXPECT_NEAR(-23.232323, result, 1e-5);
+
+  fraction = IntegerFrExp(NAN, &shift);
+  result = DoubleFromFractionAndShift(fraction, shift);
+  EXPECT_TRUE(std::isnan(result));
+
+  fraction = IntegerFrExp(INFINITY, &shift);
+  result = DoubleFromFractionAndShift(fraction, shift);
+  EXPECT_FALSE(std::isfinite(result));
+}
+
+TEST(QuantizationUtilTest, IntegerDoubleMultiply) {
+  EXPECT_NEAR(1.0, IntegerDoubleMultiply(1.0, 1.0), 1e-5);
+  EXPECT_NEAR(2.0, IntegerDoubleMultiply(1.0, 2.0), 1e-5);
+  EXPECT_NEAR(2.0, IntegerDoubleMultiply(2.0, 1.0), 1e-5);
+  EXPECT_NEAR(4.0, IntegerDoubleMultiply(2.0, 2.0), 1e-5);
+  EXPECT_NEAR(0.5, IntegerDoubleMultiply(1.0, 0.5), 1e-5);
+  EXPECT_NEAR(0.25, IntegerDoubleMultiply(0.5, 0.5), 1e-5);
+  EXPECT_NEAR(-1.0, IntegerDoubleMultiply(1.0, -1.0), 1e-5);
+  EXPECT_NEAR(-1.0, IntegerDoubleMultiply(-1.0, 1.0), 1e-5);
+  EXPECT_NEAR(1.0, IntegerDoubleMultiply(-1.0, -1.0), 1e-5);
+  EXPECT_NEAR(15000000.0, IntegerDoubleMultiply(3000.0, 5000.0), 1e-5);
+  EXPECT_TRUE(std::isnan(IntegerDoubleMultiply(NAN, 5000.0)));
+  EXPECT_TRUE(std::isnan(IntegerDoubleMultiply(3000.0, NAN)));
+}
+
+TEST(QuantizationUtilTest, IntegerDoubleCompare) {
+  EXPECT_EQ(-1, IntegerDoubleCompare(0.0, 1.0));
+  EXPECT_EQ(1, IntegerDoubleCompare(1.0, 0.0));
+  EXPECT_EQ(0, IntegerDoubleCompare(1.0, 1.0));
+  EXPECT_EQ(0, IntegerDoubleCompare(0.0, 0.0));
+  EXPECT_EQ(-1, IntegerDoubleCompare(-10.0, 10.0));
+  EXPECT_EQ(1, IntegerDoubleCompare(123.45, 10.0));
+  EXPECT_EQ(1, IntegerDoubleCompare(NAN, INFINITY));
+  EXPECT_EQ(1, IntegerDoubleCompare(INFINITY, NAN));
+}
+
+#if GTEST_HAS_DEATH_TEST
+TEST(QuantizationUtilTest, ChooseQuantizationParamsInvalidRange) {
+  EXPECT_DEATH(ChooseQuantizationParams<uint8_t>(10.0, -30.0), "");
+}
+
+TEST(QuantizationUtilTest, QuantizeMultiplierSmallerThanOneExp) {
+  auto quantize = [](double d) {
+    int32_t q;
+    int s;
+    QuantizeMultiplierSmallerThanOneExp(d, &q, &s);
+    return std::pair<int32_t, int>{q, s};
+  };
+
+  EXPECT_DEATH(quantize(-0.1), "");
+  EXPECT_DEATH(quantize(0.0), "");
+  EXPECT_THAT(quantize(0.25), Pair(1073741824, -1));
+
+  // Around 0.5 we can see the change in exponent and how we try hard to
+  // void hitting max int32.
+  EXPECT_THAT(quantize(0.50 - 5e-9), Pair(2147483627, -1));
+  EXPECT_THAT(quantize(0.50 - 1e-10), Pair(1073741824, 0));
+  EXPECT_THAT(quantize(0.50), Pair(1073741824, 0));
+
+  EXPECT_THAT(quantize(0.75), Pair(1610612736, 0));
+  EXPECT_THAT(quantize(1 - 1e-9), Pair(2147483646, 0));
+
+  // If we get close enough to 1.0 it crashes and dies in one of two ways:
+  // Either the shift becomes negative or we trigger the 'less-than-one' CHECK.
+  EXPECT_DEATH(quantize(1 - 1e-15), "");
+  EXPECT_DEATH(quantize(1 - 1e-17), "");
+  EXPECT_DEATH(quantize(1.0), "");
+}
+
+TEST(QuantizationUtilTest, QuantizeMultiplierGreaterThanOne) {
+  auto quantize = [](double d) {
+    int32_t q;
+    int s;
+    QuantizeMultiplierGreaterThanOne(d, &q, &s);
+    return std::pair<int32_t, int>{q, s};
+  };
+
+  // If we are close enough to 1.0 it crashes.
+  EXPECT_DEATH(quantize(1 + 1e-16), "");
+
+  EXPECT_THAT(quantize(1 + 1e-11), Pair(1073741824, 1));
+  EXPECT_THAT(quantize(1.25), Pair(1342177280, 1));
+  EXPECT_THAT(quantize(1.50), Pair(1610612736, 1));
+  EXPECT_THAT(quantize(1.75), Pair(1879048192, 1));
+
+  // Around the powers of two we see the change in exponent. Also,
+  // we try hard to avoid hitting max int32.
+  EXPECT_THAT(quantize(2 - 1e-9), Pair(2147483647, 1));
+  EXPECT_THAT(quantize(2 - 1e-11), Pair(1073741824, 2));
+  EXPECT_THAT(quantize(2), Pair(1073741824, 2));
+}
+
+#ifndef __APPLE__  // Some Apple toolchains don't support std::ldexp
+TEST(QuantizationUtilTest, QuantizeMultiplierUnderflow) {
+  auto quantize = [](double d) {
+    int32_t q;
+    int s;
+    QuantizeMultiplier(d, &q, &s);
+    return std::pair<int32_t, int>{q, s};
+  };
+
+  EXPECT_THAT(quantize(std::ldexp(1.0f, -31)), Pair(1073741824, -30));
+  EXPECT_THAT(quantize(std::ldexp(1.0f, -32)), Pair(1073741824, -31));
+  EXPECT_THAT(quantize(std::ldexp(0.99f, -32)), Pair(0, 0));
+  EXPECT_THAT(quantize(std::ldexp(1.0f, -33)), Pair(0, 0));
+}
+#endif
+
+TEST(QuantizationUtilTest, GetInvSqrtQuantizedMultiplierExp) {
+  auto inv_sqrt = [](std::int32_t input) {
+    int32_t output;
+    int output_shift;
+    GetInvSqrtQuantizedMultiplierExp(input, 1, &output, &output_shift);
+    return std::pair<int32_t, int>{output, output_shift};
+  };
+
+  const auto kInt32Max = std::numeric_limits<std::int32_t>::max();
+  EXPECT_THAT(inv_sqrt(0), Pair(kInt32Max, 0));
+  EXPECT_THAT(inv_sqrt(1), Pair(kInt32Max, 0));
+  EXPECT_THAT(inv_sqrt(2), Pair(1518498372, 0));
+  EXPECT_THAT(inv_sqrt(3), Pair(1239850284, 0));
+  EXPECT_THAT(inv_sqrt(4), Pair(1073741828, 0));
+  EXPECT_THAT(inv_sqrt(100), Pair(214748363, 0));
+  EXPECT_THAT(inv_sqrt(10000), Pair(343597361, 4));
+  EXPECT_THAT(inv_sqrt(1000000), Pair(274877901, 7));
+  EXPECT_THAT(inv_sqrt(100000000), Pair(219902323, 10));
+  EXPECT_THAT(inv_sqrt((1 << 30)), Pair(268435457, 12));
+  EXPECT_THAT(inv_sqrt(kInt32Max), Pair(189812531, 12));
+}
+
+TEST(QuantizationUtilTest, MultiplyByQuantizedMultiplierInt32) {
+  auto quant_and_multiply = [](int32_t x, double multiplier) {
+    int32_t quantized_multiplier;
+    int shift;
+    QuantizeMultiplier(multiplier, &quantized_multiplier, &shift);
+    return MultiplyByQuantizedMultiplier(x, quantized_multiplier, shift);
+  };
+
+  EXPECT_EQ(quant_and_multiply(0, 0.1), 0);
+  EXPECT_EQ(quant_and_multiply(1, 0), 0);
+  EXPECT_EQ(quant_and_multiply(10000, 0.00097656), 10);
+  EXPECT_EQ(quant_and_multiply(-10000, 0.00097656), -10);
+  EXPECT_EQ(quant_and_multiply(std::numeric_limits<int32_t>::min(), 0.00001),
+            -21475);
+  EXPECT_EQ(quant_and_multiply(std::numeric_limits<int32_t>::max(), 0.00001),
+            21475);
+#if !TFLITE_SINGLE_ROUNDING
+  // Single-rounding doesn't support negative multipliers, only test negative
+  // multipliers in double-rounding mode.
+  EXPECT_EQ(quant_and_multiply(10000, -0.00097656), -10);
+  EXPECT_EQ(quant_and_multiply(-10000, -0.00097656), 10);
+  EXPECT_EQ(quant_and_multiply(std::numeric_limits<int32_t>::min(), -0.00001),
+            21475);
+  EXPECT_EQ(quant_and_multiply(std::numeric_limits<int32_t>::max(), -0.00001),
+            -21475);
+#endif
+
+  // Test with maximum possible x and quantized_multiplier
+  const int32_t x = std::numeric_limits<int32_t>::max();
+  const int32_t quantized_multiplier = std::numeric_limits<int32_t>::max();
+  const int shift = -3;
+  const int32_t expected = static_cast<int32_t>(
+      TfLiteRound(static_cast<int64_t>(x) * quantized_multiplier /
+                  static_cast<double>(1LL << (31 - shift))));
+  EXPECT_EQ(MultiplyByQuantizedMultiplier(x, quantized_multiplier, shift),
+            expected);
+  EXPECT_EQ(MultiplyByQuantizedMultiplier(-x, quantized_multiplier, shift),
+            -expected);
+}
+
+TEST(QuantizationUtilTest, MultiplyByQuantizedMultiplierInt64) {
+  auto quant_and_multiply = [](int64_t x, double multiplier) {
+    int32_t quantized_multiplier;
+    int shift;
+    QuantizeMultiplier(multiplier, &quantized_multiplier, &shift);
+    return MultiplyByQuantizedMultiplier(x, quantized_multiplier, shift);
+  };
+
+  // Negative multipliers are not supported by the 64-bit
+  // MultiplyByQuantizedMultiplier, only use >= 0 multipliers.
+  EXPECT_EQ(quant_and_multiply(0, 0.1), 0);
+  EXPECT_EQ(quant_and_multiply(1, 0), 0);
+  EXPECT_EQ(quant_and_multiply(10000, 0.00097656), 10);
+  EXPECT_EQ(quant_and_multiply(-10000, 0.00097656), -10);
+  EXPECT_EQ(quant_and_multiply(-(1LL << 47), 0.00001), -1407385600);
+  EXPECT_EQ(quant_and_multiply((1LL << 47) - 1, 0.00001), 1407385600);
+
+  // Test with maximum possible x and quantized_multiplier
+  const int64_t x = (1LL << 47) - 1;
+  const int32_t quantized_multiplier = std::numeric_limits<int32_t>::max();
+  const int shift = -31;
+  // Expected is around 'x * quantized_multiplier / 2**(31 - shift)' ~= 65536
+  // As there is some rounding error, expected is a bit smaller.
+  const int32_t expected = 65534;
+  EXPECT_EQ(MultiplyByQuantizedMultiplier(x, quantized_multiplier, shift),
+            expected);
+  EXPECT_EQ(MultiplyByQuantizedMultiplier(-x, quantized_multiplier, shift),
+            -expected);
+}
+
+TEST(QuantizationUtilTest, PreprocessSoftmaxScaling) {
+  auto quantize = [](double beta, double scale, int integer_bits) {
+    int32_t q;
+    int s;
+    PreprocessSoftmaxScaling(beta, scale, integer_bits, &q, &s);
+    return std::pair<int32_t, int>{q, s};
+  };
+
+#if TFLITE_SINGLE_ROUNDING
+  // If beta * scale is greater than fits in the number of integer bits, the
+  // result is move near the maximum. Otherwise they quantize as expected.
+  // With 4 integer bits we can represent up to 8.0.
+  EXPECT_THAT(quantize(1.0, 8.0, 4), Pair(2147483646, 30));
+  EXPECT_THAT(quantize(1.0, 4.0, 4), Pair(1073741824, 30));
+  // But with 5 bits we can go further.
+  EXPECT_THAT(quantize(2.0, 8.0, 5), Pair(2147483646, 30));
+  EXPECT_THAT(quantize(2.0, 4.0, 5), Pair(1073741824, 30));
+#else
+  // If beta * scale is greater than fits in the number of integer bits, the
+  // result is move near the maximum. Otherwise they quantize as expected.
+  // With 4 integer bits we can represent up to 16.0.
+  EXPECT_THAT(quantize(1.0, 16.0, 4), Pair(2147483647, 31));
+  EXPECT_THAT(quantize(1.0, 8.0, 4), Pair(1073741824, 31));
+  // But with 5 bits we can go further.
+  EXPECT_THAT(quantize(2.0, 16.0, 5), Pair(2147483647, 31));
+  EXPECT_THAT(quantize(2.0, 8.0, 5), Pair(1073741824, 31));
+#endif
+}
+#endif  // GTEST_HAS_DEATH_TEST
+
+TEST(QuantizationUtilTest, CalculateInputRadius) {
+  EXPECT_EQ(CalculateInputRadius(4, 27), 15);
+  EXPECT_EQ(CalculateInputRadius(3, 27), 14);
+  EXPECT_EQ(CalculateInputRadius(3, 28), 7);
+  EXPECT_EQ(CalculateInputRadius(4, 2), 503316480);
+}
+
+TEST(QuantizationUtilTest, QuantizeMultiplierArray) {
+  const std::vector<double> weights = {-4,    -2,   -1,  -0.5, -0.25, -0.125, 0,
+                                       0.125, 0.25, 0.5, 1,    2,     4};
+  const int size = weights.size();
+  std::vector<int32_t> effective_scale_significand(size);
+  std::vector<int> effective_scale_shift(size);
+  QuantizeMultiplierArray(weights.data(), size,
+                          effective_scale_significand.data(),
+                          effective_scale_shift.data());
+  const std::vector<int32_t> expected_effective_scale_significand = {
+      -1073741824,  // float scale = -4
+      -1073741824,  // float scale = -2
+      -1073741824,  // float scale = -1
+      -1073741824,  // float scale = -0.5
+      -1073741824,  // float scale = -0.25
+      -1073741824,  // float scale = -0.125
+      0,            // float scale = 0
+      1073741824,   // float scale = 0.125
+      1073741824,   // float scale = 0.25
+      1073741824,   // float scale = 0.5
+      1073741824,   // float scale = 1
+      1073741824,   // float scale = 2
+      1073741824,   // float scale = 4
+  };
+
+  const std::vector<int> expected_effective_scale_shift = {
+      3,   // float scale = -4
+      2,   // float scale = -2
+      1,   // float scale = -1
+      0,   // float scale = -0.5
+      -1,  // float scale = -0.25
+      -2,  // float scale = -0.125
+      0,   // float scale = 0
+      -2,  // float scale = 0.125
+      -1,  // float scale = 0.25
+      0,   // float scale = 0.5
+      1,   // float scale = 1
+      2,   // float scale = 2
+      3,   // float scale = 4
+  };
+  EXPECT_THAT(effective_scale_significand,
+              ElementsAreArray(expected_effective_scale_significand));
+  EXPECT_THAT(effective_scale_shift,
+              ElementsAreArray(expected_effective_scale_shift));
+}
+
+}  // namespace
+}  // namespace tflite

tensorflow/lite/kernels/internal/reduce_common.h

@@ -0,0 +1,37 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REDUCE_COMMON_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_REDUCE_COMMON_H_
+
+namespace tflite {
+namespace ops {
+namespace builtin {
+namespace reduce {
+
+enum ReduceType {
+  kSum,
+  kProd,
+  kMax,
+  kMin,
+  kAny,
+  kAll,
+};
+
+}  // namespace reduce
+}  // namespace builtin
+}  // namespace ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_REDUCE_COMMON_H_

tensorflow/lite/kernels/internal/reference/add.h

@@ -0,0 +1,427 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ADD_H_
+#define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ADD_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <type_traits>
+
+#include "fixedpoint/fixedpoint.h"
+#include "tensorflow/lite/kernels/internal/common.h"
+#include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/reference/broadcast_loop.h"
+
+namespace tflite {
+
+namespace reference_ops {
+
+template <typename T>
+inline void Add(const ArithmeticParams& params,
+                const RuntimeShape& input1_shape, const T* input1_data,
+                const RuntimeShape& input2_shape, const T* input2_data,
+                const RuntimeShape& output_shape, T* output_data) {
+  T activation_min, activation_max;
+  GetActivationParams(params, &activation_min, &activation_max);
+
+  const int flat_size =
+      MatchingElementsSize(input1_shape, input2_shape, output_shape);
+  for (int i = 0; i < flat_size; ++i) {
+    output_data[i] = ActivationFunctionWithMinMax<T>(
+        input1_data[i] + input2_data[i], activation_min, activation_max);
+  }
+}
+
+// Element-wise add that can often be used for inner loop of broadcast add as
+// well as the non-broadcast add.
+
+// This function is used for 8-bit as well as for 16-bit, but the accumulator
+// is 32-bit for both cases. The overflow does not happen due to the
+// choice of the shift (20 or 15, accordingly - see add.cc for more comments).
+template <typename T>
+inline void AddElementwise(int size, const ArithmeticParams& params,
+                           const T* input1_data, const T* input2_data,
+                           T* output_data) {
+  TFLITE_DCHECK_GT(params.input1_offset, -std::numeric_limits<T>::max());
+  TFLITE_DCHECK_GT(params.input2_offset, -std::numeric_limits<T>::max());
+  TFLITE_DCHECK_LT(params.input1_offset, std::numeric_limits<T>::max());
+  TFLITE_DCHECK_LT(params.input2_offset, std::numeric_limits<T>::max());
+
+  for (int i = 0; i < size; ++i) {
+    const int32_t input1_val = params.input1_offset + input1_data[i];
+    const int32_t input2_val = params.input2_offset + input2_data[i];
+    const int32_t shifted_input1_val = input1_val * (1 << params.left_shift);
+    const int32_t shifted_input2_val = input2_val * (1 << params.left_shift);
+    const int32_t scaled_input1_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input1_val, params.input1_multiplier, params.input1_shift);
+    const int32_t scaled_input2_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input2_val, params.input2_multiplier, params.input2_shift);
+    const int32_t raw_sum = scaled_input1_val + scaled_input2_val;
+    const int32_t raw_output =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            raw_sum, params.output_multiplier, params.output_shift) +
+        params.output_offset;
+    const int32_t clamped_output =
+        std::min(params.quantized_activation_max,
+                 std::max(params.quantized_activation_min, raw_output));
+    output_data[i] = static_cast<T>(clamped_output);
+  }
+}
+
+// Scalar-broadcast add that can be used for inner loop of more general
+// broadcast add, so that, for example, scalar-broadcast with batch will still
+// be fast.
+inline void AddScalarBroadcast(int size, const ArithmeticParams& params,
+                               uint8_t input1_data, const uint8_t* input2_data,
+                               uint8_t* output_data) {
+  TFLITE_DCHECK_GT(params.input1_offset, -256);
+  TFLITE_DCHECK_GT(params.input2_offset, -256);
+  TFLITE_DCHECK_LT(params.input1_offset, 256);
+  TFLITE_DCHECK_LT(params.input2_offset, 256);
+
+  const int32_t input1_val = params.input1_offset + input1_data;
+  const int32_t shifted_input1_val = input1_val * (1 << params.left_shift);
+  const int32_t scaled_input1_val =
+      MultiplyByQuantizedMultiplierSmallerThanOneExp(
+          shifted_input1_val, params.input1_multiplier, params.input1_shift);
+  for (int i = 0; i < size; ++i) {
+    const int32_t input2_val = params.input2_offset + input2_data[i];
+    const int32_t shifted_input2_val = input2_val * (1 << params.left_shift);
+    const int32_t scaled_input2_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input2_val, params.input2_multiplier, params.input2_shift);
+    const int32_t raw_sum = scaled_input1_val + scaled_input2_val;
+    const int32_t raw_output =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            raw_sum, params.output_multiplier, params.output_shift) +
+        params.output_offset;
+    const int32_t clamped_output =
+        std::min(params.quantized_activation_max,
+                 std::max(params.quantized_activation_min, raw_output));
+    output_data[i] = static_cast<uint8_t>(clamped_output);
+  }
+}
+
+inline void Add(const ArithmeticParams& params,
+                const RuntimeShape& input1_shape, const uint8_t* input1_data,
+                const RuntimeShape& input2_shape, const uint8_t* input2_data,
+                const RuntimeShape& output_shape, uint8_t* output_data) {
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  const int flat_size =
+      MatchingElementsSize(input1_shape, input2_shape, output_shape);
+
+  TFLITE_DCHECK_GT(params.input1_offset, -256);
+  TFLITE_DCHECK_GT(params.input2_offset, -256);
+  TFLITE_DCHECK_LT(params.input1_offset, 256);
+  TFLITE_DCHECK_LT(params.input2_offset, 256);
+  AddElementwise(flat_size, params, input1_data, input2_data, output_data);
+}
+
+inline void AddGeneralParamScale(const ArithmeticParams& params,
+                                 const RuntimeShape& input1_shape,
+                                 const int16_t* input1_data,
+                                 const RuntimeShape& input2_shape,
+                                 const int16_t* input2_data,
+                                 const RuntimeShape& output_shape,
+                                 int16_t* output_data) {
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+  const int flat_size =
+      MatchingElementsSize(input1_shape, input2_shape, output_shape);
+
+  int max_value = std::numeric_limits<int16_t>::max();
+
+  TFLITE_DCHECK_GT(params.input1_offset, -max_value);
+  TFLITE_DCHECK_GT(params.input2_offset, -max_value);
+  TFLITE_DCHECK_LT(params.input1_offset, max_value);
+  TFLITE_DCHECK_LT(params.input2_offset, max_value);
+  AddElementwise(flat_size, params, input1_data, input2_data, output_data);
+}
+
+inline void Add(const ArithmeticParams& params,
+                const RuntimeShape& input1_shape, const int16_t* input1_data,
+                const RuntimeShape& input2_shape, const int16_t* input2_data,
+                const RuntimeShape& output_shape, int16_t* output_data,
+                bool pot_scale = true) {
+  if (!pot_scale) {
+    AddGeneralParamScale(params, input1_shape, input1_data, input2_shape,
+                         input2_data, output_shape, output_data);
+    return;
+  }
+
+  TFLITE_DCHECK_LE(params.quantized_activation_min,
+                   params.quantized_activation_max);
+
+  const int input1_shift = params.input1_shift;
+  const int flat_size =
+      MatchingElementsSize(input1_shape, input2_shape, output_shape);
+  const int16_t output_activation_min = params.quantized_activation_min;
+  const int16_t output_activation_max = params.quantized_activation_max;
+
+  TFLITE_DCHECK(input1_shift == 0 || params.input2_shift == 0);
+  TFLITE_DCHECK_LE(input1_shift, 0);
+  TFLITE_DCHECK_LE(params.input2_shift, 0);
+  const int16_t* not_shift_input =
+      input1_shift == 0 ? input1_data : input2_data;
+  const int16_t* shift_input = input1_shift == 0 ? input2_data : input1_data;
+  const int input_right_shift =
+      input1_shift == 0 ? -params.input2_shift : -input1_shift;
+
+  for (int i = 0; i < flat_size; i++) {
+    // F0 uses 0 integer bits, range [-1, 1].
+    using F0 = gemmlowp::FixedPoint<std::int16_t, 0>;
+
+    F0 input_ready_scaled = F0::FromRaw(not_shift_input[i]);
+    F0 scaled_input = F0::FromRaw(
+        gemmlowp::RoundingDivideByPOT(shift_input[i], input_right_shift));
+    F0 result = gemmlowp::SaturatingAdd(scaled_input, input_ready_scaled);
+    const int16_t raw_output = result.raw();
+    const int16_t clamped_output = std::min(
+        output_activation_max, std::max(output_activation_min, raw_output));
+    output_data[i] = clamped_output;
+  }
+}
+
+template <typename T>
+inline void AddBroadcast(const T* input_data, const T* broadcast_data,
+                         T* output_data, size_t size, T activation_min,
+                         T activation_max) {
+  for (size_t c = 0; c < size; ++c) {
+    output_data[c] = ActivationFunctionWithMinMax<T>(
+        input_data[c] + broadcast_data[0], activation_min, activation_max);
+  }
+}
+
+template <>
+inline void AddBroadcast<int32_t>(const int32_t* input_data,
+                                  const int32_t* broadcast_data,
+                                  int32_t* output_data, size_t size,
+                                  int32_t activation_min,
+                                  int32_t activation_max) {
+  size_t c = 0;
+#ifdef USE_NEON
+  const int32x4_t vmax = vdupq_n_s32(activation_max);
+  const int32x4_t vmin = vdupq_n_s32(activation_min);
+  const int32x4_t vb = vdupq_n_s32(broadcast_data[0]);
+  for (; c + 4 <= size; c += 4) {
+    const int32x4_t va = vld1q_s32(&input_data[c]);
+    int32x4_t vres = vaddq_s32(va, vb);
+    vres = vmaxq_s32(vmin, vres);
+    vres = vminq_s32(vmax, vres);
+    vst1q_s32(&output_data[c], vres);
+  }
+#endif
+  for (; c < size; ++c) {
+    output_data[c] = ActivationFunctionWithMinMax<int32_t>(
+        input_data[c] + broadcast_data[0], activation_min, activation_max);
+  }
+}
+
+template <typename T>
+void AddElementwise(const T* input1_data, const T* input2_data, T* output_data,
+                    size_t size, T activation_min, T activation_max) {
+  for (size_t c = 0; c < size; ++c) {
+    output_data[c] = ActivationFunctionWithMinMax<T>(
+        input1_data[c] + input2_data[c], activation_min, activation_max);
+  }
+}
+
+template <>
+inline void AddElementwise<int32_t>(const int32_t* input1_data,
+                                    const int32_t* input2_data,
+                                    int32_t* output_data, size_t size,
+                                    int32_t activation_min,
+                                    int32_t activation_max) {
+  size_t c = 0;
+#ifdef USE_NEON
+  const int32x4_t vmax = vdupq_n_s32(activation_max);
+  const int32x4_t vmin = vdupq_n_s32(activation_min);
+  for (; c + 4 <= size; c += 4) {
+    const int32x4_t va = vld1q_s32(&input1_data[c]);
+    const int32x4_t vb = vld1q_s32(&input2_data[c]);
+    int32x4_t vres = vaddq_s32(va, vb);
+    vres = vmaxq_s32(vmin, vres);
+    vres = vminq_s32(vmax, vres);
+    vst1q_s32(&output_data[c], vres);
+  }
+#endif
+  for (; c < size; ++c) {
+    output_data[c] = ActivationFunctionWithMinMax<int32_t>(
+        input1_data[c] + input2_data[c], activation_min, activation_max);
+  }
+}
+
+template <typename T,
+          // For unquantized add for small integers, explicitly set to true.
+          bool dummy = false>
+inline typename std::enable_if<!is_small_integer<T>::value || dummy, void>::type
+BroadcastAdd6DSlow(const ArithmeticParams& params,
+                   const RuntimeShape& input1_shape, const T* input1_data,
+                   const RuntimeShape& input2_shape, const T* input2_data,
+                   const RuntimeShape& output_shape, T* output_data) {
+  T activation_min, activation_max;
+  GetActivationParams(params, &activation_min, &activation_max);
+  auto op = [activation_min, activation_max](T a, T b) {
+    return ActivationFunctionWithMinMax<T>(a + b, activation_min,
+                                           activation_max);
+  };
+  BroadcastBinaryOpSimple(input1_shape, input1_data, input2_shape, input2_data,
+                          output_shape, output_data, op);
+}
+
+// This function is used for 8-bit as well as for 16-bit, but the accumulator
+// is 32-bit for both cases. The overflow does not happen due to the
+// choice of the shift (20 or 15, accordingly - see add.cc for more comments).
+template <typename T>
+inline typename std::enable_if<is_small_integer<T>::value, void>::type
+BroadcastAdd6DSlow(const ArithmeticParams& params,
+                   const RuntimeShape& input1_shape, const T* input1_data,
+                   const RuntimeShape& input2_shape, const T* input2_data,
+                   const RuntimeShape& output_shape, T* output_data) {
+  auto op = [&params](T a, T b) {
+    const int32_t input1_val = params.input1_offset + a;
+    const int32_t input2_val = params.input2_offset + b;
+    const int32_t shifted_input1_val = input1_val * (1 << params.left_shift);
+    const int32_t shifted_input2_val = input2_val * (1 << params.left_shift);
+    const int32_t scaled_input1_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input1_val, params.input1_multiplier, params.input1_shift);
+    const int32_t scaled_input2_val =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            shifted_input2_val, params.input2_multiplier, params.input2_shift);
+    const int32_t raw_sum = scaled_input1_val + scaled_input2_val;
+    const int32_t raw_output =
+        MultiplyByQuantizedMultiplierSmallerThanOneExp(
+            raw_sum, params.output_multiplier, params.output_shift) +
+        params.output_offset;
+    const int32_t clamped_output =
+        std::min(params.quantized_activation_max,
+                 std::max(params.quantized_activation_min, raw_output));
+    return static_cast<T>(clamped_output);
+  };
+  BroadcastBinaryOpSimple(input1_shape, input1_data, input2_shape, input2_data,
+                          output_shape, output_data, op);
+}
+
+template <typename T>
+inline void BroadcastAdd4DSlow(
+    const ArithmeticParams& params, const RuntimeShape& input1_shape,
+    const T* input1_data, const RuntimeShape& input2_shape,
+    const T* input2_data, const RuntimeShape& output_shape, T* output_data) {
+  return BroadcastAdd6DSlow(params, input1_shape, input1_data, input2_shape,
+                            input2_data, output_shape, output_data);
+}
+
+inline void BroadcastAddFivefold(const ArithmeticParams& unswitched_params,
+                                 const RuntimeShape& unswitched_input1_shape,
+                                 const uint8_t* unswitched_input1_data,
+                                 const RuntimeShape& unswitched_input2_shape,
+                                 const uint8_t* unswitched_input2_data,
+                                 const RuntimeShape& output_shape,
+                                 uint8_t* output_data) {
+  ArithmeticParams switched_params = unswitched_params;
+  switched_params.input1_offset = unswitched_params.input2_offset;
+  switched_params.input1_multiplier = unswitched_params.input2_multiplier;
+  switched_params.input1_shift = unswitched_params.input2_shift;
+  switched_params.input2_offset = unswitched_params.input1_offset;
+  switched_params.input2_multiplier = unswitched_params.input1_multiplier;
+  switched_params.input2_shift = unswitched_params.input1_shift;
+
+  const bool use_unswitched =
+      unswitched_params.broadcast_category ==
+      tflite::BroadcastableOpCategory::kFirstInputBroadcastsFast;
+
+  const ArithmeticParams& params =
+      use_unswitched ? unswitched_params : switched_params;
+  const uint8_t* input1_data =
+      use_unswitched ? unswitched_input1_data : unswitched_input2_data;
+  const uint8_t* input2_data =
+      use_unswitched ? unswitched_input2_data : unswitched_input1_data;
+
+  // Fivefold nested loops. The second input resets its position for each
+  // iteration of the second loop. The first input resets its position at the
+  // beginning of the fourth loop. The innermost loop is an elementwise add of
+  // sections of the arrays.
+  uint8_t* output_data_ptr = output_data;
+  const uint8_t* input1_data_ptr = input1_data;
+  const uint8_t* input2_data_reset = input2_data;
+  // In the fivefold pattern, y0, y2 and y4 are not broadcast, and so shared
+  // between input shapes. y3 for input 1 is always broadcast, and so the
+  // dimension there is 1, whereas optionally y1 might be broadcast for input 2.
+  // Put another way,
+  // input1.shape.FlatSize = y0 * y1 * y2 * y4,
+  // input2.shape.FlatSize = y0 * y2 * y3 * y4.
+  int y0 = params.broadcast_shape[0];
+  int y1 = params.broadcast_shape[1];
+  int y2 = params.broadcast_shape[2];
+  int y3 = params.broadcast_shape[3];
+  int y4 = params.broadcast_shape[4];
+  if (y4 > 1) {
+    // General fivefold pattern, with y4 > 1 so there is a non-broadcast inner
+    // dimension.
+    for (int i0 = 0; i0 < y0; ++i0) {
+      const uint8_t* input2_data_ptr;
+      for (int i1 = 0; i1 < y1; ++i1) {
+        input2_data_ptr = input2_data_reset;
+        for (int i2 = 0; i2 < y2; ++i2) {
+          for (int i3 = 0; i3 < y3; ++i3) {
+            AddElementwise(y4, params, input1_data_ptr, input2_data_ptr,
+                           output_data_ptr);
+            input2_data_ptr += y4;
+            output_data_ptr += y4;
+          }
+          // We have broadcast y4 of input1 data y3 times, and now move on.
+          input1_data_ptr += y4;
+        }
+      }
+      // We have broadcast y2*y3*y4 of input2 data y1 times, and now move on.
+      input2_data_reset = input2_data_ptr;
+    }
+  } else {
+    // Special case of y4 == 1, in which the innermost loop is a single element
+    // and can be combined with the next (y3) as an inner broadcast.
+    //
+    // Note that this handles the case of pure scalar broadcast when
+    // y0 == y1 == y2 == 1. With low overhead it handles cases such as scalar
+    // broadcast with batch (as y2 > 1).
+    //
+    // NOTE The process is the same as the above general case except simplified
+    // for y4 == 1 and the loop over y3 is contained within the
+    // AddScalarBroadcast function.
+    for (int i0 = 0; i0 < y0; ++i0) {
+      const uint8_t* input2_data_ptr;
+      for (int i1 = 0; i1 < y1; ++i1) {
+        input2_data_ptr = input2_data_reset;
+        for (int i2 = 0; i2 < y2; ++i2) {
+          AddScalarBroadcast(y3, params, *input1_data_ptr, input2_data_ptr,
+                             output_data_ptr);
+          input2_data_ptr += y3;
+          output_data_ptr += y3;
+          input1_data_ptr += 1;
+        }
+      }
+      input2_data_reset = input2_data_ptr;
+    }
+  }
+}
+
+}  // namespace reference_ops
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_ADD_H_