chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/lite/kernels/BUILD

@@ -0,0 +1,23 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+# LINT.IfChange(cc_library_padding)
+cc_library(
+    name = "padding",
+    srcs = [],
+    hdrs = ["padding.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//tensorflow/compiler/mlir/lite/kernels:__pkg__"],
+    deps = [
+        "//tensorflow/compiler/mlir/lite/core/c:tflite_common",
+    ],
+)
+# LINT.ThenChange(//tensorflow/lite/kernels/BUILD)

tensorflow/compiler/mlir/lite/kernels/internal/BUILD

@@ -0,0 +1,115 @@
+load("@bazel_skylib//lib:selects.bzl", "selects")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/compiler/mlir/lite:build_def.bzl", "tflite_copts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "compatibility_macros",
+    hdrs = ["compatibility_macros.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+)
+
+cc_library(
+    name = "runtime_shape",
+    srcs = ["runtime_shape.cc"],
+    hdrs = ["runtime_shape.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [":compatibility_macros"],
+)
+
+tf_cc_test(
+    name = "runtime_shape_test",
+    srcs = ["runtime_shape_test.cc"],
+    deps = [
+        ":runtime_shape",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+# LINT.IfChange
+
+cc_library(
+    name = "cppmath",
+    srcs = [],
+    hdrs = [
+        "cppmath.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+)
+
+cc_library(
+    name = "quantization_util",
+    srcs = ["quantization_util.cc"],
+    hdrs = ["quantization_util.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    linkopts = select({
+        "//tensorflow:windows": [],
+        "//conditions:default": ["-lm"],
+    }),
+    deps = [
+        ":compatibility_macros",
+        ":cppmath",
+    ],
+)
+
+config_setting(
+    name = "x86_32",
+    constraint_values = ["@platforms//cpu:x86_32"],
+)
+
+config_setting(
+    name = "x86_64",
+    constraint_values = ["@platforms//cpu:x86_64"],
+)
+
+selects.config_setting_group(
+    name = "x86_any",
+    match_any = [
+        ":x86_32",
+        ":x86_64",
+    ],
+)
+
+cc_library(
+    name = "cpu_check",
+    hdrs = [
+        "optimized/neon_check.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    visibility = ["//visibility:public"],
+    deps = select({
+        ":x86_any": ["@arm_neon_2_x86_sse"],
+        "//conditions:default": [],
+    }),
+)
+
+cc_library(
+    name = "common",
+    srcs = ["common.cc"],
+    hdrs = ["common.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    deps = [
+        ":compatibility_macros",
+        ":cpu_check",
+        "//tensorflow/compiler/mlir/lite/core:macros",
+        "@gemmlowp//:fixedpoint",
+    ],
+)
+
+# LINT.ThenChange(//tensorflow/lite/kernels/internal/BUILD)

tensorflow/compiler/mlir/lite/kernels/internal/README

@@ -0,0 +1,3 @@
+This folder mirrors some files in tensorflow/lite/kernels/internal that are
+needed in the converter. They mostly have the same name, although
+compatibility_macros.h mirrors compatibility.h and op_macros.h.

tensorflow/compiler/mlir/lite/kernels/internal/common.cc

@@ -0,0 +1,107 @@
+/* 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/compiler/mlir/lite/kernels/internal/common.h"
+
+// LINT.IfChange
+
+namespace tflite_migration {
+
+// 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_migration
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/common.cc)

tensorflow/compiler/mlir/lite/kernels/internal/common.h

@@ -0,0 +1,236 @@
+/* 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_COMPILER_MLIR_LITE_KERNELS_INTERNAL_COMMON_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_COMMON_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+
+#ifndef ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK
+#ifdef GEMMLOWP_ALLOW_SLOW_SCALAR_FALLBACK
+#define ALLOW_SLOW_GENERIC_DEPTHWISECONV_FALLBACK
+#endif
+#endif
+
+#include <cmath>
+#include <functional>
+
+#include "fixedpoint/fixedpoint.h"
+#include "tensorflow/compiler/mlir/lite/core/macros.h"
+#include "tensorflow/compiler/mlir/lite/kernels/internal/compatibility_macros.h"
+#include "tensorflow/compiler/mlir/lite/kernels/internal/optimized/neon_check.h"
+
+// LINT.IfChange
+
+namespace tflite_migration {
+
+constexpr int kReverseShift = -1;
+
+TFLITE_NOINLINE int32_t MultiplyByQuantizedMultiplier(
+    int32_t x, int32_t quantized_multiplier, int shift);
+
+TFLITE_NOINLINE int32_t MultiplyByQuantizedMultiplier(
+    int64_t x, int32_t quantized_multiplier, int shift);
+
+// Single-rounding MultiplyByQuantizedMultiplier
+#if TFLITE_SINGLE_ROUNDING
+inline int32_t MultiplyByQuantizedMultiplierSmallerThanOneExp(
+    int32_t x, int32_t quantized_multiplier, int shift) {
+  TFLITE_DCHECK_LE(shift, 0);
+  return MultiplyByQuantizedMultiplier(x, quantized_multiplier, shift);
+}
+
+inline int32_t MultiplyByQuantizedMultiplierGreaterThanOne(
+    int32_t x, int32_t quantized_multiplier, int shift) {
+  TFLITE_DCHECK_GE(shift, 0);
+  return MultiplyByQuantizedMultiplier(x, quantized_multiplier, shift);
+}
+
+#ifdef USE_NEON
+inline int32x4x4_t MultiplyByQuantizedMultiplier4Rows(
+    int32x4x4_t input_val, int32_t quantized_multiplier, int shift) {
+  TFLITE_DCHECK(quantized_multiplier >= 0);
+
+  const int right_shift = std::min(-1, shift);
+  const int left_shift = shift - right_shift;
+
+  const int32x4_t multiplier_dup = vdupq_n_s32(quantized_multiplier);
+  const int32x4_t left_shift_dup = vdupq_n_s32(left_shift);
+  const int32x4_t right_shift_dup = vdupq_n_s32(right_shift);
+
+  int32x4x4_t result;
+  result.val[0] = vrshlq_s32(
+      vqdmulhq_s32(vshlq_s32(input_val.val[0], left_shift_dup), multiplier_dup),
+      right_shift_dup);
+
+  result.val[1] = vrshlq_s32(
+      vqdmulhq_s32(vshlq_s32(input_val.val[1], left_shift_dup), multiplier_dup),
+      right_shift_dup);
+
+  result.val[2] = vrshlq_s32(
+      vqdmulhq_s32(vshlq_s32(input_val.val[2], left_shift_dup), multiplier_dup),
+      right_shift_dup);
+
+  result.val[3] = vrshlq_s32(
+      vqdmulhq_s32(vshlq_s32(input_val.val[3], left_shift_dup), multiplier_dup),
+      right_shift_dup);
+
+  return result;
+}
+#endif  // USE_NEON
+// Double-rounding MultiplyByQuantizedMultiplier
+#else
+inline int32_t MultiplyByQuantizedMultiplierSmallerThanOneExp(
+    int32_t x, int32_t quantized_multiplier, int left_shift) {
+  using gemmlowp::RoundingDivideByPOT;
+  using gemmlowp::SaturatingRoundingDoublingHighMul;
+  return RoundingDivideByPOT(
+      SaturatingRoundingDoublingHighMul(x, quantized_multiplier), -left_shift);
+}
+
+inline int32_t MultiplyByQuantizedMultiplierGreaterThanOne(
+    int32_t x, int32_t quantized_multiplier, int left_shift) {
+  using gemmlowp::SaturatingRoundingDoublingHighMul;
+  return SaturatingRoundingDoublingHighMul(x * (1 << left_shift),
+                                           quantized_multiplier);
+}
+
+#ifdef USE_NEON
+// Round uses ARM's rounding shift right.
+inline int32x4x4_t MultiplyByQuantizedMultiplier4Rows(
+    int32x4x4_t input_val, int32_t quantized_multiplier, int shift) {
+  const int left_shift = std::max(shift, 0);
+  const int right_shift = std::min(shift, 0);
+  int32x4x4_t result;
+
+  int32x4_t multiplier_dup = vdupq_n_s32(quantized_multiplier);
+  int32x4_t left_shift_dup = vdupq_n_s32(left_shift);
+  int32x4_t right_shift_dup = vdupq_n_s32(right_shift);
+
+  result.val[0] =
+      vrshlq_s32(vqrdmulhq_s32(vshlq_s32(input_val.val[0], left_shift_dup),
+                               multiplier_dup),
+                 right_shift_dup);
+
+  result.val[1] =
+      vrshlq_s32(vqrdmulhq_s32(vshlq_s32(input_val.val[1], left_shift_dup),
+                               multiplier_dup),
+                 right_shift_dup);
+
+  result.val[2] =
+      vrshlq_s32(vqrdmulhq_s32(vshlq_s32(input_val.val[2], left_shift_dup),
+                               multiplier_dup),
+                 right_shift_dup);
+
+  result.val[3] =
+      vrshlq_s32(vqrdmulhq_s32(vshlq_s32(input_val.val[3], left_shift_dup),
+                               multiplier_dup),
+                 right_shift_dup);
+
+  return result;
+}
+#endif  // USE_NEON
+#endif  // TFLITE_SINGLE_ROUNDING
+
+template <typename T>
+int CountLeadingZeros(T integer_input) {
+  static_assert(std::is_unsigned<T>::value,
+                "Only unsigned integer types handled.");
+  if (integer_input == 0) {
+    return std::numeric_limits<T>::digits;
+  }
+#if defined(__GNUC__)
+  if (std::is_same<T, uint32_t>::value) {
+    return __builtin_clz(integer_input);
+  } else if (std::is_same<T, uint64_t>::value) {
+    return __builtin_clzll(integer_input);
+  }
+#endif
+  const T one_in_leading_positive = static_cast<T>(1)
+                                    << (std::numeric_limits<T>::digits - 1);
+  int leading_zeros = 0;
+  while (integer_input < one_in_leading_positive) {
+    integer_input <<= 1;
+    ++leading_zeros;
+  }
+  return leading_zeros;
+}
+
+inline void GetInvSqrtQuantizedMultiplierExp(int32_t input, int reverse_shift,
+                                             int32_t* output_inv_sqrt,
+                                             int* output_shift) {
+  TFLITE_DCHECK_GE(input, 0);
+  if (input <= 1) {
+    // Handle the input value 1 separately to avoid overflow in that case
+    // in the general computation below (b/143972021). Also handle 0 as if it
+    // were a 1. 0 is an invalid input here (divide by zero) and 1 is a valid
+    // but rare/unrealistic input value. We can expect both to occur in some
+    // incompletely trained models, but probably not in fully trained models.
+    *output_inv_sqrt = std::numeric_limits<std::int32_t>::max();
+    *output_shift = 0;
+    return;
+  }
+  TFLITE_DCHECK_GT(input, 1);
+  *output_shift = 11;
+  while (input >= (1 << 29)) {
+    input /= 4;
+    ++*output_shift;
+  }
+  const unsigned max_left_shift_bits =
+      CountLeadingZeros(static_cast<uint32_t>(input)) - 1;
+  const unsigned max_left_shift_bit_pairs = max_left_shift_bits / 2;
+  const unsigned left_shift_bit_pairs = max_left_shift_bit_pairs - 1;
+  *output_shift -= left_shift_bit_pairs;
+  input <<= 2 * left_shift_bit_pairs;
+  TFLITE_DCHECK_GE(input, (1 << 27));
+  TFLITE_DCHECK_LT(input, (1 << 29));
+  using gemmlowp::FixedPoint;
+  using gemmlowp::Rescale;
+  using gemmlowp::SaturatingRoundingMultiplyByPOT;
+  // Using 3 integer bits gives us enough room for the internal arithmetic in
+  // this Newton-Raphson iteration.
+  using F3 = FixedPoint<int32_t, 3>;
+  using F0 = FixedPoint<int32_t, 0>;
+  const F3 fixedpoint_input = F3::FromRaw(input >> 1);
+  const F3 fixedpoint_half_input =
+      SaturatingRoundingMultiplyByPOT<-1>(fixedpoint_input);
+  const F3 fixedpoint_half_three =
+      GEMMLOWP_CHECKED_FIXEDPOINT_CONSTANT(F3, (1 << 28) + (1 << 27), 1.5);
+  // Newton-Raphson iteration
+  // Naive unoptimized starting guess: x = 1
+  F3 x = F3::One();
+  // Naive unoptimized number of iterations: 5
+  for (int i = 0; i < 5; i++) {
+    const F3 x3 = Rescale<3>(x * x * x);
+    x = Rescale<3>(fixedpoint_half_three * x - fixedpoint_half_input * x3);
+  }
+  const F0 fixedpoint_half_sqrt_2 =
+      GEMMLOWP_CHECKED_FIXEDPOINT_CONSTANT(F0, 1518500250, std::sqrt(2.) / 2.);
+  x = x * fixedpoint_half_sqrt_2;
+  *output_inv_sqrt = x.raw();
+  if (*output_shift < 0) {
+    *output_inv_sqrt <<= -*output_shift;
+    *output_shift = 0;
+  }
+  // Convert right shift (right is positive) to left shift.
+  *output_shift *= reverse_shift;
+}
+
+}  // namespace tflite_migration
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/common.h)
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_COMMON_H_

tensorflow/compiler/mlir/lite/kernels/internal/compatibility_macros.h

@@ -0,0 +1,63 @@
+/* Copyright 2024 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_COMPILER_MLIR_LITE_KERNELS_INTERNAL_COMPATIBILITY_MACROS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_COMPATIBILITY_MACROS_H_
+
+#ifndef TFLITE_ABORT
+#define TFLITE_ABORT abort()
+#endif
+
+#ifndef TFLITE_ASSERT_FALSE
+#if defined(NDEBUG)
+#define TFLITE_ASSERT_FALSE (static_cast<void>(0))
+#else
+#define TFLITE_ASSERT_FALSE TFLITE_ABORT
+#endif
+#endif
+
+// LINT.IfChange
+
+#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
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/compatibility.h)
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_COMPATIBILITY_MACROS_H_

tensorflow/compiler/mlir/lite/kernels/internal/cppmath.h

@@ -0,0 +1,43 @@
+/* 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_COMPILER_MLIR_LITE_KERNELS_INTERNAL_CPPMATH_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_CPPMATH_H_
+
+#include <cmath>
+
+// LINT.IfChange
+
+namespace tflite_migration {
+
+#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)
+
+}  // namespace tflite_migration
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/cppmath.h)
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_CPPMATH_H_

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

@@ -0,0 +1,32 @@
+/* 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_COMPILER_MLIR_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_
+
+// LINT.IfChange
+
+#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
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/optimized/neon_check.h)
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_OPTIMIZED_NEON_CHECK_H_

tensorflow/compiler/mlir/lite/kernels/internal/quantization_util.cc

@@ -0,0 +1,333 @@
+/* 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/compiler/mlir/lite/kernels/internal/quantization_util.h"
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+
+#include "tensorflow/compiler/mlir/lite/kernels/internal/compatibility_macros.h"
+#include "tensorflow/compiler/mlir/lite/kernels/internal/cppmath.h"
+
+namespace tflite_migration {
+
+// LINT.IfChange
+
+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_DCHECK(q_fixed <= (1LL << 31));
+  if (q_fixed == (1LL << 31)) {
+    q_fixed /= 2;
+    ++*shift;
+  }
+  TFLITE_DCHECK_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_DCHECK_GT(double_multiplier, 1.);
+  QuantizeMultiplier(double_multiplier, quantized_multiplier, left_shift);
+  TFLITE_DCHECK_GE(*left_shift, 0);
+}
+
+int64_t IntegerFrExp(double input, int* shift) {
+  // Make sure our assumptions about the double layout hold.
+  TFLITE_DCHECK_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);
+}
+
+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
+}
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/quantization_util.cc)
+
+}  // namespace tflite_migration

tensorflow/compiler/mlir/lite/kernels/internal/quantization_util.h

@@ -0,0 +1,166 @@
+/* Copyright 2024 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_COMPILER_MLIR_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_
+
+#include <cmath>
+#include <cstdint>
+#include <limits>
+
+namespace tflite_migration {
+
+// LINT.IfChange
+
+// 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.
+
+// 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);
+
+// 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/lite/kernels/internal/quantization_util.h)
+}  // namespace tflite_migration
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_QUANTIZATION_UTIL_H_

tensorflow/compiler/mlir/lite/kernels/internal/runtime_shape.cc

@@ -0,0 +1,144 @@
+/* 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.
+==============================================================================*/
+
+// This file is the MLIR copy of runtime_shape as part of the effort to
+// decouple TFLite from MLIR.
+// LINT.IfChange
+
+#include "tensorflow/compiler/mlir/lite/kernels/internal/runtime_shape.h"
+
+#include <cstdint>
+#include <cstring>
+#include <limits>
+
+#include "tensorflow/compiler/mlir/lite/kernels/internal/compatibility_macros.h"
+
+namespace mlir {
+
+namespace {
+
+bool CheckedMul(size_t a, size_t b, size_t& product) {
+  if (a != 0 && b > std::numeric_limits<size_t>::max() / a) {
+    return false;
+  }
+  product = a * b;
+  return true;
+}
+
+bool CheckedMul(size_t a, int32_t b, size_t& product) {
+  if (b < 0) return false;
+  return CheckedMul(a, static_cast<size_t>(b), product);
+}
+
+bool CheckedCast(size_t in, int& out) {
+  if (in > static_cast<size_t>(std::numeric_limits<int>::max())) return false;
+  out = static_cast<int>(in);
+  return true;
+}
+
+}  // namespace
+
+RuntimeShape::~RuntimeShape() {
+#ifndef TF_LITE_STATIC_MEMORY
+  if (size_ > kMaxSmallSize) {
+    delete[] dims_pointer_;
+  }
+#endif  // TF_LITE_STATIC_MEMORY
+}
+
+int32_t RuntimeShape::Dims(int i) const {
+  TFLITE_DCHECK_GE(i, 0);
+  TFLITE_DCHECK_LT(i, size_);
+#ifndef TF_LITE_STATIC_MEMORY
+  return size_ > kMaxSmallSize ? dims_pointer_[i] : dims_[i];
+#else
+  return dims_[i];
+#endif  // TF_LITE_STATIC_MEMORY
+}
+
+void RuntimeShape::ReplaceWith(int dimensions_count, const int32_t* dims_data) {
+  TFLITE_DCHECK_GE(dimensions_count, 0);
+#ifndef TF_LITE_STATIC_MEMORY
+  Resize(dimensions_count);
+  int32_t* dst_dims = DimsData();
+#else
+  TFLITE_DCHECK_LE(dimensions_count, kMaxSmallSize);
+  size_ = dimensions_count;
+  int32_t* dst_dims = DimsData();
+#endif  // TF_LITE_STATIC_MEMORY
+  std::memcpy(dst_dims, dims_data, dimensions_count * sizeof(int32_t));
+}
+
+int RuntimeShape::FlatSize() const {
+  int buffer_size = 1;
+  const int* dims_data = reinterpret_cast<const int*>(DimsData());
+  for (int i = 0; i < size_; i++) {
+    buffer_size *= dims_data[i];
+  }
+  return buffer_size;
+}
+
+bool RuntimeShape::CheckedNumElementsInRange(int start, int end,
+                                             size_t& out) const {
+  if (start < 0 || end < start || end > size_) return false;
+  size_t checked_out = 1;
+  const int32_t* dims_data = DimsData();
+  for (int i = start; i < end; ++i) {
+    if (!CheckedMul(checked_out, dims_data[i], checked_out)) return false;
+  }
+  out = checked_out;
+  return true;
+}
+
+bool RuntimeShape::CheckedNumElementsInRange(int start, int end,
+                                             int& out) const {
+  size_t checked_out = 0;
+  return CheckedNumElementsInRange(start, end, checked_out) &&
+         CheckedCast(checked_out, out);
+}
+
+bool RuntimeShape::CheckedSizeToDimension(int end, size_t& out) const {
+  return CheckedNumElementsInRange(0, end, out);
+}
+
+bool RuntimeShape::CheckedSizeToDimension(int end, int& out) const {
+  return CheckedNumElementsInRange(0, end, out);
+}
+
+bool RuntimeShape::CheckedSizeFromDimension(int start, size_t& out) const {
+  return CheckedNumElementsInRange(start, size_, out);
+}
+
+bool RuntimeShape::CheckedSizeFromDimension(int start, int& out) const {
+  return CheckedNumElementsInRange(start, size_, out);
+}
+
+bool RuntimeShape::CheckedFlatSize(size_t& flat_size) const {
+  return CheckedNumElementsInRange(0, size_, flat_size);
+}
+
+bool RuntimeShape::CheckedFlatSizeSkipDim(int skip_dim,
+                                          size_t& flat_size) const {
+  if (skip_dim < 0 || skip_dim >= size_) return false;
+  size_t prefix = 0;
+  if (!CheckedNumElementsInRange(0, skip_dim, prefix)) return false;
+  size_t suffix = 0;
+  if (!CheckedNumElementsInRange(skip_dim + 1, size_, suffix)) return false;
+  return CheckedMul(prefix, suffix, flat_size);
+}
+
+}  // namespace mlir
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/runtime_shape.cc)

tensorflow/compiler/mlir/lite/kernels/internal/runtime_shape.h

@@ -0,0 +1,406 @@
+/* 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_COMPILER_MLIR_LITE_KERNELS_INTERNAL_RUNTIME_SHAPE_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_RUNTIME_SHAPE_H_
+
+// This file is the MLIR copy of runtime_shape as part of the effort to
+// decouple TFLite from MLIR.
+// LINT.IfChange
+
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <initializer_list>
+#include <iterator>
+#include <memory>
+
+#include "tensorflow/compiler/mlir/lite/kernels/internal/compatibility_macros.h"
+
+namespace mlir {
+
+template <int N>
+struct Dims {
+  int sizes[N];
+  int strides[N];
+};
+
+class RuntimeShape {
+ public:
+  // Shapes with dimensions up to 6 are stored directly in the structure, while
+  // larger shapes are separately allocated.
+  static constexpr int kMaxSmallSize = 6;
+
+  RuntimeShape& operator=(RuntimeShape const&) = delete;
+
+  RuntimeShape() : size_(0) {}
+
+  explicit RuntimeShape(int dimensions_count) : size_(dimensions_count) {
+    TFLITE_DCHECK_GE(dimensions_count, 0);
+#ifndef TF_LITE_STATIC_MEMORY
+    if (dimensions_count > kMaxSmallSize) {
+      dims_pointer_ = new int32_t[dimensions_count];
+    }
+#else
+    TFLITE_DCHECK_LE(dimensions_count, kMaxSmallSize);
+#endif  // TF_LITE_STATIC_MEMORY
+  }
+
+#ifndef TF_LITE_STATIC_MEMORY
+  RuntimeShape(int shape_size, int32_t value) : size_(0) {
+    TFLITE_DCHECK_GE(shape_size, 0);
+    Resize(shape_size);
+#else
+  RuntimeShape(int shape_size, int32_t value) : size_(shape_size) {
+    TFLITE_DCHECK_GE(shape_size, 0);
+    TFLITE_DCHECK_LE(shape_size, kMaxSmallSize);
+#endif  // TF_LITE_STATIC_MEMORY
+    for (int i = 0; i < shape_size; ++i) {
+      SetDim(i, value);
+    }
+  }
+
+  RuntimeShape(int dimensions_count, const int32_t* dims_data) : size_(0) {
+    TFLITE_DCHECK_GE(dimensions_count, 0);
+    ReplaceWith(dimensions_count, dims_data);
+  }
+
+#ifndef TF_LITE_STATIC_MEMORY
+  RuntimeShape(const std::initializer_list<int> init_list) : size_(0) {
+    BuildFrom(init_list);
+  }
+
+  // Avoid using this constructor.  We should be able to delete it when C++17
+  // rolls out.
+  RuntimeShape(RuntimeShape const& other) : size_(other.DimensionsCount()) {
+    if (size_ > kMaxSmallSize) {
+      dims_pointer_ = new int32_t[size_];
+    }
+    std::memcpy(DimsData(), other.DimsData(), sizeof(int32_t) * size_);
+  }
+#endif  // TF_LITE_STATIC_MEMORY
+
+  bool operator==(const RuntimeShape& comp) const {
+    return this->size_ == comp.size_ &&
+           std::memcmp(DimsData(), comp.DimsData(), size_ * sizeof(int32_t)) ==
+               0;
+  }
+
+  ~RuntimeShape();
+
+  inline int32_t DimensionsCount() const { return size_; }
+
+  int32_t Dims(int i) const;
+
+  inline void SetDim(int i, int32_t val) {
+    TFLITE_DCHECK_GE(i, 0);
+    TFLITE_DCHECK_LT(i, size_);
+#ifndef TF_LITE_STATIC_MEMORY
+    if (size_ > kMaxSmallSize) {
+      dims_pointer_[i] = val;
+    } else {
+      dims_[i] = val;
+    }
+#else
+    dims_[i] = val;
+#endif  // TF_LITE_STATIC_MEMORY
+  }
+
+  inline int32_t* DimsData() {
+#ifndef TF_LITE_STATIC_MEMORY
+    return size_ > kMaxSmallSize ? dims_pointer_ : dims_;
+#else
+    return dims_;
+#endif  // TF_LITE_STATIC_MEMORY
+  }
+  inline const int32_t* DimsData() const {
+#ifndef TF_LITE_STATIC_MEMORY
+    return size_ > kMaxSmallSize ? dims_pointer_ : dims_;
+#else
+    return dims_;
+#endif  // TF_LITE_STATIC_MEMORY
+  }
+  // The caller must ensure that the shape is no bigger than 5-D.
+  inline const int32_t* DimsDataUpTo5D() const { return dims_; }
+
+#ifndef TF_LITE_STATIC_MEMORY
+  inline void Resize(int dimensions_count) {
+    TFLITE_DCHECK_GE(dimensions_count, 0);
+    const int32_t old_size = size_;
+    size_ = dimensions_count;
+
+    if (old_size <= kMaxSmallSize) {
+      if (dimensions_count <= kMaxSmallSize) {
+        return;
+      } else {  // Small to big.
+        int32_t* new_big_data = new int32_t[dimensions_count];
+        memcpy(new_big_data, dims_, sizeof(int32_t) * old_size);
+        dims_pointer_ = new_big_data;
+      }
+    } else {
+      if (dimensions_count > kMaxSmallSize && dimensions_count <= old_size) {
+        return;
+      }
+      std::unique_ptr<int32_t[]> old_data(dims_pointer_);
+      if (dimensions_count <= old_size) {  // Big to small.
+        memcpy(dims_, old_data.get(), sizeof(int32_t) * dimensions_count);
+      } else {  // Big to bigger.
+        dims_pointer_ = new int32_t[dimensions_count];
+        memcpy(dims_pointer_, old_data.get(), sizeof(int32_t) * old_size);
+      }
+    }
+  }
+#endif  // TF_LITE_STATIC_MEMORY
+
+  void ReplaceWith(int dimensions_count, const int32_t* dims_data);
+
+#ifndef TF_LITE_STATIC_MEMORY
+  template <typename T>
+  inline void BuildFrom(const T& src_iterable) {
+    const int dimensions_count =
+        std::distance(src_iterable.begin(), src_iterable.end());
+    Resize(dimensions_count);
+    int32_t* data = DimsData();
+    for (auto it : src_iterable) {
+      *data = it;
+      ++data;
+    }
+  }
+#endif  // TF_LITE_STATIC_MEMORY
+
+  // This will probably be factored out. Old code made substantial use of 4-D
+  // shapes, and so this function is used to extend smaller shapes. Note that
+  // (a) as Dims<4>-dependent code is eliminated, the reliance on this should be
+  // reduced, and (b) some kernels are stricly 4-D, but then the shapes of their
+  // inputs should already be 4-D, so this function should not be needed.
+  inline static RuntimeShape ExtendedShape(int new_shape_size,
+                                           const RuntimeShape& shape) {
+    TFLITE_DCHECK_GE(new_shape_size, 0);
+#ifdef TF_LITE_STATIC_MEMORY
+    TFLITE_DCHECK_LE(new_shape_size, kMaxSmallSize);
+#endif  // TF_LITE_STATIC_MEMORY
+    return RuntimeShape(new_shape_size, shape, 1);
+  }
+
+#ifndef TF_LITE_STATIC_MEMORY
+  inline void BuildFrom(const std::initializer_list<int> init_list) {
+    BuildFrom<const std::initializer_list<int>>(init_list);
+  }
+#endif  // TF_LITE_STATIC_MEMORY
+
+  // NOTE: This function does not handle potential integer overflow.
+  // The caller should verify that the size calculation won't overflow
+  // before calling this function -- for example, by calling `CheckedFlatSize`
+  // in a non-hot-path code and verifying that the return value is `true`
+  int FlatSize() const;
+
+  /**
+   * Returns false if any dimension is negative or if the product of all
+   * dimensions would overflow size_t.
+   * @param flat_size The output parameter in which to store the product of the
+   * dimensions.
+   * @return  False if any dimension is negative, or if the product would
+   * overflow size_t. Returns true otherwise. Returns 1 if the shape is empty.
+   */
+  bool CheckedFlatSize(size_t& flat_size) const;
+
+  /**
+   * Returns the checked product of dimensions in the half-open interval
+   * [start, end). It does a similar thing to FlatSize(), but it is more
+   * general and more secure.
+   * @param start The starting dimension index (inclusive).
+   * @param end The ending dimension index (exclusive).
+   * @param out The output parameter in which to store the product of the
+   * dimensions.
+   * @return  False if the range [start, end) is invalid or if any dimension is
+   * negative, or if the product would overflow size_t. Returns true otherwise.
+   * An empty range [start, start) will return 1.
+   */
+  bool CheckedNumElementsInRange(int start, int end, size_t& out) const;
+  bool CheckedNumElementsInRange(int start, int end, int& out) const;
+
+  /**
+   * Returns the checked product of dimensions in the half-open interval [0,
+   * end). It does a similar thing to FlatSize(), but it is more general and
+   * more secure.
+   * @param end The ending dimension index (exclusive).
+   * @param out The output parameter in which to store the product of the
+   * dimensions.
+   * @return  False if the index end is out of bounds or if any dimension in the
+   * interval [0, end) is negative, or if the product would overflow size_t.
+   * Returns true otherwise. Returns 1 if end is 0.
+   */
+  bool CheckedSizeToDimension(int end, size_t& out) const;
+
+  /**
+   * Returns the checked product of dimensions in the half-open interval [0,
+   * end). It does a similar thing to FlatSize(), but it is more general and
+   * more secure.
+   * @param end The ending dimension index (exclusive).
+   * @param out The output parameter in which to store the product of the
+   * dimensions.
+   * @return  False if the index end is out of bounds or if any dimension in the
+   * interval [0, end) is negative, or if the product would overflow int.
+   * Returns true otherwise. Returns 1 if end is 0.
+   */
+  bool CheckedSizeToDimension(int end, int& out) const;
+
+  /**
+   * Returns the checked product of dimensions in the half-open interval [start,
+   * DimensionsCount()).
+   * @param start The starting dimension index (inclusive).
+   * @param out The output parameter in which to store the product of the
+   * dimensions.
+   * @return  False if the index start is out of bounds or if any dimension in
+   * the interval [start, DimensionsCount()) is negative, or if the product
+   * would overflow size_t. Returns true otherwise. Returns 1 if start is
+   * DimensionsCount().
+   */
+  bool CheckedSizeFromDimension(int start, size_t& out) const;
+  /**
+   * Returns the checked product of dimensions in the half-open interval [start,
+   * DimensionsCount()).
+   * @param start The starting dimension index (inclusive).
+   * @param out The output parameter in which to store the product of the
+   * dimensions.
+   * @return  False if the index start is out of bounds or if any dimension in
+   * the interval [start, DimensionsCount()) is negative, or if the product
+   * would overflow int. Returns true otherwise. Returns 1 if start is
+   * DimensionsCount().
+   */
+  bool CheckedSizeFromDimension(int start, int& out) const;
+
+  /**
+   * Returns the checked product of dimensions in the half-open interval [0,
+   * DimensionsCount()) excluding the dimension at the given index.
+   * @param skip_dim The index of the dimension to exclude from the product.
+   * @param flat_size The output parameter in which to store the product of the
+   * dimensions.
+   * @return  False if the index skip_dim is out of bounds or if any dimension
+   * in the interval [0, DimensionsCount()) excluding the dimension at the given
+   * index is negative, or if the product would overflow size_t. Returns true
+   * otherwise. Returns 1 if the shape has only one dimension.
+   */
+  bool CheckedFlatSizeSkipDim(int skip_dim, size_t& flat_size) const;
+
+  bool operator!=(const RuntimeShape& comp) const { return !((*this) == comp); }
+
+ private:
+  // For use only by ExtendedShape(), written to guarantee (return-value) copy
+  // elision in C++17.
+  // This creates a shape padded to the desired size with the specified value.
+  RuntimeShape(int new_shape_size, const RuntimeShape& shape, int pad_value)
+#ifndef TF_LITE_STATIC_MEMORY
+      : size_(0) {
+    // If the following check fails, it is likely because a 4D-only kernel is
+    // being used with an array of larger dimension count.
+    TFLITE_DCHECK_GE(new_shape_size, shape.DimensionsCount());
+    Resize(new_shape_size);
+#else
+      : size_(new_shape_size) {
+    // If the following check fails, it is likely because a 4D-only kernel is
+    // being used with an array of larger dimension count.
+    TFLITE_DCHECK_GE(new_shape_size, shape.DimensionsCount());
+#endif  // TF_LITE_STATIC_MEMORY
+    const int size_increase = new_shape_size - shape.DimensionsCount();
+    for (int i = 0; i < size_increase; ++i) {
+      SetDim(i, pad_value);
+    }
+    std::memcpy(DimsData() + size_increase, shape.DimsData(),
+                sizeof(int32_t) * shape.DimensionsCount());
+  }
+
+  // Number of dimensions in the shape.
+  // size_t * sizeof(int32_t) is the number of bytes to allocate which should
+  // not exceed the maximum value of size_t.
+  int32_t size_;
+  union {
+    int32_t dims_[kMaxSmallSize];
+#ifndef TF_LITE_STATIC_MEMORY
+    int32_t* dims_pointer_;
+#endif  // TF_LITE_STATIC_MEMORY
+  };
+};
+
+// Converts inference-style shape to legacy tflite::Dims<4>.
+inline mlir::Dims<4> ToRuntimeDims(const mlir::RuntimeShape& array_shape) {
+  mlir::Dims<4> result;
+  const int dimensions_count = array_shape.DimensionsCount();
+  TFLITE_DCHECK_LE(dimensions_count, 4);
+  int cum_prod = 1;
+  for (int i = 0; i < 4; i++) {
+    const int new_dim =
+        (i < dimensions_count) ? array_shape.Dims(dimensions_count - 1 - i) : 1;
+    result.sizes[i] = new_dim;
+    result.strides[i] = cum_prod;
+    cum_prod *= new_dim;
+  }
+  return result;
+}
+
+#ifndef TF_LITE_STATIC_MEMORY
+// TODO(b/80418076): Move to legacy ops file, update invocations.
+inline RuntimeShape DimsToShape(const mlir::Dims<4>& dims) {
+  return RuntimeShape(
+      {dims.sizes[3], dims.sizes[2], dims.sizes[1], dims.sizes[0]});
+}
+#endif  // TF_LITE_STATIC_MEMORY
+
+// Since tensors with '0' in their shape are valid in TF, these offset functions
+// allow that as long as the corresponding index is also 0. It is upto the
+// calling ops to ensure that they perform verification checks on tensor shapes
+// if they don't support a particular behavior.
+
+inline int Offset(const RuntimeShape& shape, int i0, int i1, int i2, int i3) {
+  TFLITE_DCHECK_EQ(shape.DimensionsCount(), 4);
+  const int* dims_data = reinterpret_cast<const int*>(shape.DimsDataUpTo5D());
+  TFLITE_DCHECK((dims_data[0] == 0 && i0 == 0) ||
+                (i0 >= 0 && i0 < dims_data[0]));
+  TFLITE_DCHECK((dims_data[1] == 0 && i1 == 0) ||
+                (i1 >= 0 && i1 < dims_data[1]));
+  TFLITE_DCHECK((dims_data[2] == 0 && i2 == 0) ||
+                (i2 >= 0 && i2 < dims_data[2]));
+  TFLITE_DCHECK((dims_data[3] == 0 && i3 == 0) ||
+                (i3 >= 0 && i3 < dims_data[3]));
+  return ((i0 * dims_data[1] + i1) * dims_data[2] + i2) * dims_data[3] + i3;
+}
+
+inline int Offset(const RuntimeShape& shape, int i0, int i1, int i2, int i3,
+                  int i4) {
+  TFLITE_DCHECK_EQ(shape.DimensionsCount(), 5);
+  const int* dims_data = reinterpret_cast<const int*>(shape.DimsDataUpTo5D());
+  TFLITE_DCHECK((dims_data[0] == 0 && i0 == 0) ||
+                (i0 >= 0 && i0 < dims_data[0]));
+  TFLITE_DCHECK((dims_data[1] == 0 && i1 == 0) ||
+                (i1 >= 0 && i1 < dims_data[1]));
+  TFLITE_DCHECK((dims_data[2] == 0 && i2 == 0) ||
+                (i2 >= 0 && i2 < dims_data[2]));
+  TFLITE_DCHECK((dims_data[3] == 0 && i3 == 0) ||
+                (i3 >= 0 && i3 < dims_data[3]));
+  TFLITE_DCHECK((dims_data[4] == 0 && i4 == 0) ||
+                (i4 >= 0 && i4 < dims_data[4]));
+  return (((i0 * dims_data[1] + i1) * dims_data[2] + i2) * dims_data[3] + i3) *
+             dims_data[4] +
+         i4;
+}
+
+inline int Offset(const RuntimeShape& shape, int* index) {
+  return Offset(shape, index[0], index[1], index[2], index[3]);
+}
+
+}  // namespace mlir
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/runtime_shape.h)
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_RUNTIME_SHAPE_H_

tensorflow/compiler/mlir/lite/kernels/internal/runtime_shape_test.cc

@@ -0,0 +1,408 @@
+/* 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.
+==============================================================================*/
+// This file is the MLIR copy of runtime_shape as part of the effort to
+// decouple TFLite from MLIR.
+// LINT.IfChange
+
+#include "tensorflow/compiler/mlir/lite/kernels/internal/runtime_shape.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <initializer_list>
+#include <limits>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/algorithm/container.h"
+#include "absl/types/span.h"
+
+using testing::Each;
+using testing::ElementsAreArray;
+
+namespace mlir {
+namespace {
+
+constexpr int kSmallSize = RuntimeShape::kMaxSmallSize;
+constexpr int kBigSize = RuntimeShape::kMaxSmallSize + 1;
+
+std::vector<int32_t> IotaVector(int size, int start = 0) {
+  std::vector<int32_t> vec(size);
+  absl::c_iota(vec, start);
+  return vec;
+}
+
+absl::Span<const int32_t> AsSpan(const RuntimeShape& shape) {
+  return absl::Span<const int32_t>(shape.DimsData(), shape.DimensionsCount());
+}
+
+class RuntimeShapeTest : public testing::TestWithParam<int> {};
+
+TEST(RuntimeShapeTest, TestDefaultConstructor) {
+  const RuntimeShape shape;
+  EXPECT_EQ(shape.DimensionsCount(), 0);
+}
+
+TEST_P(RuntimeShapeTest, TestConstructorWithSize) {
+  const int size = GetParam();
+  const RuntimeShape shape(size);
+  EXPECT_EQ(shape.DimensionsCount(), size);
+}
+
+TEST_P(RuntimeShapeTest, TestConstructorWithSizeAndDefaultValue) {
+  const int size = GetParam();
+  const RuntimeShape shape(size, 34);
+  EXPECT_EQ(shape.DimensionsCount(), size);
+  EXPECT_THAT(AsSpan(shape), Each(34));
+}
+
+TEST_P(RuntimeShapeTest, TestConstructorFromCArray) {
+  const int size = GetParam();
+  const std::vector<int32_t> src = IotaVector(size);
+  const RuntimeShape shape(size, src.data());
+  EXPECT_EQ(shape.DimensionsCount(), size);
+  EXPECT_THAT(AsSpan(shape), ElementsAreArray(src));
+}
+
+TEST(RuntimeShapeTest, TestConstructorFromSmallInitList) {
+  std::initializer_list<int> init{1, 2, 3};
+  // Ensure we are testing a small initializer list.
+  ASSERT_LE(init.size(), RuntimeShape::kMaxSmallSize);
+  const RuntimeShape shape(init);
+  EXPECT_EQ(shape.DimensionsCount(), init.size());
+  EXPECT_THAT(AsSpan(shape), ElementsAreArray(init));
+}
+
+TEST(RuntimeShapeTest, TestConstructorFromBigInitList) {
+  std::initializer_list<int> init{1, 2, 3, 4, 5, 6, 7, 8, 9};
+  // Ensure we are testing a big initializer list.
+  ASSERT_GT(init.size(), RuntimeShape::kMaxSmallSize);
+  const RuntimeShape shape(init);
+  EXPECT_EQ(shape.DimensionsCount(), init.size());
+  EXPECT_THAT(AsSpan(shape), ElementsAreArray(init));
+}
+
+TEST_P(RuntimeShapeTest, TestCopyConstructorFromShape) {
+  const int size = GetParam();
+  const RuntimeShape src(size, 34);
+  const RuntimeShape dst(src);
+  EXPECT_EQ(dst.DimensionsCount(), src.DimensionsCount());
+  EXPECT_THAT(AsSpan(dst), ElementsAreArray(AsSpan(src)));
+}
+
+TEST_P(RuntimeShapeTest, TestEqualityOperator) {
+  const int size = GetParam();
+  const RuntimeShape shape1(size, 34);
+  const RuntimeShape shape2(size, 34);
+  EXPECT_TRUE(shape1 == shape2);
+  EXPECT_FALSE(shape1 != shape2);
+}
+
+TEST_P(RuntimeShapeTest, TestEqualityOperatorDifferentSizes) {
+  const int size = GetParam();
+  const RuntimeShape shape1(size, 34);
+  const RuntimeShape shape2(size + 1, 34);
+  EXPECT_FALSE(shape1 == shape2);
+  EXPECT_TRUE(shape1 != shape2);
+}
+
+TEST_P(RuntimeShapeTest, TestEqualityOperatorDifferentValues) {
+  const int size = GetParam();
+  const RuntimeShape shape1(size, 34);
+  const RuntimeShape shape2(size, 43);
+  EXPECT_FALSE(shape1 == shape2);
+  EXPECT_TRUE(shape1 != shape2);
+}
+
+TEST_P(RuntimeShapeTest, TestSetterGetter) {
+  const int size = GetParam();
+  RuntimeShape shape(size);
+  for (int i = 0; i < size; ++i) {
+    shape.SetDim(i, i);
+    EXPECT_EQ(shape.Dims(i), i);
+  }
+  EXPECT_THAT(AsSpan(shape), ElementsAreArray(IotaVector(size)));
+}
+
+TEST(RuntimeShapeTest, TestResizeSmallSmall) {
+  ASSERT_GE(kSmallSize, 1);
+  RuntimeShape shape(kSmallSize - 1, 23);
+  shape.Resize(kSmallSize);
+  EXPECT_EQ(shape.DimensionsCount(), kSmallSize);
+  EXPECT_THAT(absl::Span<const int32_t>(shape.DimsData(), kSmallSize - 1),
+              Each(23));
+}
+
+TEST(RuntimeShapeTest, TestResizeSmallBig) {
+  RuntimeShape shape(kSmallSize, 23);
+  shape.Resize(kBigSize);
+  EXPECT_EQ(shape.DimensionsCount(), kBigSize);
+  EXPECT_THAT(absl::Span<const int32_t>(shape.DimsData(), kSmallSize),
+              Each(23));
+}
+
+TEST(RuntimeShapeTest, TestResizeBigSmall) {
+  RuntimeShape shape(kBigSize, 23);
+  shape.Resize(kSmallSize);
+  EXPECT_EQ(shape.DimensionsCount(), kSmallSize);
+  EXPECT_THAT(absl::Span<const int32_t>(shape.DimsData(), kSmallSize),
+              Each(23));
+}
+
+TEST(RuntimeShapeTest, TestResizeDownBigBig) {
+  RuntimeShape shape(kBigSize + 3, 23);
+  shape.Resize(kBigSize);
+  EXPECT_EQ(shape.DimensionsCount(), kBigSize);
+  EXPECT_THAT(absl::Span<const int32_t>(shape.DimsData(), kBigSize), Each(23));
+}
+
+TEST(RuntimeShapeTest, TestResizeUpBigBig) {
+  RuntimeShape shape(kBigSize, 23);
+  shape.Resize(kBigSize + 1);
+  EXPECT_EQ(shape.DimensionsCount(), kBigSize + 1);
+  EXPECT_THAT(absl::Span<const int32_t>(shape.DimsData(), kBigSize), Each(23));
+}
+
+TEST_P(RuntimeShapeTest, TestReplaceWith) {
+  static_assert(
+      RuntimeShape::kMaxSmallSize > 2,
+      "kMaxSmallSize should be greater than 2 for this test to work.");
+  const int size = GetParam();
+  for (const int offset : {-2, 2}) {
+    const std::vector<int32_t> src =
+        IotaVector(offset + RuntimeShape::kMaxSmallSize);
+    RuntimeShape shape(size);
+    shape.ReplaceWith(src.size(), src.data());
+    EXPECT_EQ(shape.DimensionsCount(), src.size());
+    EXPECT_THAT(AsSpan(shape), testing::ElementsAreArray(src));
+  }
+}
+
+TEST_P(RuntimeShapeTest, TestBuildFrom) {
+  const int size = GetParam();
+  const std::vector<int32_t> src = IotaVector(size);
+  RuntimeShape shape;
+  shape.BuildFrom(src);
+  EXPECT_EQ(shape.DimensionsCount(), src.size());
+  EXPECT_THAT(AsSpan(shape), testing::ElementsAreArray(src));
+}
+
+TEST(RuntimeShapeTest, TestExtendedShapeSmall) {
+  ASSERT_GE(kSmallSize, 2);
+  const std::vector<int32_t> dims = IotaVector(kSmallSize - 2);
+  const RuntimeShape src(dims.size(), dims.data());
+  const RuntimeShape extended = RuntimeShape::ExtendedShape(kSmallSize, src);
+  EXPECT_EQ(extended.DimensionsCount(), kSmallSize);
+  EXPECT_EQ(extended.Dims(0), 1);
+  EXPECT_EQ(extended.Dims(1), 1);
+  EXPECT_THAT(absl::Span<const int32_t>(extended.DimsData() + 2, dims.size()),
+              ElementsAreArray(dims));
+}
+
+TEST(RuntimeShapeTest, TestExtendedShapeBig) {
+  ASSERT_GE(kSmallSize, 2);
+  const std::vector<int32_t> dims = IotaVector(kBigSize);
+  const RuntimeShape src(dims.size(), dims.data());
+  const RuntimeShape extended = RuntimeShape::ExtendedShape(kBigSize + 2, src);
+  EXPECT_EQ(extended.DimensionsCount(), kBigSize + 2);
+  EXPECT_EQ(extended.Dims(0), 1);
+  EXPECT_EQ(extended.Dims(1), 1);
+  EXPECT_THAT(absl::Span<const int32_t>(extended.DimsData() + 2, dims.size()),
+              ElementsAreArray(dims));
+}
+
+TEST(RuntimeShapeTest, TestExtendedShapeSmallToBig) {
+  const std::vector<int32_t> dims = IotaVector(kSmallSize);
+  const RuntimeShape src(dims.size(), dims.data());
+  const RuntimeShape extended = RuntimeShape::ExtendedShape(kBigSize, src);
+  EXPECT_EQ(extended.DimensionsCount(), kBigSize);
+  EXPECT_THAT(
+      absl::Span<const int32_t>(extended.DimsData(), kBigSize - kSmallSize),
+      Each(1));
+  EXPECT_THAT(absl::Span<const int32_t>(
+                  extended.DimsData() + kBigSize - kSmallSize, dims.size()),
+              ElementsAreArray(dims));
+}
+
+TEST_P(RuntimeShapeTest, TestFlatSize) {
+  const std::vector<int32_t> src = IotaVector(kSmallSize);
+  const RuntimeShape shape(src.size(), src.data());
+  int32_t flat_size = 1;
+  for (std::vector<int>::const_iterator it = src.begin(); it != src.end(); ++it)
+    flat_size *= *it;
+  EXPECT_EQ(shape.FlatSize(), flat_size);
+}
+
+TEST(RuntimeShapeTest, TestCheckedFlatSize) {
+  const RuntimeShape shape({2, 3, 4});
+  size_t flat_size = 0;
+  EXPECT_TRUE(shape.CheckedFlatSize(flat_size));
+  EXPECT_EQ(flat_size, 24u);
+}
+
+TEST(RuntimeShapeTest, TestCheckedFlatSizeScalar) {
+  const RuntimeShape shape;
+  size_t flat_size = 0;
+  EXPECT_TRUE(shape.CheckedFlatSize(flat_size));
+  EXPECT_EQ(flat_size, 1u);
+}
+
+TEST(RuntimeShapeTest, TestCheckedFlatSizeRejectsNegativeDim) {
+  const RuntimeShape shape({2, -3, 4});
+  size_t flat_size = 0;
+  EXPECT_FALSE(shape.CheckedFlatSize(flat_size));
+}
+
+TEST(RuntimeShapeTest, TestCheckedFlatSizeRejectsOverflow) {
+  const RuntimeShape shape({std::numeric_limits<int32_t>::max(),
+                            std::numeric_limits<int32_t>::max(),
+                            std::numeric_limits<int32_t>::max()});
+  size_t flat_size = 0;
+  EXPECT_FALSE(shape.CheckedFlatSize(flat_size));
+}
+
+TEST(RuntimeShapeTest, TestCheckedNumElementsInRange) {
+  const RuntimeShape shape({2, 3, 4, 5});
+  size_t flat_size = 0;
+  EXPECT_TRUE(
+      shape.CheckedNumElementsInRange(/*start=*/1, /*end=*/3, flat_size));
+  EXPECT_EQ(flat_size, 12u);
+}
+
+TEST(RuntimeShapeTest, TestCheckedNumElementsInRangeToInt) {
+  const RuntimeShape shape({2, 3, 4, 5});
+  int flat_size = 0;
+  EXPECT_TRUE(
+      shape.CheckedNumElementsInRange(/*start=*/1, /*end=*/3, flat_size));
+  EXPECT_EQ(flat_size, 12);
+}
+
+TEST(RuntimeShapeTest, TestCheckedNumElementsInRangeAllowsEmptyRange) {
+  const RuntimeShape shape({2, 3, 4});
+  size_t flat_size = 0;
+  EXPECT_TRUE(
+      shape.CheckedNumElementsInRange(/*start=*/1, /*end=*/1, flat_size));
+  EXPECT_EQ(flat_size, 1u);
+}
+
+TEST(RuntimeShapeTest, TestCheckedNumElementsInRangeRejectsInvalidRange) {
+  const RuntimeShape shape({2, 3, 4});
+  size_t flat_size = 0;
+  EXPECT_FALSE(
+      shape.CheckedNumElementsInRange(/*start=*/-1, /*end=*/2, flat_size));
+  EXPECT_FALSE(
+      shape.CheckedNumElementsInRange(/*start=*/2, /*end=*/1, flat_size));
+  EXPECT_FALSE(
+      shape.CheckedNumElementsInRange(/*start=*/0, /*end=*/4, flat_size));
+}
+
+TEST(RuntimeShapeTest, TestCheckedNumElementsInRangeRejectsNegativeDim) {
+  const RuntimeShape shape({2, -3, 4});
+  size_t flat_size = 0;
+  EXPECT_FALSE(
+      shape.CheckedNumElementsInRange(/*start=*/0, /*end=*/3, flat_size));
+}
+
+TEST(RuntimeShapeTest, TestCheckedNumElementsInRangeRejectsOverflow) {
+  const RuntimeShape shape({std::numeric_limits<int32_t>::max(),
+                            std::numeric_limits<int32_t>::max(),
+                            std::numeric_limits<int32_t>::max()});
+  size_t flat_size = 0;
+  EXPECT_FALSE(
+      shape.CheckedNumElementsInRange(/*start=*/0, /*end=*/3, flat_size));
+}
+
+TEST(RuntimeShapeTest, TestCheckedNumElementsInRangeToIntRejectsOverflow) {
+  const RuntimeShape shape({std::numeric_limits<int32_t>::max(),
+                            std::numeric_limits<int32_t>::max()});
+  int flat_size = 0;
+  EXPECT_FALSE(
+      shape.CheckedNumElementsInRange(/*start=*/0, /*end=*/2, flat_size));
+}
+
+TEST(RuntimeShapeTest, TestCheckedSizeToDimension) {
+  const RuntimeShape shape({2, 3, 4, 5});
+  size_t flat_size = 0;
+  EXPECT_TRUE(shape.CheckedSizeToDimension(/*end=*/3, flat_size));
+  EXPECT_EQ(flat_size, 24u);
+}
+
+TEST(RuntimeShapeTest, TestCheckedSizeFromDimension) {
+  const RuntimeShape shape({2, 3, 4, 5});
+  size_t flat_size = 0;
+  EXPECT_TRUE(shape.CheckedSizeFromDimension(/*start=*/1, flat_size));
+  EXPECT_EQ(flat_size, 60u);
+}
+
+TEST(RuntimeShapeTest, TestCheckedFlatSizeSkipDim) {
+  const RuntimeShape shape({2, 3, 4, 5});
+  size_t flat_size = 0;
+  EXPECT_TRUE(shape.CheckedFlatSizeSkipDim(/*skip_dim=*/2, flat_size));
+  EXPECT_EQ(flat_size, 30u);
+}
+
+TEST(RuntimeShapeTest, TestCheckedFlatSizeSkipDimRejectsInvalidSkipDim) {
+  const RuntimeShape shape({2, 3, 4});
+  size_t flat_size = 0;
+  EXPECT_FALSE(shape.CheckedFlatSizeSkipDim(/*skip_dim=*/-1, flat_size));
+  EXPECT_FALSE(shape.CheckedFlatSizeSkipDim(/*skip_dim=*/3, flat_size));
+}
+
+TEST(RuntimeShapeTest, TestCheckedFlatSizeSkipDimRejectsNegativeDim) {
+  const RuntimeShape shape({2, -3, 4});
+  size_t flat_size = 0;
+  EXPECT_FALSE(shape.CheckedFlatSizeSkipDim(/*skip_dim=*/0, flat_size));
+}
+
+TEST(RuntimeShapeTest, TestCheckedFlatSizeSkipDimRejectsOverflow) {
+  const RuntimeShape shape({std::numeric_limits<int32_t>::max(),
+                            std::numeric_limits<int32_t>::max(),
+                            std::numeric_limits<int32_t>::max(),
+                            std::numeric_limits<int32_t>::max()});
+  size_t flat_size = 0;
+  EXPECT_FALSE(shape.CheckedFlatSizeSkipDim(/*skip_dim=*/0, flat_size));
+}
+
+#ifndef NDEBUG
+TEST(RuntimeShapeTest, NegativeDimensionsCountInConstructor) {
+  EXPECT_DEATH(RuntimeShape shape(-1), "");
+}
+
+TEST(RuntimeShapeTest, NegativeDimensionsCountInConstructorWithValue) {
+  EXPECT_DEATH(RuntimeShape shape(-1, 0), "");
+}
+
+TEST(RuntimeShapeTest, NegativeDimensionsCountInResize) {
+  RuntimeShape shape;
+  EXPECT_DEATH(shape.Resize(-1), "");
+}
+
+TEST(RuntimeShapeTest, NegativeDimensionsCountInReplaceWith) {
+  RuntimeShape shape;
+  int32_t dims_data[] = {1, 2, 3};
+  EXPECT_DEATH(shape.ReplaceWith(-1, dims_data), "");
+}
+#endif
+
+INSTANTIATE_TEST_SUITE_P(BigSmall, RuntimeShapeTest,
+                         testing::Values(kSmallSize, kBigSize),
+                         [](const testing::TestParamInfo<int>& info) {
+                           return info.param == kSmallSize ? "Small" : "Big";
+                         });
+
+}  // namespace
+}  // namespace mlir
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/runtime_shape_test.cc)

tensorflow/compiler/mlir/lite/kernels/internal/utils/BUILD

@@ -0,0 +1,27 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/compiler/mlir/lite:build_def.bzl", "tflite_copts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+# LINT.IfChange(sparsity_format_converter)
+
+cc_library(
+    name = "sparsity_format_converter",
+    srcs = ["sparsity_format_converter.cc"],
+    hdrs = ["sparsity_format_converter.h"],
+    compatible_with = get_compatible_with_portable(),
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/compiler/mlir/lite/core/c:tflite_common",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+# LINT.ThenChange(//tensorflow/lite/kernels/internal/utils/BUILD)

tensorflow/compiler/mlir/lite/kernels/internal/utils/sparsity_format_converter.cc

@@ -0,0 +1,204 @@
+/* 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/compiler/mlir/lite/kernels/internal/utils/sparsity_format_converter.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/compiler/mlir/lite/core/c/tflite_types.h"
+
+namespace tflite_migration {
+namespace internal {
+namespace sparsity {
+
+// LINT.IfChange
+
+template <typename T>
+FormatConverter<T>::FormatConverter(
+    const std::vector<int>& shape, const std::vector<int>& traversal_order,
+    const std::vector<TfLiteDimensionType>& format,
+    const std::vector<int>& block_size, const std::vector<int>& block_map)
+    : dense_shape_(shape),
+      traversal_order_(traversal_order),
+      block_size_(block_size),
+      block_map_(block_map) {
+  dense_size_ = 1;
+  int block_dim = 0;
+  blocked_shape_.resize(shape.size());
+  format_.resize(shape.size() + block_map.size());
+  for (int i = 0; i < shape.size(); i++) {
+    format_[i] = format[traversal_order[i]];
+    dense_size_ *= shape[i];
+    if (block_dim < block_map.size() && block_map[block_dim] == i) {
+      blocked_shape_[i] = shape[i] / block_size[block_dim];
+      block_dim++;
+    } else {
+      blocked_shape_[i] = shape[i];
+    }
+  }
+
+  // Only dense blocks are supported.
+  for (int i = 0; i < block_map.size(); i++) {
+    format_[i + shape.size()] = kTfLiteDimDense;
+  }
+}
+
+template <typename T>
+void FormatConverter<T>::DenseToSparse(const T* src_data) {
+  int num_original_dims = dense_shape_.size();
+  int num_block_dims = block_map_.size();
+  int num_expanded_dims = num_original_dims + num_block_dims;
+  std::vector<int> expanded_shape(num_expanded_dims);
+  for (int i = 0; i < num_expanded_dims; i++) {
+    if (i < num_original_dims) {
+      expanded_shape[i] = blocked_shape_[i];
+    } else {
+      expanded_shape[i] = block_size_[i - num_original_dims];
+    }
+  }
+
+  std::vector<int> shape_offset(num_original_dims);
+  shape_offset[shape_offset.size() - 1] = 1;
+  for (int i = num_original_dims - 1; i > 0; --i) {
+    shape_offset[i - 1] = shape_offset[i] * dense_shape_[i];
+  }
+
+  std::vector<int> expanded_shape_offset(num_expanded_dims);
+  for (int i = 0; i < num_original_dims; ++i) {
+    expanded_shape_offset[i] = shape_offset[i];
+  }
+  for (int i = 0; i < num_block_dims; ++i) {
+    int mapped_dim = block_map_[i];
+    expanded_shape_offset[num_original_dims + i] = shape_offset[mapped_dim];
+    expanded_shape_offset[mapped_dim] *= block_size_[i];
+  }
+
+  std::vector<int> dst_ordered_offset(num_expanded_dims);
+  for (int i = 0; i < num_expanded_dims; ++i) {
+    dst_ordered_offset[i] = expanded_shape_offset[traversal_order_[i]];
+  }
+
+  std::vector<bool> dst_dim_has_nonzeroes(num_expanded_dims);
+  std::fill(dst_dim_has_nonzeroes.begin(), dst_dim_has_nonzeroes.end(), false);
+  std::vector<int> inner_compressed_dim(num_expanded_dims);
+  int most_recent_compressed_dim = -1;
+  std::vector<int> num_segments_of_next_compressed_dim(num_expanded_dims);
+  int segment_count = 1;
+  for (int i = num_expanded_dims - 1; i >= 0; --i) {
+    inner_compressed_dim[i] = most_recent_compressed_dim;
+    if (format_[i] == kTfLiteDimSparseCSR) {
+      most_recent_compressed_dim = i;
+      num_segments_of_next_compressed_dim[i] = segment_count;
+      segment_count = 1;
+    } else {
+      num_segments_of_next_compressed_dim[i] = -1;
+      segment_count *= expanded_shape[traversal_order_[i]];
+    }
+  }
+
+  dim_metadata_.resize(num_expanded_dims * 2);
+  std::vector<int> dst_sparse_dims;
+  dst_sparse_dims.reserve(num_expanded_dims);
+  for (int i = 0; i < num_expanded_dims; ++i) {
+    dim_metadata_[i * 2].clear();
+    dim_metadata_[i * 2 + 1].clear();
+    if (format_[i] == kTfLiteDimDense) {
+      // If dimension is dense, just store the shape.
+      dim_metadata_[i * 2].push_back(expanded_shape[traversal_order_[i]]);
+    } else {
+      dim_metadata_[i * 2].push_back(0);  // Segment array always begins with 0.
+      dst_sparse_dims.push_back(i);       // Add dimension to the sparse list.
+    }
+  }
+
+  // This algorithm assumes that the block size is small enough for all the
+  // elements to fit in cache, so the strided accesses from different traversal
+  // order and the write-first-erase-later strategy shouldn't be too slow
+  int dst_dim_idx = num_expanded_dims;
+  std::vector<int> coordinate(num_expanded_dims, 0);
+  int dense_tensor_idx = 0;
+  while (dst_dim_idx >= 0) {
+    if (dst_dim_idx == num_expanded_dims) {
+      // We have a complete coordinate. Add the element to the value array if it
+      // is not zero, or if the last dimension is dense.
+      if (!IsZero(src_data[dense_tensor_idx])) {
+        data_.push_back(src_data[dense_tensor_idx]);
+        // Mark all sparse dimensions that their current indices have nonzeroes.
+        for (auto dst_dim : dst_sparse_dims) {
+          if (!dst_dim_has_nonzeroes[dst_dim]) {
+            // Only add the index to the indices array if the current nonzero
+            // is the first nonzero of the block.
+            dim_metadata_[2 * dst_dim + 1].push_back(coordinate[dst_dim]);
+            dst_dim_has_nonzeroes[dst_dim] = true;
+          }
+        }
+      } else if (format_[num_expanded_dims - 1] == kTfLiteDimDense) {
+        data_.push_back(src_data[dense_tensor_idx]);
+      }
+      --dst_dim_idx;
+    } else {
+      int original_dim_idx = traversal_order_[dst_dim_idx];
+      int dim_size = expanded_shape[original_dim_idx];
+      if (dst_dim_has_nonzeroes[dst_dim_idx]) {
+        // If the previous block has nonzeroes, reset the flag to false since
+        // we have just moved to a new block.
+        dst_dim_has_nonzeroes[dst_dim_idx] = false;
+      } else if (format_[dst_dim_idx] == kTfLiteDimSparseCSR) {
+        // This block is empty. Delete unnecessary values if compressed.
+        int next_compressed_dim = inner_compressed_dim[dst_dim_idx];
+        int erase_offset = dim_metadata_[2 * dst_dim_idx + 1].size() *
+                           num_segments_of_next_compressed_dim[dst_dim_idx];
+        if (next_compressed_dim >= 0) {
+          auto& segments = dim_metadata_[2 * inner_compressed_dim[dst_dim_idx]];
+          segments.erase(segments.begin() + 1 + erase_offset, segments.end());
+        } else {
+          data_.erase(data_.begin() + erase_offset, data_.end());
+        }
+      }
+      if (++coordinate[dst_dim_idx] < dim_size) {
+        // The current dst_dim_idx is valid (not out of bound).
+        dense_tensor_idx += dst_ordered_offset[dst_dim_idx];
+        ++dst_dim_idx;
+      } else {
+        // dst_dim_idx has reached its dim size. Update segment array and go
+        // back to incrementing the previous dimension (dst_dim_idx - 1).
+        if (format_[dst_dim_idx] == kTfLiteDimSparseCSR) {
+          dim_metadata_[2 * dst_dim_idx].push_back(
+              dim_metadata_[2 * dst_dim_idx + 1].size());
+        }
+        coordinate[dst_dim_idx] = -1;
+        dense_tensor_idx -= dst_ordered_offset[dst_dim_idx] * dim_size;
+        --dst_dim_idx;
+      }
+    }
+  }
+}
+
+template <typename T>
+bool FormatConverter<T>::IsZero(const T val) {
+  return (val == static_cast<T>(0));
+}
+
+template class FormatConverter<int8_t>;
+template class FormatConverter<float>;
+template class FormatConverter<Eigen::half>;
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/utils/sparsity_format_converter.cc)
+
+}  // namespace sparsity
+}  // namespace internal
+}  // namespace tflite_migration

tensorflow/compiler/mlir/lite/kernels/internal/utils/sparsity_format_converter.h

@@ -0,0 +1,102 @@
+/* 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_COMPILER_MLIR_LITE_KERNELS_INTERNAL_UTILS_SPARSITY_FORMAT_CONVERTER_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_UTILS_SPARSITY_FORMAT_CONVERTER_H_
+
+#include <vector>
+
+#include "Eigen/Core"  // from @eigen_archive
+#include "tensorflow/compiler/mlir/lite/core/c/tflite_types.h"
+
+namespace tflite_migration {
+namespace internal {
+namespace sparsity {
+
+// LINT.IfChange
+
+// A converter that keeps an internal representation of sparse tensor parameters
+// and converts tensors between dense and sparse formats.
+template <typename T>
+class FormatConverter {
+ public:
+  /*
+   * Creates a dense to sparse converter.
+   * @param shape             Shape of the dense tensor.
+   * @param traversal_order   In what order to traverse all dimensions,
+   *                          including block dimensions.
+   * @param format            Whether each dimension in the dense tensor is
+   *                          dense or sparse (not in the traversal order).
+   * @param block_size        Size of each block dimension.
+   * @param block_map         Map from block dimension to original tensor
+   *                          dimension.
+   */
+  FormatConverter(const std::vector<int>& shape,
+                  const std::vector<int>& traversal_order,
+                  const std::vector<TfLiteDimensionType>& format,
+                  const std::vector<int>& block_size = {},
+                  const std::vector<int>& block_map = {});
+
+  const std::vector<T>& GetData() { return data_; }
+
+  const std::vector<std::vector<int>>& GetDimMetadata() {
+    return dim_metadata_;
+  }
+
+  // Method for dense to sparse conversion. Need to call GetData() method to get
+  // the compressed data.
+
+  void DenseToSparse(const T* src_data);
+
+  // Check if val is equal to zero.
+  bool IsZero(T val);
+
+  // Shape of the conceptual dense tensor.
+  std::vector<int> dense_shape_;
+  // Shape of the dense tensor with inner blocks reduced. For example, a (4, 4)
+  // tensor with (2, 2) block has blocked_shape (2, 2).
+  std::vector<int> blocked_shape_;
+  // Total number of elements in the dense tensor.
+  size_t dense_size_;
+  // Has n(original dimension)+k(block_dimension) elements.
+  std::vector<int> traversal_order_;
+  // Format of each dimension in the traversal order.
+  std::vector<TfLiteDimensionType> format_;
+  // Size of each block dimension, in the same order as block map.
+  std::vector<int> block_size_;
+  // Map from block dimension to the original tensor dimension.
+  std::vector<int> block_map_;
+  // Metadata of each dimension in the traversal order.
+  // Each dimension needs two vectors. For dense dimensions, the first vector
+  // stores the size of that dimension, and the second vector is empty. For
+  // sparse dimensions, the first vector stores the segments and the second one
+  // stores the indices.
+  std::vector<std::vector<int>> dim_metadata_;
+  // Actual buffer holding data after conversion. Could be sparse buffer or
+  // dense buffer.
+  std::vector<T> data_;
+};
+
+extern template class FormatConverter<int32_t>;
+extern template class FormatConverter<int8_t>;
+extern template class FormatConverter<float>;
+extern template class FormatConverter<Eigen::half>;
+
+// LINT.ThenChange(//tensorflow/lite/kernels/internal/utils/sparsity_format_converter.h)
+
+}  // namespace sparsity
+}  // namespace internal
+}  // namespace tflite_migration
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_INTERNAL_UTILS_SPARSITY_FORMAT_CONVERTER_H_

