chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/lite/utils/arithmetic_count_util.h

@@ -0,0 +1,90 @@
+/* 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_UTILS_ARITHMETIC_COUNT_UTIL_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_ARITHMETIC_COUNT_UTIL_H_
+
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+// For add/mul/div/sub and other broadcastable ops.
+class ArithmeticCountUtilHelper {
+ public:
+  static bool GetFirstOutputCount(mlir::Operation* op, int64_t* count) {
+    auto output = op->getResult(0);
+    auto output_type =
+        mlir::dyn_cast_or_null<mlir::RankedTensorType>(output.getType());
+    if (!output_type || !output_type.hasStaticShape()) return false;
+
+    *count = output_type.getNumElements();
+    return true;
+  }
+
+  static bool GetInputTensorTotalSize(mlir::Operation* op, int64_t* count) {
+    int64_t total_count = 0;
+    for (auto input : op->getOperands()) {
+      auto input_type =
+          mlir::dyn_cast_or_null<mlir::RankedTensorType>(input.getType());
+      if (!input_type || !input_type.hasStaticShape()) {
+        return false;
+      }
+      total_count += input_type.getNumElements();
+    }
+    *count = total_count;
+    return true;
+  }
+
+  // For conv2d/depthwise_conv/fully_connected ops.
+  // This algorithm actually comes from TOCO tooling_util.cc
+  static bool GetArithmeticCountForConvAndFullyconnectedOp(mlir::Operation* op,
+                                                           int64_t* count) {
+    auto weight = op->getOperand(1);
+    auto weight_type =
+        mlir::dyn_cast_or_null<mlir::RankedTensorType>(weight.getType());
+    if (weight_type == nullptr || !weight_type.hasStaticShape()) return false;
+
+    auto output = op->getResult(0);
+    auto output_type =
+        mlir::dyn_cast_or_null<mlir::RankedTensorType>(output.getType());
+    if (output_type == nullptr || !output_type.hasStaticShape()) return false;
+
+    int64_t cols = 1;
+    for (int i = 0; i < output_type.getRank() - 1; ++i) {
+      cols *= output_type.getDimSize(i);
+    }
+    const int64_t cost_per_col = 2 * weight_type.getNumElements();
+
+    *count = cost_per_col * cols;
+
+    auto bias = op->getOperand(2);
+    if (bias) {
+      auto bias_type =
+          mlir::dyn_cast_or_null<mlir::RankedTensorType>(bias.getType());
+      if (bias_type && bias_type.hasStaticShape()) {
+        *count += output_type.getNumElements();
+      }
+    }
+
+    return true;
+  }
+};
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_ARITHMETIC_COUNT_UTIL_H_

tensorflow/compiler/mlir/lite/utils/attribute_utils.cc

@@ -0,0 +1,94 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <climits>
+#include <cstddef>
+#include <cstdint>
+
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/MathExtras.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+FloatAttr ExtractSingleElementAsFloat(ElementsAttr attr) {
+  if (attr.getShapedType().getNumElements() != 1 ||
+      !mlir::isa<FloatType>(attr.getShapedType().getElementType())) {
+    return {};
+  }
+  return attr.getSplatValue<FloatAttr>();
+}
+
+FloatAttr GetSingleElementAsFloatOrSelf(Attribute attr) {
+  if (auto m = mlir::dyn_cast_or_null<ElementsAttr>(attr)) {
+    return ExtractSingleElementAsFloat(m);
+  } else {
+    return mlir::dyn_cast_or_null<FloatAttr>(attr);
+  }
+}
+
+IntegerAttr ExtractSingleElementAsInteger(ElementsAttr attr) {
+  if (attr.getShapedType().getNumElements() != 1 ||
+      !attr.getShapedType().getElementType().isSignlessInteger()) {
+    return {};
+  }
+  return attr.getSplatValue<IntegerAttr>();
+}
+
+size_t GetDenseElementBitWidth(Type elt_type) {
+  // Align the width for complex to 8 to make storage and interpretation easier.
+  if (ComplexType comp = llvm::dyn_cast<ComplexType>(elt_type))
+    return llvm::alignTo<8>(GetDenseElementBitWidth(comp.getElementType())) * 2;
+  if (elt_type.isIndex()) return IndexType::kInternalStorageBitWidth;
+  return elt_type.getIntOrFloatBitWidth();
+}
+
+bool IsValidIntOrFloat(Type type, int64_t data_element_size, bool is_int,
+                       bool is_signed) {
+  // Make sure that the data element size is the same as the type element width.
+  auto dense_elt_bit_width =
+      std::max(GetDenseElementBitWidth(type), static_cast<size_t>(CHAR_BIT));
+  auto data_size = static_cast<size_t>(data_element_size * CHAR_BIT);
+  if (dense_elt_bit_width != data_size) {
+    return false;
+  }
+
+  // Check that the element type is either float or integer or index.
+  if (!is_int) {
+    return llvm::isa<FloatType>(type);
+  }
+  if (type.isIndex()) return true;
+
+  auto int_type = llvm::dyn_cast<IntegerType>(type);
+  if (!int_type) {
+    return false;
+  }
+
+  // Make sure signedness semantics is consistent.
+  if (int_type.isSignless()) return true;
+
+  return int_type.isSigned() == is_signed;
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/attribute_utils.h

@@ -0,0 +1,127 @@
+/* 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.
+==============================================================================*/
+
+// This header file defines common utils used by TFLite transformation
+// passes to work with op attributes.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_ATTRIBUTE_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_ATTRIBUTE_UTILS_H_
+
+#include <sys/stat.h>
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/DialectResourceBlobManager.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+// Returns true if none of the three attributes are empty.
+inline bool HasAll3Attrs(Attribute a, Attribute b, Attribute c) {
+  return a != Attribute() && b != Attribute() && c != Attribute();
+}
+
+// Returns the single float element from an ElementsAttr. Returns empty
+// attribute if the number of elements in the attribute is not 1 or the
+// element isn't a float attribute.
+FloatAttr ExtractSingleElementAsFloat(ElementsAttr attr);
+
+// Returns the single float element if the input is an ElementsAttr, or return
+// itself as a float element. Returns empty attribute if the number of elements
+// in the attribute is not 1, the element or itself isn't a float attribute.
+FloatAttr GetSingleElementAsFloatOrSelf(Attribute attr);
+
+// Returns the single integer element from an ElementsAttr. Returns empty
+// attribute if the number of elements in the attribute is not 1 or the
+// element isn't a integer attribute.
+IntegerAttr ExtractSingleElementAsInteger(ElementsAttr attr);
+
+/// Check the information for a C++ data type, check if this type is valid for
+/// the current attribute. This method is used to verify specific type
+/// invariants that the templatized 'getValues' method cannot.
+bool IsValidIntOrFloat(Type type, int64_t data_elt_size, bool is_int,
+                       bool is_signed);
+
+/// Type trait used to check if the given type T is a potentially valid C++
+/// floating point type that can be used to access the underlying element
+/// types of a DenseElementsAttr.
+template <typename T>
+struct is_valid_cpp_fp_type {
+  /// The type is a valid floating point type if it is a builtin floating
+  /// point type, or is a potentially user defined floating point type. The
+  /// latter allows for supporting users that have custom types defined for
+  /// bfloat16/half/etc.
+  static constexpr bool value = llvm::is_one_of<T, float, double>::value ||
+                                (std::numeric_limits<T>::is_specialized &&
+                                 !std::numeric_limits<T>::is_integer);
+};
+
+/// Return the bit width which ElementsAttr should use for this type.
+size_t GetDenseElementBitWidth(Type elt_type);
+
+// Returns the values of the given ElementsAttr as a ArrayRef of ElementType.
+// Returns {std::nullopt} if the attribute is empty.
+// TODO(b/707702324): Add support for type like IntergAttr and APInt, etc. This
+// is a temporary solution to unblock the LiteRT. Ideally MLIR should provide
+// a common API to access the values of an DenseResourceElementsAttr.
+template <typename ElementType>
+using IntFloatValueTemplateCheckT =
+    std::enable_if_t<(!std::is_same<ElementType, bool>::value &&
+                      std::numeric_limits<ElementType>::is_integer) ||
+                     is_valid_cpp_fp_type<ElementType>::value>;
+template <typename ElementType,
+          typename = IntFloatValueTemplateCheckT<ElementType>>
+inline ArrayRef<ElementType> GetValues(ElementsAttr attr) {
+  Type element_type = attr.getElementType();
+
+  // Check if the element type is not valid for the given ElementType, return
+  // empty ArrayRef.
+  if (!IsValidIntOrFloat(element_type, sizeof(ElementType),
+                         std::numeric_limits<ElementType>::is_integer,
+                         std::numeric_limits<ElementType>::is_signed)) {
+    assert(false && "Incompatible dtype expected from the given ElementsAttr");
+  }
+
+  if (auto dense_elements_attr = dyn_cast<DenseElementsAttr>(attr)) {
+    auto raw_data = dense_elements_attr.getRawData();
+    if (raw_data.empty()) {
+      return {};
+    }
+    return llvm::ArrayRef<ElementType>(
+        reinterpret_cast<const ElementType*>(raw_data.data()),
+        raw_data.size() / sizeof(ElementType));
+  } else if (auto dense_resource_elements_attr =
+                 dyn_cast<DenseResourceElementsAttr>(attr)) {
+    if (AsmResourceBlob* blob =
+            dense_resource_elements_attr.getRawHandle().getBlob())
+      return blob->getDataAs<ElementType>();
+    return {};
+  }
+  return {};
+}
+
+}  // end namespace TFL
+}  // end namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_ATTRIBUTE_UTILS_H_

tensorflow/compiler/mlir/lite/utils/const_tensor_utils.cc

@@ -0,0 +1,529 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/const_tensor_utils.h"
+
+#include <algorithm>
+#include <cassert>
+#include <climits>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include "absl/base/casts.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "Eigen/Core"  // from @eigen_archive
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/bit.h"
+#include "llvm/Support/Endian.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
+#include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
+#include "tensorflow/compiler/mlir/lite/utils/low_bit_utils.h"
+#include "tensorflow/compiler/mlir/lite/utils/string_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tsl/platform/statusor.h"
+
+namespace mlir {
+namespace TFL {
+namespace {
+
+using ::absl::StatusOr;
+using ::mlir::Builder;
+using ::mlir::quant::QuantizedType;
+using ::tflite::TensorT;
+
+// The buffers in TFLite flatbuffers have their contents stored as a vector of
+// bytes that represent host endianness values.
+// The read_size parameter is present to allow reading both float16 and float32
+// without a case split.
+template <typename T>
+llvm::SmallVector<mlir::APInt> ReadAsHostEndian(ArrayRef<uint8_t> bytes) {
+  llvm::SmallVector<mlir::APInt> ret;
+  size_t read_size = sizeof(T);
+  size_t bytes_len = bytes.size();
+  assert(bytes_len % read_size == 0);
+
+  size_t elem_count = bytes_len / read_size;
+  ret.reserve(elem_count);
+
+  const char* data_ptr = reinterpret_cast<const char*>(bytes.data());
+  for (size_t i = 0; i < elem_count; i++) {
+    T val = llvm::support::endian::readNext<T, llvm::endianness::native,
+                                            llvm::support::unaligned>(data_ptr);
+    ret.push_back(mlir::APInt(sizeof(T) * 8, val));
+  }
+  return ret;
+}
+
+// If the values in the buffer can be clamped to a bitwidth, truncate
+// and return the new clamped integer width.
+void truncateLimitedIntegerAPInt(llvm::SmallVector<mlir::APInt>& values) {
+  mlir::APInt min = values[0];
+  mlir::APInt max = values[0];
+  for (mlir::APInt& val : values) {
+    min = llvm::APIntOps::smin(val, min);
+    max = llvm::APIntOps::smax(val, max);
+  }
+
+  for (int64_t bw = 8; bw < min.getBitWidth(); bw += bw) {
+    mlir::APInt limitMin =
+        mlir::APInt::getSignedMinValue(bw).sext(min.getBitWidth());
+    mlir::APInt limitMax =
+        mlir::APInt::getSignedMaxValue(bw).sext(min.getBitWidth());
+
+    // Skips to the next bitwidth if the min and max values are out of the range
+    // for the current bitwidth.
+    if (min.sle(limitMin) || max.sle(limitMin) || min.sge(limitMax) ||
+        max.sge(limitMax)) {
+      continue;
+    }
+
+    for (mlir::APInt& val : values) {
+      val = val.trunc(bw);
+    }
+    break;
+  }
+}
+
+}  // namespace
+
+bool IsQuantized(const TensorT& tensor) {
+  return (tensor.quantization != nullptr) &&
+         !tensor.quantization->zero_point.empty();
+}
+
+// Returns the correct type for a quantized tensor.
+// We have a special case for constants since they have a higher minimum value.
+StatusOr<QuantizedType> GetQuantizedType(const TensorT& tensor, Builder builder,
+                                         bool is_constant,
+                                         mlir::Type storage_type) {
+  tflite::QuantizationParametersT& quant_params = *tensor.quantization;
+  if (quant_params.details.AsCustomQuantization()) {
+    return absl::UnimplementedError("Cannot handle experimental quantization");
+  }
+
+  bool is_signed = true;
+  if (tensor.type == tflite::TensorType_UINT8) {
+    is_signed = false;
+    storage_type = mlir::IntegerType::get(builder.getContext(), 8);
+  } else if (tensor.type == tflite::TensorType_UINT4) {
+    is_signed = false;
+    storage_type = mlir::IntegerType::get(builder.getContext(), 4);
+  }
+
+  if (!storage_type) {
+    const mlir::Type raw_elem_type = ConvertElementType(tensor.type, builder);
+    if (!mlir::isa<mlir::IntegerType>(raw_elem_type)) {
+      return absl::InvalidArgumentError(
+          "Quantized tensors must be stored as integers");
+    }
+    storage_type = mlir::cast<mlir::IntegerType>(raw_elem_type);
+  }
+
+  // TFlite uses narrow-range [u]int8 for constant buffers of quantized weights.
+  // Since we don't know which ones are weights, we represent this optimization
+  // as a change in the storage bounds for the type for all constants of this
+  // type.
+  const int bitwidth = storage_type.getIntOrFloatBitWidth();
+  const bool is_weight_buffer = is_constant && (bitwidth == 8);
+
+  int64_t storage_min =
+      QuantizedType::getDefaultMinimumForInteger(is_signed, bitwidth) +
+      static_cast<int>(is_weight_buffer);
+  int64_t storage_max =
+      QuantizedType::getDefaultMaximumForInteger(is_signed, bitwidth);
+  uint32_t flags =
+      is_signed ? mlir::quant::QuantizationFlags::FlagValue::Signed : 0;
+
+  // Zero scales we make the minimum fp value, this is because some flatbuffers
+  // contain zero scale for zero values.
+  llvm::SmallVector<double> scales;
+  for (float scale : quant_params.scale) {
+    if (scale == 0) {
+      scales.push_back(std::numeric_limits<float>::min());
+      continue;
+    }
+    scales.push_back(scale);
+  }
+
+  // Scale size can't be zero as it is checked before.
+  if (quant_params.scale.size() != 1) {
+    return mlir::quant::UniformQuantizedPerAxisType::get(
+        flags, storage_type, builder.getF32Type(), scales,
+        quant_params.zero_point, quant_params.quantized_dimension, storage_min,
+        storage_max);
+  }
+  return mlir::quant::UniformQuantizedType::get(
+      flags, storage_type, builder.getF32Type(), scales[0],
+      quant_params.zero_point.at(0), storage_min, storage_max);
+}
+
+StatusOr<QuantizedType> GetCalibratedQuantizedType(const TensorT& tensor,
+                                                   Builder builder) {
+  if (tensor.quantization == nullptr) {
+    return absl::InvalidArgumentError("The tensor is not quantized.");
+  }
+  mlir::Type raw_elem_type = ConvertElementType(tensor.type, builder);
+  float min = tensor.quantization->min[0];
+  float max = tensor.quantization->max[0];
+  return mlir::quant::CalibratedQuantizedType::get(raw_elem_type, min, max);
+}
+
+StatusOr<mlir::TensorType> GetTensorType(const TensorT& tensor, Builder builder,
+                                         bool is_constant, bool is_intermediate,
+                                         bool get_storage) {
+  mlir::Type elem_type = ConvertElementType(tensor.type, builder);
+  if (tensor.type == tflite::TensorType_VARIANT) {
+    llvm::SmallVector<mlir::TensorType> tensor_types;
+    if (tensor.variant_tensors.size() > 1) {
+      return absl::InvalidArgumentError(
+          "Have more than one nested type in `variant_tensors`.");
+    }
+    for (const auto& nested_tensor : tensor.variant_tensors) {
+      mlir::Type nested_elem_type =
+          ConvertElementType(nested_tensor->type, builder);
+      if (nested_tensor->has_rank) {
+        llvm::SmallVector<int64_t> shape(nested_tensor->shape.begin(),
+                                         nested_tensor->shape.end());
+        tensor_types.push_back(
+            tensorflow::GetTypeFromTFTensorShape(shape, nested_elem_type));
+      } else {
+        tensor_types.push_back(UnrankedTensorType::get(nested_elem_type));
+      }
+    }
+    elem_type = mlir::TF::VariantType::get(tensor_types, builder.getContext());
+  }
+  if (IsQuantized(tensor) && !get_storage) {
+    TF_ASSIGN_OR_RETURN(elem_type,
+                        GetQuantizedType(tensor, builder, is_constant));
+  } else if (IsQuantized(tensor) && get_storage) {
+    // If the type is quantized we strip the signedness from the storage type.
+    elem_type = mlir::IntegerType::get(elem_type.getContext(),
+                                       elem_type.getIntOrFloatBitWidth());
+  }
+
+  // Intermediate tensors with calibration value (but not scale and zero points)
+  // should return calibrated quantized type.
+  if (is_intermediate && tensor.quantization != nullptr &&
+      !IsQuantized(tensor)) {
+    TF_ASSIGN_OR_RETURN(elem_type, GetCalibratedQuantizedType(tensor, builder));
+  }
+
+  if (tensor.shape.empty() && (is_constant || tensor.has_rank)) {
+    return RankedTensorType::get({}, elem_type);
+  }
+
+  if (!tensor.shape_signature.empty()) {
+    llvm::SmallVector<int64_t, 4> shape(tensor.shape_signature.begin(),
+                                        tensor.shape_signature.end());
+    return tensorflow::GetTypeFromTFTensorShape(shape, elem_type);
+  }
+
+  if (!tensor.shape.empty()) {
+    llvm::SmallVector<int64_t, 4> shape(tensor.shape.begin(),
+                                        tensor.shape.end());
+    return tensorflow::GetTypeFromTFTensorShape(shape, elem_type);
+  }
+
+  return UnrankedTensorType::get(elem_type);
+}
+
+mlir::ElementsAttr GetSplat(RankedTensorType type, int unique_index,
+                            Builder builder) {
+  mlir::Type element_ty = getElementTypeOrSelf(type);
+
+  if (element_ty.isSignlessInteger())
+    return DenseElementsAttr::get(
+        type, builder.getIntegerAttr(element_ty, unique_index));
+
+  if (mlir::isa<mlir::FloatType>(element_ty))
+    return DenseElementsAttr::get(
+        type, builder.getFloatAttr(element_ty, unique_index));
+
+  if (auto qtype = mlir::dyn_cast<QuantizedType>(element_ty)) {
+    mlir::RankedTensorType new_type = tensorflow::GetTypeFromTFTensorShape(
+        type.getShape(), qtype.getStorageType());
+    return DenseElementsAttr::get(
+        new_type, builder.getIntegerAttr(qtype.getStorageType(), unique_index));
+  }
+  llvm_unreachable("unhandled element type");
+}
+
+StatusOr<mlir::ElementsAttr> ConvertIntBuffer(
+    mlir::RankedTensorType shaped_type, const std::vector<uint8_t>& buffer,
+    bool truncate) {
+  mlir::Type elem_type = shaped_type.getElementType();
+  unsigned bit_width;
+  bool is_signed;
+  if (auto itype = mlir::dyn_cast<mlir::IntegerType>(elem_type)) {
+    bit_width = itype.getWidth();
+    is_signed = !itype.isUnsigned();
+  } else if (auto qtype =
+                 mlir::dyn_cast<mlir::quant::QuantizedType>(elem_type)) {
+    bit_width = qtype.getStorageTypeIntegralWidth();
+    is_signed = qtype.isSigned();
+    shaped_type = tensorflow::GetTypeFromTFTensorShape(shaped_type.getShape(),
+                                                       qtype.getStorageType());
+  } else {
+    return absl::InvalidArgumentError("unsupported integer constant type");
+  }
+
+  llvm::SmallVector<mlir::APInt> values;
+  switch (bit_width) {
+    case 1: {
+      // vector<bool> doesn't convert to an ArrayRef
+      llvm::SmallVector<bool, 8> boolValues;
+      boolValues.reserve(buffer.size());
+      for (auto b : buffer) {
+        boolValues.emplace_back(b != 0);
+      }
+      return mlir::ElementsAttr(
+          DenseElementsAttr::get(shaped_type, ArrayRef<bool>(boolValues)));
+    }
+    case 2: {
+      auto i2Values = tflite::UnpackDenseLowBitIntoInt8(
+          buffer, shaped_type.getNumElements(), /*bit_width=*/2);
+      // Use `getFromRawBuffer()` instead of `get()` to bypass a templated size
+      // check which doesn't work with int2 because int2_t doesn't exist.
+      return mlir::ElementsAttr(DenseElementsAttr::getFromRawBuffer(
+          shaped_type, ArrayRef<char>(i2Values)));
+    }
+    case 4: {
+      if (is_signed) {
+        auto i4Values = tflite::UnpackDenseLowBitIntoInt8(
+            buffer, shaped_type.getNumElements(), /*bit_width=*/4);
+        // Use `getFromRawBuffer()` instead of `get()` to bypass a templated
+        // size check which doesn't work with int4 because int4_t doesn't
+        // exist.
+        return mlir::ElementsAttr(DenseElementsAttr::getFromRawBuffer(
+            shaped_type, ArrayRef<char>(i4Values)));
+      } else {
+        auto ui4Values = tflite::UnpackDenseLowBitIntoUint8(
+            buffer, shaped_type.getNumElements(), /*bit_width=*/4);
+        return mlir::ElementsAttr(DenseElementsAttr::getFromRawBuffer(
+            shaped_type, ArrayRef<char>(ui4Values)));
+      }
+    }
+    case 8: {
+      return mlir::ElementsAttr(
+          DenseElementsAttr::get(shaped_type, ArrayRef<uint8_t>(buffer)));
+    }
+    case 16: {
+      values = ReadAsHostEndian<uint16_t>(buffer);
+      break;
+    }
+    case 32: {
+      values = ReadAsHostEndian<uint32_t>(buffer);
+      break;
+    }
+    case 64: {
+      values = ReadAsHostEndian<uint64_t>(buffer);
+      break;
+    }
+    default:
+      return absl::UnimplementedError(
+          absl::StrCat("Cannot handle bit width ", bit_width));
+  }
+
+  if (truncate) {
+    truncateLimitedIntegerAPInt(values);
+    auto sign = mlir::cast<mlir::IntegerType>(shaped_type.getElementType())
+                    .getSignedness();
+    auto ety = mlir::IntegerType::get(shaped_type.getContext(),
+                                      values[0].getBitWidth(), sign);
+    shaped_type =
+        tensorflow::GetTypeFromTFTensorShape(shaped_type.getShape(), ety);
+  }
+
+  return mlir::ElementsAttr(DenseElementsAttr::get(shaped_type, values));
+}
+
+StatusOr<mlir::ElementsAttr> ConvertFloatBuffer(
+    mlir::RankedTensorType shaped_type, const std::vector<uint8_t>& buffer) {
+  size_t bytes_len = buffer.size();
+  mlir::Type elem_type = shaped_type.getElementType();
+
+  // The bytes of floats are stored little-endian.
+  switch (elem_type.getIntOrFloatBitWidth()) {
+    case 8: {
+      assert(bytes_len == shaped_type.getNumElements());
+      assert(mlir::isa<mlir::Float8E4M3FNType>(elem_type) ||
+             mlir::isa<mlir::Float8E5M2Type>(elem_type));
+      return mlir::ElementsAttr(DenseElementsAttr::getFromRawBuffer(
+          shaped_type,
+          llvm::ArrayRef<char>(reinterpret_cast<const char*>(buffer.data()),
+                               bytes_len)));
+    }
+    case 16: {
+      assert(bytes_len % 2 == 0);
+      // Supports both BF16 and F16.
+      assert(elem_type.isF16() || elem_type.isBF16());
+      size_t elem_count = bytes_len / 2;
+
+      if (elem_type.isF16()) {
+        std::vector<Eigen::half> values;
+        values.reserve(elem_count);
+
+        const char* data = reinterpret_cast<const char*>(buffer.data());
+
+        for (size_t i = 0; i < elem_count; i++) {
+          uint16_t bit_repr = llvm::support::endian::readNext<
+              uint16_t, llvm::endianness::native, llvm::support::unaligned>(
+              data);
+          values.push_back(Eigen::numext::bit_cast<Eigen::half>(bit_repr));
+        }
+
+        return mlir::ElementsAttr(
+            DenseElementsAttr::get(shaped_type, ArrayRef<Eigen::half>(values)));
+      } else {
+        std::vector<Eigen::bfloat16> values;
+        values.reserve(elem_count);
+
+        const char* data = reinterpret_cast<const char*>(buffer.data());
+
+        for (size_t i = 0; i < elem_count; i++) {
+          uint16_t bit_repr = llvm::support::endian::readNext<
+              uint16_t, llvm::endianness::native, llvm::support::unaligned>(
+              data);
+          values.push_back(Eigen::numext::bit_cast<Eigen::bfloat16>(bit_repr));
+        }
+
+        return mlir::ElementsAttr(DenseElementsAttr::get(
+            shaped_type, ArrayRef<Eigen::bfloat16>(values)));
+      }
+    }
+    case 32: {
+      assert(bytes_len % 4 == 0);
+      size_t elem_count = bytes_len / 4;
+      std::vector<float> values;
+      values.reserve(elem_count);
+
+      const char* data = reinterpret_cast<const char*>(buffer.data());
+
+      for (size_t i = 0; i < elem_count; i++) {
+        uint32_t bit_repr =
+            llvm::support::endian::readNext<uint32_t, llvm::endianness::native,
+                                            llvm::support::unaligned>(data);
+        values.push_back(absl::bit_cast<float>(bit_repr));
+      }
+      return mlir::ElementsAttr(
+          DenseElementsAttr::get(shaped_type, ArrayRef<float>(values)));
+    }
+    case 64: {
+      assert(bytes_len % 8 == 0);
+      size_t elem_count = bytes_len / 8;
+      std::vector<double> values;
+      values.reserve(elem_count);
+
+      const char* data = reinterpret_cast<const char*>(buffer.data());
+
+      for (size_t i = 0; i < elem_count; i++) {
+        uint64_t bit_repr =
+            llvm::support::endian::readNext<uint64_t, llvm::endianness::native,
+                                            llvm::support::unaligned>(data);
+        values.push_back(absl::bit_cast<double>(bit_repr));
+      }
+      return mlir::ElementsAttr(
+          DenseElementsAttr::get(shaped_type, ArrayRef<double>(values)));
+    }
+  }
+  return absl::InvalidArgumentError(absl::StrCat(
+      "unsupported bit width ", elem_type.getIntOrFloatBitWidth()));
+}
+
+tensorflow::TensorProto ConvertTfliteConstTensor(
+    const TensorT& tensor, const std::vector<uint8_t>& buffer) {
+  tensorflow::TensorProto ret;
+  ret.set_dtype(TflTypeToTfType(tensor.type));
+
+  tensorflow::TensorShapeProto* shape = ret.mutable_tensor_shape();
+  shape->set_unknown_rank(false);
+  for (auto dim : tensor.shape) {
+    shape->add_dim()->set_size(int64_t{dim});
+  }
+  // TensorFlow Lite uses tflite::DynamicBufer to encode vector of strings.
+  if (tensor.type == tflite::TensorType_STRING) {
+    for (int i = 0; i < mlir::TFL::GetStringCount(buffer.data()); ++i) {
+      mlir::TFL::StringRef str = mlir::TFL::GetString(buffer.data(), i);
+      ret.add_string_val(str.str, str.len);
+    }
+    return ret;
+  }
+  std::string content;
+  content.assign(reinterpret_cast<const char*>(buffer.data()), buffer.size());
+  ret.set_tensor_content(content);
+  return ret;
+}
+
+int64_t GetSizeInBits(mlir::ShapedType shaped_type) {
+  if (!shaped_type.hasStaticShape()) return 0;
+  return GetSizeInBits(shaped_type.getElementType()) *
+         shaped_type.getNumElements();
+}
+
+int64_t GetSizeInBits(mlir::quant::QuantizedType quant_type) {
+  const int64_t bits = std::max(quant_type.getStorageTypeIntegralWidth(),
+                                static_cast<uint32_t>(CHAR_BIT));
+  assert(IsPowerOfTwo(bits));
+  return bits;
+}
+
+int64_t GetSizeInBits(mlir::Type type) {
+  if (type.isIntOrFloat()) {
+    const int64_t bits =
+        std::max(type.getIntOrFloatBitWidth(), static_cast<uint32_t>(CHAR_BIT));
+    assert(IsPowerOfTwo(bits));
+    return bits;
+  }
+  if (mlir::isa<mlir::ShapedType>(type)) {
+    auto shaped_type = mlir::cast<mlir::ShapedType>(type);
+    if (mlir::isa<mlir::ComplexType>(shaped_type.getElementType())) {
+      auto complex_type =
+          mlir::cast<mlir::ComplexType>(shaped_type.getElementType());
+      return GetSizeInBits(complex_type.getElementType()) * 2;
+    } else if (mlir::isa<mlir::quant::QuantizedType>(
+                   shaped_type.getElementType())) {
+      auto quant_type =
+          mlir::cast<mlir::quant::QuantizedType>(shaped_type.getElementType());
+      return GetSizeInBits(quant_type);
+    } else {
+      return GetSizeInBits(shaped_type);
+    }
+  }
+
+  return 0;
+}
+
+int64_t GetSizeInBytes(mlir::Type type) {
+  return ExactIntegerDivide(GetSizeInBits(type), CHAR_BIT);
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/const_tensor_utils.h

@@ -0,0 +1,111 @@
+/* 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_UTILS_CONST_TENSOR_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONST_TENSOR_UTILS_H_
+
+#include <stdbool.h>
+
+#include <cassert>
+#include <cstdint>
+#include <type_traits>
+#include <vector>
+
+#include "absl/base/attributes.h"
+#include "absl/meta/type_traits.h"
+#include "absl/status/statusor.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/AffineMap.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+
+namespace mlir {
+namespace TFL {
+
+bool IsQuantized(const tflite::TensorT& tensor);
+
+absl::StatusOr<mlir::quant::QuantizedType> GetQuantizedType(
+    const tflite::TensorT& tensor, mlir::Builder builder,
+    bool is_constant = false, mlir::Type storage_type = {});
+
+// Imports float tensor with calibration value into calibrated quantized type.
+absl::StatusOr<mlir::quant::QuantizedType> GetCalibratedQuantizedType(
+    const tflite::TensorT& tensor, mlir::Builder builder);
+
+absl::StatusOr<mlir::TensorType> GetTensorType(const tflite::TensorT& tensor,
+                                               mlir::Builder builder,
+                                               bool is_constant = false,
+                                               bool is_intermediate = false,
+                                               bool get_storage = false);
+
+// Gets a constant splat for the given value of type. Requires value to be of
+// type static shaped RankedTensorType. `unique_index` is used to get the unique
+// value for the attribute.
+mlir::ElementsAttr GetSplat(mlir::RankedTensorType type, int unique_index,
+                            mlir::Builder builder);
+
+absl::StatusOr<mlir::ElementsAttr> ConvertIntBuffer(
+    mlir::RankedTensorType shaped_type, const std::vector<uint8_t>& buffer,
+    bool truncate = false);
+
+absl::StatusOr<mlir::ElementsAttr> ConvertFloatBuffer(
+    mlir::RankedTensorType shaped_type, const std::vector<uint8_t>& buffer);
+
+tensorflow::TensorProto ConvertTfliteConstTensor(
+    const tflite::TensorT& tensor, const std::vector<uint8_t>& buffer);
+
+// Get the size of the type in bits. The type can be ComplexType, FloatType,
+// IntegerType, QuantizedType, or ShapeType of other supported types.
+//
+// Sub-byte types, e.g. qu4 and i2, are treated as a full i8.
+int64_t GetSizeInBits(mlir::ShapedType shaped_type);
+int64_t GetSizeInBits(mlir::Type type);
+int64_t GetSizeInBits(mlir::quant::QuantizedType quant_type);
+
+// Get the size of the type in bytes.
+//
+// Sub-byte element types, e.g. qu4 and i2, are treated as a full i8.
+// e.g. GetSizeInBytes(tensor<4xi2>) == 4, instead of 1.
+int64_t GetSizeInBytes(mlir::Type type);
+
+// Performs an integer divide and checks that the remainder is zero.
+// It supports int64 version as well.
+template <typename Integer,
+          typename = std::enable_if_t<std::is_same<int32_t, Integer>::value ||
+                                      std::is_same<uint32_t, Integer>::value ||
+                                      std::is_same<int64_t, Integer>::value ||
+                                      std::is_same<uint64_t, Integer>::value>>
+ABSL_ATTRIBUTE_ALWAYS_INLINE Integer ExactIntegerDivide(Integer numerator,
+                                                        int64_t denominator) {
+  const Integer ratio = numerator / denominator;
+  assert((numerator % denominator) == 0);
+  return ratio;
+}
+
+template <typename IntType,
+          std::enable_if_t<!std::is_unsigned<IntType>::value, int> = 0>
+ABSL_ATTRIBUTE_ALWAYS_INLINE bool IsPowerOfTwo(IntType n) {
+  static_assert(std::is_integral<IntType>::value, "");
+  return n > 0 && (n & (n - 1)) == 0;
+}
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONST_TENSOR_UTILS_H_

tensorflow/compiler/mlir/lite/utils/constant_utils.cc

@@ -0,0 +1,157 @@
+/* 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/utils/constant_utils.h"
+
+#include <complex>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/mangling_util.h"
+#include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace mlir {
+namespace TFL {
+
+absl::StatusOr<TypedAttr> CreateTypedAttr(ShapedType shaped_type, int value) {
+  Type element_type = shaped_type.getElementType();
+  if (element_type.isF16()) {
+    auto floatType = mlir::Float16Type::get(element_type.getContext());
+    auto floatAttr = mlir::FloatAttr::get(floatType, static_cast<float>(value));
+    std::vector<Attribute> floatValues({floatAttr});
+    return DenseElementsAttr::get(shaped_type, floatValues);
+  } else if (element_type.isBF16()) {
+    auto floatType = mlir::BFloat16Type::get(element_type.getContext());
+    auto floatAttr = mlir::FloatAttr::get(floatType, static_cast<float>(value));
+    std::vector<Attribute> floatValues({floatAttr});
+    return DenseElementsAttr::get(shaped_type, floatValues);
+  } else if (element_type.isF32()) {
+    return DenseElementsAttr::get<float>(shaped_type,
+                                         static_cast<float>(value));
+  } else if (auto complex_type =
+                 mlir::dyn_cast<mlir::ComplexType>(element_type)) {
+    auto etype = complex_type.getElementType();
+    if (etype.isF32()) {
+      tensorflow::TensorProto repr;
+      repr.set_dtype(tensorflow::DT_COMPLEX64);
+
+      tensorflow::TensorShapeProto* shape = repr.mutable_tensor_shape();
+      shape->set_unknown_rank(false);
+      shape->add_dim()->set_size(int64_t{1});
+      std::string content;
+      auto complex_value = std::complex<float>(static_cast<float>(value), 0.0f);
+      content.assign(reinterpret_cast<const char*>(&complex_value),
+                     sizeof(complex_value));
+      repr.set_tensor_content(content);
+      std::string mangled = tensorflow::mangling_util::MangleTensor(repr);
+
+      return mlir::TF::TensorProtoAttr::get(shaped_type, mangled);
+    } else {
+      return absl::Status(absl::StatusCode::kInvalidArgument,
+                          "Unsupported type");
+    }
+  } else if (auto itype = mlir::dyn_cast<mlir::IntegerType>(element_type)) {
+    if (element_type.isSignedInteger()) {
+      switch (itype.getWidth()) {
+        case 8:
+          return DenseElementsAttr::get<int8_t>(shaped_type,
+                                                static_cast<int8_t>(value));
+          break;
+        case 16:
+          return DenseElementsAttr::get<int16_t>(shaped_type,
+                                                 static_cast<int16_t>(value));
+          break;
+        case 32:
+          return DenseElementsAttr::get<int32_t>(shaped_type,
+                                                 static_cast<int32_t>(value));
+          break;
+        case 64:
+          return DenseElementsAttr::get<int64_t>(shaped_type,
+                                                 static_cast<int64_t>(value));
+          break;
+        default:
+          return absl::Status(absl::StatusCode::kInvalidArgument,
+                              "Unsupported type");
+      }
+    } else {
+      switch (itype.getWidth()) {
+        case 8:
+          return DenseElementsAttr::get<uint8_t>(shaped_type,
+                                                 static_cast<uint8_t>(value));
+          break;
+        case 16:
+          return DenseElementsAttr::get<uint16_t>(shaped_type,
+                                                  static_cast<uint16_t>(value));
+          break;
+        case 32:
+          return DenseElementsAttr::get<uint32_t>(shaped_type,
+                                                  static_cast<uint32_t>(value));
+          break;
+        case 64:
+          return DenseElementsAttr::get<uint64_t>(shaped_type,
+                                                  static_cast<uint64_t>(value));
+          break;
+        default:
+          return absl::Status(absl::StatusCode::kInvalidArgument,
+                              "Unsupported type");
+      }
+    }
+  } else {
+    return absl::Status(absl::StatusCode::kInvalidArgument, "Unsupported type");
+  }
+}
+
+// Returns a Constant op with a splat vector value.
+absl::StatusOr<arith::ConstantOp> CreateConstOpWithVectorValue(
+    PatternRewriter* rewriter, Location loc, ShapedType shaped_type,
+    int value) {
+  ShapedType dense_type = RankedTensorType::get(shaped_type.getShape(),
+                                                shaped_type.getElementType());
+  auto attr = CreateTypedAttr(dense_type, value);
+
+  return arith::ConstantOp::create(*rewriter, loc, dense_type,
+                                   cast<TypedAttr>(*attr));
+}
+
+absl::StatusOr<arith::ConstantOp> CreateConstOpWithSingleValue(
+    PatternRewriter* rewriter, Location loc, ShapedType shaped_type,
+    int value) {
+  ShapedType scalar_type =
+      RankedTensorType::get({}, shaped_type.getElementType());
+  auto attr = CreateTypedAttr(scalar_type, value);
+
+  return arith::ConstantOp::create(*rewriter, loc, scalar_type,
+                                   cast<TypedAttr>(*attr));
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/constant_utils.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_UTILS_CONSTANT_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONSTANT_UTILS_H_
+
+#include "absl/status/statusor.h"
+#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/AffineMap.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "tsl/platform/statusor.h"
+
+namespace mlir {
+namespace TFL {
+
+// Returns a Constant op with a single value.
+absl::StatusOr<arith::ConstantOp> CreateConstOpWithSingleValue(
+    PatternRewriter* rewriter, Location loc, ShapedType shaped_type, int value);
+
+// Returns a Constant op with a splat vector value.
+absl::StatusOr<arith::ConstantOp> CreateConstOpWithVectorValue(
+    PatternRewriter* rewriter, Location loc, ShapedType shaped_type, int value);
+
+}  // namespace TFL
+}  // namespace mlir
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONSTANT_UTILS_H_

tensorflow/compiler/mlir/lite/utils/control_edges.h

@@ -0,0 +1,34 @@
+/* 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_UTILS_CONTROL_EDGES_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONTROL_EDGES_H_
+
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+namespace tflite {
+
+// LINT.IfChange
+
+using ControlEdge = std::pair<int32_t, int32_t>;
+using ControlEdges = std::vector<ControlEdge>;
+
+// LINT.ThenChange(//tensorflow/lite/graph_info.h)
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONTROL_EDGES_H_

tensorflow/compiler/mlir/lite/utils/convert_type.cc

@@ -0,0 +1,265 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/convert_type.h"
+
+#include "absl/status/statusor.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tflite {
+
+using absl::StatusOr;
+
+namespace errors = tensorflow::errors;
+
+tflite::TensorType ConvertTypeToTensorType(mlir::Type type) {
+  if (type.isF16()) {
+    return tflite::TensorType_FLOAT16;
+  } else if (type.isBF16()) {
+    return tflite::TensorType_BFLOAT16;
+  } else if (type.isF32()) {
+    return tflite::TensorType_FLOAT32;
+  } else if (type.isF64()) {
+    return tflite::TensorType_FLOAT64;
+  } else if (mlir::isa<mlir::Float8E4M3FNType>(type)) {
+    return tflite::TensorType_FLOAT8_E4M3FN;
+  } else if (mlir::isa<mlir::Float8E5M2Type>(type)) {
+    return tflite::TensorType_FLOAT8_E5M2;
+  } else if (mlir::isa<mlir::TF::StringType>(type)) {
+    return tflite::TensorType_STRING;
+  } else if (auto complex_type = mlir::dyn_cast<mlir::ComplexType>(type)) {
+    if (complex_type.getElementType().isF32()) {
+      return tflite::TensorType_COMPLEX64;
+    } else if (complex_type.getElementType().isF64()) {
+      return tflite::TensorType_COMPLEX128;
+    }
+    llvm_unreachable("invalid complex Type in conversion");
+  } else if (auto itype = mlir::dyn_cast<mlir::IntegerType>(type)) {
+    switch (itype.getWidth()) {
+      case 1:
+        return tflite::TensorType_BOOL;
+      case 4:
+        if (itype.isUnsigned())
+          return tflite::TensorType_UINT4;
+        else
+          return tflite::TensorType_INT4;
+      case 8:
+        if (itype.isUnsigned())
+          return tflite::TensorType_UINT8;
+        else
+          return tflite::TensorType_INT8;
+      case 16:
+        return tflite::TensorType_INT16;
+      case 32:
+        return tflite::TensorType_INT32;
+      case 64:
+        if (itype.isUnsigned())
+          return tflite::TensorType_UINT64;
+        else
+          return tflite::TensorType_INT64;
+      default:
+        llvm_unreachable("invalid integer Type in conversion");
+    }
+  }
+  llvm_unreachable("invalid Type in conversion");
+}
+
+mlir::Type ConvertElementType(tflite::TensorType type, mlir::Builder builder) {
+  switch (type) {
+    case tflite::TensorType_FLOAT16:
+      return builder.getF16Type();
+    case tflite::TensorType_BFLOAT16:
+      return builder.getBF16Type();
+    case tflite::TensorType_FLOAT32:
+      return builder.getF32Type();
+    case tflite::TensorType_FLOAT64:
+      return builder.getF64Type();
+    case tflite::TensorType_FLOAT8_E4M3FN:
+      return mlir::Float8E4M3FNType::get(builder.getContext());
+    case tflite::TensorType_FLOAT8_E5M2:
+      return mlir::Float8E5M2Type::get(builder.getContext());
+    case tflite::TensorType_INT32:
+      return builder.getIntegerType(32);
+    case tflite::TensorType_UINT16:
+      return builder.getIntegerType(16, /*isSigned=*/false);
+    case tflite::TensorType_UINT32:
+      return builder.getIntegerType(32, /*isSigned=*/false);
+    case tflite::TensorType_UINT8:
+      return builder.getIntegerType(8, /*isSigned=*/false);
+    case tflite::TensorType_INT64:
+      return builder.getIntegerType(64);
+    case tflite::TensorType_STRING:
+      return mlir::TF::StringType::get(builder.getContext());
+    case tflite::TensorType_BOOL:
+      return builder.getI1Type();
+    case tflite::TensorType_INT16:
+      return builder.getIntegerType(16);
+    case tflite::TensorType_COMPLEX64:
+      return mlir::ComplexType::get(builder.getF32Type());
+    case tflite::TensorType_COMPLEX128:
+      return mlir::ComplexType::get(builder.getF64Type());
+    case tflite::TensorType_INT2:
+      return builder.getIntegerType(2);
+    case tflite::TensorType_INT4:
+      return builder.getIntegerType(4);
+    case tflite::TensorType_UINT4:
+      return builder.getIntegerType(4, /*isSigned=*/false);
+    case tflite::TensorType_INT8:
+      return builder.getIntegerType(8);
+    case tflite::TensorType_UINT64:
+      return builder.getIntegerType(64, /*isSigned=*/false);
+    case tflite::TensorType_RESOURCE:
+      return mlir::TF::ResourceType::get(builder.getContext());
+    case tflite::TensorType_VARIANT:
+      return mlir::TF::VariantType::get(builder.getContext());
+  }
+}
+
+tensorflow::DataType TflTypeToTfType(tflite::TensorType type) {
+  switch (type) {
+    case tflite::TensorType_BOOL:
+      return tensorflow::DT_BOOL;
+    case tflite::TensorType_COMPLEX64:
+      return tensorflow::DT_COMPLEX64;
+    case tflite::TensorType_COMPLEX128:
+      return tensorflow::DT_COMPLEX128;
+    case tflite::TensorType_FLOAT16:
+      return tensorflow::DT_HALF;
+    case tflite::TensorType_BFLOAT16:
+      return tensorflow::DT_BFLOAT16;
+    case tflite::TensorType_FLOAT32:
+      return tensorflow::DT_FLOAT;
+    case tflite::TensorType_FLOAT64:
+      return tensorflow::DT_DOUBLE;
+    case tflite::TensorType_FLOAT8_E4M3FN:
+      return tensorflow::DT_FLOAT8_E4M3FN;
+    case tflite::TensorType_FLOAT8_E5M2:
+      return tensorflow::DT_FLOAT8_E5M2;
+    // TODO(b/246806634): Tensorflow DT_INT2/4 type doesn't exist yet
+    case tflite::TensorType_INT2:
+      return tensorflow::DT_INT8;
+    case tflite::TensorType_INT4:
+      return tensorflow::DT_INT8;
+    case tflite::TensorType_UINT4:
+      return tensorflow::DT_UINT4;
+    case tflite::TensorType_INT8:
+      return tensorflow::DT_INT8;
+    case tflite::TensorType_INT16:
+      return tensorflow::DT_INT16;
+    case tflite::TensorType_INT32:
+      return tensorflow::DT_INT32;
+    case tflite::TensorType_UINT32:
+      return tensorflow::DT_UINT32;
+    case tflite::TensorType_INT64:
+      return tensorflow::DT_INT64;
+    case tflite::TensorType_STRING:
+      return tensorflow::DT_STRING;
+    case tflite::TensorType_UINT8:
+      return tensorflow::DT_UINT8;
+    case tflite::TensorType_UINT16:
+      return tensorflow::DT_UINT16;
+    case tflite::TensorType_UINT64:
+      return tensorflow::DT_UINT64;
+    case tflite::TensorType_RESOURCE:
+      return tensorflow::DT_RESOURCE;
+    case tflite::TensorType_VARIANT:
+      return tensorflow::DT_VARIANT;
+  }
+}
+
+absl::StatusOr<tflite::TensorType> TfTypeToTflType(tensorflow::DataType type) {
+  switch (type) {
+    case tensorflow::DT_BOOL:
+      return tflite::TensorType_BOOL;
+    case tensorflow::DT_COMPLEX64:
+      return tflite::TensorType_COMPLEX64;
+    case tensorflow::DT_COMPLEX128:
+      return tflite::TensorType_COMPLEX128;
+    case tensorflow::DT_HALF:
+      return tflite::TensorType_FLOAT16;
+    case tensorflow::DT_BFLOAT16:
+      return tflite::TensorType_BFLOAT16;
+    case tensorflow::DT_FLOAT:
+      return tflite::TensorType_FLOAT32;
+    case tensorflow::DT_DOUBLE:
+      return tflite::TensorType_FLOAT64;
+    case tensorflow::DT_FLOAT8_E4M3FN:
+      return tflite::TensorType_FLOAT8_E4M3FN;
+    case tensorflow::DT_FLOAT8_E5M2:
+      return tflite::TensorType_FLOAT8_E5M2;
+    case tensorflow::DT_UINT4:
+      return tflite::TensorType_UINT4;
+    case tensorflow::DT_INT8:
+      return tflite::TensorType_INT8;
+    case tensorflow::DT_INT16:
+      return tflite::TensorType_INT16;
+    case tensorflow::DT_INT32:
+      return tflite::TensorType_INT32;
+    case tensorflow::DT_UINT32:
+      return tflite::TensorType_UINT32;
+    case tensorflow::DT_INT64:
+      return tflite::TensorType_INT64;
+    case tensorflow::DT_UINT64:
+      return tflite::TensorType_UINT64;
+    case tensorflow::DT_STRING:
+      return tflite::TensorType_STRING;
+    case tensorflow::DT_UINT8:
+      return tflite::TensorType_UINT8;
+    case tensorflow::DT_RESOURCE:
+      return tflite::TensorType_RESOURCE;
+    case tensorflow::DT_VARIANT:
+      return tflite::TensorType_VARIANT;
+    default:
+      return absl::InvalidArgumentError(
+          absl::StrCat("unsupported tensor data type", type));
+  }
+}
+
+mlir::Type GetShapeStrippedType(mlir::TypeAttr type_attr) {
+  auto type = type_attr.getValue();
+  auto shaped_type = mlir::dyn_cast<mlir::ShapedType>(type);
+  if (shaped_type) {
+    return shaped_type.getElementType();
+  } else {
+    return type;
+  }
+}
+
+bool NotFromQuantOpOrSameQuantType(mlir::Value val, mlir::TypeAttr qtype_attr) {
+  auto val_defn_op = val.getDefiningOp();
+  mlir::TFL::QuantizeOp q_op =
+      llvm::dyn_cast_or_null<mlir::TFL::QuantizeOp>(val_defn_op);
+  if (!q_op) return true;
+
+  // Ignore shape details - we're really only trying to
+  // check if quantization is the same.
+  auto stripped_src_qtype = GetShapeStrippedType(q_op.getQtypeAttr());
+  auto stripped_qtype = GetShapeStrippedType(qtype_attr);
+  return stripped_src_qtype == stripped_qtype;
+}
+
+}  // namespace tflite

tensorflow/compiler/mlir/lite/utils/convert_type.h

@@ -0,0 +1,52 @@
+/* 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_UTILS_CONVERT_TYPE_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONVERT_TYPE_H_
+
+#include "absl/status/statusor.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/schema/schema_generated.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace mlir {
+class Builder;
+}  // namespace mlir
+
+namespace tflite {
+// Convert the MLIR type to the corresponding TFLite tensor.
+tflite::TensorType ConvertTypeToTensorType(mlir::Type type);
+
+// Convert the scalar type of a TFlite tensor to the corresponding MLIR type.
+mlir::Type ConvertElementType(tflite::TensorType type, mlir::Builder builder);
+
+// Convert the scalar type of a TFLite tensor to the corresponding
+// Tensorflow type
+tensorflow::DataType TflTypeToTfType(tflite::TensorType type);
+
+// Convert the Tensorflow scalar type to the corresponding TFLite type
+absl::StatusOr<tflite::TensorType> TfTypeToTflType(tensorflow::DataType type);
+
+// Returns element type from attribute Type 'type_attr'.
+mlir::Type GetShapeStrippedType(mlir::TypeAttr type_attr);
+
+// Returns true if 'val' is not from Quantize op or
+// from Quantize Op with same quant type as 'qtype_attr'
+bool NotFromQuantOpOrSameQuantType(mlir::Value val, mlir::TypeAttr qtype_attr);
+
+}  // namespace tflite
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_CONVERT_TYPE_H_

tensorflow/compiler/mlir/lite/utils/fake_quant_utils.cc

@@ -0,0 +1,116 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/fake_quant_utils.h"
+
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/utils/utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
+
+namespace mlir {
+namespace TFL {
+
+// Moves the TF operations out from the tfl.TFCustomOps wrappers inside the
+// function. This is a no-op for the ops which are not wrapped.
+LogicalResult UnwrapTFCustomOps(func::FuncOp fn, OpBuilder& builder) {
+  llvm::SmallVector<Operation*, 4> wrapped_ops;
+  fn.walk([&](TFL::CustomTfOp custom_op) {
+    auto* real_op = &custom_op.getBody().front().front();
+    wrapped_ops.push_back(real_op);
+  });
+
+  for (auto* op : wrapped_ops) {
+    auto parent_op = op->getParentOfType<TFL::CustomTfOp>();
+    if (!parent_op) continue;
+    builder.setInsertionPoint(parent_op);
+
+    // Recreate the operation by using the wrapper's operands and return types.
+    // TODO(fengliuai): copy the regions.
+    OperationState state(op->getLoc(), op->getName().getStringRef(),
+                         parent_op->getOperands(), parent_op->getResultTypes(),
+                         op->getAttrs(), op->getSuccessors());
+    Operation* inlined = builder.create(state);
+
+    parent_op->replaceAllUsesWith(inlined);
+    parent_op->erase();
+  }
+  return success();
+}
+
+// Three instances of the rule to cover the three different types of
+// TF::FakeQuant operators
+using PreparePerTensorFakeQuant = InsertTFLQuantOpsAfterTFFakeQuantOp<
+    TF::FakeQuantWithMinMaxVarsOp, /*PerAxis=*/false,
+    FetchConstantMinMaxInputs<TF::FakeQuantWithMinMaxVarsOp>>;
+
+using PreparePerChannelFakeQuant = InsertTFLQuantOpsAfterTFFakeQuantOp<
+    TF::FakeQuantWithMinMaxVarsPerChannelOp, /*PerAxis=*/true,
+    FetchConstantMinMaxInputs<TF::FakeQuantWithMinMaxVarsPerChannelOp>>;
+
+using PreparePerTensorFakeQuantWithMinMaxArgs =
+    InsertTFLQuantOpsAfterTFFakeQuantOp<
+        TF::FakeQuantWithMinMaxArgsOp, /*PerAxis=*/false,
+        FetchMinMaxAttrs<TF::FakeQuantWithMinMaxArgsOp>>;
+
+// Removes the wrapper of the tf.FakeQuant* ops and creates the tfl.quantize
+// and tfl.dequantize pairs before tf.FakeQuant* being foled.
+LogicalResult ConvertFakeQuantOps(func::FuncOp func, MLIRContext* ctx,
+                                  bool use_fake_quant_num_bits) {
+  OpBuilder builder(func);
+  if (failed(UnwrapTFCustomOps(func, builder))) {
+    return failure();
+  }
+
+  // Insert the tfl.quantize/tfl.dequantize ops after the tf.FakeQuant* ops to
+  // preserve the quantization parameters.
+  func.walk([&](Operation* op) {
+    if (auto fake_quant = llvm::dyn_cast<TF::FakeQuantWithMinMaxArgsOp>(op)) {
+      (void)PreparePerTensorFakeQuantWithMinMaxArgs(use_fake_quant_num_bits)
+          .matchAndRewrite(fake_quant, builder);
+    } else if (auto fake_quant =
+                   llvm::dyn_cast<TF::FakeQuantWithMinMaxVarsOp>(op)) {
+      (void)PreparePerTensorFakeQuant(use_fake_quant_num_bits)
+          .matchAndRewrite(fake_quant, builder);
+    } else if (auto fake_quant =
+                   llvm::dyn_cast<TF::FakeQuantWithMinMaxVarsPerChannelOp>(
+                       op)) {
+      (void)PreparePerChannelFakeQuant(use_fake_quant_num_bits)
+          .matchAndRewrite(fake_quant, builder);
+    }
+  });
+
+  return success();
+}
+
+std::vector<std::string> AllTfFakeQuantOps() {
+  return {
+      mlir::TF::FakeQuantWithMinMaxVarsOp::getOperationName().str(),
+      mlir::TF::FakeQuantWithMinMaxVarsPerChannelOp::getOperationName().str(),
+      mlir::TF::FakeQuantWithMinMaxArgsOp::getOperationName().str()};
+}
+
+}  // end namespace TFL
+}  // end namespace mlir

tensorflow/compiler/mlir/lite/utils/fake_quant_utils.h