tensorflow/compiler/mlir/lite/kernels/padding.h

@@ -0,0 +1,59 @@
+/* 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_COMPILER_MLIR_LITE_KERNELS_PADDING_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_PADDING_H_
+
+// LINT.IfChange
+#include "tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h"
+
+namespace tflite_migration {
+
+// Matching GetWindowedOutputSize in TensorFlow.
+inline int ComputeOutSize(TfLitePadding padding, int image_size,
+                          int filter_size, int stride, int dilation_rate = 1) {
+  int effective_filter_size = (filter_size - 1) * dilation_rate + 1;
+
+  // TODO(b/186448822): This uses 0 since the function has no other way to
+  // report error case
+  if (stride == 0) return 0;
+
+  switch (padding) {
+    case kTfLitePaddingSame:
+      return (image_size + stride - 1) / stride;
+    case kTfLitePaddingValid:
+      return (image_size + stride - effective_filter_size) / stride;
+    default:
+      return 0;
+  }
+}
+
+// It's not guaranteed that padding is symmetric. It's important to keep
+// offset for algorithms need all paddings.
+inline int ComputePaddingWithOffset(int stride, int dilation_rate, int in_size,
+                                    int filter_size, int out_size,
+                                    int* offset) {
+  int effective_filter_size = (filter_size - 1) * dilation_rate + 1;
+  int total_padding =
+      ((out_size - 1) * stride + effective_filter_size - in_size);
+  total_padding = total_padding > 0 ? total_padding : 0;
+  *offset = total_padding % 2;
+  return total_padding / 2;
+}
+
+}  // namespace tflite_migration
+
+// LINT.ThenChange(//tensorflow/lite/kernels/padding.h)
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_KERNELS_PADDING_H_