@@ -0,0 +1,177 @@
+/* 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.
+==============================================================================*/
+
+// This header file defines common utils used by TFLite transformation
+// passes to work with tf.FakeQuant* ops.
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_FAKE_QUANT_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_FAKE_QUANT_UTILS_H_
+
+#include <string>
+#include <vector>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/quantization/common/quantization_lib/quantization_utils.h"
+#include "tensorflow/compiler/mlir/lite/utils/utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
+
+namespace mlir {
+namespace TFL {
+
+template <class TFFakeQuantOp>
+struct FetchMinMaxAttrs {
+  using AttrType = FloatAttr;
+  bool operator()(TFFakeQuantOp tf_op, AttrType &min_value,
+                  AttrType &max_value) const {
+    min_value = tf_op.getMinAttr();
+    max_value = tf_op.getMaxAttr();
+    return true;  // Successfully matched and fetched.
+  }
+};
+
+template <class TFFakeQuantOp>
+struct FetchConstantMinMaxInputs {
+  using AttrType = DenseFPElementsAttr;
+  bool operator()(TFFakeQuantOp tf_op, AttrType &min_value,
+                  AttrType &max_value) const {
+    Value min = tf_op.getMin(), max = tf_op.getMax();
+    if (!matchPattern(min, m_Constant(&min_value))) {
+      return false;
+    }
+    if (!matchPattern(max, m_Constant(&max_value))) {
+      return false;
+    }
+    return true;  // Successfully matched and fetched.
+  }
+};
+
+// Inserts a "tfl.quantize" and "tfl.dequantize" op pair (QDQs) after the
+// tf.FakeQyantWithMinMax{Vars|VarsPerChannel|Args}Op
+// before the op being constant folded. Since the constant
+// folding logic will use a "arith.constant" op to replace the
+// "tf.FakeQuantWithMinMaxVarsOp", the "tfl.quantize" op is used to preserve
+// the quantization parameters as a TypeAttr and "tfl.dequantize" op used to
+// convert the output type to the next op. Here are the transformations:
+//
+// input   min cst       max cst          input   min cst       max cst
+//  \       |             |                \       |             |
+//   \  (tf.Identity) (tf.Identity)   =>    \  (tf.Identity) (tf.Identity)
+//    \     |             |                  \     |             |
+//       tf.FakeQuantWithMinMaxVars       tf.FakeQuantWithMinMaxVars
+//                   |                                 |
+//                                                tfl.quantize
+//                                                     |
+//                                                tfl.dequantize
+//                                                     |
+// If the input is a constant, the result pattern will eventually converted to
+//
+//            quant-emulated input
+//                   |
+//               tfl.quantize
+//                   |
+//              tfl.dequantize
+//                   |
+//
+//
+// Warns if the (most likely unwanted, currently not quite correctly handled)
+// case of back-to-back tf.FakeQuant occurs
+//
+//             tf.FakeQuant*
+//                   |
+//             tf.FakeQuant*
+//
+template <typename TFFakeQuantOp, bool PerAxis, class FetchMinMax>
+class InsertTFLQuantOpsAfterTFFakeQuantOp {
+ public:
+  explicit InsertTFLQuantOpsAfterTFFakeQuantOp(bool use_fake_quant_num_bits)
+      : use_fake_quant_num_bits_(use_fake_quant_num_bits) {}
+
+  FetchMinMax fetch_min_max_;
+
+  using FetchAttrType = typename FetchMinMax::AttrType;
+  LogicalResult matchAndRewrite(TFFakeQuantOp tf_op,
+                                OpBuilder &rewriter) const {
+    // We don't want to insert quantize/dequantize if the quantize op exists.
+    auto res = tf_op.getOutputs();
+    if (!res.hasOneUse() || isa<QuantizeOp>(*res.user_begin())) {
+      return failure();
+    }
+
+    // Extract the min/max constant values from the operands. We also consider
+    // a special case that there are tf.Identity ops between the min/max
+    // constants and the tf.FakeQuantWithMinMaxVarsOp.
+
+    FetchAttrType min_value, max_value;
+    if (!fetch_min_max_(tf_op, min_value, max_value)) {
+      return failure();
+    }
+
+    int quant_dim = -1;
+    if (PerAxis) {
+      // This is a special case that the quant_dim is the last dimensions.
+      quant_dim = mlir::cast<ShapedType>(res.getType()).getRank() - 1;
+    }
+    // Use the min/max from the operands and the num_bits and narrow_range
+    // attribute to create the quantization parameter for the new quantize op.
+    rewriter.setInsertionPointAfter(tf_op.getOperation());
+    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
+    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
+    Type res_type = tf_op.getType();
+    TypeAttr qtype = GetQuantizedTypeAttr(
+        rewriter, res_type, min_value, max_value, quant_dim, num_bits,
+        narrow_range, /*is_signed=*/false, /*legacy_float_scale=*/false,
+        use_fake_quant_num_bits_);
+    if (!qtype) {
+      return failure();
+    }
+
+    // Finally, use the quantization parameter to create the quantize and
+    // dequantize ops, and insert them between the tf.FakeQuantWithMinMaxVarsOp
+    // and its users.
+    Value value = tf_op.getOutputs();
+    auto quantize = TFL::QuantizeOp::create(rewriter, tf_op.getLoc(),
+                                            qtype.getValue(), value, qtype);
+    auto dequantize = TFL::DequantizeOp::create(rewriter, tf_op.getLoc(),
+                                                res_type, quantize.getOutput());
+    value.replaceAllUsesWith(dequantize);
+    quantize.getOperation()->replaceUsesOfWith(dequantize, value);
+
+    return success();
+  }
+
+  bool use_fake_quant_num_bits_;
+};
+
+// Removes the wrapper of the tf.FakeQuant* ops and creates the tfl.quantize
+// and tfl.dequantize pairs before tf.FakeQuant* being foled.
+LogicalResult ConvertFakeQuantOps(func::FuncOp func, MLIRContext *ctx,
+                                  bool use_fake_quant_num_bits = false);
+
+// Returns the names of all the considered tf.FakeQuant* ops.
+std::vector<std::string> AllTfFakeQuantOps();
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_FAKE_QUANT_UTILS_H_

tensorflow/compiler/mlir/lite/utils/low_bit_utils.cc

@@ -0,0 +1,67 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/low_bit_utils.h"
+
+#include <cassert>
+#include <cstdint>
+#include <vector>
+
+namespace tflite {
+
+std::vector<char> UnpackDenseLowBitIntoInt8(
+    const std::vector<uint8_t>& src_buffer, int64_t num_elements,
+    int bit_width) {
+  std::vector<char> unpacked_buffer;
+  unpacked_buffer.reserve(num_elements);
+  const int elements_per_byte = 8 / bit_width;
+  const int sign_bit_shift = 8 - bit_width;
+  const uint8_t mask = (1 << bit_width) - 1;
+
+  for (uint8_t value : src_buffer) {
+    for (int i = 0; i < elements_per_byte; ++i) {
+      if (unpacked_buffer.size() == num_elements) break;
+      int bit_offset = i * bit_width;
+      uint8_t extracted_value = (value >> bit_offset) & mask;
+      // Sign extend
+      unpacked_buffer.push_back(
+          static_cast<int8_t>(extracted_value << sign_bit_shift) >>
+          sign_bit_shift);
+    }
+  }
+
+  return unpacked_buffer;
+}
+
+std::vector<char> UnpackDenseLowBitIntoUint8(
+    const std::vector<uint8_t>& src_buffer, int64_t num_elements,
+    int bit_width) {
+  std::vector<char> unpacked_buffer;
+  unpacked_buffer.reserve(num_elements);
+  const int elements_per_byte = 8 / bit_width;
+  const uint8_t mask = (1 << bit_width) - 1;
+
+  for (uint8_t value : src_buffer) {
+    for (int i = 0; i < elements_per_byte; ++i) {
+      if (unpacked_buffer.size() == num_elements) break;
+      int bit_offset = i * bit_width;
+      unpacked_buffer.push_back((value >> bit_offset) & mask);
+    }
+  }
+
+  return unpacked_buffer;
+}
+
+}  // namespace tflite

tensorflow/compiler/mlir/lite/utils/low_bit_utils.h

@@ -0,0 +1,161 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_LOW_BIT_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_LOW_BIT_UTILS_H_
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <vector>
+
+#include "absl/functional/function_ref.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace tflite {
+// Packs raw byte inputs (chemically containing 2-bit or 4-bit values) into a
+// dense bit-packed representation using Little-Endian (low-bits-first) order.
+// The input buffer must contain unpacked values (1 value per byte).
+//
+// Requirement: 8 must be evenly divisible by kBitWidth.
+//
+// Template Args:
+//   kBitWidth: The logical bit width of the values (must be 2 or 4).
+template <int kBitWidth>
+absl::Status StreamPackLowBitValues8Bit(
+    llvm::ArrayRef<char> values,
+    absl::FunctionRef<absl::Status(absl::string_view)> apply_chunk) {
+  static_assert(kBitWidth == 2 || kBitWidth == 4);
+  // Use a chunk size that balances overhead with cache locality.
+  constexpr size_t kChunkSize = 1024 * 1024;
+
+  std::vector<char> packed_buffer;
+
+  constexpr size_t kElementsPerByte = 8 / kBitWidth;
+  constexpr size_t kPackedBufferSize =
+      (kChunkSize + kElementsPerByte - 1) / kElementsPerByte;
+  packed_buffer.resize(kPackedBufferSize);
+
+  // Outer Loop: Iterate over large chunks of the input
+  for (size_t offset = 0; offset < values.size(); offset += kChunkSize) {
+    llvm::ArrayRef<char> chunk_buffer =
+        values.slice(offset, std::min(kChunkSize, values.size() - offset));
+
+    const size_t num_elements = chunk_buffer.size();
+    constexpr uint8_t mask = (1 << kBitWidth) - 1;
+    size_t packed_idx = 0;
+
+    // Inner Loop: Iterate through the chunk to pack bits
+    for (size_t local_idx = 0; local_idx < num_elements;
+         local_idx += kElementsPerByte) {
+      uint8_t out_byte = 0;
+      for (int j = 0; j < kElementsPerByte; ++j) {
+        if (local_idx + j < num_elements) {
+          out_byte |= (chunk_buffer[local_idx + j] & mask) << (j * kBitWidth);
+        }
+      }
+      packed_buffer[packed_idx++] = static_cast<char>(out_byte);
+    }
+
+    auto status =
+        apply_chunk(absl::string_view(packed_buffer.data(), packed_idx));
+    if (!status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+// Packs elements from a generic range (e.g., mlir::APInt from
+// DenseElementsAttr) into a dense representation using little-endian
+// (low-bits-first) bit-packing.
+//
+// Requirement: 8 must be evenly divisible by kBitWidth.
+//
+// Template Args:
+//   kBitWidth: The logical bit width of the values (must be 2 or 4).
+//   Range: A range type providing `begin()` and `end()` returning `APInt`.
+template <int kBitWidth, typename Range>
+absl::Status StreamPackLowBitValues(
+    Range&& values,
+    absl::FunctionRef<absl::Status(absl::string_view)> apply_chunk) {
+  static_assert(kBitWidth == 2 || kBitWidth == 4);
+  constexpr size_t kChunkSize = 1024 * 1024;
+
+  std::vector<uint8_t> chunk_buffer;
+  std::vector<char> packed_buffer;
+
+  chunk_buffer.reserve(kChunkSize);
+
+  constexpr size_t kElementsPerByte = 8 / kBitWidth;
+  constexpr size_t kPackedBufferSize =
+      (kChunkSize + kElementsPerByte - 1) / kElementsPerByte;
+  packed_buffer.resize(kPackedBufferSize);
+
+  auto it = values.begin();
+  auto end = values.end();
+
+  while (it != end) {
+    // Phase 1: Unroll complex iterator into a simple contiguous buffer
+    chunk_buffer.clear();
+    for (size_t i = 0; i < kChunkSize && it != end; ++it, ++i) {
+      chunk_buffer.push_back(static_cast<uint8_t>(*((*it).getRawData())));
+    }
+
+    const size_t num_elements = chunk_buffer.size();
+    constexpr uint8_t mask = (1 << kBitWidth) - 1;
+    size_t packed_idx = 0;
+
+    // Phase 2: Pack bits (Auto-vectorized by compiler)
+    for (size_t i = 0; i < num_elements; i += kElementsPerByte) {
+      uint8_t out_byte = 0;
+      for (int j = 0; j < kElementsPerByte; ++j) {
+        if (i + j < num_elements) {
+          out_byte |= (chunk_buffer[i + j] & mask) << (j * kBitWidth);
+        }
+      }
+      packed_buffer[packed_idx++] = static_cast<char>(out_byte);
+    }
+
+    auto status =
+        apply_chunk(absl::string_view(packed_buffer.data(), packed_idx));
+    if (!status.ok()) {
+      return status;
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+// Assumes `src_buffer` contains densely packed low bit elements.
+// Returns a vector where each int8 element contains a sign-extended value.
+std::vector<char> UnpackDenseLowBitIntoInt8(
+    const std::vector<uint8_t>& src_buffer, int64_t num_elements,
+    int bit_width);
+
+// Assumes `src_buffer` contains densely packed low bit elements.
+// Returns a vector where each uint8 element contains an unpacked value.
+std::vector<char> UnpackDenseLowBitIntoUint8(
+    const std::vector<uint8_t>& src_buffer, int64_t num_elements,
+    int bit_width);
+}  // namespace tflite
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_LOW_BIT_UTILS_H_

tensorflow/compiler/mlir/lite/utils/low_bit_utils_test.cc

@@ -0,0 +1,143 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/low_bit_utils.h"
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir {
+namespace TFL {
+
+namespace {
+
+#ifndef EXPECT_OK
+#define EXPECT_OK(x) EXPECT_TRUE(x.ok());
+#endif
+
+class LowBitUtilsTest : public ::testing::Test {
+ protected:
+  LowBitUtilsTest() = default;
+
+  void SetUp() override {
+    context_ = std::make_unique<mlir::MLIRContext>();
+    builder_ = std::make_unique<mlir::Builder>(context_.get());
+  }
+
+  void TearDown() override { builder_.reset(); }
+
+  std::unique_ptr<mlir::MLIRContext> context_;
+  std::unique_ptr<mlir::Builder> builder_;
+};
+
+TEST_F(LowBitUtilsTest, Stream4BitValues) {
+  auto type = mlir::RankedTensorType::get({3}, builder_->getIntegerType(32));
+  auto raw_values = std::vector<int32_t>{1, 2, 3};
+  auto attr = mlir::DenseElementsAttr::get(type, llvm::ArrayRef(raw_values));
+
+  std::vector<uint8_t> packed_values;
+  auto apply_chunk_fn = [&](absl::string_view chunk) {
+    const uint8_t* data = reinterpret_cast<const uint8_t*>(chunk.data());
+    packed_values.insert(packed_values.end(), data, data + chunk.size());
+    return absl::OkStatus();
+  };
+
+  EXPECT_OK(tflite::StreamPackLowBitValues</*kBitWidth=*/4>(
+      attr.getValues<mlir::APInt>(), apply_chunk_fn));
+  EXPECT_EQ(packed_values.size(), 2);
+  EXPECT_EQ(packed_values[0], 0x21);
+  EXPECT_EQ(packed_values[1], 0x03);
+}
+
+TEST_F(LowBitUtilsTest, Stream4BitValues8Bit) {
+  auto type = mlir::RankedTensorType::get({3}, builder_->getIntegerType(8));
+  auto raw_values = std::vector<uint8_t>{1, 2, 3};
+  auto attr =
+      mlir::DenseElementsAttr::get(type, llvm::ArrayRef<uint8_t>(raw_values));
+
+  std::vector<uint8_t> packed_values;
+  auto apply_chunk_fn = [&](absl::string_view chunk) {
+    const uint8_t* data = reinterpret_cast<const uint8_t*>(chunk.data());
+    packed_values.insert(packed_values.end(), data, data + chunk.size());
+    return absl::OkStatus();
+  };
+
+  EXPECT_EQ(attr.getRawData().size(), 3);
+  EXPECT_OK(tflite::StreamPackLowBitValues8Bit</*kBitWidth=*/4>(
+      attr.getRawData(), apply_chunk_fn));
+  EXPECT_EQ(packed_values.size(), 2);
+  EXPECT_EQ(packed_values[0], 0x21);
+  EXPECT_EQ(packed_values[1], 0x03);
+}
+
+TEST_F(LowBitUtilsTest, Stream2BitValues) {
+  auto type = mlir::RankedTensorType::get({7}, builder_->getIntegerType(32));
+  auto raw_values = std::vector<int32_t>{1, 1, 2, 3, 3, 2, 1};
+  auto attr = mlir::DenseElementsAttr::get(type, llvm::ArrayRef(raw_values));
+
+  std::vector<uint8_t> packed_values;
+  auto apply_chunk_fn = [&](absl::string_view chunk) {
+    const uint8_t* data = reinterpret_cast<const uint8_t*>(chunk.data());
+    packed_values.insert(packed_values.end(), data, data + chunk.size());
+    return absl::OkStatus();
+  };
+
+  EXPECT_OK(tflite::StreamPackLowBitValues</*kBitWidth=*/2>(
+      attr.getValues<mlir::APInt>(), apply_chunk_fn));
+  EXPECT_EQ(packed_values.size(), 2);
+  EXPECT_EQ(packed_values[0], 0xE5);  // 1110 0101
+  EXPECT_EQ(packed_values[1], 0x1B);  // 0001 1011
+}
+
+TEST_F(LowBitUtilsTest, Stream2BitValues8Bit) {
+  auto type = mlir::RankedTensorType::get({7}, builder_->getIntegerType(8));
+  auto raw_values = std::vector<uint8_t>{1, 1, 2, 3, 3, 2, 1};
+  auto attr =
+      mlir::DenseElementsAttr::get(type, llvm::ArrayRef<uint8_t>(raw_values));
+
+  std::vector<uint8_t> packed_values;
+  auto apply_chunk_fn = [&](absl::string_view chunk) {
+    const uint8_t* data = reinterpret_cast<const uint8_t*>(chunk.data());
+    packed_values.insert(packed_values.end(), data, data + chunk.size());
+    return absl::OkStatus();
+  };
+
+  EXPECT_EQ(attr.getRawData().size(), 7);
+  EXPECT_OK(tflite::StreamPackLowBitValues8Bit</*kBitWidth=*/2>(
+      attr.getRawData(), apply_chunk_fn));
+  EXPECT_EQ(packed_values.size(), 2);
+  EXPECT_EQ(packed_values[0], 0xE5);  // 1110 0101
+  EXPECT_EQ(packed_values[1], 0x1B);  // 0001 1011
+}
+
+}  // namespace
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/lstm_utils.cc

@@ -0,0 +1,888 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/lstm_utils.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/utils/utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h"
+
+namespace mlir {
+namespace TFL {
+
+namespace {
+
+Value CreateI32SplatConst(OpBuilder* builder, ArrayRef<int64_t> shape,
+                          int32_t val, mlir::Location location) {
+  auto type = RankedTensorType::get(shape, builder->getIntegerType(32));
+  auto attr = DenseElementsAttr::get(type, val);
+  return arith::ConstantOp::create(*builder, location, type, attr);
+}
+
+Value CreateF32SplatConst(OpBuilder* builder, ArrayRef<int64_t> shape,
+                          float val, mlir::Location location) {
+  auto type = RankedTensorType::get(shape, builder->getF32Type());
+  auto attr = DenseElementsAttr::get(type, val);
+  return arith::ConstantOp::create(*builder, location, type, attr);
+}
+
+Value CreatTfF32ConstOp(OpBuilder* builder, ArrayRef<int64_t> shape, float val,
+                        mlir::Location location) {
+  auto type = RankedTensorType::get(shape, builder->getF32Type());
+  auto ele_type = RankedTensorType::get({1}, builder->getF32Type());
+  auto attr = DenseElementsAttr::get(ele_type, val);
+  return TF::ConstOp::create(*builder, location, type, attr);
+}
+
+Value CreateI64DenseConst(OpBuilder* builder, ArrayRef<int64_t> shape,
+                          ArrayRef<int64_t> values, mlir::Location location) {
+  auto type = RankedTensorType::get(static_cast<int>(shape.size()),
+                                    builder->getIntegerType(64));
+  auto attr = DenseElementsAttr::get(type, values);
+  return arith::ConstantOp::create(*builder, location, type, attr);
+}
+
+Value CreateI32DenseConst(OpBuilder* builder, ArrayRef<int32_t> values,
+                          mlir::Location location) {
+  auto type = RankedTensorType::get(static_cast<int>(values.size()),
+                                    builder->getIntegerType(32));
+  auto attr = DenseElementsAttr::get(type, values);
+  return arith::ConstantOp::create(*builder, location, type, attr);
+}
+
+Value CreateNoneValue(OpBuilder* builder, mlir::Location location) {
+  return TFL::NoValueOp::create(*builder, location, builder->getNoneType(),
+                                builder->getUnitAttr());
+}
+
+Value Transpose(OpBuilder* builder, Value value_to_transpose,
+                SmallVector<int32_t, 4> perm, RankedTensorType original_type,
+                mlir::Location location) {
+  // Create a constant op for transpose permutation.
+  auto perm_op = CreateI32DenseConst(builder, perm, location);
+
+  // Create tensor type for the transpose result.
+  auto transpose_type = original_type;
+  auto transpose_shape =
+      llvm::to_vector<8>(llvm::map_range(perm, [transpose_type](int32_t dim) {
+        return transpose_type.getDimSize(dim);
+      }));
+  auto elem_type = transpose_type.getElementType();
+  auto result_type = RankedTensorType::get(transpose_shape, elem_type);
+
+  return TF::TransposeOp::create(*builder, location, result_type,
+                                 value_to_transpose, perm_op);
+}
+
+Value Transpose2D(OpBuilder* builder, Value value_to_transpose,
+                  RankedTensorType type, mlir::Location location) {
+  // Create a constant op for transpose permutation.
+  SmallVector<int32_t, 4> perm = {1, 0};
+  return Transpose(builder, value_to_transpose, perm, type, location);
+}
+
+Value Reverse(OpBuilder* builder, Value value_to_reverse, int axis,
+              RankedTensorType type, mlir::Location location) {
+  auto axis_op = CreateI32SplatConst(builder, {1}, axis, location);
+  // The result type will be the same as the input.
+  return TF::ReverseV2Op::create(*builder, location, type, value_to_reverse,
+                                 axis_op);
+}
+
+ArrayRef<int64_t> GetRankedTensorShape(Value value) {
+  return mlir::cast<RankedTensorType>(value.getType()).getShape();
+}
+
+Value SliceRankedTensor(OpBuilder* builder, Value input,
+                        ArrayRef<int64_t> begin_shape,
+                        ArrayRef<int64_t> begin_values,
+                        ArrayRef<int64_t> size_shape,
+                        ArrayRef<int64_t> size_values,
+                        mlir::Location location) {
+  // If the size of the tensor to be sliced from the input overflows
+  // the input tensor's dimensions, return 0-valued tensor of the requested
+  // shape.
+  ArrayRef<int64_t> input_shape = GetRankedTensorShape(input);
+  for (int i = 0, end = input_shape.size(); i < end; i++) {
+    if (begin_values[i] < 0 ||
+        (begin_values[i] + size_values[i] > input_shape[i])) {
+      return CreateF32SplatConst(builder, size_shape, 0, location);
+    }
+  }
+
+  // Create a dense constant op for slice's begin
+  auto slice_i2c_begin =
+      CreateI64DenseConst(builder, begin_shape, begin_values, location);
+
+  // Create a dense constant op for slice's size
+  auto slice_i2c_size =
+      CreateI64DenseConst(builder, size_shape, size_values, location);
+
+  return TF::SliceOp::create(
+      *builder, location,
+      RankedTensorType::get(
+          size_values,
+          mlir::cast<RankedTensorType>(input.getType()).getElementType()),
+      input, slice_i2c_begin, slice_i2c_size);
+}
+
+Value CreateStridedSliceOp(mlir::Location loc, ArrayRef<int64_t> output_shape,
+                           Value input, ArrayRef<int32_t> begin,
+                           ArrayRef<int32_t> end, ArrayRef<int32_t> strides,
+                           int64_t begin_mask, int64_t end_mask,
+                           int64_t ellipsis_mask, int64_t new_axis_mask,
+                           int64_t shrink_axis_mask, OpBuilder* builder) {
+  auto output_type = RankedTensorType::get(
+      output_shape,
+      mlir::cast<RankedTensorType>(input.getType()).getElementType());
+  auto begin_tensor = CreateI32DenseConst(builder, begin, loc);
+  auto end_tensor = CreateI32DenseConst(builder, end, loc);
+  auto strides_tensor = CreateI32DenseConst(builder, strides, loc);
+
+  return TF::StridedSliceOp::create(
+      *builder, loc, output_type, input, begin_tensor, end_tensor,
+      strides_tensor, builder->getI64IntegerAttr(begin_mask),
+      builder->getI64IntegerAttr(end_mask),
+      builder->getI64IntegerAttr(ellipsis_mask),
+      builder->getI64IntegerAttr(new_axis_mask),
+      builder->getI64IntegerAttr(shrink_axis_mask));
+}
+
+}  // namespace
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetWeightForInputToCellGate() {
+  SmallVector<int64_t, 2> begin_i2c_values = {0, 0};
+  input2cell_ = SliceRankedTensor(
+      &builder_, weight_transposed_, weight_slice_shape_, begin_i2c_values,
+      weight_slice_shape_, weight_slice_size_input_values_,
+      fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetWeightForInputToInputGate() {
+  SmallVector<int64_t, 2> begin_i2i_values = {n_cell_, 0};
+  input2input_ = couple_input_forget_gates_
+                     ? none_
+                     : SliceRankedTensor(&builder_, weight_transposed_,
+                                         weight_slice_shape_, begin_i2i_values,
+                                         weight_slice_shape_,
+                                         weight_slice_size_input_values_,
+                                         fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetWeightForInputToForgetGate() {
+  int input_forget_start = couple_input_forget_gates_ ? n_cell_ : 2 * n_cell_;
+  SmallVector<int64_t, 2> begin_i2f_values = {input_forget_start, 0};
+  input2forget_ = SliceRankedTensor(
+      &builder_, weight_transposed_, weight_slice_shape_, begin_i2f_values,
+      weight_slice_shape_, weight_slice_size_input_values_,
+      fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetWeightForInputToOutputGate() {
+  int input_output_start =
+      couple_input_forget_gates_ ? 2 * n_cell_ : 3 * n_cell_;
+  SmallVector<int64_t, 2> begin_i2o_values = {input_output_start, 0};
+  input2output_ = SliceRankedTensor(
+      &builder_, weight_transposed_, weight_slice_shape_, begin_i2o_values,
+      weight_slice_shape_, weight_slice_size_input_values_,
+      fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetWeightForRecurrentToCellGate() {
+  SmallVector<int64_t, 2> begin_rec2c_values = {0, n_input_};
+  rec2cell_ = SliceRankedTensor(
+      &builder_, weight_transposed_, weight_slice_shape_, begin_rec2c_values,
+      weight_slice_shape_, weight_slice_size_recurrent_values_,
+      fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetWeightForRecurrentToInputGate() {
+  SmallVector<int64_t, 2> begin_rec2i_values = {n_cell_, n_input_};
+  rec2input_ = couple_input_forget_gates_
+                   ? none_
+                   : SliceRankedTensor(&builder_, weight_transposed_,
+                                       weight_slice_shape_, begin_rec2i_values,
+                                       weight_slice_shape_,
+                                       weight_slice_size_recurrent_values_,
+                                       fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetWeightForRecurrentToForgetGate() {
+  int rec_forget_start = couple_input_forget_gates_ ? n_cell_ : 2 * n_cell_;
+  SmallVector<int64_t, 2> begin_rec2f_values = {rec_forget_start, n_input_};
+  rec2forget_ = SliceRankedTensor(
+      &builder_, weight_transposed_, weight_slice_shape_, begin_rec2f_values,
+      weight_slice_shape_, weight_slice_size_recurrent_values_,
+      fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetWeightForRecurrentToOutputGate() {
+  int rec_output_start = couple_input_forget_gates_ ? 2 * n_cell_ : 3 * n_cell_;
+  SmallVector<int64_t, 2> begin_rec2o_values = {rec_output_start, n_input_};
+  rec2output_ = SliceRankedTensor(
+      &builder_, weight_transposed_, weight_slice_shape_, begin_rec2o_values,
+      weight_slice_shape_, weight_slice_size_recurrent_values_,
+      fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetBiasToCellGate() {
+  SmallVector<int64_t, 1> begin_bias2c_values = {0};
+  bias2cell_ = SliceRankedTensor(&builder_, bias_, bias_slice_shape_,
+                                 begin_bias2c_values, bias_slice_shape_,
+                                 bias_size_values_, fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetBiasToInputGate() {
+  SmallVector<int64_t, 1> begin_bias2i_values = {n_cell_};
+  bias2input_ =
+      couple_input_forget_gates_
+          ? none_
+          : SliceRankedTensor(&builder_, bias_, bias_slice_shape_,
+                              begin_bias2i_values, bias_slice_shape_,
+                              bias_size_values_, fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetBiasToForgetGate() {
+  int bias_forget_start = couple_input_forget_gates_ ? n_cell_ : 2 * n_cell_;
+  SmallVector<int64_t, 1> begin_bias2f_values = {bias_forget_start};
+  bias2forget_ = SliceRankedTensor(&builder_, bias_, bias_slice_shape_,
+                                   begin_bias2f_values, bias_slice_shape_,
+                                   bias_size_values_, fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetBiasToOutputGate() {
+  int bias_output_start =
+      couple_input_forget_gates_ ? 2 * n_cell_ : 3 * n_cell_;
+  SmallVector<int64_t, 1> begin_bias2o_values = {bias_output_start};
+  bias2output_ = SliceRankedTensor(&builder_, bias_, bias_slice_shape_,
+                                   begin_bias2o_values, bias_slice_shape_,
+                                   bias_size_values_, fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetProjection() {
+  SmallVector<int64_t, 2> projection_slice_shape = {
+      1, num_cols_projection_transposed_};
+  SmallVector<int64_t, 2> projection_slice_size_values = {n_output_, n_cell_};
+  SmallVector<int64_t, 2> projection_slice_begin_values = {0, 0};
+  proj_weight_ =
+      !projection_
+          ? none_
+          : SliceRankedTensor(
+                &builder_, projection_transposed_, projection_slice_shape,
+                projection_slice_begin_values, projection_slice_shape,
+                projection_slice_size_values, fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetProjectionBias() {
+  proj_bias_ = !projection_type_
+                   ? none_
+                   : CreateF32SplatConst(&builder_, {n_output_}, 0,
+                                         fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetInputActivationState() {
+  input_activation_state_ = CreateF32SplatConst(&builder_, {1, n_output_}, 0,
+                                                fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetInputCellState() {
+  input_cell_state_ =
+      CreateF32SplatConst(&builder_, {1, n_cell_}, 0, fused_func_op_.getLoc());
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetCellLayerNormCoefficients() {
+  cell_layer_norm_coefficients_ = none_;
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetInputLayerNormCoefficients() {
+  input_layer_norm_coefficients_ = none_;
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::SetForgetLayerNormCoefficients() {
+  forget_layer_norm_coefficients_ = none_;
+}
+void ConvertLSTMCellSimpleToFusedLSTM::SetOutputLayerNormCoefficients() {
+  output_layer_norm_coefficients_ = none_;
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::GenerateFusedOpOperands() {
+  // Transpose both weight and projection.
+  weight_transposed_ =
+      Transpose2D(&builder_, weight_, weight_type_, fused_func_op_.getLoc());
+  projection_transposed_ = Transpose2D(&builder_, projection_, projection_type_,
+                                       fused_func_op_.getLoc());
+
+  none_ = CreateNoneValue(&builder_, fused_func_op_.getLoc());
+  // Extract input to cifg gates via slicing the weight tensor
+  SetWeightForInputToCellGate();
+  SetWeightForInputToInputGate();
+  SetWeightForInputToForgetGate();
+  SetWeightForInputToOutputGate();
+
+  // Extract recurrent to cifg gates via slicing the weight tensor
+  SetWeightForRecurrentToCellGate();
+  SetWeightForRecurrentToInputGate();
+  SetWeightForRecurrentToForgetGate();
+  SetWeightForRecurrentToOutputGate();
+
+  // Extract bias to cifg gates via slicing the bias tensor
+  SetBiasToCellGate();
+  SetBiasToInputGate();
+  SetBiasToForgetGate();
+  SetBiasToOutputGate();
+
+  // Extract projection and set an empty projection bias
+  SetProjection();
+  SetProjectionBias();
+
+  // Set the variable tensors
+  SetInputActivationState();
+  SetInputCellState();
+
+  // Extract the layer norm coefficients
+  SetCellLayerNormCoefficients();
+  SetInputLayerNormCoefficients();
+  SetForgetLayerNormCoefficients();
+  SetOutputLayerNormCoefficients();
+}
+
+void ConvertLSTMCellSimpleToFusedLSTM::UpdateFuncSignature() {
+  // https://github.com/tensorflow/community/pull/113
+  SmallVector<int64_t, 2> output_shape{1, tensorflow::kTFDynamicSize};
+  auto input_types = fused_func_op_.getFunctionType().getInputs();
+  auto output_type = tensorflow::GetTypeFromTFTensorShape(
+      output_shape,
+      mlir::cast<RankedTensorType>(input_.getType()).getElementType());
+  fused_func_op_.setType(mlir::FunctionType::get(fused_func_op_.getContext(),
+                                                 input_types, output_type));
+}
+
+LogicalResult ConvertLSTMCellSimpleToFusedLSTM::RewriteFunc() {
+  LogicalResult result = Initialize();
+  if (failed(result)) {
+    return result;
+  }
+
+  // Update the func signature, based on output shape.
+  // The func will ultimately return the output of the fused
+  // LSTM op.
+  UpdateFuncSignature();
+
+  // Transform the weights, projection, bias and layer norm coefficients
+  // to generate operands for the TFL fused LSTM op.
+  GenerateFusedOpOperands();
+
+  // Create the fused LSTM op.
+  SmallVector<int64_t, 2> output_shape = {1, n_output_};
+  auto result_type = mlir::RankedTensorType::get(
+      output_shape,
+      mlir::cast<RankedTensorType>(input_.getType()).getElementType());
+  lstm_ = mlir::TFL::LSTMOp::create(
+      builder_, fused_func_op_.getLoc(), result_type, input_, input2input_,
+      input2forget_, input2cell_, input2output_, rec2input_, rec2forget_,
+      rec2cell_, rec2output_, /*cell_to_input_weights*/ none_,
+      /*cell_to_forget_weights*/ none_,
+      /*cell_to_output_weights*/ none_, bias2input_, bias2forget_, bias2cell_,
+      bias2output_, proj_weight_, proj_bias_, input_activation_state_,
+      input_cell_state_, input_layer_norm_coefficients_,
+      forget_layer_norm_coefficients_, cell_layer_norm_coefficients_,
+      output_layer_norm_coefficients_, builder_.getStringAttr("TANH"),
+      builder_.getF32FloatAttr(10.0), builder_.getF32FloatAttr(0.0),
+      mlir::TFL::LSTMKernelTypeAttr::get(builder_.getContext(),
+                                         mlir::TFL::LSTMKernelType::FULL),
+      /*asymmetric_quantize_inputs=*/mlir::BoolAttr(),
+      /*input_to_input_intermediate=*/mlir::TypeAttr(),
+      /*input_to_forget_intermediate=*/mlir::TypeAttr(),
+      /*input_to_cell_intermediate=*/mlir::TypeAttr(),
+      /*input_to_output_intermediate=*/mlir::TypeAttr(),
+      /*effective_hidden_scale_intermediate=*/mlir::TypeAttr());
+
+  // Cast the static shaped lstm result to FuncOp's signature -
+  // Ranked but unknown 2nd dimension to support stacking these.
+  SmallVector<int64_t, 2> func_output_shape = {1, tensorflow::kTFDynamicSize};
+  auto func_result_type = tensorflow::GetTypeFromTFTensorShape(
+      func_output_shape,
+      mlir::cast<RankedTensorType>(input_.getType()).getElementType());
+
+  auto tensor_cast = mlir::tensor::CastOp::create(
+      builder_, fused_func_op_.getLoc(), func_result_type, lstm_.getResult());
+  mlir::func::ReturnOp::create(builder_, fused_func_op_.getLoc(),
+                               tensor_cast.getResult());
+  return success();
+}
+
+LogicalResult ConvertLSTMCellSimpleToFusedLSTM::InitializeFromFuncAttributes() {
+  auto attr = fused_func_op_->getAttrOfType<StringAttr>(kTFImplements);
+  if (!attr) {
+    return fused_func_op_.emitError()
+           << "Invalid function attribute, expected " << kTFImplements
+           << " attribute "
+              "not found";
+  }
+
+  // TODO(ashwinm, b/144775479): Make these NamedAttribute on TF import
+  // once tf.function can support this.
+  llvm::SmallVector<llvm::StringRef, 4> attr_tokens;
+  attr.getValue().split(attr_tokens, ",");
+  if (attr_tokens.empty()) {
+    return fused_func_op_.emitError()
+           << kTFImplements << " attribute should be set";
+  }
+
+  // Check if the interface matches.
+  if (GetCompositeOpName().str() != attr_tokens[0]) {
+    return fused_func_op_.emitError()
+           << "Unexpected interface for the composite op. Expected: "
+           << GetCompositeOpName() << " Actual: " << attr_tokens[0];
+  }
+
+  // Extract other interface attributes, for now cifg.
+  couple_input_forget_gates_ =
+      std::find(attr_tokens.begin() + 1, attr_tokens.end(),
+                kCoupleInputForgetGates) != attr_tokens.end();
+
+  return success();
+}
+
+LogicalResult ConvertLSTMCellSimpleToFusedLSTM::Initialize() {
+  if (failed(InitializeFromFuncAttributes())) {
+    return fused_func_op_.emitError()
+           << "Expected function attributes were not set on the function "
+              "encapsulating the composite op";
+  }
+
+  num_gates_ = couple_input_forget_gates_ ? 3 : 4;
+
+  input_ = fused_func_op_.getArgument(0);
+  bias_ = fused_func_op_.getArgument(2);
+
+  weight_ = fused_func_op_.getArgument(1);
+  weight_type_ = mlir::cast<RankedTensorType>(weight_.getType());
+
+  if (weight_type_.getRank() != 2) {
+    return fused_func_op_.emitError() << "The weight tensor was not of rank 2";
+  }
+
+  if (weight_type_.getDimSize(1) % num_gates_ != 0) {
+    return fused_func_op_.emitError()
+           << "Invalid dimension 1 of weight tensor, "
+              "should be divisible by the number of gates";
+  }
+  n_cell_ = weight_type_.getDimSize(1) / num_gates_;
+
+  projection_ = fused_func_op_.getArgument(3);
+  projection_type_ = mlir::cast<RankedTensorType>(projection_.getType());
+  if (projection_type_.getRank() != 2) {
+    n_output_ = n_cell_;
+  } else {
+    n_output_ = projection_type_.getDimSize(1);
+  }
+  n_input_ = weight_type_.getDimSize(0) - n_output_;
+  num_cols_weight_transposed_ = weight_type_.getDimSize(0);
+  num_cols_projection_transposed_ = projection_type_.getDimSize(0);
+
+  bias_slice_shape_ = {n_cell_};
+  bias_size_values_ = {n_cell_};
+  weight_slice_shape_ = {1, num_cols_weight_transposed_};
+  weight_slice_size_input_values_ = {n_cell_, n_input_};
+  weight_slice_size_recurrent_values_ = {n_cell_, n_output_};
+
+  return success();
+}
+
+LogicalResult ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM::Initialize() {
+  if (failed(ConvertLSTMCellSimpleToFusedLSTM::Initialize())) {
+    return fused_func_op_.emitError()
+           << "Specified LayerNormalizedLSTMCellSimple was not of the expected "
+              "interface and cannot not be converted to the fused LSTM op";
+  }
+
+  layer_norm_scale_ = fused_func_op_.getArgument(4);
+  layer_norm_scale_type_ =
+      mlir::cast<RankedTensorType>(layer_norm_scale_.getType());
+  if (layer_norm_scale_type_.getRank() != 1) {
+    return fused_func_op_.emitError()
+           << "The layer_norm_scale tensor was not of rank 1";
+  }
+  layer_norm_slice_shape_ = {n_cell_};
+  layer_norm_size_values_ = {n_cell_};
+
+  return success();
+}
+
+void ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM::
+    SetCellLayerNormCoefficients() {
+  SmallVector<int64_t, 1> begin_cell_layer_norm_values = {0};
+  cell_layer_norm_coefficients_ =
+      SliceRankedTensor(&builder_, layer_norm_scale_, layer_norm_slice_shape_,
+                        begin_cell_layer_norm_values, layer_norm_slice_shape_,
+                        layer_norm_size_values_, fused_func_op_.getLoc());
+}
+
+void ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM::
+    SetInputLayerNormCoefficients() {
+  SmallVector<int64_t, 1> begin_input_layer_norm_values = {n_cell_};
+  input_layer_norm_coefficients_ =
+      couple_input_forget_gates_
+          ? none_
+          : SliceRankedTensor(
+                &builder_, layer_norm_scale_, layer_norm_slice_shape_,
+                begin_input_layer_norm_values, layer_norm_slice_shape_,
+                layer_norm_size_values_, fused_func_op_.getLoc());
+}
+
+void ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM::
+    SetForgetLayerNormCoefficients() {
+  SmallVector<int64_t, 1> begin_forget_layer_norm_values = {2 * n_cell_};
+  forget_layer_norm_coefficients_ =
+      SliceRankedTensor(&builder_, layer_norm_scale_, layer_norm_slice_shape_,
+                        begin_forget_layer_norm_values, layer_norm_slice_shape_,
+                        layer_norm_size_values_, fused_func_op_.getLoc());
+}
+
+void ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM::
+    SetOutputLayerNormCoefficients() {
+  SmallVector<int64_t, 1> begin_output_layer_norm_values = {3 * n_cell_};
+  output_layer_norm_coefficients_ =
+      SliceRankedTensor(&builder_, layer_norm_scale_, layer_norm_slice_shape_,
+                        begin_output_layer_norm_values, layer_norm_slice_shape_,
+                        layer_norm_size_values_, fused_func_op_.getLoc());
+}
+
+TF::ConstOp Create1DConstantOp(const std::vector<int>& value, Location loc,
+                               OpBuilder* builder) {
+  auto type =
+      mlir::RankedTensorType::get(value.size(), builder->getIntegerType(32));
+  auto dense_values = mlir::DenseIntElementsAttr::get(type, value);
+  return TF::ConstOp::create(*builder, loc, dense_values);
+}
+
+TF::ConstOp CreateScalarConstantOp(int value, Location loc,
+                                   OpBuilder* builder) {
+  return TF::ConstOp::create(*builder, loc, builder->getI32IntegerAttr(value));
+}
+
+TF::ReshapeOp CreateFlattenOP(const Value& input, Location loc,
+                              OpBuilder* builder) {
+  auto output_shape = Create1DConstantOp({-1}, loc, builder);
+  return mlir::TF::ReshapeOp::create(*builder, loc,
+                                     /*tensor=*/input,
+                                     /*shape=*/output_shape.getResult());
+}
+
+LogicalResult CreateEqualSizeSplitVOp(Value input, int axis, int splits,
+                                      Location loc, OpBuilder* builder,
+                                      Operation** result) {
+  auto input_type = mlir::cast<RankedTensorType>(input.getType());
+  SmallVector<int64_t, 4> output_shape;
+  int size_of_splits;
+  if (input_type.getRank() < axis || axis < 0) return failure();
+  for (int i = 0; i < input_type.getRank(); ++i) {
+    int64_t dim = input_type.getDimSize(i);
+    if (i == axis) {
+      if (dim % splits != 0) {
+        return failure();
+      }
+      size_of_splits = dim / splits;
+      output_shape.push_back(size_of_splits);
+    } else {
+      output_shape.push_back(dim);
+    }
+  }
+
+  SmallVector<mlir::Type, 4> output_types;
+  for (int i = 0; i < splits; ++i) {
+    output_types.push_back(
+        mlir::RankedTensorType::get(output_shape, input_type.getElementType()));
+  }
+  auto size_of_splits_op = Create1DConstantOp(
+      {size_of_splits, size_of_splits, size_of_splits, size_of_splits}, loc,
+      builder);
+
+  auto axis_op = CreateScalarConstantOp(axis, loc, builder);
+  *result =
+      TF::SplitVOp::create(*builder, loc, output_types, input,
+                           size_of_splits_op.getResult(), axis_op.getResult());
+  return success();
+}
+
+// TODO(b/147436982): Consider refactoring these to be more general.
+LogicalResult ConvertKerasLSTMLayer(mlir::func::FuncOp func_op,
+                                    OpBuilder* builder) {
+  return ConvertKerasLSTMLayer(func_op, builder, false);
+}
+
+LogicalResult ConvertKerasLSTMLayer(mlir::func::FuncOp func_op,
+                                    OpBuilder* builder, bool indy) {
+  // For argument order, please check out standard_lstm under
+  // tensorflow/python/keras/layers/recurrent_v2.py
+  Value input = func_op.getArgument(0);
+  Value output_init_state = func_op.getArgument(1);
+  Value hidden_init_state = func_op.getArgument(2);
+  Value weight_kernel = func_op.getArgument(3);
+  Value recurrent_kernel = func_op.getArgument(4);
+  Value bias = func_op.getArgument(5);
+
+  // The func op should have 5 outputs.
+  if (func_op.getNumResults() != 5) return failure();
+
+  // TFL lstm only supports time-majored inputs, so if it's not time-majored,
+  // we will transpose the inputs and outputs.
+  auto time_major_attr = func_op->getAttrOfType<BoolAttr>("tf.time_major");
+  if (time_major_attr == nullptr) return failure();
+
+  bool time_majored = time_major_attr.getValue();
+  auto input_type = mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
+  if (!input_type) {
+    func_op.emitError() << "Input type is not a ranked tensor type";
+    return failure();
+  }
+
+  auto final_inputs = input;
+  auto final_input_type = input_type;
+
+  // Handle go_backwards:
+  // LSTM in Keras semantic will reverse the input sequence if it's go_backwards
+  auto go_backwards_attr = func_op->getAttrOfType<BoolAttr>("tf.go_backwards");
+
+  if (go_backwards_attr != nullptr && go_backwards_attr.getValue()) {
+    int time_dim = time_majored ? 0 : 1;
+    final_inputs = Reverse(builder, final_inputs, time_dim, final_input_type,
+                           func_op.getLoc());
+  }
+
+  int64_t batch = time_majored ? final_input_type.getDimSize(1)
+                               : final_input_type.getDimSize(0);
+  int64_t time = time_majored ? final_input_type.getDimSize(0)
+                              : final_input_type.getDimSize(1);
+
+  // Setup correct weights.
+  RankedTensorType weight_type =
+      mlir::cast<RankedTensorType>(weight_kernel.getType());
+  if (weight_type.getRank() != 2)
+    return func_op.emitError() << "The weight should be rank of 2";
+
+  Value transposed_weight_kernel =
+      Transpose2D(builder, weight_kernel, weight_type, func_op.getLoc());
+
+  RankedTensorType recurrent_kernel_type =
+      mlir::cast<RankedTensorType>(recurrent_kernel.getType());
+  const int64_t n_output = recurrent_kernel_type.getDimSize(0);
+
+  Value transpose_recurrent_kernel = Transpose2D(
+      builder, recurrent_kernel, recurrent_kernel_type, func_op.getLoc());
+
+  // Splits the weights into 4: i, f, c, o.
+  const int splits = 4;
+
+  Operation* weights_array;
+  if (failed(CreateEqualSizeSplitVOp(transposed_weight_kernel, 0, splits,
+                                     func_op.getLoc(), builder,
+                                     &weights_array)))
+    return failure();
+
+  // Splits the recurrent_weights into 4:
+  Operation* recurrent_weights_array;
+  if (failed(CreateEqualSizeSplitVOp(transpose_recurrent_kernel, 0, splits,
+                                     func_op.getLoc(), builder,
+                                     &recurrent_weights_array)))
+    return failure();
+
+  // Reshape recurrent weights to vectors if indy behaviour is enabled.
+  // IndyLSTMs are a LSTM variant with diagonal recurrent weight
+  // matrices. For optimization purposes these are provided as vectors.
+  Value recurrent_to_input_weights =
+      indy ? mlir::cast<Value>(
+                 CreateFlattenOP(recurrent_weights_array->getResult(0),
+                                 func_op.getLoc(), builder)
+                     .getResult())
+           : recurrent_weights_array->getResult(0);
+  Value recurrent_to_forget_weights =
+      indy ? mlir::cast<Value>(
+                 CreateFlattenOP(recurrent_weights_array->getResult(1),
+                                 func_op.getLoc(), builder)
+                     .getResult())
+           : recurrent_weights_array->getResult(1);
+  Value recurrent_to_cell_weights =
+      indy ? mlir::cast<Value>(
+                 CreateFlattenOP(recurrent_weights_array->getResult(2),
+                                 func_op.getLoc(), builder)
+                     .getResult())
+           : recurrent_weights_array->getResult(2);
+  Value recurrent_to_output_weights =
+      indy ? mlir::cast<Value>(
+                 CreateFlattenOP(recurrent_weights_array->getResult(3),
+                                 func_op.getLoc(), builder)
+                     .getResult())
+           : recurrent_weights_array->getResult(3);
+
+  // Splits the bias into 4:
+  Operation* bias_array;
+  if (failed(CreateEqualSizeSplitVOp(bias, 0, splits, func_op.getLoc(), builder,
+                                     &bias_array)))
+    return failure();
+
+  // Build the lstm op.
+  SmallVector<int64_t, 3> output_shape;
+  if (time_majored) {
+    output_shape = {time, batch, n_output};
+  } else {
+    output_shape = {batch, time, n_output};
+  }
+  auto result_type = mlir::RankedTensorType::get(
+      output_shape,
+      mlir::cast<RankedTensorType>(final_inputs.getType()).getElementType());
+
+  Value none = CreateNoneValue(builder, func_op.getLoc());
+  auto lstm = mlir::TFL::UnidirectionalSequenceLSTMOp::create(
+      *builder, func_op.getLoc(), result_type, /*input=*/final_inputs,
+      /*input_to_input_weights=*/weights_array->getResult(0),
+      /*input_to_forget_weights=*/weights_array->getResult(1),
+      /*input_to_cell_weights=*/weights_array->getResult(2),
+      /*input_to_output_weights=*/weights_array->getResult(3),
+      /*recurrent_to_input_weights=*/recurrent_to_input_weights,
+      /*recurrent_to_forget_weights=*/recurrent_to_forget_weights,
+      /*recurrent_to_cell_weights=*/recurrent_to_cell_weights,
+      /*recurrent_to_output_weights=*/recurrent_to_output_weights,
+      /*cell_to_input_weights=*/none,
+      /*cell_to_forget_weights=*/none,
+      /*cell_to_output_weights=*/none,
+      /*input_gate_bias=*/bias_array->getResult(0),
+      /*forget_gate_bias=*/bias_array->getResult(1),
+      /*cell_bias=*/bias_array->getResult(2),
+      /*output_gate_bias=*/bias_array->getResult(3),
+      /*projection_weights=*/none,
+      /*projection_bias=*/none,
+      /*input_activation_state=*/output_init_state,
+      /*input_cell_state=*/hidden_init_state,
+      /*input_layer_norm_coefficients=*/none,
+      /*forget_layer_norm_coefficients=*/none,
+      /*cell_layer_norm_coefficients=*/none,
+      /*output_layer_norm_coefficients=*/none,
+      /*fused_activation_function*/ builder->getStringAttr("TANH"),
+      /*cell_clip*/ builder->getF32FloatAttr(10.0),
+      /*proj_clip*/ builder->getF32FloatAttr(0.0),
+      /*time_major*/ builder->getBoolAttr(time_majored),
+      /*asymmetric_quantize_inputs=*/mlir::BoolAttr(),
+      /*diagonal_recurrent_tensors=*/builder->getBoolAttr(indy),
+      /*input_to_input_intermediate=*/mlir::TypeAttr(),
+      /*input_to_forget_intermediate=*/mlir::TypeAttr(),
+      /*input_to_cell_intermediate=*/mlir::TypeAttr(),
+      /*input_to_output_intermediate=*/mlir::TypeAttr(),
+      /*effective_hidden_scale_intermediate=*/mlir::TypeAttr());
+
+  auto final_output_full_sequences = lstm.getResult();
+
+  // Populate the last output: last output is sliced from the full sequences.
+  // If time_major: last_output = outputs[-1, :, :]
+  // else: last_output = outputs[:, -1, :]
+  //
+  // As we are creating the strided_slice op, we need to populate the following
+  // fields:
+  // end: should always be (0, 0, 0)
+  // strides: should always be (1, 1, 1)
+  // begin: should be (0, -1, 0) or (-1, 0, 0) if it's time-majored.
+  // new_axis_mask: should always be 0.
+  // ellipsis_mask: should always be 0.
+  // begin_mask & end_mask: should be 0b101 = 5 or 0b110 = 4 if it's
+  // time-majored. shrink_axis_mask: should be 0b010 = 2 or 0b001 = 1 if it's
+  // time-majored.
+  SmallVector<int64_t, 2> last_output_shape({batch, n_output});
+
+  SmallVector<int32_t, 3> end({0, 0, 0});
+  SmallVector<int32_t, 3> strides({1, 1, 1});
+  SmallVector<int32_t, 3> begin;
+
+  int64_t new_axis_mask = 0;
+  int64_t ellipsis_mask = 0;
+  int64_t begin_mask;
+  int64_t end_mask;
+  int64_t shrink_axis_mask;
+  if (time_majored) {
+    begin_mask = 6;
+    end_mask = 6;
+    shrink_axis_mask = 1;
+    begin = {-1, 0, 0};
+  } else {
+    begin_mask = 5;
+    end_mask = 5;
+    shrink_axis_mask = 2;
+    begin = {0, -1, 0};
+  }
+
+  auto last_output = CreateStridedSliceOp(
+      func_op.getLoc(), last_output_shape, final_output_full_sequences, begin,
+      end, strides, begin_mask, end_mask, ellipsis_mask, new_axis_mask,
+      shrink_axis_mask, builder);
+
+  SmallVector<Value, 5> outputs;
+  SmallVector<Type, 5> output_types;
+
+  // Due to the existence of the while loop, the timestamp may be unknown
+  // for the signature, for us, since we know the inputs, we can infer the time
+  // steps.
+
+  // Last output.
+  outputs.push_back(last_output);
+  output_types.push_back(last_output.getType());
+
+  // Full sequences.
+  outputs.push_back(final_output_full_sequences);
+  output_types.push_back(final_output_full_sequences.getType());
+
+  // All the rest: states, device.
+  for (int i = 2; i < 5; ++i) {
+    auto result_type =
+        mlir::dyn_cast<RankedTensorType>(func_op.getResultTypes()[i]);
+    outputs.push_back(CreatTfF32ConstOp(builder, result_type.getShape(), 0.0f,
+                                        func_op.getLoc()));
+    output_types.push_back(result_type);
+  }
+
+  // Update function signatures.
+  func_op.setType(mlir::FunctionType::get(func_op.getContext(),
+                                          func_op.getFunctionType().getInputs(),
+                                          output_types));
+
+  mlir::func::ReturnOp::create(*builder, func_op.getLoc(), outputs);
+  return success();
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/lstm_utils.h

@@ -0,0 +1,224 @@
+/* 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.
+==============================================================================*/
+
+// This header file defines common utils used by TFLite transformation
+// passes to work with op attributes.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_LSTM_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_LSTM_UTILS_H_
+
+#include <cstdint>
+
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/utils/utils.h"
+
+namespace mlir {
+namespace TFL {
+
+constexpr char kTFImplements[] = "tf._implements";
+constexpr char kLstmCellSimple[] = "LSTMCellSimple";
+constexpr char kLayerNormalizedLstmCellSimple[] =
+    "LayerNormalizedLstmCellSimple";
+constexpr char kCoupleInputForgetGates[] = "CoupleInputForgetGates";
+
+// A utility class that enables the conversion of the LSTMCellSimple composite
+// op into a fused TFL LSTM op. The fused op is contained within a FuncOp
+// that also contains other supporting ops needed to construct the operands for
+// the fused op. The caller provides the containing FuncOp as input with
+// arguments specifying the input, weight, projection and bias.
+// The weight, projection, bias and layer norm scale all need to be
+// RankedTensorType.
+// This class sets the layer norm coefficients to NoneType.
+class ConvertLSTMCellSimpleToFusedLSTM {
+ public:
+  explicit ConvertLSTMCellSimpleToFusedLSTM(mlir::func::FuncOp fused_func_op)
+      : fused_func_op_(fused_func_op),
+        couple_input_forget_gates_(false),
+        builder_(fused_func_op.getBody()) {}
+
+  // not copyable.
+  ConvertLSTMCellSimpleToFusedLSTM(const ConvertLSTMCellSimpleToFusedLSTM&) =
+      delete;
+  ConvertLSTMCellSimpleToFusedLSTM& operator=(
+      const ConvertLSTMCellSimpleToFusedLSTM&) = delete;
+  virtual ~ConvertLSTMCellSimpleToFusedLSTM() = default;
+
+  virtual llvm::StringRef GetCompositeOpName() { return kLstmCellSimple; }
+
+  // Rewrite the func body with constructed fused lstm.
+  LogicalResult RewriteFunc();
+
+  int GetNumInputs() { return n_input_; }
+
+ protected:
+  // verify input func op arguments/attributes and initialize internal state.
+  virtual LogicalResult InitializeFromFuncAttributes();
+  virtual LogicalResult Initialize();
+
+  void UpdateFuncSignature();
+  void GenerateFusedOpOperands();
+
+  void SetWeightForInputToCellGate();
+  void SetWeightForInputToInputGate();
+  void SetWeightForInputToForgetGate();
+  void SetWeightForInputToOutputGate();
+
+  void SetWeightForRecurrentToCellGate();
+  void SetWeightForRecurrentToInputGate();
+  void SetWeightForRecurrentToForgetGate();
+  void SetWeightForRecurrentToOutputGate();
+
+  void SetBiasToCellGate();
+  void SetBiasToInputGate();
+  void SetBiasToForgetGate();
+  void SetBiasToOutputGate();
+
+  void SetProjection();
+  void SetProjectionBias();
+
+  void SetInputActivationState();
+  void SetInputCellState();
+
+  virtual void SetCellLayerNormCoefficients();
+  virtual void SetInputLayerNormCoefficients();
+  virtual void SetForgetLayerNormCoefficients();
+  virtual void SetOutputLayerNormCoefficients();
+
+  // specified state
+  func::FuncOp fused_func_op_;
+  Value input_;
+  Value weight_;
+  Value bias_;
+  Value projection_;
+  bool couple_input_forget_gates_;
+
+  // internal state
+  Value weight_transposed_;
+  Value projection_transposed_;
+  RankedTensorType weight_type_;
+  RankedTensorType projection_type_;
+  int num_gates_;
+  int n_cell_;
+  int n_output_;
+  int n_input_;
+  int num_cols_weight_transposed_;
+  int num_cols_projection_transposed_;
+
+  // input -> cifg
+  Value input2input_;
+  Value input2forget_;
+  Value input2cell_;
+  Value input2output_;
+
+  // recurrent -> cifg
+  Value rec2input_;
+  Value rec2forget_;
+  Value rec2cell_;
+  Value rec2output_;
+
+  // bias -> cifg
+  Value bias2input_;
+  Value bias2forget_;
+  Value bias2cell_;
+  Value bias2output_;
+
+  // projection
+  Value proj_weight_;
+  Value proj_bias_;
+
+  // state
+  Value input_activation_state_;
+  Value input_cell_state_;
+
+  // layer norm coefficients
+  Value input_layer_norm_coefficients_;
+  Value forget_layer_norm_coefficients_;
+  Value cell_layer_norm_coefficients_;
+  Value output_layer_norm_coefficients_;
+
+  mlir::TFL::LSTMOp lstm_;
+
+  Value none_;
+  SmallVector<int64_t, 1> bias_slice_shape_;
+  SmallVector<int64_t, 1> bias_size_values_;
+  SmallVector<int64_t, 2> weight_slice_shape_;
+  SmallVector<int64_t, 2> weight_slice_size_input_values_;
+  SmallVector<int64_t, 2> weight_slice_size_recurrent_values_;
+  OpBuilder builder_;
+};
+
+// A utility class that enables the conversion of the
+// LayerNormalizedLSTMCellSimple composite op into a fused TFL LSTM op. The
+// fused op is contained within a FuncOp that also contains other supporting ops
+// needed to construct the operands for the fused op. The caller provides the
+// containing FuncOp as input with arguments specifying the input, weight,
+// projection, bias and layer norm scale. The weight, projection, bias and
+// layer norm scale all need to be RankedTensorType.
+// This class overrides the layer norm coefficient setters from the base class.
+class ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM
+    : public ConvertLSTMCellSimpleToFusedLSTM {
+ public:
+  explicit ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM(
+      mlir::func::FuncOp fused_func_op)
+      : ConvertLSTMCellSimpleToFusedLSTM(fused_func_op) {}
+
+  // not copyable.
+  ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM(
+      const ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM&) = delete;
+  ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM& operator=(
+      const ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM&) = delete;
+  ~ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM() override = default;
+
+  llvm::StringRef GetCompositeOpName() override {
+    return kLayerNormalizedLstmCellSimple;
+  }
+
+ protected:
+  LogicalResult Initialize() override;
+
+  void SetCellLayerNormCoefficients() override;
+  void SetInputLayerNormCoefficients() override;
+  void SetForgetLayerNormCoefficients() override;
+  void SetOutputLayerNormCoefficients() override;
+
+ private:
+  // specified state
+  Value layer_norm_scale_;
+
+  // internal state
+  RankedTensorType layer_norm_scale_type_;
+  SmallVector<int64_t, 1> layer_norm_slice_shape_;
+  SmallVector<int64_t, 1> layer_norm_size_values_;
+};
+
+LogicalResult ConvertKerasLSTMLayer(mlir::func::FuncOp func_op,
+                                    OpBuilder* builder);
+
+LogicalResult ConvertKerasLSTMLayer(mlir::func::FuncOp func_op,
+                                    OpBuilder* builder, bool indy);
+
+}  // end namespace TFL
+}  // end namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_LSTM_UTILS_H_

tensorflow/compiler/mlir/lite/utils/lstm_utils_test.cc

@@ -0,0 +1,262 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/lstm_utils.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringExtras.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir {
+namespace TFL {
+
+func::FuncOp createLstmCompositeFunc(mlir::Builder* builder, bool ln,
+                                     bool cifg) {
+  SmallVector<int64_t, 2> input_shape{1, 2};
+  SmallVector<int64_t, 2> weight_shape{3, 12};
+  SmallVector<int64_t, 1> bias_shape{2};
+  SmallVector<int64_t, 2> projection_shape{1, 2};
+  SmallVector<int64_t, 1> layer_norm_scale{4};
+  SmallVector<int64_t, 2> output_shape{1, 2};
+  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());
+  auto weight_type = RankedTensorType::get(weight_shape, builder->getF32Type());
+  auto bias_type = RankedTensorType::get(bias_shape, builder->getF32Type());
+  auto projection_type =
+      RankedTensorType::get(projection_shape, builder->getF32Type());
+  auto layer_norm_scale_type =
+      RankedTensorType::get(layer_norm_scale, builder->getF32Type());
+  auto output_type = RankedTensorType::get(output_shape, builder->getF32Type());
+  SmallVector<mlir::Type, 4> input_types{input_type, weight_type, bias_type,
+                                         projection_type,
+                                         layer_norm_scale_type};
+  auto func_type = builder->getFunctionType(input_types, output_type);
+
+  auto func = func::FuncOp::create(
+      mlir::NameLoc::get(builder->getStringAttr("fused_func")), "fused_func",
+      func_type, {});
+  func.addEntryBlock();
+
+  std::vector<std::string> attributes;
+  if (ln) {
+    attributes.push_back(kLayerNormalizedLstmCellSimple);
+  } else {
+    attributes.push_back(kLstmCellSimple);
+  }
+
+  if (cifg) {
+    attributes.push_back(kCoupleInputForgetGates);
+  }
+
+  mlir::StringAttr attr_values =
+      builder->getStringAttr(llvm::join(attributes, ","));
+
+  func->setAttr(kTFImplements, attr_values);
+  return func;
+}
+
+class LstmUtilsTest : public ::testing::Test {
+ protected:
+  LstmUtilsTest() = default;
+
+  void SetUp() override {
+    context_ = std::make_unique<mlir::MLIRContext>();
+    context_->loadDialect<arith::ArithDialect, mlir::func::FuncDialect,
+                          tensor::TensorDialect, mlir::TF::TensorFlowDialect,
+                          TensorFlowLiteDialect>();
+    builder_ = std::make_unique<mlir::Builder>(context_.get());
+    fused_lstm_func_ = createLstmCompositeFunc(builder_.get(), false, false);
+    fused_lstm_func_cifg_ =
+        createLstmCompositeFunc(builder_.get(), false, true);
+    fused_ln_lstm_func_ = createLstmCompositeFunc(builder_.get(), true, false);
+  }
+
+  void TearDown() override {
+    fused_lstm_func_.erase();
+    fused_lstm_func_cifg_.erase();
+    fused_ln_lstm_func_.erase();
+    builder_.reset();
+  }
+
+  func::FuncOp fused_lstm_func_;
+  func::FuncOp fused_lstm_func_cifg_;
+  func::FuncOp fused_ln_lstm_func_;
+  std::unique_ptr<mlir::MLIRContext> context_;
+  std::unique_ptr<mlir::Builder> builder_;
+};
+
+TEST_F(LstmUtilsTest, ConvertLSTMCellSimple) {
+  mlir::TFL::ConvertLSTMCellSimpleToFusedLSTM convert(fused_lstm_func_);
+
+  auto result = convert.RewriteFunc();
+  EXPECT_FALSE(failed(result));
+  fused_lstm_func_.dump();
+
+  // verify transpose
+  EXPECT_EQ(
+      fused_lstm_func_->getAttrOfType<StringAttr>(kTFImplements).getValue(),
+      convert.GetCompositeOpName());
+  EXPECT_EQ(fused_lstm_func_.getNumArguments(), 5);
+  EXPECT_EQ(fused_lstm_func_.getFunctionType().getNumResults(), 1);
+
+  auto transpose_op = fused_lstm_func_.getBody().front().begin();
+  transpose_op++;
+  EXPECT_EQ(mlir::cast<RankedTensorType>(transpose_op->getOperand(0).getType())
+                .getDimSize(0),
+            3);
+  EXPECT_EQ(mlir::cast<RankedTensorType>(transpose_op->getOperand(0).getType())
+                .getDimSize(1),
+            12);
+  EXPECT_EQ(mlir::cast<RankedTensorType>(transpose_op->getResult(0).getType())
+                .getDimSize(0),
+            12);
+  EXPECT_EQ(mlir::cast<RankedTensorType>(transpose_op->getResult(0).getType())
+                .getDimSize(1),
+            3);
+
+  auto it = fused_lstm_func_.getBody().back().rbegin();
+  EXPECT_EQ(it->getName().getStringRef(),
+            mlir::func::ReturnOp::getOperationName());
+  it++;  // tensor_cast
+  it++;  // lstm
+  EXPECT_EQ(it->getName().getStringRef(),
+            mlir::TFL::LSTMOp::getOperationName());
+  EXPECT_EQ(it->getNumOperands(), 24);
+  EXPECT_EQ(it->getNumResults(), 1);
+  // cifg = false, so input2input is not None.
+  EXPECT_FALSE(mlir::isa<NoneType>(it->getOperand(1).getType()));
+  // input layer norm is None
+  EXPECT_TRUE(mlir::isa<NoneType>(it->getOperand(20).getType()));
+  // proj_bias is F32
+  EXPECT_TRUE(mlir::cast<RankedTensorType>(it->getOperand(17).getType())
+                  .getElementType()
+                  .isF32());
+
+  // output gate bias is 0 since it is out of bounds of the bias tensor, so
+  // we set its value as a const tensor of specified size and value 0.
+  EXPECT_TRUE(
+      mlir::cast<ElementsAttr>(mlir::cast<mlir::arith::ConstantOp>(
+                                   it->getOpOperand(15).get().getDefiningOp())
+                                   .getValue())
+          .getValues<FloatAttr>()[0]
+          .getValue()
+          .isExactlyValue(0.0f));
+
+  EXPECT_EQ(fused_lstm_func_.getFunctionType().getNumResults(), 1);
+  auto output_types = fused_lstm_func_.getFunctionType().getResults();
+  SmallVector<int64_t, 2> output_shape{1, mlir::ShapedType::kDynamic};
+  EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getShape().size(),
+            output_shape.size());
+  for (int i = 0; i < output_shape.size(); i++) {
+    EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getDimSize(i),
+              output_shape[i]);
+  }
+}
+
+TEST_F(LstmUtilsTest, ConvertLSTMCellSimpleToFusedLSTMCoupleInputForget) {
+  mlir::TFL::ConvertLSTMCellSimpleToFusedLSTM convert(fused_lstm_func_cifg_);
+
+  auto result = convert.RewriteFunc();
+  EXPECT_FALSE(failed(result));
+  fused_lstm_func_cifg_.dump();
+
+  llvm::SmallVector<std::string, 2> attributes{kLstmCellSimple,
+                                               kCoupleInputForgetGates};
+  EXPECT_EQ(fused_lstm_func_cifg_->getAttrOfType<StringAttr>(kTFImplements)
+                .getValue(),
+            llvm::join(attributes, ","));
+
+  auto it = fused_lstm_func_cifg_.getBody().back().rbegin();
+  EXPECT_EQ(it->getName().getStringRef(),
+            mlir::func::ReturnOp::getOperationName());
+  it++;
+  it++;
+  EXPECT_EQ(it->getName().getStringRef(),
+            mlir::TFL::LSTMOp::getOperationName());
+  EXPECT_EQ(it->getNumOperands(), 24);
+  EXPECT_EQ(it->getNumResults(), 1);
+  // cifg = true, so input2input is None.
+  EXPECT_TRUE(mlir::isa<NoneType>(it->getOperand(1).getType()));
+}
+
+TEST_F(LstmUtilsTest, ConvertLayerNormLSTMCellSimpleToFusedLSTM) {
+  mlir::TFL::ConvertLayerNormalizedLSTMCellSimpleToFusedLSTM convert(
+      fused_ln_lstm_func_);
+
+  auto result = convert.RewriteFunc();
+  EXPECT_FALSE(failed(result));
+  fused_ln_lstm_func_.dump();
+
+  EXPECT_EQ(
+      fused_ln_lstm_func_->getAttrOfType<StringAttr>(kTFImplements).getValue(),
+      convert.GetCompositeOpName());
+  EXPECT_EQ(fused_ln_lstm_func_.getNumArguments(), 5);
+  EXPECT_EQ(fused_ln_lstm_func_.getFunctionType().getNumResults(), 1);
+
+  auto it = fused_ln_lstm_func_.getBody().back().rbegin();
+  EXPECT_EQ(it->getName().getStringRef(),
+            mlir::func::ReturnOp::getOperationName());
+  it++;
+  it++;
+  EXPECT_EQ(it->getName().getStringRef(),
+            mlir::TFL::LSTMOp::getOperationName());
+  EXPECT_EQ(it->getNumOperands(), 24);
+  EXPECT_EQ(it->getNumResults(), 1);
+  // cifg = false, so input2input is not None.
+  EXPECT_FALSE(mlir::isa<NoneType>(it->getOperand(1).getType()));
+
+  // input layer norm
+  EXPECT_FALSE(mlir::isa<NoneType>(it->getOperand(20).getType()));
+  EXPECT_EQ(mlir::cast<RankedTensorType>(it->getOperand(20).getType())
+                .getShape()
+                .size(),
+            1);
+  EXPECT_EQ(
+      mlir::cast<RankedTensorType>(it->getOperand(20).getType()).getDimSize(0),
+      3);
+
+  EXPECT_EQ(fused_ln_lstm_func_.getFunctionType().getNumResults(), 1);
+  auto output_types = fused_ln_lstm_func_.getFunctionType().getResults();
+  SmallVector<int64_t, 2> output_shape{1, mlir::ShapedType::kDynamic};
+  EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getShape().size(),
+            output_shape.size());
+  for (int i = 0; i < output_shape.size(); i++) {
+    EXPECT_EQ(mlir::cast<RankedTensorType>(output_types[0]).getDimSize(i),
+              output_shape[i]);
+  }
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/metadata_utils.h

@@ -0,0 +1,29 @@
+/* Copyright 2025 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_UTILS_METADATA_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_METADATA_UTILS_H_
+
+namespace mlir {
+namespace TFL {
+
+// This constant serves as a marker on a FusedLoc to identify it as a
+// temporary wrapper added by the TFLite importer for roundtrip purposes.
+inline constexpr char kImporterWrapper[] = "tflite.importer_wrapper";
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_METADATA_UTILS_H_

tensorflow/compiler/mlir/lite/utils/mlir_module_utils.h

@@ -0,0 +1,76 @@
+/* Copyright 2025 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_UTILS_MLIR_MODULE_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_MLIR_MODULE_UTILS_H_
+
+#include <stdlib.h>
+
+#include <cstdint>
+
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/utils/const_tensor_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
+
+namespace mlir {
+namespace TFL {
+
+// This function estimates the size of the module in mega bytes. It does so by
+// iterating through all the constant-like attributes and tensors in the module
+// and summing up their sizes.
+//
+// This function is used to reserve space in the buffer before serializing the
+// module to avoid reallocating the buffer during serialization.
+//
+// This function may need to be improved to give more accurate size of the
+// module if the current estimate is not good enough and causes huge
+// reallocations during serialization.
+inline uint64_t GetApproximateModuleSize(mlir::ModuleOp module) {
+  uint64_t module_size_estimate = 0;
+  mlir::DenseSet<mlir::Attribute> unique_tensors;
+
+  for (auto global_tensor_op :
+       module.getOps<mlir::tf_saved_model::GlobalTensorOp>()) {
+    mlir::ElementsAttr elements_attr = global_tensor_op.getValueAttr();
+    uint64_t tensor_size =
+        mlir::TFL::GetSizeInBytes(global_tensor_op.getType());
+    unique_tensors.insert(elements_attr);
+    module_size_estimate += tensor_size;
+  }
+
+  module.walk([&](Operation* op) {
+    mlir::ElementsAttr attr;
+    if (mlir::detail::constant_op_binder<mlir::ElementsAttr>(&attr).match(op)) {
+      // If the tensor hasn't been seen before
+      if (!unique_tensors.contains(attr)) {
+        uint64_t tensor_size =
+            mlir::TFL::GetSizeInBytes(op->getResult(0).getType());
+        unique_tensors.insert(attr);  // Store the size in the map
+        module_size_estimate += tensor_size;
+      }
+    }
+  });
+  return module_size_estimate;
+}
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_MLIR_MODULE_UTILS_H_

tensorflow/compiler/mlir/lite/utils/nms_utils.cc

@@ -0,0 +1,217 @@
+/* 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/utils/nms_utils.h"
+
+#include <cstddef>
+#include <string>
+
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/utils/utils.h"
+
+namespace mlir {
+namespace TFL {
+
+namespace {
+
+// TODO(b/162842801): Consolidate all util definitions of kTFImplements.
+constexpr char kTFImplements[] = "tf._implements";
+constexpr char kCustomSSDPostprocessing[] = "TFLite_Detection_PostProcess";
+constexpr char kTfNMSPadded[] = "non_max_suppression_padded_v2";
+
+inline ConstBytesAttr CustomOption(OpBuilder* builder,
+                                   const std::string& content) {
+  return ConstBytesAttr::get(builder->getContext(),
+                             StringRef(content.data(), content.size()));
+}
+
+}  // namespace
+
+void ConvertNMSPaddedFunc::RewriteFunc() {
+  func_->setAttr(kTFImplements,
+                 StringAttr::get(func_.getContext(), kTfNMSPadded));
+  Value boxes = func_.getArgument(0);
+  Value scores = func_.getArgument(1);
+  Value max_output_size = func_.getArgument(2);
+  Value iou_threshold = func_.getArgument(3);
+  Value score_threshold = func_.getArgument(4);
+  auto output_type0 = func_.getFunctionType().getResult(0);
+  auto output_type1 = func_.getFunctionType().getResult(1);
+
+  OpBuilder builder(func_.getBody());
+  auto op = mlir::TFL::NonMaxSuppressionV4Op::create(
+      builder, func_.getLoc(), output_type0, output_type1, boxes, scores,
+      max_output_size, iou_threshold, score_threshold);
+
+  mlir::func::ReturnOp::create(builder, func_.getLoc(), op.getResults());
+}
+
+LogicalResult ConvertNMSPaddedFunc::VerifySignature() {
+  // Verify high-level function signature.
+  // Relevant argument characteristics are checked by the TFL op definition.
+  if (func_.getNumArguments() < 5) {
+    return func_.emitWarning()
+           << "Invalid number of arguments to "
+              "non_max_suppression_padded_v2 (need at least 5): "
+           << func_.getNumArguments();
+  }
+  if (func_.getFunctionType().getNumResults() != 2) {
+    return func_.emitWarning() << "Invalid number of results from "
+                                  "non_max_suppression_padded_v2 (need 2): "
+                               << func_.getFunctionType().getNumResults();
+  }
+  // The TFLite fused op does not support batching yet.
+  // TODO(b/158709815): Add support for batches with padded NMS.
+  auto boxes_type =
+      mlir::dyn_cast<RankedTensorType>(func_.getFunctionType().getInput(0));
+  if (boxes_type == nullptr || !boxes_type.hasRank() ||
+      boxes_type.getRank() != 2) {
+    return func_.emitWarning() << "TFLite does not support batched input for "
+                                  "non_max_suppression_padded";
+  }
+  return success();
+}
+
+LogicalResult ConvertSSDPostProcessFunc::RewriteFunc() {
+  func_.eraseBody();
+  func_.addEntryBlock();
+  func_->setAttr(kTFImplements,
+                 StringAttr::get(func_.getContext(), kCustomSSDPostprocessing));
+
+  OpBuilder builder(func_.getBody());
+  std::string custom_option_buffer;
+  if (failed(CreateNMSCustomOptions(func_, attr_.getAttrs(),
+                                    custom_option_buffer))) {
+    return failure();
+  }
+  auto op = CustomOp::create(builder, func_.getLoc(),
+                             func_.getFunctionType().getResults(),
+                             func_.getArguments(), kCustomSSDPostprocessing,
+                             CustomOption(&builder, custom_option_buffer));
+  func::ReturnOp::create(builder, func_.getLoc(), op.getResults());
+
+  return success();
+}
+
+LogicalResult ConvertSSDPostProcessFunc::CreateNMSCustomOptions(
+    func::FuncOp func, DictionaryAttr attrs,
+    std::string& custom_option_buffer) {
+  flexbuffers::Builder fbb;
+  size_t start_map = fbb.StartMap();
+
+  if (failed(AddIntAttr(func, attrs, "max_detections", &fbb)) ||
+      failed(AddIntAttr(func, attrs, "max_classes_per_detection", &fbb)) ||
+      failed(AddIntAttr(func, attrs, "num_classes", &fbb)) ||
+      failed(AddFloatAttr(func, attrs, "nms_score_threshold", &fbb)) ||
+      failed(AddFloatAttr(func, attrs, "nms_iou_threshold", &fbb)) ||
+      failed(AddFloatAttr(func, attrs, "y_scale", &fbb)) ||
+      failed(AddFloatAttr(func, attrs, "x_scale", &fbb)) ||
+      failed(AddFloatAttr(func, attrs, "h_scale", &fbb)) ||
+      failed(AddFloatAttr(func, attrs, "w_scale", &fbb)))
+    return failure();
+  auto use_regular_nms =
+      mlir::dyn_cast_or_null<BoolAttr>(attrs.get("use_regular_nms"));
+  if (!use_regular_nms) {
+    return func.emitError()
+           << "use_regular_nms attribute is not set or not a bool";
+  }
+  fbb.Int("use_regular_nms", use_regular_nms.getValue());
+
+  fbb.EndMap(start_map);
+  fbb.Finish();
+  custom_option_buffer.assign(fbb.GetBuffer().begin(), fbb.GetBuffer().end());
+  return success();
+}
+
+LogicalResult ConvertSSDPostProcessFunc::AddIntAttr(
+    func::FuncOp func, DictionaryAttr attrs, const std::string& attribute,
+    flexbuffers::Builder* builder) {
+  auto int_attr = mlir::dyn_cast_or_null<IntegerAttr>(attrs.get(attribute));
+  if (!int_attr) {
+    return func.emitError()
+           << attribute.c_str() << " attribute is not set or not an integer";
+  }
+  builder->Int(attribute.c_str(), int_attr.getInt());
+  return success();
+}
+
+LogicalResult ConvertSSDPostProcessFunc::AddFloatAttr(
+    func::FuncOp func, DictionaryAttr attrs, const std::string& attribute,
+    flexbuffers::Builder* builder) {
+  auto float_attr = mlir::dyn_cast_or_null<FloatAttr>(attrs.get(attribute));
+  if (!float_attr) {
+    return func.emitError()
+           << attribute.c_str() << " attribute is not set or not a float";
+  }
+  builder->Float(attribute.c_str(), float_attr.getValue().convertToFloat());
+  return success();
+}
+
+LogicalResult ConvertSSDPostProcessFunc::HasIntAttr(
+    func::FuncOp func, DictionaryAttr attrs, const std::string& attribute) {
+  auto int_attr = mlir::dyn_cast_or_null<IntegerAttr>(attrs.get(attribute));
+  if (!int_attr) {
+    return func.emitWarning()
+           << attribute.c_str() << " attribute is not set or not an integer";
+  }
+  return success();
+}
+
+LogicalResult ConvertSSDPostProcessFunc::HasFloatAttr(
+    func::FuncOp func, DictionaryAttr attrs, const std::string& attribute) {
+  auto float_attr = mlir::dyn_cast_or_null<FloatAttr>(attrs.get(attribute));
+  if (!float_attr) {
+    return func.emitWarning()
+           << attribute.c_str() << " attribute is not set or not a float";
+  }
+  return success();
+}
+
+LogicalResult ConvertSSDPostProcessFunc::VerifySignature() {
+  // Verify high-level function signature.
+  if (func_.getNumArguments() != 3) {
+    return func_.emitWarning()
+           << "Invalid number of arguments to " << kCustomSSDPostprocessing
+           << ": " << func_.getNumArguments();
+  }
+  if (func_.getFunctionType().getNumResults() != 4) {
+    return func_.emitWarning()
+           << "Invalid number of results from " << kCustomSSDPostprocessing
+           << ": " << func_.getFunctionType().getNumResults();
+  }
+
+  auto attrs = attr_.getAttrs();
+  if (failed(HasIntAttr(func_, attrs, "max_detections")) ||
+      failed(HasIntAttr(func_, attrs, "max_classes_per_detection")) ||
+      failed(HasIntAttr(func_, attrs, "num_classes")) ||
+      failed(HasFloatAttr(func_, attrs, "nms_score_threshold")) ||
+      failed(HasFloatAttr(func_, attrs, "nms_iou_threshold")) ||
+      failed(HasFloatAttr(func_, attrs, "y_scale")) ||
+      failed(HasFloatAttr(func_, attrs, "x_scale")) ||
+      failed(HasFloatAttr(func_, attrs, "h_scale")) ||
+      failed(HasFloatAttr(func_, attrs, "w_scale"))) {
+    return failure();
+  }
+  return success();
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/nms_utils.h

@@ -0,0 +1,84 @@
+/* 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.
+==============================================================================*/
+
+// This header file defines common utils used by TFLite transformation
+// passes to work with NMS ops in TFLite.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_NMS_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_NMS_UTILS_H_
+
+#include <string>
+
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+
+namespace mlir {
+namespace TFL {
+
+// Abstracts the conversion of the padded NMS composite function.
+class ConvertNMSPaddedFunc {
+ public:
+  explicit ConvertNMSPaddedFunc(func::FuncOp func) : func_(func) {}
+
+  void RewriteFunc();
+
+  LogicalResult VerifySignature();
+
+ private:
+  func::FuncOp func_;
+};
+
+// Abstracts the conversion of the SSD post-processing composite function to
+// TFLite.
+class ConvertSSDPostProcessFunc {
+ public:
+  explicit ConvertSSDPostProcessFunc(func::FuncOp func, mlir::TF::FuncAttr attr)
+      : func_(func), attr_(attr) {}
+
+  LogicalResult RewriteFunc();
+
+  LogicalResult VerifySignature();
+
+ private:
+  LogicalResult CreateNMSCustomOptions(func::FuncOp func, DictionaryAttr attrs,
+                                       std::string& custom_option_buffer);
+
+  LogicalResult AddIntAttr(func::FuncOp func, DictionaryAttr attrs,
+                           const std::string& attribute,
+                           flexbuffers::Builder* builder);
+
+  LogicalResult AddFloatAttr(func::FuncOp func, DictionaryAttr attrs,
+                             const std::string& attribute,
+                             flexbuffers::Builder* builder);
+
+  LogicalResult HasIntAttr(func::FuncOp func, DictionaryAttr attrs,
+                           const std::string& attribute);
+
+  LogicalResult HasFloatAttr(func::FuncOp func, DictionaryAttr attrs,
+                             const std::string& attribute);
+
+  func::FuncOp func_;
+  mlir::TF::FuncAttr attr_;
+};
+
+}  // end namespace TFL
+}  // end namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_TFTEXT_UTILS_H_

tensorflow/compiler/mlir/lite/utils/perception_ops_utils.cc

@@ -0,0 +1,256 @@
+/* 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/utils/perception_ops_utils.h"
+
+#include <string>
+
+#include "flatbuffers/base.h"  // from @flatbuffers
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/core/c/builtin_op_data.h"
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+
+namespace mlir {
+namespace TFL {
+
+namespace {
+
+constexpr char kTFImplements[] = "tf._implements";
+constexpr char kMaxUnpooling[] = "MaxUnpooling2D";
+constexpr char kImageWarping[] = "DenseImageWarp";
+
+inline ConstBytesAttr CustomOption(OpBuilder* builder,
+                                   const std::string& content) {
+  return ConstBytesAttr::get(builder->getContext(),
+                             StringRef(content.data(), content.size()));
+}
+
+inline LogicalResult HasIntegerArrayWithSize(func::FuncOp* func,
+                                             const DictionaryAttr& attrs,
+                                             const std::string& attr_name,
+                                             int N) {
+  ArrayAttr array_attr =
+      mlir::dyn_cast_or_null<ArrayAttr>(attrs.get(attr_name));
+  if (array_attr == nullptr || array_attr.size() != N) {
+    return func->emitWarning()
+           << "'" << attr_name << "' attribute for " << kMaxUnpooling
+           << " must be set and has size of " << N;
+  }
+  for (Attribute integer_attr : array_attr.getValue()) {
+    IntegerAttr value = mlir::dyn_cast<IntegerAttr>(integer_attr);
+    if (!value) {
+      return func->emitWarning()
+             << "'" << attr_name << "' attribute for " << kMaxUnpooling
+             << " does not contain integer values";
+    }
+  }
+  return success();
+}
+
+inline LogicalResult GetIntegerArraySafe(
+    func::FuncOp* func, const DictionaryAttr& attrs,
+    const std::string& attr_name, llvm::SmallVectorImpl<int32_t>* results,
+    int N) {
+  ArrayAttr array_attr =
+      mlir::dyn_cast_or_null<ArrayAttr>(attrs.get(attr_name));
+  if (array_attr == nullptr || array_attr.size() != N) {
+    return func->emitError()
+           << "'" << attr_name << "' attribute for " << kMaxUnpooling
+           << " must be set and has size of " << N;
+  }
+  results->reserve(N);
+
+  for (Attribute integer_attr : array_attr.getValue()) {
+    IntegerAttr value = mlir::dyn_cast<IntegerAttr>(integer_attr);
+    if (!value) {
+      return func->emitError()
+             << "'" << attr_name << "' attribute for " << kMaxUnpooling
+             << " does not contain integer values";
+    }
+    results->push_back(value.getInt());
+  }
+  return success();
+}
+
+}  // namespace
+
+LogicalResult ConvertMaxUnpoolingFunc::RewriteFunc() {
+  func_.eraseBody();
+  func_.addEntryBlock();
+  func_->setAttr(kTFImplements,
+                 StringAttr::get(func_.getContext(), kMaxUnpooling));
+
+  OpBuilder builder(func_.getBody());
+  std::string custom_option_buffer;
+  if (failed(CreateCustomOptions(custom_option_buffer))) {
+    return failure();
+  }
+  auto op = CustomOp::create(builder, func_.getLoc(),
+                             func_.getFunctionType().getResults(),
+                             func_.getArguments(), kMaxUnpooling,
+                             CustomOption(&builder, custom_option_buffer));
+  func::ReturnOp::create(builder, func_.getLoc(), op.getResults());
+
+  return success();
+}
+
+LogicalResult ConvertMaxUnpoolingFunc::VerifySignature() {
+  // Verify high-level function signature.
+  if (func_.getNumArguments() != 2) {
+    return func_.emitWarning()
+           << "Invalid number of arguments to " << kMaxUnpooling << ": "
+           << func_.getNumArguments();
+  }
+  if (func_.getFunctionType().getNumResults() != 1) {
+    return func_.emitWarning()
+           << "Invalid number of results from " << kMaxUnpooling << ": "
+           << func_.getFunctionType().getNumResults();
+  }
+
+  auto attrs = attr_.getAttrs();
+
+  if (failed(HasIntegerArrayWithSize(&func_, attrs, "pool_size", 2))) {
+    return failure();
+  }
+
+  if (failed(HasIntegerArrayWithSize(&func_, attrs, "strides", 2))) {
+    return failure();
+  }
+
+  // Retrieves padding.
+  auto padding = mlir::dyn_cast_or_null<StringAttr>(attrs.get("padding"));
+  if (!padding) {
+    return func_.emitWarning() << "'padding' attribute for " << kMaxUnpooling
+                               << " is not set or not a string";
+  }
+  if (padding.getValue() != "VALID" && padding.getValue() != "SAME") {
+    return func_.emitWarning()
+           << "Padding for " << kMaxUnpooling << " must be 'SAME' or 'VALID'";
+  }
+  return success();
+}
+
+LogicalResult ConvertMaxUnpoolingFunc::CreateCustomOptions(
+    std::string& custom_option_buffer) {
+  auto attrs = attr_.getAttrs();
+  TfLitePoolParams pool_params;
+
+  llvm::SmallVector<int32_t, 2> pool_size;
+  if (failed(GetIntegerArraySafe(&func_, attrs, "pool_size", &pool_size, 2))) {
+    return failure();
+  }
+  pool_params.filter_height = pool_size[0];
+  pool_params.filter_width = pool_size[1];
+
+  // Retrieve strides.
+  llvm::SmallVector<int32_t, 2> strides;
+  if (failed(GetIntegerArraySafe(&func_, attrs, "strides", &strides, 2))) {
+    return failure();
+  }
+  pool_params.stride_height = strides[0];
+  pool_params.stride_width = strides[1];
+
+  // Retrieves padding.
+  auto padding = mlir::dyn_cast_or_null<StringAttr>(attrs.get("padding"));
+  if (!padding) {
+    return func_.emitError() << "'padding' attribute for " << kMaxUnpooling
+                             << " is not set or not a string";
+  }
+  if (padding.getValue() == "VALID") {
+    pool_params.padding = kTfLitePaddingValid;
+  } else if (padding.getValue() == "SAME") {
+    pool_params.padding = kTfLitePaddingSame;
+  } else {
+    return func_.emitError()
+           << "Padding for " << kMaxUnpooling << " must be 'SAME' or 'VALID'";
+  }
+
+  pool_params.activation = kTfLiteActNone;
+  pool_params.computed.padding = TfLitePaddingValues{0, 0, 0, 0};
+
+#if FLATBUFFERS_LITTLEENDIAN == 0
+  int32_t* p = reinterpret_cast<int32_t*>(&pool_params);
+  for (size_t i = 0; i < sizeof(TfLitePoolParams) / 4; i++, p++)
+    *p = flatbuffers::EndianSwap(*p);
+#endif
+
+  custom_option_buffer.assign(reinterpret_cast<char*>(&pool_params),
+                              sizeof(TfLitePoolParams));
+  return success();
+}
+
+LogicalResult ConvertDenseImageWarpFunc::RewriteFunc() {
+  func_.eraseBody();
+  func_.addEntryBlock();
+  func_->setAttr(kTFImplements,
+                 StringAttr::get(func_.getContext(), kImageWarping));
+
+  OpBuilder builder(func_.getBody());
+  auto op = CustomOp::create(builder, func_.getLoc(),
+                             func_.getFunctionType().getResults(),
+                             func_.getArguments(), kImageWarping,
+                             CustomOption(&builder, /*content=*/""));
+  func::ReturnOp::create(builder, func_.getLoc(), op.getResults());
+
+  return success();
+}
+
+LogicalResult ConvertDenseImageWarpFunc::VerifySignature() {
+  // Verify high-level function signature.
+  if (func_.getNumArguments() != 2) {
+    return func_.emitWarning()
+           << "Invalid number of arguments to " << kImageWarping << ": "
+           << func_.getNumArguments();
+  }
+  if (func_.getFunctionType().getNumResults() != 1) {
+    return func_.emitWarning()
+           << "Invalid number of results from " << kImageWarping << ": "
+           << func_.getFunctionType().getNumResults();
+  }
+
+  // Check types and shapes.
+  auto image_type = mlir::dyn_cast_or_null<RankedTensorType>(
+      func_.getFunctionType().getInput(0));
+  if (!image_type || !image_type.getElementType().isF32() ||
+      image_type.getRank() != 4) {
+    return func_.emitWarning() << "Image should be a 4D float tensor";
+  }
+
+  auto flow_type = mlir::dyn_cast_or_null<RankedTensorType>(
+      func_.getFunctionType().getInput(1));
+  if (!flow_type || !flow_type.getElementType().isF32() ||
+      flow_type.getRank() != 4) {
+    return func_.emitWarning() << "Flow should be a 4D float tensor";
+  }
+
+  auto output_type = mlir::dyn_cast_or_null<RankedTensorType>(
+      func_.getFunctionType().getResult(0));
+  if (!output_type || !output_type.getElementType().isF32() ||
+      output_type.getRank() != 4) {
+    return func_.emitWarning() << "Output should be a 4D float tensor";
+  }
+
+  return success();
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/perception_ops_utils.h

@@ -0,0 +1,63 @@
+/* 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_UTILS_PERCEPTION_OPS_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_PERCEPTION_OPS_UTILS_H_
+
+#include <string>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+
+namespace mlir {
+namespace TFL {
+
+// Fuse MaxUnpooling2D ops annotated by tf.function to a TFLite custom op.
+class ConvertMaxUnpoolingFunc {
+ public:
+  explicit ConvertMaxUnpoolingFunc(func::FuncOp func, mlir::TF::FuncAttr attr)
+      : func_(func), attr_(attr) {}
+
+  LogicalResult RewriteFunc();
+
+  LogicalResult VerifySignature();
+
+ private:
+  LogicalResult CreateCustomOptions(std::string& custom_option_buffer);
+
+  func::FuncOp func_;
+  mlir::TF::FuncAttr attr_;
+};
+
+// Fuse DenseImageWarp ops annotated by tf.function to a TFLite custom op.
+class ConvertDenseImageWarpFunc {
+ public:
+  explicit ConvertDenseImageWarpFunc(func::FuncOp func) : func_(func) {}
+
+  LogicalResult RewriteFunc();
+
+  LogicalResult VerifySignature();
+
+ private:
+  func::FuncOp func_;
+};
+
+}  // end namespace TFL
+}  // end namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_PERCEPTION_OPS_UTILS_H_

tensorflow/compiler/mlir/lite/utils/perception_ops_utils_test.cc

@@ -0,0 +1,199 @@
+/* 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/utils/perception_ops_utils.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir {
+namespace TFL {
+namespace {
+
+template <int NInput, int NOutput>
+func::FuncOp createMaxUnpoolingFunc(
+    mlir::Builder* builder, const SmallVector<mlir::Type, NInput>& input_types,
+    const SmallVector<mlir::Type, NOutput>& output_types) {
+  auto func_type = builder->getFunctionType(input_types, output_types);
+  auto func = func::FuncOp::create(
+      mlir::NameLoc::get(builder->getStringAttr("fused_func")), "fused_func",
+      func_type, {});
+
+  func.addEntryBlock();
+  mlir::StringAttr attr_value = builder->getStringAttr("MaxUnpooling2D");
+  func->setAttr("tf._implements", attr_value);
+  return func;
+}
+
+func::FuncOp createMaxUnpoolingFunc(
+    mlir::Builder* builder, const SmallVector<int64_t, 4>& input_shape,
+    const SmallVector<int64_t, 4>& output_shape) {
+  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());
+  auto indices_type = RankedTensorType::get(input_shape, builder->getI64Type());
+  auto output_type = RankedTensorType::get(output_shape, builder->getF32Type());
+  SmallVector<mlir::Type, 2> input_types{input_type, indices_type};
+  SmallVector<mlir::Type, 1> output_types{output_type};
+  return createMaxUnpoolingFunc<2, 1>(builder, input_types, output_types);
+}
+
+template <int N>
+ArrayAttr createInt32Array(mlir::Builder* builder, mlir::MLIRContext* context,
+                           const SmallVector<int32_t, N>& values) {
+  SmallVector<Attribute, N> ret;
+  for (int32_t value : values) {
+    ret.push_back(builder->getI32IntegerAttr(value));
+  }
+  return ArrayAttr::get(context, ret);
+}
+
+template <int N>
+ArrayAttr createInt64Array(mlir::Builder* builder, mlir::MLIRContext* context,
+                           const SmallVector<int64_t, N>& values) {
+  SmallVector<Attribute, N> ret;
+  for (int64_t value : values) {
+    ret.push_back(builder->getI64IntegerAttr(value));
+  }
+  return ArrayAttr::get(context, ret);
+}
+
+mlir::TF::FuncAttr createMaxUnpoolingAttr(mlir::MLIRContext* context,
+                                          const std::string& padding,
+                                          const ArrayAttr& pool_size,
+                                          const ArrayAttr& strides) {
+  SmallVector<::mlir::NamedAttribute, 3> fields;
+
+  auto padding_id = ::mlir::StringAttr::get(context, "padding");
+  fields.emplace_back(padding_id, StringAttr::get(context, padding));
+
+  auto pool_size_id = ::mlir::StringAttr::get(context, "pool_size");
+  fields.emplace_back(pool_size_id, pool_size);
+
+  auto strides_id = ::mlir::StringAttr::get(context, "strides");
+  fields.emplace_back(strides_id, strides);
+
+  DictionaryAttr dict = DictionaryAttr::get(context, fields);
+  return TF::FuncAttr::get(context, "MaxUnpooling2D", dict);
+}
+
+}  // namespace
+
+class PerceptionUtilsTest : public ::testing::Test {
+ protected:
+  PerceptionUtilsTest() = default;
+
+  void SetUp() override {
+    context_ = std::make_unique<mlir::MLIRContext>();
+    context_->loadDialect<mlir::arith::ArithDialect, mlir::func::FuncDialect,
+                          mlir::TF::TensorFlowDialect, TensorFlowLiteDialect>();
+    builder_ = std::make_unique<mlir::Builder>(context_.get());
+
+    fused_max_unpooling_func_ =
+        createMaxUnpoolingFunc(builder_.get(), {2, 4, 4, 2}, {2, 2, 2, 2});
+
+    func_attr_ = createMaxUnpoolingAttr(
+        context_.get(), "SAME",
+        createInt32Array<2>(builder_.get(), context_.get(), {2, 2}),
+        createInt32Array<2>(builder_.get(), context_.get(), {2, 2}));
+  }
+
+  void TearDown() override {
+    fused_max_unpooling_func_.erase();
+    builder_.reset();
+  }
+
+  func::FuncOp fused_max_unpooling_func_;
+  mlir::TF::FuncAttr func_attr_;
+  std::unique_ptr<mlir::MLIRContext> context_;
+  std::unique_ptr<mlir::Builder> builder_;
+};
+
+TEST_F(PerceptionUtilsTest, VerifySignatureValid) {
+  mlir::TFL::ConvertMaxUnpoolingFunc convert(fused_max_unpooling_func_,
+                                             func_attr_);
+
+  EXPECT_FALSE(failed(convert.VerifySignature()));
+}
+
+TEST_F(PerceptionUtilsTest, VerifySignatureInvalid) {
+  auto input_type = RankedTensorType::get({1, 2, 2, 1}, builder_->getF32Type());
+  auto output_type =
+      RankedTensorType::get({1, 2, 1, 1}, builder_->getF32Type());
+  SmallVector<mlir::Type, 1> input_types{input_type};
+  SmallVector<mlir::Type, 1> output_types{output_type};
+
+  auto max_unpooling_func =
+      createMaxUnpoolingFunc<1, 1>(builder_.get(), input_types, output_types);
+  mlir::TFL::ConvertMaxUnpoolingFunc convert(max_unpooling_func, func_attr_);
+
+  EXPECT_TRUE(failed(convert.VerifySignature()));
+  max_unpooling_func->erase();
+}
+
+TEST_F(PerceptionUtilsTest, RewriteValid) {
+  mlir::TFL::ConvertMaxUnpoolingFunc convert(fused_max_unpooling_func_,
+                                             func_attr_);
+
+  EXPECT_FALSE(failed(convert.RewriteFunc()));
+}
+
+TEST_F(PerceptionUtilsTest, RewriteWrongPadding) {
+  auto func_attr = createMaxUnpoolingAttr(
+      context_.get(), "INVALID",
+      createInt32Array<2>(builder_.get(), context_.get(), {2, 2}),
+      createInt32Array<2>(builder_.get(), context_.get(), {2, 2}));
+  mlir::TFL::ConvertMaxUnpoolingFunc convert(fused_max_unpooling_func_,
+                                             func_attr);
+
+  EXPECT_TRUE(failed(convert.RewriteFunc()));
+}
+
+TEST_F(PerceptionUtilsTest, RewriteWrongFilter) {
+  auto func_attr = createMaxUnpoolingAttr(
+      context_.get(), "VALID",
+      createInt32Array<2>(builder_.get(), context_.get(), {2, 2, 2}),
+      createInt32Array<2>(builder_.get(), context_.get(), {2, 2}));
+  mlir::TFL::ConvertMaxUnpoolingFunc convert(fused_max_unpooling_func_,
+                                             func_attr);
+
+  EXPECT_TRUE(failed(convert.RewriteFunc()));
+}
+
+TEST_F(PerceptionUtilsTest, RewriteWrongStrides) {
+  auto func_attr = createMaxUnpoolingAttr(
+      context_.get(), "VALID",
+      createInt32Array<2>(builder_.get(), context_.get(), {2, 2}),
+      createInt32Array<2>(builder_.get(), context_.get(), {2, 2, 0}));
+  mlir::TFL::ConvertMaxUnpoolingFunc convert(fused_max_unpooling_func_,
+                                             func_attr);
+
+  EXPECT_TRUE(failed(convert.RewriteFunc()));
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/region_isolation.cc

@@ -0,0 +1,70 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/region_isolation.h"
+
+#define DEBUG_TYPE "tfl_isolate_regions"
+
+#include <optional>
+
+#include "absl/strings/str_format.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/RegionUtils.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+std::optional<llvm::SetVector<Value>> IsolateRegions(Operation* op_with_regions,
+                                                     OpBuilder& b) {
+  LLVM_DEBUG(
+      llvm::dbgs() << absl::StrFormat("Isolating Op with %u regions...\n",
+                                      op_with_regions->getNumRegions()));
+  LLVM_DEBUG(op_with_regions->print(llvm::dbgs()));
+  LLVM_DEBUG(llvm::dbgs() << "\n");
+
+  if (op_with_regions->getNumRegions() == 0) {
+    return {};
+  }
+
+  llvm::SetVector<Value> shared_signature;
+  getUsedValuesDefinedAbove(op_with_regions->getRegions(), shared_signature);
+
+  for (auto& reg : op_with_regions->getRegions()) {
+    if (!reg.hasOneBlock()) {
+      LLVM_DEBUG(
+          llvm::dbgs()
+          << "Region isolation only supports regions with a single block\n");
+      return {};
+    }
+    auto& block = reg.getBlocks().front();
+    if (block.getNumArguments() != 0) {
+      LLVM_DEBUG(llvm::dbgs() << "Region isolation reguires empty blargs\n");
+    }
+    for (auto val : shared_signature) {
+      auto blarg = block.addArgument(val.getType(), b.getUnknownLoc());
+      replaceAllUsesInRegionWith(val, blarg, reg);
+    }
+  }
+
+  return shared_signature;
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/region_isolation.h

@@ -0,0 +1,41 @@
+/* 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_UTILS_REGION_ISOLATION_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_REGION_ISOLATION_H_
+
+#include <optional>
+
+#include "llvm/ADT/SetVector.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+// Isolates op's contained regions. Replaces all references to values defined
+// above these (single block) regions with a block argument. The union of all
+// values referenced this way is returned. Each region will have an identical
+// signature, which is the types of the returned vector in the same order.
+// NOTE: Critically, llvm::SetVector iterates deterministically in order of
+// insertion.
+std::optional<llvm::SetVector<Value>> IsolateRegions(Operation* op_with_regions,
+                                                     OpBuilder& b);
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_REGION_ISOLATION_H_

tensorflow/compiler/mlir/lite/utils/region_isolation_test.cc

@@ -0,0 +1,163 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/region_isolation.h"
+
+#include <cstdint>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinDialect.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+
+namespace mlir::TFL {
+namespace {
+
+using ::testing::UnorderedElementsAreArray;
+
+static constexpr int64_t kValShape[] = {2, 2};
+
+RankedTensorType TestValType(OpBuilder& b) {
+  return RankedTensorType::get(kValShape, b.getI32Type());
+}
+
+DenseElementsAttr TestValData(int fill, OpBuilder& b) {
+  return b
+      .getI32TensorAttr(
+          llvm::SmallVector<int32_t, 4>(TestValType(b).getNumElements(), fill))
+      .reshape(TestValType(b));
+}
+
+TEST(RegionIsolationTest, CaseOp) {
+  //
+  // Make Test Model
+  //
+
+  //   %0 = const : tensor<2x2xi32>
+  //   %1 = const : tensor<2x2xi32>
+  //   %2 = const : tensor<2x2xi32>
+  //   %3 = const : tensor<i32>
+  //   %4 = "stablehlo.case"(%3) ({
+  //     %7 = add(%0, %0)
+  //   }, {
+  //     %6 = add(%0, %1)
+  //   }, {
+  //     %5 = add(%0, %2)
+  //   }
+
+  // Setup context, empty module, etc.
+
+  MLIRContext ctx;
+
+  {
+    DialectRegistry reg;
+    reg.insert<arith::ArithDialect, BuiltinDialect,
+               stablehlo::StablehloDialect>();
+    ctx.appendDialectRegistry(reg);
+    ctx.loadAllAvailableDialects();
+  }
+
+  OpBuilder b(&ctx);
+
+  OwningOpRef<ModuleOp> root(ModuleOp::create(b, b.getUnknownLoc()));
+
+  {
+    auto& block = root->getBodyRegion().front();
+    b.setInsertionPointToStart(&block);
+  }
+
+  // Add values to be referenced within later regions.
+
+  auto root_val_1 =
+      arith::ConstantOp::create(b, b.getUnknownLoc(), TestValData(1, b));
+
+  auto root_val_2 =
+      arith::ConstantOp::create(b, b.getUnknownLoc(), TestValData(2, b));
+
+  auto root_val_3 =
+      arith::ConstantOp::create(b, b.getUnknownLoc(), TestValData(3, b));
+
+  // Iteration convenience.
+  llvm::SmallVector<arith::ConstantOp> root_vals(
+      {root_val_1, root_val_2, root_val_3});
+
+  // Make a regioned op with computations that reference defined above vals.
+
+  auto root_ind = arith::ConstantOp::create(
+      b, b.getUnknownLoc(),
+      DenseIntElementsAttr::get(RankedTensorType::get({}, b.getI32Type()),
+                                {0}));
+
+  auto regioned_op = stablehlo::CaseOp::create(
+      b, root_val_1.getLoc(), llvm::SmallVector<Type>({TestValType(b)}),
+      root_ind,
+      /*branch_count=*/3);
+
+  // Populate each branch with a computation that references the
+  // above values.
+
+  for (auto [reg, val] : llvm::zip(regioned_op.getBranches(), root_vals)) {
+    auto& block = reg.emplaceBlock();
+    b.setInsertionPointToStart(&block);
+    auto res = stablehlo::AddOp::create(b, b.getUnknownLoc(), root_val_1, val);
+    llvm::SmallVector<Value> rets({res});
+    stablehlo::ReturnOp::create(b, b.getUnknownLoc(), rets);
+  }
+
+  //
+  // Isolate Regions in Place
+  //
+
+  {
+    auto result = IsolateRegions(regioned_op.getOperation(), b);
+    ASSERT_TRUE(result.has_value());
+
+    auto& vals_defined_above = result.value();
+    EXPECT_THAT(vals_defined_above,
+                UnorderedElementsAreArray(llvm::to_vector(llvm::map_range(
+                    root_vals, [](auto op) { return op.getResult(); }))));
+
+    auto& branch_front = regioned_op.getBranches().front();
+    ASSERT_EQ(branch_front.getArgumentTypes().size(),
+              vals_defined_above.size());
+
+    for (int i = 0; i < vals_defined_above.size(); ++i) {
+      EXPECT_EQ(branch_front.getArgumentTypes()[i],
+                vals_defined_above[i].getType());
+    }
+
+    for (auto& reg : regioned_op.getBranches()) {
+      EXPECT_EQ(reg.getArgumentTypes(), branch_front.getArgumentTypes());
+    }
+  }
+}
+
+}  // namespace
+
+}  // namespace mlir::TFL

tensorflow/compiler/mlir/lite/utils/shape_and_size_utils.cc

@@ -0,0 +1,96 @@
+/* Copyright 2025 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/utils/shape_and_size_utils.h"
+
+#include <cstddef>
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+int32_t ConvertToTfliteSize(int64_t size) {
+  return mlir::ShapedType::isDynamic(size) ? -1 : static_cast<int32_t>(size);
+}
+
+absl::StatusOr<int32_t> GetQuantDimensionAfterReshape(
+    const mlir::ArrayRef<int64_t> input_shape,
+    const mlir::ArrayRef<int64_t> output_shape, int32_t input_quant_dim) {
+  if (input_quant_dim >= input_shape.size()) {
+    return absl::InvalidArgumentError(
+        "Input quantization dimension is not found in the output shape.");
+  }
+
+  // We might be able to handle the folded case but not the split case because
+  // MLIR Quant only supports single axis quantization. So, leaving both the
+  // cases out-of-scope/unsupported for now.
+  int32_t input_dim = -1, output_dim = -1;
+  while (input_quant_dim > input_dim) {
+    ++input_dim;
+    ++output_dim;
+    if (input_shape[input_dim] > output_shape[output_dim]) {
+      // This means that the input dimension is being reduced/split. Multiply
+      // the output dims until it matches the input dim.
+      size_t split_output_shape_product = output_shape[output_dim];
+      while (split_output_shape_product < input_shape[input_dim]) {
+        split_output_shape_product *= output_shape[++output_dim];
+      }
+      // If the split output shape product doesn't match the input shape, we
+      // return an error. This is a safety check.
+      if (split_output_shape_product != input_shape[input_dim]) {
+        return absl::InvalidArgumentError(
+            "Input quantization dimension is not found in the output shape.");
+      }
+    } else if (input_shape[input_dim] < output_shape[output_dim]) {
+      // This means that the input dimension is being folded. We need to
+      // multiply the input dims until it matches the output dim.
+      size_t folded_input_shape_product = input_shape[input_dim];
+      while (folded_input_shape_product < output_shape[output_dim]) {
+        folded_input_shape_product *= input_shape[++input_dim];
+      }
+      // If the folded input shape product doesn't match the output shape, we
+      // return an error. This is a safety check.
+      if (folded_input_shape_product != output_shape[output_dim]) {
+        return absl::InvalidArgumentError(
+            "Input quantization dimension is not found in the output shape.");
+      }
+    }
+  }
+
+  // Safety check to ensure that the input quantization dimension was not folded
+  // or split. We should land exactly at input_quant_dim == input_dim, here.
+  if (input_quant_dim < input_dim) {
+    return absl::InvalidArgumentError(
+        "Input quantization dimension is not found in the output shape.");
+  }
+
+  // If the input quantization dimension is the same as the input dimension,
+  // but the input and output shapes are different, we return an error.
+  if (input_quant_dim == input_dim &&
+      input_shape[input_dim] != output_shape[output_dim]) {
+    return absl::InvalidArgumentError(
+        "Input quantization dimension is not found in the output shape.");
+  }
+
+  return output_dim;
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/shape_and_size_utils.h

@@ -0,0 +1,44 @@
+/* Copyright 2025 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_UTILS_SHAPE_AND_SIZE_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_SHAPE_AND_SIZE_UTILS_H_
+
+#include <cstdint>
+
+#include "mlir/Support/LLVM.h"
+#include "absl/status/statusor.h"
+
+namespace mlir {
+namespace TFL {
+
+// Converts a TF size (64-bit) to TFLite (32-bit) and properly converts TF's
+// value for dynamic size (`std::numeric_limits<int64_t>::min()`) to the
+// TFLite-specific value.
+int32_t ConvertToTfliteSize(int64_t size);
+
+// Returns the quantization dimension after a reshape operation.
+//
+// TFL Reshape Op can fold multiple dimensions into one, or split one
+// dimension into multiple dimensions. This function will return an error if the
+// input quantization dimension is part of the folded/split dimensions.
+absl::StatusOr<int32_t> GetQuantDimensionAfterReshape(
+    mlir::ArrayRef<int64_t> input_shape, mlir::ArrayRef<int64_t> output_shape,
+    int32_t input_quant_dim);
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_SHAPE_AND_SIZE_UTILS_H_

tensorflow/compiler/mlir/lite/utils/shape_and_size_utils_test.cc

@@ -0,0 +1,95 @@
+/* Copyright 2025 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/utils/shape_and_size_utils.h"
+
+#include <cstdint>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "tensorflow/core/platform/test.h"
+
+namespace mlir {
+namespace TFL {
+namespace {
+
+TEST(SizeAndShapeUtilTest, ConvertsSize) {
+  EXPECT_EQ(ConvertToTfliteSize(1), 1);
+  EXPECT_EQ(ConvertToTfliteSize(-1), -1);
+  EXPECT_EQ(ConvertToTfliteSize(mlir::ShapedType::kDynamic), -1);
+}
+
+TEST(SizeAndShapeUtilTest, GetQuantDimensionAfterReshape) {
+  // Test that the input quantization dimension is part of the folded/split
+  // dimensions.
+  absl::StatusOr<int32_t> result =
+      GetQuantDimensionAfterReshape({1, 2, 3, 4, 5}, {1, 24, 5}, 4);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 2);
+
+  result = GetQuantDimensionAfterReshape({1, 2, 3, 4, 5}, {1, 2, 12, 5}, 4);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 3);
+
+  result = GetQuantDimensionAfterReshape({1, 2, 3, 4, 5}, {2, 12, 5}, 4);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 2);
+
+  result = GetQuantDimensionAfterReshape({1, 2, 3, 5, 5}, {30, 5}, 4);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 1);
+
+  result = GetQuantDimensionAfterReshape({1, 2, 3, 4, 5}, {1, 6, 2, 2, 5}, 4);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 4);
+
+  result = GetQuantDimensionAfterReshape({1, 2, 3, 4, 5}, {6, 2, 2, 5}, 4);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 3);
+
+  result =
+      GetQuantDimensionAfterReshape({1, 2, 3, 4, 5}, {1, 2, 3, 2, 2, 5}, 4);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 5);
+
+  result = GetQuantDimensionAfterReshape({5, 1, 1, 8}, {1, 5, 1, 8}, 3);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 3);
+
+  result = GetQuantDimensionAfterReshape({5, 1, 1, 8}, {1, 1, 5, 8}, 3);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 3);
+
+  result = GetQuantDimensionAfterReshape({5, 1, 1, 8}, {1, 1, 5, 8}, 0);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 2);
+
+  result = GetQuantDimensionAfterReshape({5, 20}, {5, 1, 20}, 1);
+  ASSERT_EQ(result.status(), absl::OkStatus());
+  EXPECT_EQ(*result, 2);
+
+  EXPECT_NE(GetQuantDimensionAfterReshape({5, 10, 2}, {5, 5, 2, 2}, 1).status(),
+            absl::OkStatus());
+  EXPECT_NE(GetQuantDimensionAfterReshape({5, 2, 3, 5}, {5, 15, 5}, 1).status(),
+            absl::OkStatus());
+  // TODO(b/396164748): Fix the test case. Supports the case where the input
+  // dimensions are folded and split into new dimension values.
+  // EXPECT_NE(GetQuantDimensionAfterReshape({5, 2, 3}, {2, 5, 3}, 2).status(),
+  // absl::OkStatus());
+}
+}  // namespace
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/stateful_ops_utils.cc

@@ -0,0 +1,37 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/stateful_ops_utils.h"
+
+#include <vector>
+
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/lite/utils/utils.h"
+
+namespace mlir {
+namespace TFL {
+
+bool IsStatefulOp(Operation* op, std::vector<int>* stateful_operand_indices) {
+  if (auto stateful_op = dyn_cast_or_null<StatefulOpInterface>(op)) {
+    *stateful_operand_indices = stateful_op.GetStatefulOperands();
+    return true;
+  }
+
+  return false;
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/stateful_ops_utils.h

@@ -0,0 +1,34 @@
+/* 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_UTILS_STATEFUL_OPS_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_STATEFUL_OPS_UTILS_H_
+
+#include <vector>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+// Check if the given op has stateful operands and return their stateful
+// operand indices.
+bool IsStatefulOp(Operation* op, std::vector<int>* stateful_operand_indices);
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_STATEFUL_OPS_UTILS_H_

tensorflow/compiler/mlir/lite/utils/string_utils.cc

@@ -0,0 +1,85 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/string_utils.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <vector>
+
+namespace mlir::TFL {
+
+bool SimpleDynamicBuffer::AddString(const char* str, size_t len) {
+  // If `data_.size() + len` is greater than `SIZE_MAX` then the left hand side
+  // will overflow to something less than max_length_. After checking `len <=
+  // max_length_` we can use this subtraction to check for overflow.
+  if (len > max_length_ || data_.size() >= max_length_ - len) return false;
+  // It is undefined to call memcpy with `nullptr`, and that will be the case if
+  // `len` is 0, and `data_` is empty. An empty record is created though.
+  if (len == 0) {
+    offset_.push_back(offset_.back());
+    return true;
+  }
+  data_.resize(data_.size() + len);
+  memcpy(data_.data() + offset_.back(), str, len);
+  offset_.push_back(offset_.back() + len);
+  return true;
+}
+
+int SimpleDynamicBuffer::WriteToBuffer(char** buffer) {
+  // Allocate sufficient memory to tensor buffer.
+  int32_t num_strings = offset_.size() - 1;
+  // Total bytes include:
+  //   * size of content (data_.size)
+  //   * offset of each tensor (sizeof(int32_t) * num_strings)
+  //   * length of whole buffer (int32_t)
+  //   * num of strings (int32_t).
+  int32_t bytes = data_.size()                            // size of content
+                  + sizeof(int32_t) * (num_strings + 2);  // size of header
+  // Caller will take ownership of buffer.
+  *buffer = reinterpret_cast<char*>(malloc(bytes));
+
+  if (*buffer == nullptr) {
+    return -1;
+  }
+
+  // Set num of string
+  //
+  // NOTE: The string buffer is accessed here as if it's native endian (instead
+  // of small endian, as documented in the header). This will protentially break
+  // when TFLite is ported to big endian platforms.
+  // TODO(b/165919229): This code will need changing if/when we port to a
+  // big-endian platform.
+  memcpy(*buffer, &num_strings, sizeof(int32_t));
+
+  // Set offset of strings.
+  int32_t start = sizeof(int32_t) * (num_strings + 2);
+  for (size_t i = 0; i < offset_.size(); i++) {
+    // TODO(b/165919229): This code will need changing if/when we port to a
+    // big-endian platform.
+    int32_t offset = start + offset_[i];
+    memcpy(*buffer + sizeof(int32_t) * (i + 1), &offset, sizeof(int32_t));
+  }
+
+  // Copy data of strings.
+  if (!data_.empty()) {
+    memcpy(*buffer + start, data_.data(), data_.size());
+  }
+  return bytes;
+}
+
+}  // namespace mlir::TFL

tensorflow/compiler/mlir/lite/utils/string_utils.h

@@ -0,0 +1,110 @@
+/* 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.
+==============================================================================*/
+
+// Util methods to store a an ordered collection of strings in a char buffer.
+// The format of the char buffer is:
+//   [0, 3] 4 bytes: N, num of strings in the collection.
+//   [(i+1)*4, (i+1)*4+3] 4 bytes: offset of i-th string in little endian,
+//                                 for i from 0 to N-1.
+//   [(N+1)*4, (N+1)*4+3] 4 bytes: length of the whole char buffer.
+//   [offset(i), offset(i+1) - 1] : content of i-th string.
+//
+// A typical usage:
+//   SimpleDynamicBuffer buf;
+//   char* buffer;
+//   # Add string "AB", string is stored in dynamic buffer.
+//   buf.AddString("AB", 2);
+//   # Write content of SimpleDynamicBuffer to buffer in format described above.
+//   buf.WriteToBuffer(&buffer)
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_STRING_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_STRING_UTILS_H_
+
+#include <stdint.h>
+
+#include <cstddef>
+#include <limits>
+#include <vector>
+
+namespace mlir::TFL {
+
+constexpr uint64_t kDefaultMaxLength = std::numeric_limits<int>::max();
+
+class SimpleDynamicBuffer {
+ public:
+  explicit SimpleDynamicBuffer(size_t max_length = kDefaultMaxLength)
+      : offset_({0}), max_length_(max_length) {}
+
+  // Add string to dynamic buffer by resizing the buffer and copying the data.
+  bool AddString(const char* str, size_t len);
+
+  // Fill content into a buffer and returns the number of bytes stored.
+  // The function allocates space for the buffer but does NOT take ownership.
+  int WriteToBuffer(char** buffer);
+
+ protected:
+  // Data buffer to store contents of strings, not including headers.
+  std::vector<char> data_;
+  // Offset of the starting index of each string in data buffer.
+  std::vector<size_t> offset_;
+  // Max length in number of characters that we permit the total
+  // buffer containing the concatenation of all added strings to be.
+  // For historical reasons this is limited to 32bit length. At this files
+  // inception, sizes were represented using 32bit which forced an implicit cap
+  // on the size of the buffer. When this was refactored to use size_t (which
+  // could be 64bit) we enforce that the buffer remains at most 32bit length to
+  // avoid a change in behavior.
+  const size_t max_length_;
+};
+
+// Convenient structure to store string pointer and length. Note that
+// methods on SimpleDynamicBuffer enforce that the whole buffer (and by
+// extension every contained string) is of max length (2ul << 30) - 1. See
+// string_util.cc for more info.
+typedef struct {
+  const char* str;
+  size_t len;
+} StringRef;
+
+// Return num of strings in a String tensor.
+inline int GetStringCount(const void* raw_buffer) {
+  // The first integers in the raw buffer is the number of strings.
+  //
+  // NOTE: The string buffer is accessed here as if it's native endian (instead
+  // of small endian, as documented in the header). This will protentially break
+  // when TFLite is ported to big endian platforms.
+  // TODO(b/165919229): This code will need changing if/when we port to a
+  // big-endian platform.
+  return *static_cast<const int32_t*>(raw_buffer);
+}
+
+// Get String pointer and length of index-th string in tensor.
+// NOTE: This will not create a copy of string data.
+inline StringRef GetString(const void* raw_buffer, int string_index) {
+  // NOTE: The string buffer is accessed here as if it's native endian (instead
+  // of small endian, as documented in the header). This will protentially break
+  // when TFLite is ported to big endian platforms.
+  // TODO(b/165919229): This code will need changing if/when we port to a
+  // big-endian platform.
+  const int32_t* offset =
+      static_cast<const int32_t*>(raw_buffer) + (string_index + 1);
+  const size_t string_len = (*(offset + 1)) - (*offset);
+  return StringRef{static_cast<const char*>(raw_buffer) + (*offset),
+                   string_len};
+}
+
+}  // namespace mlir::TFL
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_STRING_UTILS_H_

tensorflow/compiler/mlir/lite/utils/tftext_utils.cc

@@ -0,0 +1,392 @@
+/* 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/utils/tftext_utils.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/ir/types/dialect.h"
+
+namespace mlir {
+namespace TFL {
+
+namespace {
+
+constexpr char kNgrams[] = "tftext:Ngrams";
+constexpr char kWhitespaceTokenizer[] = "tftext:WhitespaceTokenizer";
+constexpr char kCustomSgnnProjection[] = "tftext:custom:SgnnProjection";
+constexpr char kTFImplements[] = "tf._implements";
+
+using mlir::TF::FuncAttr;
+using mlir::TF::StringType;
+
+inline ConstBytesAttr CustomOption(OpBuilder* builder,
+                                   const std::string& content) {
+  return ConstBytesAttr::get(builder->getContext(),
+                             StringRef(content.data(), content.size()));
+}
+
+inline TensorType GetInputType(func::FuncOp func, int idx) {
+  return mlir::dyn_cast_or_null<TensorType>(
+      func.getFunctionType().getInput(idx));
+}
+
+inline TensorType GetResultType(func::FuncOp func, int idx) {
+  return mlir::dyn_cast_or_null<TensorType>(
+      func.getFunctionType().getResult(idx));
+}
+
+inline bool RankEquals(const TensorType& type, int rank) {
+  return type && type.hasRank() && type.getRank() == rank;
+}
+
+LogicalResult VerifyWhitespaceTokenizer(func::FuncOp func) {
+  // In the case of input tensor with 0 rank.
+  // Whitespace tokenizer generates 1 output:
+  // * String tensor for tokens.
+  //
+  // In the case of 1-D input tensor,
+  // Whitespace tokenizer generates 2 outputs to make up a ragged tensor:
+  // * 1st output is the value of ragged tensor;
+  // * 2nd output is the offset.
+  //
+  // In the case of batched input tesnor,
+  // Whitespace tokenizer has 3 outputs to make up a nested ragged tensor:
+  // * 1st output is the value of ragged tensor;
+  // * 2nd output is the inner offset;
+  // * 3rd output is the outer offset.
+  auto input_type = GetInputType(func, 0);
+  if (!input_type || !mlir::isa<StringType>(input_type.getElementType()) ||
+      !input_type.hasRank()) {
+    return func.emitError() << "Input should be a string tensor";
+  }
+
+  const std::vector<int> kValidNumOfOutput = {1, 2, 3};
+  if (input_type.getRank() >= kValidNumOfOutput.size()) {
+    return func.emitError()
+           << "Unrecognized input rank: " << input_type.getRank();
+  }
+  if (func.getNumResults() != kValidNumOfOutput[input_type.getRank()]) {
+    return func.emitError()
+           << "Expect " << kValidNumOfOutput[input_type.getRank()]
+           << "output(s) when input has rank " << input_type.getRank();
+  }
+
+  auto value_type = GetResultType(func, 0);
+  if (!RankEquals(value_type, 1) ||
+      !mlir::isa<StringType>(value_type.getElementType())) {
+    return func.emitError() << "1st output should be string tensor";
+  }
+  if (func.getNumResults() > 1) {
+    auto offset_type = GetResultType(func, 1);
+    if (!RankEquals(offset_type, 1) ||
+        !offset_type.getElementType().isInteger(64)) {
+      return func.emitError() << "2nd output should be int64 tensor";
+    }
+  }
+  if (func.getNumResults() > 2) {
+    auto offset_type = GetResultType(func, 2);
+    if (!RankEquals(offset_type, 1) ||
+        !offset_type.getElementType().isInteger(64)) {
+      return func.emitError() << "3rd output should be int64 tensor";
+    }
+  }
+
+  return success();
+}
+
+LogicalResult ConvertWhitespaceTokenizer(func::FuncOp func, llvm::StringRef api,
+                                         FuncAttr attr) {
+  func.eraseBody();
+  func.addEntryBlock();
+  func->setAttr(kTFImplements, attr);
+  OpBuilder builder(func.getBody());
+  std::string empty_option_buffer;
+  auto op = CustomOp::create(
+      builder, func.getLoc(), func.getFunctionType().getResults(),
+      func.getArguments(), api, CustomOption(&builder, empty_option_buffer));
+  func::ReturnOp::create(builder, func.getLoc(), op.getResults());
+  return success();
+}
+
+LogicalResult VerifyNgrams(func::FuncOp func) {
+  // The inputs and outputs should be the same:
+  // * A string tensor for tokens/ragged tensor values.
+  // * Zero or more row_split tensors.
+  constexpr int kValues = 0;
+  constexpr int kRowSplits = 1;
+
+  if (func.getFunctionType().getInputs().size() !=
+      func.getFunctionType().getResults().size()) {
+    return func.emitError() << "Mismatched number of inputs and outputs.";
+  }
+
+  int row_splits = func.getFunctionType().getInputs().size() - kRowSplits;
+  if (row_splits == 0) {
+    auto input_values = GetInputType(func, kValues);
+    if (!input_values ||
+        !mlir::isa<StringType>(input_values.getElementType())) {
+      return func.emitError()
+             << "Input " << kValues << " should be a string tensor";
+    }
+    auto output_values = GetResultType(func, kValues);
+    if (!output_values ||
+        !mlir::isa<StringType>(output_values.getElementType())) {
+      return func.emitError()
+             << "Output " << kValues << " should be a string tensor";
+    }
+
+    if (input_values.hasRank() && output_values.hasRank() &&
+        input_values.getRank() != output_values.getRank()) {
+      return func.emitError() << "Input " << kValues << " and output "
+                              << kValues << " should have the same rank";
+    }
+  } else {
+    auto input_values = GetInputType(func, kValues);
+    if (!RankEquals(input_values, 1) ||
+        !mlir::isa<StringType>(input_values.getElementType())) {
+      return func.emitError()
+             << "Input " << kValues << " should be a 1D string tensor";
+    }
+    auto output_values = GetResultType(func, kValues);
+    if (!RankEquals(output_values, 1) ||
+        !mlir::isa<StringType>(output_values.getElementType())) {
+      return func.emitError()
+             << "Output " << kValues << " should be a 1D string tensor";
+    }
+
+    for (int i = 0; i < row_splits; ++i) {
+      const int row_index = i + kRowSplits;
+      auto input_row_splits = GetInputType(func, row_index);
+      if (!RankEquals(input_row_splits, 1) ||
+          !input_row_splits.getElementType().isInteger(64)) {
+        return func.emitError()
+               << "Input " << row_index << " should be a 1D int64 tensor";
+      }
+      auto output_row_splits = GetResultType(func, row_index);
+      if (!RankEquals(output_row_splits, 1) ||
+          !output_row_splits.getElementType().isInteger(64)) {
+        return func.emitError()
+               << "Output " << row_index << " should be a 1D int64 tensor";
+      }
+    }
+  }
+
+  return success();
+}
+
+LogicalResult CreateNgramsCustomOption(func::FuncOp func, DictionaryAttr attrs,
+                                       std::string& custom_option_buffer) {
+  flexbuffers::Builder fbb;
+  size_t start_map = fbb.StartMap();
+
+  auto width = mlir::dyn_cast_or_null<IntegerAttr>(attrs.get("width"));
+  if (!width) {
+    return func.emitError() << "'width' attribute is not set or not an integer";
+  }
+  fbb.Int("width", width.getInt());
+
+  auto string_separator =
+      mlir::dyn_cast_or_null<StringAttr>(attrs.get("string_separator"));
+  if (!string_separator) {
+    return func.emitError()
+           << "'string_separator' attribute is not set or not a string";
+  }
+  // StringAttrs are not guaranteed to be NUL terminated, but flexbuffers
+  // strings expect NUL terminated strings.
+  std::string string_separator_str(string_separator.getValue().data(),
+                                   string_separator.getValue().size());
+  fbb.String("string_separator", string_separator_str);
+
+  auto axis = mlir::dyn_cast_or_null<IntegerAttr>(attrs.get("axis"));
+  if (!axis) {
+    return func.emitError() << "'axis' attribute is not set or not an integer";
+  }
+  fbb.Int("axis", axis.getInt());
+
+  auto reduction_type =
+      mlir::dyn_cast_or_null<StringAttr>(attrs.get("reduction_type"));
+  if (!reduction_type) {
+    return func.emitError()
+           << "'reduction_type' attribute is not set or not a string";
+  }
+  // StringAttrs are not guaranteed to be NUL terminated, but flexbuffers
+  // strings expect NUL terminated strings.
+  std::string reduction_type_str(reduction_type.getValue().data(),
+                                 reduction_type.getValue().size());
+  fbb.String("reduction_type", reduction_type_str);
+
+  fbb.EndMap(start_map);
+  fbb.Finish();
+  custom_option_buffer.assign(fbb.GetBuffer().begin(), fbb.GetBuffer().end());
+  return success();
+}
+
+LogicalResult ConvertNgrams(func::FuncOp func, llvm::StringRef api,
+                            FuncAttr attr) {
+  func.eraseBody();
+  func.addEntryBlock();
+  func->setAttr(kTFImplements, attr);
+  OpBuilder builder(func.getBody());
+  std::string custom_option_buffer;
+  if (failed(CreateNgramsCustomOption(func, attr.getAttrs(),
+                                      custom_option_buffer))) {
+    return failure();
+  }
+  auto op = CustomOp::create(
+      builder, func.getLoc(), func.getFunctionType().getResults(),
+      func.getArguments(), api, CustomOption(&builder, custom_option_buffer));
+  func::ReturnOp::create(builder, func.getLoc(), op.getResults());
+  return success();
+}
+
+LogicalResult VerifySgnnProjection(func::FuncOp func, FuncAttr attr) {
+  if (func.getFunctionType().getNumInputs() != 2 ||
+      func.getFunctionType().getNumResults() != 1) {
+    return func.emitError() << "Mismatched number of inputs and outputs.";
+  }
+  auto values_type = GetInputType(func, 0);
+  if (!values_type || !mlir::isa<StringType>(values_type.getElementType())) {
+    return func.emitError() << "First input should be a string tensor";
+  }
+  auto row_splits_type = GetInputType(func, 1);
+  if (!row_splits_type ||
+      !mlir::isa<IntegerType>(row_splits_type.getElementType())) {
+    return func.emitError() << "Second input should be an integer tensor";
+  }
+
+  auto hash_seed =
+      mlir::dyn_cast_or_null<ArrayAttr>(attr.getAttrs().get("hash_seed"));
+  if (!hash_seed) {
+    return func.emitError()
+           << "'hash_seed' attribute is not set or not an array";
+  }
+  auto output_type = GetResultType(func, 0);
+  if (!output_type || !mlir::isa<FloatType>(output_type.getElementType()) ||
+      !RankEquals(output_type, 2)) {
+    return func.emitError() << "Output should be a 2D float tensor.";
+  }
+  if (output_type.getDimSize(1) != hash_seed.size()) {
+    return func.emitError()
+           << "Output 2nd dimension should be the num of hash seeds.";
+  }
+
+  auto buckets =
+      mlir::dyn_cast_or_null<IntegerAttr>(attr.getAttrs().get("buckets"));
+  if (!buckets) {
+    return func.emitError() << "'buckets' attribute is not set or not int";
+  }
+
+  return success();
+}
+
+LogicalResult CreateSgnnProjectionCustomOption(
+    func::FuncOp func, DictionaryAttr attrs,
+    std::string& custom_option_buffer) {
+  flexbuffers::Builder fbb;
+  size_t start_map = fbb.StartMap();
+
+  auto hash_seed = mlir::dyn_cast_or_null<ArrayAttr>(attrs.get("hash_seed"));
+  auto vector_start = fbb.StartVector("hash_seed");
+  for (int i = 0; i < hash_seed.size(); i++) {
+    fbb.Add(static_cast<int32_t>(
+        mlir::dyn_cast<IntegerAttr>(*(hash_seed.getValue().data() + i))
+            .getInt()));
+  }
+  fbb.EndVector(vector_start, /*typed=*/true, /*fixed=*/false);
+
+  auto buckets = mlir::dyn_cast_or_null<IntegerAttr>(attrs.get("buckets"));
+  fbb.Int("buckets", buckets.getInt());
+
+  fbb.EndMap(start_map);
+  fbb.Finish();
+  custom_option_buffer.assign(fbb.GetBuffer().begin(), fbb.GetBuffer().end());
+  return success();
+}
+
+LogicalResult ConvertSgnnProjection(func::FuncOp func, llvm::StringRef api,
+                                    FuncAttr attr) {
+  // See more details in tensorflow_models/sequence_projection/sgnn/sgnn.py
+  func.eraseBody();
+  func.addEntryBlock();
+  func->setAttr(kTFImplements, attr);
+  OpBuilder builder(func.getBody());
+  std::string custom_option_buffer;
+  if (failed(CreateSgnnProjectionCustomOption(func, attr.getAttrs(),
+                                              custom_option_buffer))) {
+    return failure();
+  }
+  auto op = CustomOp::create(
+      builder, func.getLoc(), func.getFunctionType().getResults(),
+      func.getArguments(), api, CustomOption(&builder, custom_option_buffer));
+  func::ReturnOp::create(builder, func.getLoc(), op.getResults());
+  return success();
+}
+}  // namespace
+
+LogicalResult ConvertTFTextAPI(func::FuncOp func, llvm::StringRef api,
+                               FuncAttr attr) {
+  if (api.str() == kWhitespaceTokenizer) {
+    if (succeeded(VerifyWhitespaceTokenizer(func))) {
+      return ConvertWhitespaceTokenizer(func, api, attr);
+    }
+  } else if (api.str() == kNgrams) {
+    if (succeeded(VerifyNgrams(func))) {
+      return ConvertNgrams(func, api, attr);
+    }
+  } else if (api.str() == kCustomSgnnProjection) {
+    if (succeeded(VerifySgnnProjection(func, attr))) {
+      return ConvertSgnnProjection(func, api, attr);
+    }
+  }
+  return failure();
+}
+
+bool IsTFTextRegistered(const tensorflow::OpRegistry* op_registery) {
+  const std::vector<std::string> kTFTextOps = {
+      "WhitespaceTokenizeWithOffsets",
+  };
+  for (const auto& iter : kTFTextOps) {
+    if (op_registery->LookUp(iter)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/tftext_utils.h

@@ -0,0 +1,47 @@
+/* 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.
+==============================================================================*/
+
+// This header file defines common utils used by TFLite transformation
+// passes to work with op attributes.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_TFTEXT_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_TFTEXT_UTILS_H_
+
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/core/framework/op.h"
+
+namespace mlir {
+namespace TFL {
+
+// Fuse TF.Text APIs annotated by tf.function to a TFLite custom op.
+LogicalResult ConvertTFTextAPI(mlir::func::FuncOp func, llvm::StringRef api,
+                               mlir::TF::FuncAttr attr);
+
+// Check if TF.Text Tensorflow ops are registered.
+bool IsTFTextRegistered(const tensorflow::OpRegistry* op_registery);
+
+}  // end namespace TFL
+}  // end namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_TFTEXT_UTILS_H_

tensorflow/compiler/mlir/lite/utils/tftext_utils_test.cc

@@ -0,0 +1,57 @@
+/* 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/utils/tftext_utils.h"
+
+#include <memory>
+#include <string>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/platform/test.h"
+#include "tsl/platform/status.h"
+
+namespace mlir {
+namespace TFL {
+
+using tensorflow::OpRegistrationData;
+using tensorflow::OpRegistry;
+using tensorflow::Status;
+
+namespace {
+
+void Register(const std::string& op_name, OpRegistry* registry) {
+  registry->Register([op_name](OpRegistrationData* op_reg_data) -> Status {
+    op_reg_data->op_def.set_name(op_name);
+    return absl::OkStatus();
+  });
+}
+
+}  // namespace
+
+TEST(TfTextUtilsTest, TestTfTextRegistered) {
+  std::unique_ptr<OpRegistry> registry = std::make_unique<OpRegistry>();
+  Register("WhitespaceTokenizeWithOffsets", registry.get());
+  EXPECT_TRUE(IsTFTextRegistered(registry.get()));
+}
+
+TEST(TfTextUtilsTest, TestTfTextNotRegistered) {
+  std::unique_ptr<OpRegistry> registry = std::make_unique<OpRegistry>();
+  Register("Test", registry.get());
+  EXPECT_FALSE(IsTFTextRegistered(registry.get()));
+}
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/utils.h

@@ -0,0 +1,607 @@
+/* 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_UTILS_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_UTILS_H_
+
+#include <algorithm>
+#include <complex>
+#include <cstddef>
+#include <cstdint>
+#include <set>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/Dialect/Traits.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+using llvm::ArrayRef;
+using mlir::Operation;
+using mlir::ShapedType;
+using mlir::Value;
+
+// Returns true if the value is the min float value.
+inline bool IsNegInfiniteValue(APFloat value) {
+  if (!value.isNegative()) return false;
+  return value.isInfinity();
+}
+
+// Returns true if the value is the max float value.
+inline bool IsPosInfiniteValue(APFloat value) {
+  if (value.isNegative()) return false;
+  return value.isInfinity();
+}
+
+// Returns 1D 32-bit dense elements attribute with the given values.
+inline DenseIntElementsAttr GetI32ElementsAttr(ArrayRef<int32_t> values,
+                                               Builder* builder) {
+  RankedTensorType ty = mlir::RankedTensorType::get(
+      {static_cast<int32_t>(values.size())}, builder->getIntegerType(32));
+  return DenseIntElementsAttr::get(ty, values);
+}
+
+inline DenseIntElementsAttr GetI64ElementsAttr(ArrayRef<int64_t> values,
+                                               Builder* builder) {
+  RankedTensorType ty = RankedTensorType::get(
+      {static_cast<int64_t>(values.size())}, builder->getIntegerType(64));
+  return DenseIntElementsAttr::get(ty, values);
+}
+
+// Returns true if all tensor value in `values` has static shape and same shape.
+inline bool OpHasSameStaticShapes(Operation* op) {
+  auto values = op->getOperands();
+  int operand_num = 0;
+  ArrayRef<int64_t> shape;
+  for (Value value : values) {
+    auto shaped_type = mlir::dyn_cast<ShapedType>(value.getType());
+    if (!shaped_type || !shaped_type.hasStaticShape()) {
+      return false;
+    }
+    if (operand_num == 0) {
+      shape = shaped_type.getShape();
+    } else {
+      if (shape != shaped_type.getShape()) {
+        return false;
+      }
+    }
+    ++operand_num;
+  }
+  return true;
+}
+
+// Utility function to map final permutation to initial permutation
+// initial -> permutation1 -> permutation2 -> final
+inline DenseElementsAttr RemapPermutation(Value permutation1,
+                                          DenseElementsAttr perm2_const) {
+  SmallVector<int32_t> initial_permutation;
+  DenseElementsAttr perm1_const;
+
+  SmallVector<int32_t> new_permutation;
+  if (matchPattern(permutation1, m_Constant(&perm1_const))) {
+    for (int32_t idx = 0; idx < perm1_const.getNumElements(); ++idx) {
+      initial_permutation.push_back(idx);
+    }
+    for (auto perm : perm2_const.getValues<APInt>()) {
+      new_permutation.push_back(
+          initial_permutation[perm1_const
+                                  .getValues<APInt>()[perm.getSExtValue()]
+                                  .getSExtValue()]);
+    }
+  }
+
+  return mlir::DenseElementsAttr::get(
+      RankedTensorType::get(
+          {static_cast<int>(new_permutation.size())},
+          mlir::IntegerType::get(permutation1.getContext(), 32)),
+      llvm::ArrayRef(new_permutation));
+}
+
+// Utility function to map final permutation to initial permutation
+// initial -> permutation1 -> permutation2 -> final
+inline DenseElementsAttr RemapPermutation(Value permutation1,
+                                          Value permutation2) {
+  DenseElementsAttr perm2_const;
+  (void)matchPattern(permutation2, m_Constant(&perm2_const));
+
+  return RemapPermutation(permutation1, perm2_const);
+}
+
+inline bool IsTransposeNoop(Value permutation) {
+  DenseElementsAttr perm_values_attr;
+  if (!matchPattern(permutation, m_Constant(&perm_values_attr))) return false;
+
+  for (const auto& [idx, perm_value] :
+       llvm::enumerate(perm_values_attr.getValues<APInt>())) {
+    if (perm_value.getSExtValue() != idx) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// Returns true if the transpose op is trivial. Trivial means that
+// the permutation is a cyclic permutation of the original shape with only the
+// identity dimensions permuted.
+inline bool IsTransposeTrivial(llvm::ArrayRef<int64_t> input_shape,
+                               Value perm) {
+  DenseElementsAttr perm_values_attr;
+  if (!matchPattern(perm, m_Constant(&perm_values_attr))) return false;
+
+  SmallVector<int64_t, 8> perm_values;
+  for (const auto& dim : perm_values_attr.getValues<APInt>()) {
+    // Valid range is [-input_shape.size(), input_shape.size()).
+    int64_t p = dim.getSExtValue();
+    if (p < 0) {
+      p += input_shape.size();
+    }
+    if (p < 0 || p >= input_shape.size()) {
+      return false;
+    }
+    perm_values.push_back(p);
+  }
+
+  // This should never happen unless the input graph is malformed.
+  if (input_shape.size() != perm_values.size()) {
+    return false;
+  }
+
+  SmallVector<int, 8> old_major_index_ordering;
+  SmallVector<int, 8> new_major_index_ordering;
+  for (int i = 0, end = input_shape.size(); i < end; i++) {
+    if (input_shape[i] != 1) {
+      old_major_index_ordering.push_back(i);
+    }
+
+    if (input_shape[perm_values[i]] != 1) {
+      new_major_index_ordering.push_back(perm_values[i]);
+    }
+  }
+  return (old_major_index_ordering == new_major_index_ordering);
+}
+
+// Returns the permutation that maps the input shape to the output shape.
+// This is only valid for trivial reshape ops.
+inline DenseElementsAttr GetPermutationFromTrivialReshape(
+    mlir::ShapedType input_type, mlir::ShapedType output_type) {
+  ArrayRef<int64_t> in_shape = input_type.getShape();
+  ArrayRef<int64_t> out_shape = output_type.getShape();
+
+  // Get the indexes of the non-identity dimensions and the identity dimensions
+  // in the input shape.
+  SmallVector<int32_t> input_nonidentity_dims_index_array;
+  SmallVector<int32_t> input_identity_dims_index_array;
+
+  // Since the reshape is trivial, the input and output shapes should have the
+  // same number of dimensions. And the non-identity dimensions must be in the
+  // same cyclic order.
+  for (size_t idx = 0; idx < in_shape.size(); ++idx) {
+    if (in_shape[idx] != 1) {
+      input_nonidentity_dims_index_array.push_back(idx);
+    } else {
+      input_identity_dims_index_array.push_back(idx);
+    }
+  }
+
+  // Get the permutation that maps the input shape to the output shape.
+  SmallVector<int32_t> permutation;
+  size_t nonidentity_dims_index_poiter = 0;
+  size_t identity_dims_index_pointer = 0;
+  for (auto out_dim : out_shape) {
+    if (out_dim != 1) {
+      permutation.push_back(
+          input_nonidentity_dims_index_array[nonidentity_dims_index_poiter++]);
+    } else {
+      permutation.push_back(
+          input_identity_dims_index_array[identity_dims_index_pointer++]);
+    }
+  }
+
+  return mlir::DenseElementsAttr::get(
+      RankedTensorType::get(
+          {static_cast<int>(permutation.size())},
+          mlir::IntegerType::get(input_type.getContext(), 32)),
+      llvm::ArrayRef(permutation));
+}
+
+// Returns true if the reshape op is equivalent to a transpose op.
+// This is true if the reshape op is a trivial reshape op, meaning no change in
+// the order of non-identity dimensions.
+inline bool IsReshapeEquivalentToTranspose(mlir::ShapedType input_type,
+                                           mlir::ShapedType output_type) {
+  std::vector<int64_t> in_shape{input_type.getShape().vec()};
+  std::vector<int64_t> out_shape{output_type.getShape().vec()};
+
+  // If the reshape changes the number of dimensions so it cannot be interpreted
+  // as a transpose.
+  if (in_shape.size() != out_shape.size()) {
+    return false;
+  }
+
+  in_shape.erase(std::remove(in_shape.begin(), in_shape.end(), 1),
+                 in_shape.end());
+  out_shape.erase(std::remove(out_shape.begin(), out_shape.end(), 1),
+                  out_shape.end());
+  return in_shape == out_shape;
+}
+
+// Checks if all elements in the constant attribute value are 1.
+inline bool IsAllOnesConstant(Attribute value) {
+  auto values = mlir::cast<DenseElementsAttr>(value).getValues<int32_t>();
+  return !std::any_of(values.begin(), values.end(),
+                      [](int32_t element_value) { return element_value != 1; });
+}
+
+// Checks if all elements in the constant attribute value are non-negative.
+inline bool HasNonNegativeValues(Attribute value) {
+  auto values = mlir::cast<DenseElementsAttr>(value).getValues<APInt>();
+  return !std::any_of(
+      values.begin(), values.end(),
+      [](const APInt& element_value) { return element_value.isNegative(); });
+}
+
+// Utility function to get the offset between two dense attribute values.
+inline TypedAttr GetOffSet(Attribute begin, Attribute end) {
+  auto begin_values = mlir::cast<DenseElementsAttr>(begin).getValues<int32_t>();
+  auto end_values = mlir::cast<DenseElementsAttr>(end).getValues<int32_t>();
+
+  SmallVector<int32_t> offsets;
+  if (begin_values.size() == end_values.size()) {
+    for (size_t i = 0; i < begin_values.size(); ++i) {
+      offsets.push_back(end_values[i] - begin_values[i]);
+    }
+  }
+
+  return mlir::DenseElementsAttr::get(
+      RankedTensorType::get({static_cast<int>(offsets.size())},
+                            mlir::IntegerType::get(begin.getContext(), 32)),
+      llvm::ArrayRef(offsets));
+}
+
+// Check if the offset between two dense attribute values is non-negative.
+inline bool HasNonNegativeOffset(Attribute begin, Attribute end) {
+  return HasNonNegativeValues(GetOffSet(begin, end));
+}
+
+// Return true if the permutation value only swaps the last two dimensions
+inline bool AreLastTwoDimsTransposed(Value permutation) {
+  if (!permutation) return false;
+  DenseElementsAttr perm_values_attr;
+
+  if (!matchPattern(permutation, m_Constant(&perm_values_attr))) return false;
+  auto perm_values = perm_values_attr.getValues<APInt>();
+  size_t idx = 0;
+  for (; idx < perm_values_attr.size() - 2; ++idx) {
+    if (perm_values[idx].getSExtValue() != idx) return false;
+  }
+
+  return (perm_values[idx].getSExtValue() == perm_values_attr.size() - 1) &&
+         (perm_values[idx + 1].getSExtValue() == idx);
+}
+
+// Gets the new type after transposing the last 2 dimensions.
+inline Type TransposeLastTwoDims(Type type) {
+  auto shaped_type = mlir::dyn_cast<ShapedType>(type);
+  if (!shaped_type.hasStaticShape() || shaped_type.getRank() < 2) {
+    return nullptr;
+  }
+  int rank = shaped_type.getRank();
+  if (rank < 2) {
+    return nullptr;
+  }
+  SmallVector<int64_t> new_shape(shaped_type.getShape().begin(),
+                                 shaped_type.getShape().end());
+  std::swap(new_shape[rank - 1], new_shape[rank - 2]);
+  return shaped_type.clone(new_shape);
+}
+
+// Returns a ShapedType for a permutation and the shape of input after
+// applying the permutation to the given shape through a transpose.
+inline mlir::ShapedType GetTransposedType(
+    Value input, llvm::ArrayRef<int64_t> permutation_array) {
+  auto input_type = mlir::cast<ShapedType>(input.getType());
+  if (permutation_array.size() != input_type.getRank()) {
+    return nullptr;
+  }
+  llvm::SmallVector<int64_t> transposed_shape(permutation_array.size());
+  for (int64_t i = 0; i < permutation_array.size(); ++i) {
+    transposed_shape[i] = input_type.getDimSize(permutation_array[i]);
+  }
+  auto transposed_type =
+      RankedTensorType::get(transposed_shape, input_type.getElementType());
+  return transposed_type;
+}
+
+// Return the resultant shape if the shape of the supplied attribute/value is
+// expanded by n leading 1s'.
+inline SmallVector<int32_t> GetExpandedShape(Value input_val, int n) {
+  auto input_shape = mlir::cast<ShapedType>(input_val.getType()).getShape();
+  SmallVector<int32_t> expanded_shape;
+  expanded_shape.reserve(input_shape.size() + n);
+  for (int i = 0; i < n; ++i) {
+    expanded_shape.push_back(1);
+  }
+  expanded_shape.insert(expanded_shape.end(), input_shape.begin(),
+                        input_shape.end());
+  return expanded_shape;
+}
+
+// Return the resultant shape as a DenseElementsAttr if the shape of the
+// supplied attribute/value is expanded by n leading 1s'.
+inline DenseElementsAttr GetExpandedShapeAttr(Value input_val, int n) {
+  auto expanded_shape = GetExpandedShape(input_val, n);
+
+  return mlir::DenseElementsAttr::get(
+      RankedTensorType::get({static_cast<int>(expanded_shape.size())},
+                            mlir::IntegerType::get(input_val.getContext(), 32)),
+      llvm::ArrayRef(expanded_shape));
+}
+
+// Return the resultant shape type if the shape of the supplied attribute/value
+// is expanded by n leading 1s'.
+inline mlir::ShapedType GetExpandedShapeType(Value input_val, int n) {
+  auto expanded_shape = GetExpandedShape(input_val, n);
+  return RankedTensorType::get(
+      SmallVector<int64_t>{expanded_shape.begin(), expanded_shape.end()},
+      mlir::cast<ShapedType>(input_val.getType()).getElementType());
+}
+
+// Returns shape of a ranked tensor as a SmallVector.
+// Precondition: input_value's is ranked tensor.
+// Returns a squeezed shape when `squeeze_leading_ones` is set to true.
+inline SmallVector<int32_t> GetShape(Value input_value,
+                                     bool squeeze_leading_ones = false) {
+  auto output_shape =
+      mlir::dyn_cast<ShapedType>(input_value.getType()).getShape();
+
+  SmallVector<int32_t> shape;
+  shape.reserve(output_shape.size());
+
+  bool can_squeeze = true;
+  for (size_t dim_idx = 0; dim_idx < output_shape.size(); ++dim_idx) {
+    int64_t dim = output_shape[dim_idx];
+    if (squeeze_leading_ones && can_squeeze && dim == 1) {
+      continue;
+    } else if (can_squeeze && dim != 1) {
+      can_squeeze = false;
+    }
+    shape.push_back(ShapedType::isDynamic(dim) ? -1
+                                               : static_cast<int32_t>(dim));
+  }
+  return shape;
+}
+
+// Returns shape of a ranked tensor as a DenseElementsAttr.
+// Precondition: input_value's is ranked tensor.
+// Returns a squeezed shape when `squeeze_leading_ones` is set to true.
+inline DenseElementsAttr GetShapeAttr(Value input_value,
+                                      bool squeeze_leading_ones = false) {
+  SmallVector<int32_t> shape = GetShape(input_value, squeeze_leading_ones);
+
+  return mlir::DenseElementsAttr::get(
+      RankedTensorType::get(
+          {static_cast<int>(shape.size())},
+          mlir::IntegerType::get(input_value.getContext(), 32)),
+      llvm::ArrayRef(shape));
+}
+
+// Returns the value of a constant attribute as an int array, if the value is
+// not a constant, returns an error status.
+inline absl::StatusOr<SmallVector<int32_t>> GetValueAsIntArray(Value value) {
+  DenseElementsAttr values_const_attr;
+  if (!matchPattern(value, m_Constant(&values_const_attr))) {
+    return absl::InvalidArgumentError("Value is not a constant.");
+  }
+
+  SmallVector<int32_t> values;
+  for (const auto& value : values_const_attr.getValues<APInt>()) {
+    values.push_back(value.getSExtValue());
+  }
+  return values;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+///////////////// OP BROADCASTING UTILITIES ////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+// Returns whether the resultant type of any broadcastable operation with
+// operands `a` and `b` matches `expected_output`. Returns false if `a` is not
+// broadcast-compatible with `b`.
+inline bool OperandsBroadcastToOutputType(Type a, Type b,
+                                          Type expected_output) {
+  Type output_element_type =
+      mlir::cast<ShapedType>(expected_output).getElementType();
+  Type broadcasted_type =
+      OpTrait::util::getBroadcastedType(a, b, output_element_type);
+  return broadcasted_type != Type() && broadcasted_type == expected_output;
+}
+
+// Returns int, float or complex DenseElementsAttr with scalar shape with the
+// given element type and the integer value.
+template <typename T>
+DenseElementsAttr GetScalarOfType(Type ty, T raw_value) {
+  RankedTensorType scalar_ty = RankedTensorType::get({}, ty);
+  if (auto float_ty = mlir::dyn_cast<FloatType>(ty)) {
+    FloatAttr attr = FloatAttr::get(float_ty, raw_value);
+    return DenseElementsAttr::get(scalar_ty, attr);
+  } else if (auto int_ty = mlir::dyn_cast<IntegerType>(ty)) {
+    IntegerAttr attr = IntegerAttr::get(int_ty, raw_value);
+    return DenseElementsAttr::get(scalar_ty, attr);
+  } else if (auto complex_ty = mlir::dyn_cast<ComplexType>(ty)) {
+    Type complex_element_ty = complex_ty.getElementType();
+    if (complex_element_ty.isF32()) {
+      return DenseElementsAttr::get(
+          scalar_ty, static_cast<std::complex<float>>(raw_value));
+    } else if (complex_element_ty.isF64()) {
+      return DenseElementsAttr::get(
+          scalar_ty, static_cast<std::complex<double>>(raw_value));
+    }
+  }
+  llvm_unreachable("unsupported type");
+}
+
+// Checks if reduction axes and broadcast axes are disjoint.
+// Broadcast axes are derived by comparing the shape of `input_val` to the shape
+// represented by `target_shape_attr` according to standard broadcasting rules.
+// Returns true if the sets of axes are disjoint, false otherwise or on error.
+inline bool AreBroadcastAndReductionAxesIndependent(
+    mlir::Value input_val, const mlir::Attribute& indices_attr,
+    const mlir::Attribute& target_shape_attr) {
+  // 1. Get input type and shape.
+  // Use llvm::dyn_cast for safer casting.
+  auto ranked_input_type =
+      llvm::dyn_cast<mlir::RankedTensorType>(input_val.getType());
+  if (!ranked_input_type) {
+    // Consider logging or error emission if builder context is
+    // available/needed.
+    return false;  // Expect ranked type.
+  }
+  llvm::ArrayRef<int64_t> input_shape = ranked_input_type.getShape();
+  const int64_t input_rank = ranked_input_type.getRank();
+
+  // 2. Validate and extract reduction axes.
+  // Use llvm::dyn_cast for safer casting.
+  auto indices = llvm::dyn_cast<mlir::DenseElementsAttr>(indices_attr);
+  if (!indices || !indices.getElementType().isIntOrIndex()) {
+    return false;  // Invalid indices attribute.
+  }
+
+  // Use std::set for efficient storage and lookup of axes.
+  std::set<int64_t> reduction_axes_set;
+  if (!indices.empty()) {  // Only process if there are reduction axes.
+    if (input_rank == 0) {
+      // It's invalid to specify reduction axes for a scalar (rank 0) input.
+      return false;
+    }
+
+    // Iterate using range-based for loop and structured binding (if applicable)
+    // or direct value access.
+    for (const mlir::APInt& axis_val : indices.getValues<mlir::APInt>()) {
+      int64_t axis =
+          axis_val.getSExtValue();  // Use sign extension for neg axes.
+
+      // Normalize axis and check bounds.
+      if (axis < -input_rank || axis >= input_rank) {
+        return false;  // Axis out of bounds.
+      }
+      if (axis < 0) {
+        axis += input_rank;  // Convert negative axis to positive.
+      }
+      reduction_axes_set.insert(axis);
+    }
+  }
+
+  // If there are no reduction axes, they are trivially independent of any
+  // broadcast axes.
+  if (reduction_axes_set.empty()) {
+    return true;
+  }
+
+  // 3. Validate and extract target shape for broadcast.
+  // Use llvm::dyn_cast for safer casting.
+  auto target_shape_value_attr =
+      llvm::dyn_cast<mlir::DenseElementsAttr>(target_shape_attr);
+  if (!target_shape_value_attr ||
+      !target_shape_value_attr.getElementType().isIntOrIndex()) {
+    return false;  // Invalid target shape attribute.
+  }
+
+  // Use llvm::SmallVector for efficient shape storage.
+  llvm::SmallVector<int64_t, 4> target_shape_vec;
+  target_shape_vec.reserve(
+      target_shape_value_attr.getNumElements());  // Pre-allocate
+  for (const mlir::APInt& shape_val :
+       target_shape_value_attr.getValues<mlir::APInt>()) {
+    // Assuming shape dimensions should be non-negative, consider getZExtValue.
+    // However, getSExtValue is safe if intermediate calculations handle signs.
+    target_shape_vec.push_back(shape_val.getSExtValue());
+  }
+  // Use llvm::ArrayRef for safe, non-owning view of the shape vector.
+  llvm::ArrayRef<int64_t> target_shape = target_shape_vec;
+  const int64_t target_rank = target_shape.size();
+
+  // 4. Determine broadcast axes based on standard broadcasting rules.
+  std::set<int64_t> broadcast_axes_set;
+  const int64_t max_rank = std::max(input_rank, target_rank);
+
+  // Iterate through dimensions, aligning from the right (trailing dimensions).
+  for (int64_t i = 0; i < max_rank; ++i) {
+    // Calculate indices relative to the end of the shape arrays.
+    const int64_t input_dim_idx = input_rank - 1 - i;
+    const int64_t target_dim_idx = target_rank - 1 - i;
+
+    // Treat dimensions missing due to lower rank as having size 1.
+    const int64_t input_dim =
+        (input_dim_idx >= 0) ? input_shape[input_dim_idx] : 1;
+    const int64_t target_dim =
+        (target_dim_idx >= 0) ? target_shape[target_dim_idx] : 1;
+
+    // Check for incompatible shapes (dimensions differ and neither is 1).
+    // This indicates an invalid broadcast according to NumPy rules.
+    if (input_dim != target_dim && input_dim != 1 && target_dim != 1) {
+      // Consider if the specific broadcast op allows other behaviors (e.g.,
+      // -1). For standard rules, this is an incompatibility.
+      return false;
+    }
+
+    // An axis in the *input* tensor is involved in broadcasting if its size is
+    // 1 and the corresponding target dimension size is greater than 1.
+    if (input_dim == 1 && target_dim > 1) {
+      // Ensure the axis index is valid for the input tensor's rank.
+      if (input_dim_idx >= 0) {
+        broadcast_axes_set.insert(input_dim_idx);
+      }
+      // Note: If input_dim_idx < 0, broadcasting occurs due to rank difference,
+      // but it doesn't correspond to an axis *within* the original input
+      // tensor.
+    }
+  }
+
+  // 5. Check for intersection between the set of reduction axes and the set of
+  //    broadcast axes derived above.
+  for (int64_t reduction_axis : reduction_axes_set) {
+    if (broadcast_axes_set.count(reduction_axis)) {
+      // Found an axis that is present in both sets.
+      return false;
+    }
+  }
+
+  // 6. No overlapping axes were found.
+  return true;
+}
+
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_UTILS_H_

tensorflow/compiler/mlir/lite/utils/utils.td

@@ -0,0 +1,239 @@
+/* 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.
+==============================================================================*/
+
+// Utility predicates that are shared by multiple passes.
+
+include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "mlir/IR/PatternBase.td"
+
+
+////////////////////////////////////////////////////////////////////////////////
+///////////////// TENSOR TYPE UTILITIES ////////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+def IsQuantized : Constraint<CPred<
+  "llvm::dyn_cast<ShapedType>($0.getType()) && "
+  "llvm::isa<quant::UniformQuantizedType>("
+  "llvm::dyn_cast<ShapedType>($0.getType()).getElementType())">>;
+
+def IsNotQuantized : Constraint<Neg<IsQuantized.predicate>>;
+
+////////////////////////////////////////////////////////////////////////////////
+///////////////// TENSOR RANK UTILITIES ////////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+// Checks if the rank of the value is less than or equal to the rank of the
+// other value.
+def IsRankLessThanEqualTo : Constraint<CPred<
+  "llvm::cast<ShapedType>($0.getType()).getRank() <= "
+  "llvm::cast<ShapedType>($1.getType()).getRank()">>;
+
+// Checks if the value has rank at most 'n'.
+class HasRankAtMost<int n> : Constraint<
+    CPred<"llvm::cast<ShapedType>($0.getType()).hasRank() && "
+          "llvm::cast<ShapedType>($0.getType()).getRank() <= " # n>>;
+
+////////////////////////////////////////////////////////////////////////////////
+///////////////// DENSE UTILITIES /////////////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+def DenseFPElementsAttrPred : CPred<"llvm::isa<DenseFPElementsAttr>($_self)">;
+def DenseIntElementsAttrPred : CPred<"llvm::isa<DenseIntElementsAttr>($_self)">;
+
+////////////////////////////////////////////////////////////////////////////////
+///////////////// SPLAT CONSTANT UTILITIES /////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+def DenseElementsAttrIsSplatPred
+    : CPred<"llvm::cast<DenseElementsAttr>($_self).isSplat()">;
+
+class DenseFPElementsAttrSplatValueEqualToPred<string val>
+    : CPred<"llvm::cast<DenseFPElementsAttr>($_self).getSplatValue<FloatAttr>()"
+            ".getValueAsDouble() == " # val>;
+
+class DenseFPElementsAttrSplatValueEqualToPredWithTolerance<string val, string tolerance>
+    : CPred<"std::abs(llvm::cast<DenseFPElementsAttr>($_self).getSplatValue<FloatAttr>()"
+            ".getValueAsDouble() - " # val # ") <= "#tolerance>;
+
+class DenseIntElementsAttrSplatValueEqualToPred<string val>
+    : CPred<"llvm::isa<DenseIntElementsAttr>($_self) && "
+            "llvm::isa<IntegerType>("
+	    "llvm::cast<DenseIntElementsAttr>($_self).getElementType()) && "
+            "llvm::cast<DenseIntElementsAttr>($_self).isSplat() && "
+            "llvm::cast<DenseIntElementsAttr>($_self).getSplatValue<IntegerAttr>()"
+            "  .getValue().getSExtValue() == " # val>;
+
+// AttrConstraint to match a floating point dense elements attribute with a
+// splat value equals to `Value`.
+class FPSplatConstAttr<string Value> : AttrConstraint<And<[
+    DenseFPElementsAttrPred, DenseElementsAttrIsSplatPred,
+    DenseFPElementsAttrSplatValueEqualToPred<Value>]>>;
+
+class FPSplatConstAttrWithTolerance<string Value, string Tolerance> :
+    AttrConstraint<And<[
+        DenseFPElementsAttrPred, DenseElementsAttrIsSplatPred,
+        DenseFPElementsAttrSplatValueEqualToPredWithTolerance<Value, Tolerance>]>>;
+
+// AttrConstraint to match a dense elements attribute with a splat value equal
+// to `Value`.
+class IntSplatConstAttr<string value> : AttrConstraint<
+    DenseIntElementsAttrSplatValueEqualToPred<value>>;
+
+// A constant tensor that has elements with the same value.
+def SplatElementsAttr : ElementsAttrBase<DenseElementsAttrIsSplatPred,
+  "splat tensor attr">;
+
+// A constant tensor that has elements with the same floating point value.
+def SplatFPElementsAttr : ElementsAttrBase<
+  And<[DenseFPElementsAttrPred, DenseElementsAttrIsSplatPred]>,
+  "f32 splat tensor attr">;
+
+// A constant tensor that has elements with the same integer value.
+def SplatIntElementsAttr : ElementsAttrBase<
+  And<[DenseIntElementsAttrPred, DenseElementsAttrIsSplatPred]>,
+  "integer splat tensor attr">;
+
+// Extracts the scalar constant value.
+def GetScalarElementsAttrFromSplat : NativeCodeCall<
+    "DenseElementsAttr::get("
+    " RankedTensorType::get({},"
+        " llvm::cast<mlir::DenseElementsAttr>($0).getType().getElementType()),"
+    " llvm::cast<mlir::DenseElementsAttr>($0).getSplatValue<mlir::Attribute>())">;
+
+////////////////////////////////////////////////////////////////////////////////
+///////////////// OP BROADCASTING UTILITIES ////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+def OperandsBroadcastToOutputType : Constraint<CPred<
+  "TFL::OperandsBroadcastToOutputType($0.getType(), $1.getType(), "
+                                     "$2.getType())">>;
+
+def OperandsDontBroadcastToOutputType : Constraint<Neg<
+  OperandsBroadcastToOutputType.predicate>>;
+
+////////////////////////////////////////////////////////////////////////////////
+///////////////// TENSOR SHAPE UTILITIES ///////////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+def HasSameStaticShapes : Constraint<
+    CPred<"llvm::cast<ShapedType>($0.getType()).hasStaticShape() && "
+          "llvm::cast<ShapedType>($1.getType()).hasStaticShape() && "
+          "llvm::cast<ShapedType>($0.getType()).getShape() =="
+          "llvm::cast<ShapedType>($1.getType()).getShape()">,
+    "have the same static shape">;
+
+def CreateNoneValue : NativeCodeCall<
+  "TFL::NoValueOp::create($_builder, $0.getLoc(), $_builder.getUnitAttr())">;
+
+// Returns shape of a ranked tensor.
+// if called without a ranked tensor it will fail.
+def GetShapeAttr: NativeCodeCall<"GetShapeAttr($0)">;
+
+// Return the resultant shape if the shape of the supplied attribute/value is
+// expanded by n leading 1s'.
+class GetExpandedShapeAttr<int n> : NativeCodeCall<
+  "GetExpandedShapeAttr($0, " # n # ")">;
+
+// Return the resultant shape type if the shape of the supplied attribute/value
+// is expanded by n leading 1s'.
+class GetExpandedShapeType<int n> : NativeCodeCall<
+  "GetExpandedShapeType($0, " # n # ")">;
+
+// Constraint that values in list attribute are all ones.
+def IsAllOnesConstant : Constraint<CPred<"TFL::IsAllOnesConstant($0)">>;
+
+// Constraint that checks if the transpose op is trivial. Trivial means that
+// the permutation is a cyclic permutation of the original shape with only the
+// identity dimensions permuted.
+def IsTransposeTrivial : Constraint<CPred<
+  "TFL::IsTransposeTrivial(llvm::cast<ShapedType>($0.getType()).getShape(), $1)">>;
+
+// Constraint that checks if the transpose op is a no-op.
+def IsTransposeNoop : Constraint<CPred<"TFL::IsTransposeNoop($0)">>;
+
+// Constraint that checks if the reshape op is equivalent to a transpose op.
+// This is true if the reshape op is a trivial reshape op, meaning no change in
+// the order of non-identity dimensions.
+def IsReshapeEquivalentToTranspose : Constraint<CPred<
+  "TFL::IsReshapeEquivalentToTranspose("
+    "llvm::cast<ShapedType>($0.getType()),"
+    "llvm::cast<ShapedType>($1.getType()))">>;
+
+// Returns the permutation of the trivial reshape op, this will be used to
+// construct the transpose op.
+def GetPermutationFromTrivialReshape : NativeCodeCall<
+  "TFL::GetPermutationFromTrivialReshape("
+    "llvm::cast<ShapedType>($0.getType()),"
+    "llvm::cast<ShapedType>($1.getType()))">;
+
+// Constraint that checks if all values in offset between two
+// attributes are non-negative.
+def HasNonNegativeOffset : Constraint<CPred<"TFL::HasNonNegativeOffset($0, $1)">>;
+
+// Constraint that checks if all values in list attribute are non-negative.
+def HasNonNegativeValues : Constraint<CPred<"TFL::HasNonNegativeValues($0)">>;
+
+// Utility function to get the offset between two dense attribute values.
+def GetOffSet : NativeCodeCall<"TFL::GetOffSet($0, $1)">;
+
+// Attribute Constraint that checks if the attribute value is zero.
+def ZeroIntAttr
+  : AttrConstraint<CPred<"llvm::cast<::mlir::IntegerAttr>($_self).getInt() == 0">>;
+
+// Checks if the value has rank at most 'n'.
+class HasRankAtLeast<int n> : Constraint<
+    CPred<"llvm::cast<ShapedType>($0.getType()).hasRank() && "
+          "llvm::cast<ShapedType>($0.getType()).getRank() >= " # n>>;
+
+// Accepts two inputs and check if both have the same element type.
+def SameElementType : Constraint<
+  CPred<"getElementTypeOrSelf($0) == getElementTypeOrSelf($1)">>;
+
+// Returns a ShapedType for a permutation and the shape of input after
+// applying the permutation to the given shape through a transpose.
+class GetTransposedType<string perm> : NativeCodeCall<
+  "GetTransposedType($0, " # perm # ")">;
+
+// Function to map final permutation to initial permutation
+// initial -> permutation1 -> permutation2 -> final
+def RemapPermutation: NativeCodeCall<"RemapPermutation($0, $1)">;
+
+// Checks if all of an ops inputs are the same static shape.
+// BUILD NOTE: "OpHasSameStaticShapes" here refers to the C++ function defined 
+// in `utils/utils.h`. The `utils.h` header is included in `tfl_ops.h` so all
+// of our files will have access to `OpHasSameStaticShapes` when including files
+// generated from table-gen.
+def OpHasSameStaticShapesPred : CPred<"OpHasSameStaticShapes($0.getDefiningOp())">;
+def OpHasSameStaticShapes : Constraint<OpHasSameStaticShapesPred, "op must have static same input shapes">;
+def OpHasNotSameStaticShapes : Constraint<Neg<OpHasSameStaticShapesPred>, "op must have not static same input shapes">;
+
+def TransposeFCLastTwoDims:
+  NativeCodeCall<"TransposeLastTwoDims($0[0].getType())">;
+
+def AreLastTwoDimsTransposed : Constraint<CPred<
+  "TFL::AreLastTwoDimsTransposed($0)">>;
+
+// Checks if the param passed is of NoneType.
+def IsNoneType : Constraint<CPred<"llvm::isa<NoneType>($0.getType())">>;
+
+def ConstantLikePred : CPred<"::mlir::matchPattern($0, ::mlir::m_Constant())">;
+def IsConstantLike : Constraint<ConstantLikePred>;
+def NotConstantLike : Constraint<Neg<ConstantLikePred>>;
+
+// Here, the element type can be any integer or float type. But, note that only
+// 32 bit integers are supported for the values.
+class GetScalarOfType<int value> : NativeCodeCall<
+  "GetScalarOfType(getElementTypeOrSelf($0)," # value # ")">;

tensorflow/compiler/mlir/lite/utils/utils_test.cc

@@ -0,0 +1,128 @@
+/* Copyright 2025 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/utils/utils.h"
+
+#include <cstdint>
+
+#include <gtest/gtest.h>
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+namespace {
+
+// Test fixture for AreBroadcastAndReductionAxesIndependent function.
+class BroadcastAndReductionAxesIndependentTest : public ::testing::Test {
+ protected:
+  BroadcastAndReductionAxesIndependentTest() : builder_(&context_) {
+    context_.loadDialect<arith::ArithDialect>();
+  }
+
+  // Builds an mlir::Value representing a tensor with the given shape.
+  Value BuildTensor(ArrayRef<int64_t> shape) {
+    return arith::ConstantOp::create(
+        builder_, builder_.getUnknownLoc(),
+        RankedTensorType::get(shape, builder_.getF32Type()),
+        builder_.getZeroAttr(
+            RankedTensorType::get(shape, builder_.getF32Type())));
+  }
+
+  // Builds a DenseElementsAttr representing an integer array.
+  DenseElementsAttr BuildIntArrayAttr(ArrayRef<int32_t> values) {
+    return DenseElementsAttr::get(
+        RankedTensorType::get({static_cast<int32_t>(values.size())},
+                              builder_.getI32Type()),
+        values);
+  }
+
+  MLIRContext context_;
+  OpBuilder builder_;
+};
+
+TEST_F(BroadcastAndReductionAxesIndependentTest, IndependentAxes) {
+  Value input_tensor = BuildTensor({2, 1, 4, 1});
+  DenseElementsAttr reduction_axes = BuildIntArrayAttr({0, 2});
+  DenseElementsAttr target_shape = BuildIntArrayAttr({2, 3, 4, 5});
+
+  EXPECT_TRUE(AreBroadcastAndReductionAxesIndependent(
+      input_tensor, reduction_axes, target_shape));
+  input_tensor.getDefiningOp()->destroy();
+}
+
+TEST_F(BroadcastAndReductionAxesIndependentTest, OverlappingAxes) {
+  Value input_tensor = BuildTensor({1, 3, 4, 5});
+  DenseElementsAttr reduction_axes = BuildIntArrayAttr({0, 2});
+  DenseElementsAttr target_shape = BuildIntArrayAttr({2, 3, 4, 5});
+
+  EXPECT_FALSE(AreBroadcastAndReductionAxesIndependent(
+      input_tensor, reduction_axes, target_shape));
+  input_tensor.getDefiningOp()->destroy();
+}
+
+TEST_F(BroadcastAndReductionAxesIndependentTest, EmptyReductionAxes) {
+  Value input_tensor = BuildTensor({1, 3, 1, 5});
+  DenseElementsAttr reduction_axes = BuildIntArrayAttr({});
+  DenseElementsAttr target_shape = BuildIntArrayAttr({2, 3, 4, 5});
+
+  EXPECT_TRUE(AreBroadcastAndReductionAxesIndependent(
+      input_tensor, reduction_axes, target_shape));
+  input_tensor.getDefiningOp()->destroy();
+}
+
+TEST_F(BroadcastAndReductionAxesIndependentTest, UnrankedInput) {
+  Value input_tensor = arith::ConstantOp::create(
+      builder_, builder_.getUnknownLoc(), builder_.getF32Type(),
+      builder_.getZeroAttr(builder_.getF32Type()));
+  DenseElementsAttr reduction_axes = BuildIntArrayAttr({0, 2});
+  DenseElementsAttr target_shape = BuildIntArrayAttr({2, 3, 4, 5});
+
+  EXPECT_FALSE(AreBroadcastAndReductionAxesIndependent(
+      input_tensor, reduction_axes, target_shape));
+  input_tensor.getDefiningOp()->destroy();
+}
+
+TEST_F(BroadcastAndReductionAxesIndependentTest, InvalidReductionAxesType) {
+  Value input_tensor = BuildTensor({2, 3, 4, 5});
+  DenseElementsAttr reduction_axes = DenseElementsAttr::get(
+      RankedTensorType::get({2}, builder_.getF32Type()), {1.0f, 2.0f});
+  DenseElementsAttr target_shape = BuildIntArrayAttr({1, 3, 1, 5});
+
+  EXPECT_FALSE(AreBroadcastAndReductionAxesIndependent(
+      input_tensor, reduction_axes, target_shape));
+  input_tensor.getDefiningOp()->destroy();
+}
+
+TEST_F(BroadcastAndReductionAxesIndependentTest, InvalidTargetShapeType) {
+  Value input_tensor = BuildTensor({2, 3, 4, 5});
+  DenseElementsAttr reduction_axes = BuildIntArrayAttr({0, 2});
+  DenseElementsAttr target_shape = DenseElementsAttr::get(
+      RankedTensorType::get({2}, builder_.getF32Type()), {1.0f, 2.0f});
+
+  EXPECT_FALSE(AreBroadcastAndReductionAxesIndependent(
+      input_tensor, reduction_axes, target_shape));
+  input_tensor.getDefiningOp()->destroy();
+}
+
+}  // namespace
+}  // namespace TFL
+
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/validators.cc

@@ -0,0 +1,147 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/lite/utils/validators.h"
+
+#include <algorithm>
+#include <cstdint>
+
+#include "mlir/Dialect/Traits.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+// Returns true if the given `op`
+//   * has an attribute with the given `name`,
+//   * and the attribute is an integer list of the form [1, X, Y, 1],
+// and writes X, Y as 32-bit integer attribute to `x`, `y`.
+bool TFIntListIs1XY1(Operation *op, StringRef name, IntegerAttr *x,
+                     IntegerAttr *y) {
+  auto attr = op->getAttrOfType<ArrayAttr>(name);
+  if (!attr) return false;
+
+  auto elements = attr.getValue();
+  if (elements.size() != 4 ||
+      std::any_of(elements.begin(), elements.end(),
+                  [](Attribute e) { return !mlir::isa<IntegerAttr>(e); }))
+    return false;
+
+  if (mlir::cast<IntegerAttr>(elements.front()).getInt() != 1 ||
+      mlir::cast<IntegerAttr>(elements.back()).getInt() != 1)
+    return false;
+
+  Builder b(op->getContext());
+  *x = b.getI32IntegerAttr(mlir::cast<IntegerAttr>(elements[1]).getInt());
+  *y = b.getI32IntegerAttr(mlir::cast<IntegerAttr>(elements[2]).getInt());
+
+  return true;
+}
+
+// Returns true if the attribute is an integer list of the form [1, X, Y, 1].
+bool TFIntListIs1XY1(const Attribute attr) {
+  const auto &elements = mlir::cast<ArrayAttr>(attr).getValue();
+  if (elements.size() != 4 ||
+      std::any_of(elements.begin(), elements.end(),
+                  [](Attribute e) { return !mlir::isa<IntegerAttr>(e); }))
+    return false;
+
+  if (mlir::cast<IntegerAttr>(elements.front()).getValue() != 1 ||
+      mlir::cast<IntegerAttr>(elements.back()).getValue() != 1)
+    return false;
+  return true;
+}
+
+// Returns true if the attribute is an integer list of the form [1, 1, X, Y].
+bool TFIntListIs11XY(const Attribute attr) {
+  const auto &elements = mlir::cast<ArrayAttr>(attr).getValue();
+  if (elements.size() != 4 ||
+      std::any_of(elements.begin(), elements.end(),
+                  [](Attribute e) { return !mlir::isa<IntegerAttr>(e); }))
+    return false;
+
+  const Attribute *data = elements.data();
+  if (mlir::cast<IntegerAttr>(data[0]).getValue() != 1 ||
+      mlir::cast<IntegerAttr>(data[1]).getValue() != 1)
+    return false;
+  return true;
+}
+
+// Returns true if the given `op`
+//   * has an attribute with the given `name`,
+//   * and the attribute is an integer list of the form [1, X, Y, Z, 1],
+// and writes X, Y as 32-bit integer attribute to `x`, `y`, z.
+bool TFIntListIs1XYZ1(Operation *op, StringRef name, IntegerAttr *x,
+                      IntegerAttr *y, IntegerAttr *z) {
+  auto attr = op->getAttrOfType<ArrayAttr>(name);
+  if (!attr) return false;
+
+  auto elements = attr.getValue();
+  if (elements.size() != 5 ||
+      std::any_of(elements.begin(), elements.end(),
+                  [](Attribute e) { return !mlir::isa<IntegerAttr>(e); }))
+    return false;
+
+  if (mlir::cast<IntegerAttr>(elements.front()).getInt() != 1 ||
+      mlir::cast<IntegerAttr>(elements.back()).getInt() != 1)
+    return false;
+
+  Builder b(op->getContext());
+  *x = b.getI32IntegerAttr(mlir::cast<IntegerAttr>(elements[1]).getInt());
+  *y = b.getI32IntegerAttr(mlir::cast<IntegerAttr>(elements[2]).getInt());
+  *z = b.getI32IntegerAttr(mlir::cast<IntegerAttr>(elements[3]).getInt());
+
+  return true;
+}
+
+// Returns true if every element of the attribute is 1. All elements of `attr`
+// must be `IntegerAttr`.
+bool TFIntListIsAllOnes(const Attribute attr) {
+  const auto &elements = mlir::cast<ArrayAttr>(attr).getValue();
+
+  return !std::any_of(elements.begin(), elements.end(), [](Attribute e) {
+    return mlir::cast<IntegerAttr>(e).getValue() != 1;
+  });
+}
+
+bool IsBroadcastableElementsAttrs(mlir::TypedAttr a, mlir::TypedAttr b) {
+  // This would return false if we had unranked tensors (where they should
+  // probably be considered as broadcastable), but given we are working with
+  // attributes here that shouldn't be an issue,
+  return OpTrait::util::getBroadcastedType(a.getType(), b.getType()) != Type();
+}
+
+bool IsDimensionsDegenerateExceptLastOne(ArrayRef<int64_t> elements_shape) {
+  if (elements_shape.empty()) return true;
+
+  for (auto dim : elements_shape.drop_back(1)) {
+    if (dim != 1) return false;
+  }
+  return true;
+}
+
+bool IsDimensionsDegenerateExceptLastOne(TypedAttr val) {
+  if (auto ranked_type = mlir::dyn_cast<RankedTensorType>(val.getType())) {
+    return IsDimensionsDegenerateExceptLastOne(ranked_type.getShape());
+  }
+  return false;
+}
+
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/validators.h

@@ -0,0 +1,126 @@
+/* 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.
+==============================================================================*/
+
+// This header file defines common validators used by TFLite transformation
+// passes to validate op attributes or values.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_LITE_UTILS_VALIDATORS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_VALIDATORS_H_
+
+#include <cstdint>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+
+// TODO(jpienaar): Change these to being one of these variants and/or generate
+// these predicates.
+
+// Returns true if the given TensorFlow op does not have a `data_format`
+// attribute (then default to "NHWC"), or its `data_format` attribute is "NHWC".
+inline bool TFDataFormatIsNHWC(Operation *op) {
+  auto attr = op->getAttrOfType<StringAttr>("data_format");
+  return !attr || attr.getValue() == "NHWC";
+}
+
+// Returns true if the given TensorFlow op does not have a `data_format`
+// attribute (then default to "NDHWC"), or its `data_format` attribute is
+// "NDHWC".
+inline bool TFDataFormatIsNDHWC(Operation *op) {
+  auto attr = op->getAttrOfType<StringAttr>("data_format");
+  return !attr || attr.getValue() == "NDHWC";
+}
+
+// Returns true if the given `op`
+//   * has an attribute with the given `name`,
+//   * and the attribute is an integer list of the form [1, X, Y, 1],
+// and writes X, Y as 32-bit integer attribute to `x`, `y`.
+bool TFIntListIs1XY1(Operation *op, StringRef name, IntegerAttr *x,
+                     IntegerAttr *y);
+
+// Returns true if the attribute is an integer list of the form [1, X, Y, 1].
+bool TFIntListIs1XY1(Attribute attr);
+
+// Returns true if the attribute is an integer list of the form [1, 1, X, Y].
+bool TFIntListIs11XY(Attribute attr);
+
+// Returns true if the given `op`
+//   * has an attribute with the given `name`,
+//   * and the attribute is an integer list of the form [1, X, Y, Z, 1],
+// and writes X, Y as 32-bit integer attribute to `x`, `y`, z.
+bool TFIntListIs1XYZ1(Operation *op, StringRef name, IntegerAttr *x,
+                      IntegerAttr *y, IntegerAttr *z);
+
+// Returns true if every element of the attribute is 1. All elements of `attr`
+// must be `IntegerAttr`.
+bool TFIntListIsAllOnes(Attribute attr);
+
+// Returns true iff the given value is a float32 tensor.
+// is "DT_FLOAT".
+inline bool TFTypeIsFloat32Tensor(Value value) {
+  auto tensorType = mlir::dyn_cast<TensorType>(value.getType());
+  if (!tensorType) return false;
+  return tensorType.getElementType().isF32();
+}
+
+// Returns true iff the given value is a bf16 tensor.
+inline bool TFTypeIsBFloat16Tensor(Value value) {
+  auto tensorType = mlir::dyn_cast<TensorType>(value.getType());
+  if (!tensorType) return false;
+  return tensorType.getElementType().isBF16();
+}
+
+// Returns true iff the given value is a f16 tensor.
+inline bool TFTypeIsHalfTensor(Value value) {
+  auto tensorType = mlir::dyn_cast<TensorType>(value.getType());
+  if (!tensorType) return false;
+  return tensorType.getElementType().isF16();
+}
+
+// Returns true iff the given value is a f16 or bf16 tensor.
+inline bool TFTypeIsBFloat16OrHalfTensor(Value value) {
+  return TFTypeIsBFloat16Tensor(value) || TFTypeIsHalfTensor(value);
+}
+
+// Returns true iff the given TensorFlow op has a `padding` attribute whose
+// value is "SAME" or "VALID", and writes the attribute to `padding`.
+inline bool TFPaddingIsSameOrValid(Operation *op, StringAttr *padding) {
+  auto padding_attr = op->getAttrOfType<StringAttr>("padding");
+  if (padding_attr.getValue() != "SAME" && padding_attr.getValue() != "VALID")
+    return false;
+  *padding = padding_attr;
+  return true;
+}
+
+/// Returns whether the given `a` and `b` have broadcast-compatible
+/// types.
+bool IsBroadcastableElementsAttrs(mlir::TypedAttr a, mlir::TypedAttr b);
+// Returns true if every dimension of the attribute is 1 except the last one.
+bool IsDimensionsDegenerateExceptLastOne(mlir::TypedAttr val);
+// Returns true if every element is 1 except the last one.
+bool IsDimensionsDegenerateExceptLastOne(ArrayRef<int64_t> elements_shape);
+
+}  // end namespace TFL
+}  // end namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_VALIDATORS_H_

tensorflow/compiler/mlir/lite/utils/variables_utils.cc

@@ -0,0 +1,70 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/lite/utils/variables_utils.h"
+
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+
+namespace mlir {
+namespace TFL {
+namespace utils {
+
+bool IsSupportedVariableType(Operation* op) {
+  ShapedType type;
+  if (llvm::isa<TF::ReadVariableOp>(op)) {
+    type = llvm::cast<ShapedType>(op->getResult(0).getType());
+  } else if (llvm::isa<TF::AssignVariableOp>(op)) {
+    type = llvm::cast<ShapedType>(op->getOperand(1).getType());
+  } else if (llvm::isa<TF::VarHandleOp>(op)) {
+    type =
+        llvm::cast<tf_type::TensorFlowTypeWithSubtype>(
+            llvm::cast<ShapedType>(op->getResult(0).getType()).getElementType())
+            .GetSubtypes()
+            .back();
+  }
+  return IsSupportedVariableType(type);
+}
+
+bool IsSupportedVariableType(ShapedType type) {
+  auto element_type = type.getElementType();
+  // Check complex types.
+  if (auto complex_type = llvm::dyn_cast<ComplexType>(element_type)) {
+    auto complex_element_type = complex_type.getElementType();
+    if (complex_element_type.isF32() || complex_element_type.isF64())
+      return true;
+  }
+  // Check quantized types.
+  if (auto quant_type = llvm::dyn_cast<quant::QuantizedType>(element_type)) {
+    // TFLite supports QI16, QI32, QI8, and QUI8
+    if ((quant_type.getStorageTypeIntegralWidth() == 16 &&
+         quant_type.isSigned()) ||
+        quant_type.getStorageTypeIntegralWidth() == 8 ||
+        (quant_type.getStorageTypeIntegralWidth() == 32 &&
+         quant_type.isSigned()))
+      return true;
+  }
+  return element_type.isF32() || element_type.isF64() ||
+         element_type.isInteger(1) || element_type.isInteger(8) ||
+         element_type.isInteger(32) || element_type.isSignlessInteger(64);
+}
+
+}  // namespace utils
+}  // namespace TFL
+}  // namespace mlir

tensorflow/compiler/mlir/lite/utils/variables_utils.h

@@ -0,0 +1,36 @@
+/* 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_UTILS_VARIABLES_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_LITE_UTILS_VARIABLES_UTILS_H_
+
+#include "mlir/IR/AffineMap.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFL {
+namespace utils {
+
+// Returns true if 'op' has type that is supported by native TFLite
+// variables.
+bool IsSupportedVariableType(Operation* op);
+
+// Returns true if 'type' is supported by native tflite variables.
+bool IsSupportedVariableType(ShapedType type);
+
+}  // namespace utils
+}  // namespace TFL
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_LITE_UTILS_VARIABLES_UTILS_H_