chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tosa/transforms/convert_metadata.cc

@@ -0,0 +1,114 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <string>
+
+#include "llvm/ADT/StringExtras.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir::tosa {
+
+#define GEN_PASS_DEF_CONVERTFUNCTIONMETADATA
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+namespace {
+
+// Extract the input and output names
+static void splitFunctionIONames(StringAttr namesAttr,
+                                 llvm::SmallVectorImpl<std::string> &names) {
+  SmallVector<StringRef, 4> namesRef;
+  llvm::SplitString(namesAttr.getValue(), namesRef, ",");
+  for (auto nameRef : namesRef) {
+    names.push_back(nameRef.str());
+  }
+}
+
+class ConvertFunctionMetadataPass
+    : public impl::ConvertFunctionMetadataBase<ConvertFunctionMetadataPass> {
+ public:
+  void runOnOperation() override {
+    auto funcOp = getOperation();
+
+    // Setup entry functions for compilation and preserve the
+    // associated metadata. Note that TFLite uses `tf.entry_function`.
+    auto entryFunctionAttr =
+        funcOp->getAttrOfType<DictionaryAttr>("tf.entry_function");
+    if (entryFunctionAttr) {
+      setupEntryPointAttrs(funcOp, entryFunctionAttr);
+    }
+  }
+
+ private:
+  // TF/TFL pack their I/O names in a dictionary, convert into arg attributes.
+  void setupEntryPointAttrs(func::FuncOp funcOp,
+                            DictionaryAttr entryFunctionAttr) {
+    funcOp.setPublic();
+
+    if (funcOp.getNumArguments() > 0) {
+      auto inputsAttr =
+          dyn_cast_or_null<StringAttr>(entryFunctionAttr.get("inputs"));
+      if (!inputsAttr) {
+        funcOp.emitError() << "functions with tf.entry_function must have "
+                              "input names to be handled by backend";
+        return signalPassFailure();
+      }
+      SmallVector<std::string, 4> inputNames;
+      splitFunctionIONames(inputsAttr, inputNames);
+      if (inputNames.size() != funcOp.getNumArguments()) {
+        funcOp.emitError()
+            << "tf.entry_function attribute malformed: inputs don't "
+               "match the function signature";
+        return signalPassFailure();
+      }
+      for (auto [i, name] : llvm::enumerate(inputNames)) {
+        funcOp.setArgAttr(i, "ml_program.identifier",
+                          StringAttr::get(&getContext(), name));
+      }
+    }
+    if (funcOp.getNumResults() > 0) {
+      auto outputsAttr =
+          dyn_cast_or_null<StringAttr>(entryFunctionAttr.get("outputs"));
+      if (!outputsAttr) {
+        funcOp.emitError() << "functions with tf.entry_function must have "
+                              "output names to be handled by backend";
+        return signalPassFailure();
+      }
+      SmallVector<std::string, 4> outputNames;
+      splitFunctionIONames(outputsAttr, outputNames);
+      if (outputNames.size() != funcOp.getNumResults()) {
+        funcOp.emitError()
+            << "tf.entry_function attribute malformed: outputs don't "
+               "match the function signature";
+        return signalPassFailure();
+      }
+      for (auto [i, name] : llvm::enumerate(outputNames)) {
+        funcOp.setResultAttr(i, "ml_program.identifier",
+                             StringAttr::get(&getContext(), name));
+      }
+    }
+  }
+};
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertFunctionMetadataPass() {
+  return std::make_unique<ConvertFunctionMetadataPass>();
+}
+
+}  // namespace mlir::tosa

tensorflow/compiler/mlir/tosa/transforms/convert_tfl_uint.cc

@@ -0,0 +1,547 @@
+/* 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 pass converts a TFLite unsigned integer graph to the signed integer
+// domain, with adaptors at input and output tensors. This is needed because
+// TOSA precision is implemented in the signed integer domain. This pass:
+// 1. Matches TFL::QConst with unsigned integers, generates TFL::QConst with
+// signed integers with values remapped.
+// 2. Inserts tosa.RESCALE unsigned -> signed at graph inputs
+// 3. Inserts tosa.RESCALE signed -> unsigned at graph outputs
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"                // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"                // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/QuantUtils.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/PatternMatch.h"                        // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"                      // from @llvm-project
+#include "mlir/Support/LLVM.h"                           // from @llvm-project
+#include "mlir/Support/LogicalResult.h"                  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-convert-tfl-uint-to-int"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSACONVERTTFLUNSIGNEDINTTOSIGNEDPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+// Validate that zero point follows TOSA requirements:
+// - int8/uint8 can have zero point within their valid range
+// - uint16 zero point must be either 0 or 32768
+// - All other types must have zero point equal to 0
+// Ref:
+// https://github.com/arm/tosa-specification/blob/9fe5e964e2193f0e345670f7f4098beecd7fd6eb/tosa.xml#L2479
+bool isValidTOSAZeroPoint(int32_t zp, unsigned bit_width, bool is_unsigned) {
+  if (bit_width == 8) {
+    // int8/uint8 can have any zero point in their valid range
+    if (is_unsigned) {
+      return zp >= 0 && zp <= 255;
+    } else {
+      return zp >= -128 && zp <= 127;
+    }
+  } else if (bit_width == 16 && is_unsigned) {
+    // uint16 must have zp of 0 or 32768
+    return (zp == 0 || zp == 32768);
+  } else {
+    // All other types (int16, int32, uint32) must have zero point = 0
+    return zp == 0;
+  }
+}
+
+// Adjust zero points according to TOSA requirements
+int32_t adjustZeroPointForTOSA(int32_t zp, unsigned bit_width, bool is_unsigned) {
+  if (bit_width == 16 && is_unsigned) {
+    // uint16 must have zp of 0 or 32768
+    if (zp != 0 && zp != 32768) {
+      // TOSA doesn't support other zero points for uint16. Round to nearest
+      // valid value. This leads to precision loss in the rescaled type.
+      return (zp < 16384) ? 0 : 32768;
+    }
+    return zp;
+  } else if (bit_width == 8) {
+    // int8/uint8 can have any zero point in their valid range
+    return zp;
+  } else {
+    // All other types must have zero point = 0. This also leads to precision loss.
+    return 0;
+  }
+}
+
+// Only support 8, 16, and 32-bit integers as tosa.rescale op only supports
+// these bitWidths as both input and output types
+bool isUnsupportedBitWidth(unsigned bit_width) {
+  return (bit_width != 8 && bit_width != 16 && bit_width != 32);
+}
+
+// Create a signed quantized type from an unsigned one
+// Returns the signed quantized type and the adjusted zero point
+std::pair<mlir::quant::UniformQuantizedType, int32_t>
+createSignedQuantizedTypeFromUnsigned(
+    OpBuilder& builder, mlir::quant::UniformQuantizedType unsigned_quant_type) {
+  unsigned bit_width = unsigned_quant_type.getStorageTypeIntegralWidth();
+
+  int32_t unsigned_zp = unsigned_quant_type.getZeroPoint();
+
+  // Calculate real value range
+  double real_min = static_cast<double>(
+                        unsigned_quant_type.getStorageTypeMin() - unsigned_zp) *
+                    unsigned_quant_type.getScale();
+  double real_max = static_cast<double>(
+                        unsigned_quant_type.getStorageTypeMax() - unsigned_zp) *
+                    unsigned_quant_type.getScale();
+
+  // Determine signed storage range
+  int64_t signed_min = quant::AnyQuantizedType::getDefaultMinimumForInteger(
+      /* isSignedInteger = */ true, bit_width);
+  int64_t signed_max = quant::AnyQuantizedType::getDefaultMaximumForInteger(
+      /* isSignedInteger = */ true, bit_width);
+
+  // Adjust for narrow range
+  if (unsigned_quant_type.getStorageTypeMin() == 1) {
+    signed_min++;
+  }
+
+  // Compute optimal scale and zero point
+  double real_range = real_max - real_min;
+  double storage_range = static_cast<double>(signed_max - signed_min);
+  double new_scale = real_range / storage_range;
+
+  // Compute zero point to align real_min with signed_min
+  int32_t signed_zp =
+      static_cast<int32_t>(std::round(signed_min - real_min / new_scale));
+
+  int32_t adjusted_zp =
+      adjustZeroPointForTOSA(signed_zp, bit_width, /*is_unsigned=*/false);
+
+  // If zero point was changed by TOSA constraints, recompute scale
+  if (adjusted_zp != signed_zp) {
+    // With the adjusted zp, recompute scale to ensure coverage of real range
+    // real_min = new_scale * (signed_min - adjusted_zp)
+    // real_max = new_scale * (signed_max - adjusted_zp)
+    double scale_for_min =
+        real_min / static_cast<double>(signed_min - adjusted_zp);
+    double scale_for_max =
+        real_max / static_cast<double>(signed_max - adjusted_zp);
+    new_scale = std::max(scale_for_min, scale_for_max);
+  }
+
+  signed_zp = adjusted_zp;
+
+  auto signed_quant_type = quant::UniformQuantizedType::getChecked(
+      builder.getUnknownLoc(), quant::QuantizationFlags::Signed,
+      builder.getIntegerType(bit_width), unsigned_quant_type.getExpressedType(),
+      new_scale, signed_zp, signed_min, signed_max);
+
+  return {signed_quant_type, signed_zp};
+}
+
+// Pattern for converting unsigned QConst ops
+struct ConvertUnsignedQConstOp : public RewritePattern {
+  explicit ConvertUnsignedQConstOp(MLIRContext* context)
+      : RewritePattern(TFL::QConstOp::getOperationName(), 1, context) {}
+
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& builder) const override {
+    auto tfl_qconst_op = cast<TFL::QConstOp>(op);
+
+    // Skip if it's not ranked tensor type.
+    auto output_type =
+        dyn_cast<mlir::RankedTensorType>(tfl_qconst_op.getResult().getType());
+    if (!output_type)
+      return builder.notifyMatchFailure(op, "not ranked tensor");
+
+    // Skip if output is not per-tensor quantized type.
+    auto output_element_type = dyn_cast<mlir::quant::UniformQuantizedType>(
+        output_type.getElementType());
+    if (!output_element_type)
+      return builder.notifyMatchFailure(
+          op, "not per-tensor quantized tfl.qconst op");
+
+    // Skip if output is already signed.
+    if (output_element_type.isSigned()) {
+      return builder.notifyMatchFailure(
+          op, "expected output type to be unsigned quantized type");
+    }
+
+    // Get the bit_width from the quantized type itself
+    unsigned bit_width = output_element_type.getStorageTypeIntegralWidth();
+
+    if (isUnsupportedBitWidth(bit_width)) {
+      return builder.notifyMatchFailure(
+          op, "only 8, 16, and 32-bit integers are supported");
+    }
+
+    int32_t unsigned_zp = output_element_type.getZeroPoint();
+    if (!isValidTOSAZeroPoint(unsigned_zp, bit_width, /*is_unsigned=*/true)) {
+      return builder.notifyMatchFailure(op,
+                                        "Zeropoint is not supported by TOSA.");
+    }
+
+    // Calculate the zero-point offset for this bit_width
+    auto [signed_quant_type, signed_zp] =
+        createSignedQuantizedTypeFromUnsigned(builder, output_element_type);
+
+    auto dst_qconst_type = TypeAttr::get(
+        RankedTensorType::get(output_type.getShape(), signed_quant_type));
+
+    Type dst_dense_element_type = builder.getIntegerType(bit_width);
+
+    mlir::DenseElementsAttr src_dense_attr =
+        mlir::cast<DenseElementsAttr>(tfl_qconst_op.getValue());
+
+    int64_t value_offset = unsigned_zp - signed_zp;
+    auto dst_dense_attr = src_dense_attr.mapValues(
+        dst_dense_element_type,
+        [bit_width, value_offset](const APInt& in) -> APInt {
+          int64_t in_i64 = in.getLimitedValue();
+          int64_t out_i64 = in_i64 - value_offset;
+          return APInt(bit_width, out_i64, true);
+        });
+
+    builder.replaceOpWithNewOp<TFL::QConstOp>(op, dst_qconst_type,
+                                              dst_dense_attr);
+
+    return success();
+  }
+};
+
+// returns true iff @a type is a shaped type with element type that is unsigned
+// integer. If it is, then update the rescaled type, input_zp, and output_zp to
+// use to rescale type to signed type with adjusted zero point.
+bool ConvertUnsignedToSignedRescale(OpBuilder& builder, const Type type,
+                          Type& rescaled_type, int32_t& unsigned_zp,
+                          int32_t& output_zp, unsigned& bit_width) {
+  auto shaped_type = dyn_cast<mlir::ShapedType>(type);
+  if (!shaped_type) return false;
+
+  auto element_type = shaped_type.getElementType();
+  auto unsigned_element_quant_type =
+      dyn_cast<mlir::quant::UniformQuantizedType>(element_type);
+
+  // Check if it's unsigned quantized type
+  if (unsigned_element_quant_type && !unsigned_element_quant_type.isSigned()) {
+    bit_width = unsigned_element_quant_type.getStorageTypeIntegralWidth();
+    if (isUnsupportedBitWidth(bit_width)) {
+      return false;
+    }
+
+    unsigned_zp = unsigned_element_quant_type.getZeroPoint();
+    // Use helper function to create signed quantized type
+    auto [signed_quant_type, signed_zp] = createSignedQuantizedTypeFromUnsigned(
+        builder, unsigned_element_quant_type);
+
+    output_zp = signed_zp;
+    rescaled_type = shaped_type.clone(signed_quant_type);
+    return true;
+  }
+
+  // Check for plain unsigned integer types
+  if (element_type.isUnsignedInteger()) {
+    bit_width = element_type.getIntOrFloatBitWidth();
+    if (isUnsupportedBitWidth(bit_width)) {
+      return false;
+    }
+
+    int64_t zp_offset = 1LL << (bit_width - 1);
+    unsigned_zp = 0;
+    output_zp = adjustZeroPointForTOSA(unsigned_zp - zp_offset, bit_width,
+                                       /*is_unsigned=*/false);
+    // Determine signed storage range
+    int64_t signed_min = quant::AnyQuantizedType::getDefaultMinimumForInteger(
+        /* isSignedInteger = */ true, bit_width);
+    int64_t signed_max = quant::AnyQuantizedType::getDefaultMaximumForInteger(
+        /* isSignedInteger = */ true, bit_width);
+
+    rescaled_type = shaped_type.clone(quant::UniformQuantizedType::getChecked(
+        builder.getUnknownLoc(), quant::QuantizationFlags::Signed,
+        builder.getIntegerType(bit_width), builder.getF32Type(),
+        /* scale = */ 1.0,
+        /* zeroPoint = */ output_zp,
+        /* storageTypeMin = */ signed_min,
+        /* storageTypeMax = */ signed_max));
+
+    return true;
+  }
+
+  return false;
+}
+
+LogicalResult convert_graph_unsigned_tensor(mlir::MLIRContext& context,
+                                            mlir::func::FuncOp& function) {
+  size_t num_blocks_in_main = 0;
+  mlir::Region* region = function.getCallableRegion();
+  OpBuilder builder(&context);
+  auto loc = function.getLoc();
+
+  for (mlir::Block& bb : region->getBlocks()) {
+    // Always have one block for each region right now
+    num_blocks_in_main++;
+    if (num_blocks_in_main > 1) {
+      return function.emitError("Invalid MLIR: multiple blocks in a region");
+    }
+
+    if (!bb.isEntryBlock()) {
+      return function.emitError("Invalid MLIR: block must be entry block");
+    }
+
+    // Create multiplier and shift tensors (these are always i32 and i8
+    // respectively)
+    auto multiplier = tosa::getConstTensorInt<int32_t>(builder, loc, {1 << 30});
+    auto shift = tosa::getConstTensorInt<int8_t>(builder, loc, {30});
+
+    // Insert rescale unsigned->signed for input arguments
+    for (Value arg : bb.getArguments()) {
+      auto shaped_type = dyn_cast<ShapedType>(arg.getType());
+      if (!shaped_type) continue;
+
+      // Check if zeropoint is supported by TOSA if quantized unsigned type
+      auto element_type = shaped_type.getElementType();
+      auto unsigned_element_quant_type =
+          dyn_cast<mlir::quant::UniformQuantizedType>(element_type);
+
+      if (unsigned_element_quant_type &&
+          !unsigned_element_quant_type.isSigned()) {
+        unsigned bit_width =
+            unsigned_element_quant_type.getStorageTypeIntegralWidth();
+        int32_t unsigned_zp = unsigned_element_quant_type.getZeroPoint();
+
+        if (!isValidTOSAZeroPoint(unsigned_zp, bit_width, /*is_unsigned=*/true)) {
+          return function.emitError()
+                 << "Input argument has unsigned quantized type with zero "
+                    "point "
+                 << unsigned_zp
+                 << " which is not supported by TOSA for bitwidth " << bit_width
+                 << ".";
+        }
+      }
+
+      Type rescaled_type;
+      int32_t rescale_input_zp_val, rescale_output_zp_val;
+      unsigned bit_width;
+
+      if (!ConvertUnsignedToSignedRescale(builder, arg.getType(), rescaled_type,
+                                rescale_input_zp_val, rescale_output_zp_val,
+                                bit_width))
+        continue;
+
+      // Keep original input_val use with tmp_val.
+      auto tmp_const_type =
+          RankedTensorType::get({1}, builder.getIntegerType(bit_width));
+      auto tmp_const_attr =
+          DenseElementsAttr::get(tmp_const_type, {APInt(bit_width, 0)});
+      Value tmp_val =
+          builder.create<TFL::ConstOp>(loc, tmp_const_type, tmp_const_attr);
+      arg.replaceAllUsesWith(tmp_val);
+
+      // Create zero point constants with appropriate bit_width
+      // input_zp must match input bit_width, output_zp must match
+      // output bit_width
+      auto rescale_input_zp = createZeroPointTensor(
+          builder, loc, tmp_const_type, rescale_input_zp_val);
+      auto rescale_output_zp = createZeroPointTensor(
+          builder, loc, tmp_const_type, rescale_output_zp_val);
+
+      if (!rescale_input_zp || !rescale_output_zp) {
+        return function.emitError("Failed to create input zero point tensors.");
+      }
+
+      const auto rounding_mode_attr = tosa::RoundingModeAttr::get(
+          builder.getContext(), tosa::RoundingMode::SINGLE_ROUND);
+      auto rescale_op = builder.create<tosa::RescaleOp>(
+          loc, rescaled_type, arg, multiplier, shift, rescale_input_zp.value(),
+          rescale_output_zp.value(),
+          /* scale32 = */ builder.getBoolAttr(true),
+          /* rounding_mode = */ rounding_mode_attr,
+          /* per_channel = */ builder.getBoolAttr(false),
+          /* input_unsigned = */ builder.getBoolAttr(true),     // unsigned ->
+          /* output_unsigned = */ builder.getBoolAttr(false));  // signed
+
+      Operation* op_rescale_op = static_cast<Operation*>(rescale_op);
+      bb.push_front(op_rescale_op);
+      tmp_val.replaceAllUsesWith(rescale_op.getResult());
+      tmp_val.getDefiningOp()->erase();
+      bb.push_front(rescale_output_zp.value().getDefiningOp());
+      bb.push_front(rescale_input_zp.value().getDefiningOp());
+    }
+
+    bb.push_front(shift.getDefiningOp());
+    bb.push_front(multiplier.getDefiningOp());
+
+    // Record types of original graph output before we convert intermediate
+    // tensors.
+    auto terminator = bb.getTerminator();
+    SmallVector<Type, 4> output_types;
+    for (Value val : terminator->getOperands()) {
+      output_types.push_back(val.getType());
+    }
+
+    // Convert intermediate tensors.
+    for (auto& op : bb) {
+      if (llvm::dyn_cast<tosa::ConstOp>(&op)) {
+        // Skip tosa const ops created during rescaling.
+        continue;
+      }
+
+      for (Value output_val : op.getResults()) {
+        auto shaped_type = dyn_cast<ShapedType>(output_val.getType());
+        if (!shaped_type) continue;
+        Type new_type;
+        int32_t unused_input_zp, unused_output_zp;
+        unsigned bit_width;
+
+        if (ConvertUnsignedToSignedRescale(builder, output_val.getType(), new_type,
+                                 unused_input_zp, unused_output_zp, bit_width)) {
+          output_val.setType(new_type);
+        }
+      }
+    }
+
+    if (terminator->getNumOperands() != output_types.size()) {
+      return function.emitError(
+          "Terminator's operand mismatch with number of outputs in graph");
+    }
+
+    // Insert signed->unsigned rescale before all terminator's operands
+    for (int32_t i = 0; i < terminator->getNumOperands(); i++) {
+      auto defining_op = terminator->getOperand(i).getDefiningOp();
+      // Skip if operand of terminator is block arg (nullptr in this case)
+      if (!defining_op) continue;
+      Value input_val = defining_op->getResult(0);
+
+      // Check if graph output is unsigned type.
+      auto unsigned_output_type = dyn_cast<mlir::ShapedType>(output_types[i]);
+      if (!unsigned_output_type) continue;
+
+      // Check if graph output is unsigned type.
+      Type rescaled_type;
+      int32_t unsigned_zp_val, unused_output_zp_val;
+      unsigned bit_width;
+
+      if (!ConvertUnsignedToSignedRescale(builder, output_types[i], rescaled_type,
+                                unsigned_zp_val, unused_output_zp_val,
+                                bit_width))
+        continue;
+
+      // convert terminator operand type back to original output_type.
+      auto terminator_operand_type =
+          dyn_cast<mlir::ShapedType>(terminator->getOperand(i).getType());
+      if (!terminator_operand_type) continue;
+      int operand_zp_val = 0;
+      auto quantized_type = dyn_cast<mlir::quant::UniformQuantizedType>(
+          terminator_operand_type.getElementType());
+      if (quantized_type) {
+        operand_zp_val = adjustZeroPointForTOSA(quantized_type.getZeroPoint(),
+                                                bit_width, /*is_unsigned=*/true);
+      }
+
+      // Keep original input_val use with tmp_val.
+      auto tmp_const_type =
+          RankedTensorType::get({1}, builder.getIntegerType(bit_width));
+      auto tmp_const_attr =
+          DenseElementsAttr::get(tmp_const_type, {APInt(bit_width, 0)});
+      Value tmp_val =
+          builder.create<TFL::ConstOp>(loc, tmp_const_type, tmp_const_attr);
+      input_val.replaceUsesWithIf(tmp_val, [&terminator](OpOperand& use) {
+        return use.getOwner() == terminator;
+      });
+
+      // Create zero point constants with appropriate bit_width
+      // input_zp must match input bit_width (signed), output_zp must match
+      // output bit_width (unsigned)
+      auto rescale_input_zp =
+          createZeroPointTensor(builder, loc, tmp_const_type, operand_zp_val);
+      auto rescale_output_zp =
+          createZeroPointTensor(builder, loc, tmp_const_type, unsigned_zp_val);
+
+      if (!rescale_input_zp || !rescale_output_zp) {
+        return function.emitError("Failed to create input zero point tensors.");
+      }
+
+      const auto rounding_mode_attr = tosa::RoundingModeAttr::get(
+          builder.getContext(), tosa::RoundingMode::SINGLE_ROUND);
+      auto rescale_op = builder.create<tosa::RescaleOp>(
+          loc, unsigned_output_type, input_val, multiplier, shift,
+          rescale_input_zp.value(), rescale_output_zp.value(),
+          /* scale32 = */ builder.getBoolAttr(true),
+          /* rounding_mode = */ rounding_mode_attr,
+          /* per_channel = */ builder.getBoolAttr(false),
+          /* input_unsigned = */ builder.getBoolAttr(false),   // signed ->
+          /* output_unsigned = */ builder.getBoolAttr(true));  // unsigned
+
+      Operation* op_rescale_op = static_cast<Operation*>(rescale_op);
+      bb.push_back(op_rescale_op);
+      op_rescale_op->moveBefore(terminator);
+      tmp_val.replaceAllUsesWith(rescale_op.getResult());
+      tmp_val.getDefiningOp()->erase();
+      bb.push_front(rescale_output_zp.value().getDefiningOp());
+      bb.push_front(rescale_input_zp.value().getDefiningOp());
+    }
+  }
+
+  return success();
+}
+
+class ConvertTFLUnsignedToSigned
+    : public impl::TosaConvertTFLUnsignedIntToSignedPassBase<
+          ConvertTFLUnsignedToSigned> {
+ public:
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    auto& ctx = getContext();
+    mlir::func::FuncOp func = getOperation();
+
+    func.walk([&](Operation* op) {
+      if (isa<TosaOp>(op)) {
+        // Run this before calling convert_graph_uint8_tensor as rescaling
+        // introduces tosa ops
+        op->emitError(
+            "tosa operations are not expected in this pass. Run "
+            "tosa-convert-tfl-uint-to-int before tosa-legalize-tfl");
+      }
+    });
+
+    // Convert uint const tensor. const needs to be handled specifically.
+    patterns.add<ConvertUnsignedQConstOp>(&ctx);
+    (void)applyPatternsGreedily(func, std::move(patterns));
+
+    // Replace uint tensor in the graph and insert rescale as needed.
+    (void)convert_graph_unsigned_tensor(ctx, func);
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTFLUnsignedIntToSignedPass() {
+  return std::make_unique<ConvertTFLUnsignedToSigned>();
+}
+
+}  // namespace tosa
+}  // namespace mlir

tensorflow/compiler/mlir/tosa/transforms/dequantize_tfl_softmax.cc

@@ -0,0 +1,90 @@
+/* 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 <memory>
+#include <utility>
+
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-dequantize_tfl-softmax"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSADEQUANTIZETFLSOFTMAXPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+class TosaDequantizeTFLSoftmax
+    : public impl::TosaDequantizeTFLSoftmaxPassBase<TosaDequantizeTFLSoftmax> {
+ public:
+  explicit TosaDequantizeTFLSoftmax() = default;
+  void runOnOperation() override;
+};
+
+struct TosaDequantizeTFLSoftmaxPattern : public RewritePattern {
+  explicit TosaDequantizeTFLSoftmaxPattern(MLIRContext* context)
+      : RewritePattern(TFL::SoftmaxOp::getOperationName(), 1, context) {}
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override;
+};
+
+LogicalResult TosaDequantizeTFLSoftmaxPattern::matchAndRewrite(
+    Operation* op, PatternRewriter& rewriter) const {
+  TFL::SoftmaxOp tfl_softmax_op = cast<TFL::SoftmaxOp>(op);
+  RankedTensorType input_type =
+      mlir::cast<RankedTensorType>(tfl_softmax_op.getInput().getType());
+  if (!mlir::isa<mlir::quant::QuantizedType>(input_type.getElementType())) {
+    return failure();
+  }
+  Location loc = tfl_softmax_op.getLoc();
+  RankedTensorType dequantized_input_type =
+      RankedTensorType::get(input_type.getShape(), rewriter.getF32Type());
+  Value dequantized_input = rewriter.create<TFL::DequantizeOp>(
+      loc, dequantized_input_type, tfl_softmax_op.getInput());
+  Value dequantized_softmax_output = rewriter.create<TFL::SoftmaxOp>(
+      loc, dequantized_input_type, dequantized_input, tfl_softmax_op.getBeta());
+  Type qtype = tfl_softmax_op.getOutput().getType();
+  rewriter.replaceOpWithNewOp<TFL::QuantizeOp>(tfl_softmax_op, qtype,
+                                               dequantized_softmax_output,
+                                               mlir::TypeAttr::get(qtype));
+  return success();
+}
+
+void TosaDequantizeTFLSoftmax::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  patterns.add<TosaDequantizeTFLSoftmaxPattern>(&getContext());
+  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createDequantizeTFLSoftmaxPass() {
+  return std::make_unique<TosaDequantizeTFLSoftmax>();
+}
+
+}  // namespace tosa
+
+}  // namespace mlir

tensorflow/compiler/mlir/tosa/transforms/fuse_bias_tf.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.
+==============================================================================*/
+
+// Fuse tf.Op + tf.BiasAdd and legalized to TOSA
+
+#include <memory>
+#include <optional>
+#include <utility>
+
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-fuse-bias-tf"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSAFUSEBIASTFPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+class FuseBiasTF : public impl::TosaFusebiasTFPassBase<FuseBiasTF> {
+ public:
+  explicit FuseBiasTF() = default;
+  void runOnOperation() override;
+};
+
+struct ConvertTFBiasAddOp : public RewritePattern {
+  explicit ConvertTFBiasAddOp(MLIRContext* context)
+      : RewritePattern(TF::BiasAddOp::getOperationName(), 1, context) {}
+  LogicalResult matchAndRewrite(Operation* op,
+                                PatternRewriter& rewriter) const override;
+};
+
+// Replaces the following pattern, take conv2d as an example:
+//   %1 = tf.Conv2D (%ifm, %filter)
+//   %2 = tf.BiasAdd(%1, %bias)
+//   with
+//   %1 = tosa.conv2d(%ifm, %filter, %bias)
+//   This can also be done using the pair ot Pat<> options in
+//   tf_optimize_patterns.td
+//   However, this explicit code can handle both when the LHS or RHS is the
+//   defining conv2d op.
+// TODO: support other pattern. e.g. tf.DepthwiseConv2DNative
+
+LogicalResult ConvertTFBiasAddOp::matchAndRewrite(
+    Operation* op, PatternRewriter& rewriter) const {
+  auto tf_biasadd_op = cast<TF::BiasAddOp>(op);
+  auto output_type =
+      dyn_cast<RankedTensorType>(tf_biasadd_op.getResult().getType());
+
+  if (!output_type) {
+    return rewriter.notifyMatchFailure(op, "output not a ranked tensor");
+  }
+
+  auto value = tf_biasadd_op.getValue();
+  auto bias = tf_biasadd_op.getBias();
+  auto bias_shape = mlir::cast<RankedTensorType>(bias.getType()).getShape();
+  if (bias_shape.size() != 1) {
+    return rewriter.notifyMatchFailure(op, "bias tensor must be rank 1");
+  }
+
+  if (TF::Conv2DOp tf_conv2d_op =
+          llvm::dyn_cast_if_present<TF::Conv2DOp>(value.getDefiningOp())) {
+    // Sanity check to confirm rhs() has the expected shape of bias
+    auto filter_shape =
+        mlir::cast<RankedTensorType>(tf_conv2d_op.getFilter().getType())
+            .getShape();
+
+    // Assume the filter shape is [H, W, I, O]
+    if (filter_shape.back() != bias_shape.back()) {
+      return rewriter.notifyMatchFailure(
+          op, "bias dimension must match filter output channels");
+    }
+
+    auto result = convertTFConv2DCommon(
+        rewriter, op, output_type, tf_conv2d_op.getInput(),
+        tf_conv2d_op.getFilter(), bias, tf_conv2d_op.getStrides(),
+        tf_conv2d_op.getDilations(), tf_conv2d_op.getExplicitPaddings(),
+        tf_conv2d_op.getPadding(), tf_conv2d_op.getDataFormat());
+
+    if (!result) return failure();
+
+    rewriter.replaceOp(op, {result.value()});
+
+    return success();
+  }
+
+  if (TF::Conv3DOp tf_conv3d_op =
+          llvm::dyn_cast_if_present<TF::Conv3DOp>(value.getDefiningOp())) {
+    // Sanity check to confirm rhs() has the expected shape of bias
+    auto filter_shape =
+        mlir::cast<RankedTensorType>(tf_conv3d_op.getFilter().getType())
+            .getShape();
+
+    // Assume the filter shape is [D, H, W, I, O]
+    if (filter_shape.back() != bias_shape.back()) {
+      return rewriter.notifyMatchFailure(
+          op, "bias dimension must match filter output channels");
+    }
+
+    std::optional<Value> result = convertTFConv3DCommon(
+        rewriter, op, output_type, tf_conv3d_op.getInput(),
+        tf_conv3d_op.getFilter(), bias, tf_conv3d_op.getStrides(),
+        tf_conv3d_op.getDilations(), tf_conv3d_op.getPadding(),
+        tf_conv3d_op.getDataFormat());
+
+    if (!result) return failure();
+
+    rewriter.replaceOp(op, {result.value()});
+
+    return success();
+  }
+
+  return failure();
+}
+
+void FuseBiasTF::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  auto* ctx = &getContext();
+  auto func = getOperation();
+
+  // Add the generated patterns to the list.
+  patterns.add<ConvertTFBiasAddOp>(ctx);
+  (void)applyPatternsGreedily(func, std::move(patterns));
+}
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createFuseBiasTFPass() {
+  return std::make_unique<FuseBiasTF>();
+}
+
+}  // namespace tosa
+
+}  // namespace mlir

tensorflow/compiler/mlir/tosa/transforms/legalize_common.h

@@ -0,0 +1,337 @@
+/* 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_TOSA_TRANSFORMS_LEGALIZE_COMMON_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_COMMON_H_
+
+#include <cstdint>
+#include <optional>
+
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"     // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+
+// This file contains legalizations common to mapping both TensorFlow and
+// TensorFlow Lite to TOSA.
+//
+// Conversion functions return None on a failure or result value on success.
+// Callers must check and return a LogicalResult failure on nullptr.
+//
+// For these functions, the framework-specific operands/attributes/defaults
+// are already extracted and placed in a common form for lowering.
+
+namespace mlir {
+namespace tosa {
+
+// Lowers the Pack operator to TOSA.
+std::optional<Value> convertPackOp(PatternRewriter& rewriter, Operation* op,
+                                   Value result_value,
+                                   SmallVectorImpl<Value>& inputs,
+                                   int32_t axis);
+
+// Lowers the Unpack operator to TOSA.
+std::optional<SmallVector<Value>> convertUnpackOp(PatternRewriter& rewriter,
+                                                  Operation* op,
+                                                  Value input_value,
+                                                  int32_t axis);
+
+// Lowers the Select operator to TOSA.
+std::optional<Value> convertSelectOp(PatternRewriter& rewriter, Operation* op,
+                                     Value result_value, Value condition_value,
+                                     Value x_value, Value y_value);
+
+// Lowers the ZerosLike operator to TOSA by creating a constant
+// of the desired type and shape.
+std::optional<Value> convertZerosLikeOp(PatternRewriter& rewriter,
+                                        Operation* op, Value result,
+                                        Value input);
+
+// Lowers the Mul operator to TOSA.  For quantized types, this requires
+// inserting rescale operators before and after the operation.
+std::optional<Value> convertMultiplyOp(PatternRewriter& rewriter, Operation* op,
+                                       Value output_val, Value input_lhs_val,
+                                       Value input_rhs_val);
+
+// Lowers the SquaredDifference operator to TOSA.
+std::optional<Value> convertSquaredDifferenceOp(PatternRewriter& rewriter,
+                                                Operation* op, Value result,
+                                                Value x, Value y);
+
+// Lowers the Round operator to TOSA.
+std::optional<Value> convertRoundOp(PatternRewriter& rewriter, Operation* op,
+                                    Value result, Value input);
+
+// Lowers ConcatV2 to TOSA.
+std::optional<Value> convertConcatV2Op(PatternRewriter& rewriter, Operation* op,
+                                       ShapedType result_type,
+                                       SmallVectorImpl<Value>& values,
+                                       int32_t axis);
+
+// Lowers SpaceToBatchND to TOSA.
+std::optional<Value> convertSpaceToBatchNDOp(PatternRewriter& rewriter,
+                                             Operation* op, Value result_value,
+                                             Value input_value,
+                                             Value block_shape_value,
+                                             Value paddings_value);
+
+// Lowers BatchToSpaceND to TOSA.
+std::optional<Value> convertBatchToSpaceNDOp(PatternRewriter& rewriter,
+                                             Operation* op, Value result_value,
+                                             Value input_value,
+                                             Value block_shape_value,
+                                             Value crops_value);
+
+// Lowers ExpandDims to TOSA.
+std::optional<Value> convertExpandDimsOp(PatternRewriter& rewriter,
+                                         Operation* op, Value result_value,
+                                         Value input_value, Value dim_value);
+
+// Lowers Squeeze to TOSA.
+std::optional<Value> convertSqueezeOp(PatternRewriter& rewriter, Operation* op,
+                                      Value result_value, Value input_value,
+                                      SmallVectorImpl<int32_t>& squeeze_dims);
+
+// Lowers ELU to a sequence of TOSA ops.
+std::optional<Value> convertEluOp(PatternRewriter& rewriter, Operation* op,
+                                  Value result_value, Value features_value);
+
+// Lowers Softmax to a sequence of TOSA ops.
+std::optional<Value> convertSoftmaxOp(PatternRewriter& rewriter, Operation* op,
+                                      Value result_value, Value logits_value,
+                                      double beta);
+
+// Lowers LogSoftmax to a sequence of TOSA ops.
+std::optional<Value> convertLogSoftmaxOp(PatternRewriter& rewriter,
+                                         Operation* op, Value result_value,
+                                         Value logits_value);
+
+// Lowers SpaceToDepth to a sequence of TOSA ops.  Supports NHWC.
+std::optional<Value> convertSpaceToDepthOp(PatternRewriter& rewriter,
+                                           Operation* op, Value result_value,
+                                           Value input_value,
+                                           IntegerAttr block_size_attr,
+                                           StringAttr data_format);
+
+// Lowers DepthToSpace to a sequence of TOSA ops.  Supports NHWC.
+std::optional<Value> convertDepthToSpaceOp(PatternRewriter& rewriter,
+                                           Operation* op, Value result_value,
+                                           Value input_value,
+                                           IntegerAttr block_size_attr,
+                                           StringAttr data_format);
+
+// Lowers Split to a sequence of TOSA ops.
+std::optional<SmallVector<Value>> convertSplitOp(
+    PatternRewriter& rewriter, Operation* op, Value result_value,
+    Value input_value, int32_t num_split, int32_t axis);
+
+// Lowers SplitV to a sequence of TOSA ops.
+std::optional<SmallVector<Value>> convertSplitVOp(
+    PatternRewriter& rewriter, Operation* op, Value result_value,
+    Value input_value, SmallVectorImpl<int32_t>& size_split, int32_t axis);
+
+// Lowers StridedSlice to a sequence of TOSA ops.
+std::optional<Value> convertStridedSliceOp(
+    PatternRewriter& rewriter, Operation* op, Value result_value,
+    Value input_value, Value begin_value, Value end_value, Value strides_value,
+    int32_t begin_mask, int32_t end_mask, int32_t ellipsis_mask,
+    int32_t new_axis_mask, int32_t shrink_axis_mask);
+
+// Helper function to perform division with floor rounding mode (rounding result
+// down) for integer type inputs.
+Value floorIntDiv(PatternRewriter& rewriter, Operation* op, ShapedType outType,
+                  Value lhs, Value rhs);
+
+// Lowers FloorDiv to a sequence of TOSA operators.
+std::optional<Value> convertFloorDivOp(PatternRewriter& rewriter, Operation* op,
+                                       Value result_value, Value lhs_value,
+                                       Value rhs_value);
+
+// Lowers FloorMod to a sequence of TOSA operators.
+std::optional<Value> convertFloorModOp(PatternRewriter& rewriter, Operation* op,
+                                       Value result_value, Value lhs_value,
+                                       Value rhs_value);
+
+// Lowers FusedActivation to a sequence of TOSA ops.
+std::optional<Value> convertFusedActivation(PatternRewriter& rewriter,
+                                            Operation* op, Value input_value,
+                                            StringAttr fused_activation_fn);
+
+// Helper function for implementing quantized divide by power-of-two in TOSA
+// ops.
+std::optional<Value> convertRoundingDivideByPOT(PatternRewriter& rewriter,
+                                                Operation* op,
+                                                Value input_value,
+                                                Value rshift_value);
+
+// Lowers ReduceAll to a sequence of TOSA ops.
+std::optional<Value> convertReduceAllOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims);
+
+// Lowers ReduceAny to a sequence of TOSA ops.
+std::optional<Value> convertReduceAnyOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims);
+
+// Lowers ReduceMin to a sequence of TOSA ops.
+std::optional<Value> convertReduceMinOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims,
+                                        std::optional<tosa::NanPropagationMode> nan_mode = tosa::NanPropagationMode::PROPAGATE);
+
+// Lowers ReduceMax to a sequence of TOSA ops.
+std::optional<Value> convertReduceMaxOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims,
+                                        std::optional<tosa::NanPropagationMode> nan_mode = tosa::NanPropagationMode::PROPAGATE);
+
+// Lowers ReduceProd to a sequence of TOSA ops.
+std::optional<Value> convertReduceProdOp(PatternRewriter& rewriter,
+                                         Operation* op,
+                                         RankedTensorType output_type,
+                                         Value input_value,
+                                         ElementsAttr axes_elems,
+                                         bool keep_dims);
+
+// Lowers ReduceSum to a sequence of TOSA ops.
+std::optional<Value> convertReduceSumOp(PatternRewriter& rewriter,
+                                        Operation* op,
+                                        RankedTensorType output_type,
+                                        Value input_value,
+                                        ElementsAttr axes_elems,
+                                        bool keep_dims);
+
+// Lowers ReduceMean to a sequence of TOSA ops.
+std::optional<Value> convertReduceMeanOp(PatternRewriter& rewriter,
+                                         Operation* op,
+                                         RankedTensorType output_type,
+                                         Value input_value,
+                                         ElementsAttr axes_elem,
+                                         bool keep_dims);
+
+// Lowers ResizeBilinear and ResizeNearestNeighbor to TOSA resize.
+std::optional<Value> convertResizeOp(PatternRewriter& rewriter, Operation* op,
+                                     RankedTensorType output_type,
+                                     Value input_value, tosa::ResizeMode mode,
+                                     bool align_corners,
+                                     bool half_pixel_centers);
+
+// Lowers Quantize to a sequence of TOSA quantization ops.
+std::optional<Value> convertQuantizeOp(PatternRewriter& rewriter, Operation* op,
+                                       ShapedType output_type,
+                                       Value input_value, double scale,
+                                       int64_t zeropoint);
+
+// Lowers Dequantize to a sequence of TOSA dequantization ops.
+std::optional<Value> convertDequantizeOp(PatternRewriter& rewriter,
+                                         Operation* op, ShapedType output_type,
+                                         Value input_value,
+                                         ArrayRef<float> scale,
+                                         ArrayRef<float> zeropoint,
+                                         int64_t dim);
+
+// Lowers FakeQuant to a sequence of TOSA quantization ops.
+std::optional<Value> convertFakeQuantOp(PatternRewriter& rewriter,
+                                        Operation* op, ShapedType output_type,
+                                        Value input_value, double min,
+                                        double max, int64_t num_bits,
+                                        bool narrow_range);
+
+// Align to TF_MirrorPadOp::mode and TFL_MirrorPadOp::mode
+enum class TFTFLMirrorPaddingType : uint32_t {
+  REFLECT = 0,
+  SYMMETRIC = 1,
+};
+
+std::optional<Value> convertMirrorPadCommon(PatternRewriter& rewriter,
+                                            Operation* op,
+                                            RankedTensorType output_type,
+                                            Value input, Value pad,
+                                            TFTFLMirrorPaddingType mode);
+
+// Lowers TensorFlow Conv2D to a sequence of TOSA quantization ops.
+std::optional<Value> convertTFConv2DCommon(
+    PatternRewriter& rewriter, Operation* op, RankedTensorType output_type,
+    Value input, Value filter, Value bias, ArrayAttr strides_attr,
+    ArrayAttr dilations_attr, ArrayAttr explicit_padding_attr,
+    StringRef padding_ref, StringRef data_format_ref);
+
+// Lowers TensorFlow and TensorFlow Lite Conv3D to a sequence of TOSA
+// quantization ops.
+std::optional<Value> convertConv3DCommon(
+    PatternRewriter& rewriter, Operation* op, ShapedType output_type,
+    Value input, Value filter, Value bias, DenseI64ArrayAttr pads,
+    DenseI64ArrayAttr strides, DenseI64ArrayAttr dilations, TypeAttr acc_type,
+    StringRef data_format_ref);
+
+// Preprocess TensorFlow Conv3D attributes prior to calling
+// `convertConv3DCommon`
+std::optional<Value> convertTFConv3DCommon(
+    PatternRewriter& rewriter, Operation* op, ShapedType output_type,
+    Value input, Value filter, Value bias, ArrayAttr strides_attr,
+    ArrayAttr dilations_attr, StringRef padding_ref, StringRef data_format_ref);
+
+// Lowers Gather operator to a sequence of TOSA ops.
+std::optional<Value> convertGatherOp(PatternRewriter& rewriter, Operation* op,
+                                     Value params_value, Value indices_value,
+                                     int32_t batch_dims, int32_t axis,
+                                     bool tosaOnly = true);
+
+// Lowers GatherNd operator to a sequence of TOSA ops.
+std::optional<Value> convertGatherNdOp(PatternRewriter& rewriter, Operation* op,
+                                       Value result_value, Value params_value,
+                                       Value indices_value);
+
+// Lowers ScatterNd operator to a sequence of TOSA ops.
+std::optional<Value> convertScatterNdOp(PatternRewriter& rewriter,
+                                        Operation* op, Value result_value,
+                                        Value indices_value,
+                                        Value updates_value, Value shape_value);
+
+// Lowers OneHot operator to a sequence of TOSA ops.
+std::optional<Value> convertOneHotOp(PatternRewriter& rewriter, Operation* op,
+                                     Value result_value, Value indices_value,
+                                     Value on_value, Value off_value,
+                                     int32_t depth, int32_t axis);
+
+// Lowers cast operator to a sequence of TOSA ops.
+std::optional<Value> convertCastOp(PatternRewriter& rewriter, Operation* op,
+                                   Value input, RankedTensorType output_type);
+
+// Lowers Sign operator to a sequence of TOSA ops.
+std::optional<Value> convertSignOp(PatternRewriter& rewriter, Operation* op,
+                                   Value input, RankedTensorType output_type);
+
+// Lowers BroadcastTo operator to a sequence of TOSA ops.
+std::optional<Value> convertBroadcastToOp(PatternRewriter& rewriter,
+                                          Operation* op, Value input,
+                                          Value shape);
+
+};  // namespace tosa
+};  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_COMMON_H_

tensorflow/compiler/mlir/tosa/transforms/legalize_tf_tfl.cc

@@ -0,0 +1,68 @@
+/* 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.
+==============================================================================*/
+
+// Legalize TensorFlow and TensorFlow Lite to TOSA
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+namespace mlir {
+namespace tosa {
+
+namespace {
+
+#define GEN_PASS_DEF_TOSALEGALIZETFTFLPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+// Performs lowering to TOSA dialect
+class LegalizeTFTFL : public impl::TosaLegalizeTFTFLPassBase<LegalizeTFTFL> {
+ public:
+  explicit LegalizeTFTFL() = default;
+  void runOnOperation() override;
+};
+
+void LegalizeTFTFL::runOnOperation() {
+  MLIRContext* ctx = &getContext();
+  RewritePatternSet patterns(ctx);
+  populateLegalizeTFPatterns(ctx, patterns);
+  populateLegalizeTFLPatterns(ctx, patterns);
+
+  func::FuncOp func = getOperation();
+  if (ApplyPatternsWithShapeResolution(func, std::move(patterns)).failed()) {
+    signalPassFailure();
+  }
+}
+
+}  // anonymous namespace
+
+// Creates an instance of the TensorFlow Lite dialect LegalizeTFL pass.
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeTFTFLPass() {
+  return std::make_unique<LegalizeTFTFL>();
+}
+
+}  // namespace tosa
+}  // namespace mlir

tensorflow/compiler/mlir/tosa/transforms/legalize_tfl_stateful.cc

@@ -0,0 +1,189 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Legalize TensorFlow Lite StatefulOps to TOSA
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-legalize-tfl-stateful"
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSALEGALIZETFLSTATEFULPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+// Performs lowering tfl stateful operators to TOSA
+class TosaLegalizeTFLStateful
+    : public impl::TosaLegalizeTFLStatefulPassBase<TosaLegalizeTFLStateful> {
+ public:
+  explicit TosaLegalizeTFLStateful() = default;
+  void runOnOperation() override;
+};
+
+void TosaLegalizeTFLStateful::runOnOperation() {
+  auto moduleOp = getOperation();
+  mlir::OpBuilder builder(moduleOp.getBodyRegion());
+
+  DenseMap<StringRef, func::FuncOp> symNameToFunction;
+  for (auto func : moduleOp.getOps<func::FuncOp>()) {
+    symNameToFunction[func.getSymName()] = func;
+  }
+
+  llvm::SmallVector<mlir::TFL::VarHandleOp, 6> handleOps;
+  llvm::SmallVector<mlir::TFL::AssignVariableOp, 6> assignOps;
+  llvm::SmallVector<mlir::TFL::ReadVariableOp, 6> readOps;
+  SmallVector<mlir::TFL::CallOnceOp> callOnceOps;
+  DenseMap<StringRef, mlir::tosa::VariableOp> symbolRefMap;
+
+  for (auto it : symNameToFunction) {
+    auto func = std::get<1>(it);
+    // We also want to grab the list of operations to replace.
+    for (auto& op : func.getOps()) {
+      if (auto handle = dyn_cast<mlir::TFL::VarHandleOp>(op))
+        handleOps.push_back(handle);
+      if (auto assign = dyn_cast<mlir::TFL::AssignVariableOp>(op))
+        assignOps.push_back(assign);
+      if (auto read = dyn_cast<mlir::TFL::ReadVariableOp>(op))
+        readOps.push_back(read);
+    }
+  }
+
+  for (auto func : moduleOp.getOps<func::FuncOp>()) {
+    for (auto init : func.getOps<mlir::TFL::CallOnceOp>()) {
+      callOnceOps.push_back(init);
+    }
+  }
+
+  // Look through the initialization functions and find the assigned values
+  // for each handle, save out the constant value.
+  for (auto init : callOnceOps) {
+    auto findInitFunc =
+        symNameToFunction.find(init.getSessionInitFunctionAttr());
+    if (findInitFunc == symNameToFunction.end()) {
+      init.emitError("unable to find initialization function: ");
+      continue;
+    }
+    func::FuncOp initFunc = std::get<1>(*findInitFunc);
+    for (auto assign : initFunc.getOps<mlir::TFL::AssignVariableOp>()) {
+      // 1. var_handle part
+      auto handle = dyn_cast<mlir::TFL::VarHandleOp>(
+          assign.getResourceId().getDefiningOp());
+      if (!handle) continue;
+
+      // 2. pseudo_const part
+      DenseElementsAttr constant;
+      if (!matchPattern(assign.getValue(), m_Constant(&constant))) {
+        // Quantized types we can not use the m_Constant matcher.
+        if (auto constOp = dyn_cast<mlir::TFL::QConstOp>(
+                assign.getValue().getDefiningOp())) {
+          constant = cast<DenseElementsAttr>(constOp.getValue());
+        }
+      }
+      if (!constant) continue;
+
+      // Create TOSA VariableOps
+      auto name = handle.getSharedName();
+      auto global = builder.create<mlir::tosa::VariableOp>(
+          handle.getLoc(), name, constant.getType(), constant);
+      symbolRefMap[name] = global;
+    }
+  }
+  // TF::CallOnceOps are no longer needed as we have already extracted their
+  // state.
+  for (auto op : callOnceOps) op.erase();
+
+  // Replace the assign ops with a tosa store operation.
+  for (auto assign : assignOps) {
+    auto handle = dyn_cast<mlir::TFL::VarHandleOp>(
+        assign.getResourceId().getDefiningOp());
+    if (!handle) continue;
+
+    Value value = assign.getValue();
+    auto globalOpIt = symbolRefMap.find(handle.getSharedName());
+    if (globalOpIt == symbolRefMap.end()) {
+      assign->emitError(
+          "unable to find corresponding TosaOp for op's VarHandle");
+      continue;
+    }
+    auto globalOp = std::get<1>(*globalOpIt);
+
+    builder.setInsertionPoint(assign);
+    auto variableType = mlir::tosa::getVariableType(globalOp);
+    if (variableType != value.getType()) {
+      value = builder
+                  .create<UnrealizedConversionCastOp>(assign.getLoc(),
+                                                      variableType, value)
+                  .getResult(0);
+    }
+
+    builder.create<mlir::tosa::VariableWriteOp>(
+        assign.getLoc(), llvm::StringRef(globalOp.getName()), value);
+    assign.erase();
+  }
+
+  for (auto read : readOps) {
+    auto handle =
+        dyn_cast<mlir::TFL::VarHandleOp>(read.getResourceId().getDefiningOp());
+    if (!handle) continue;
+
+    auto globalOpIt = symbolRefMap.find(handle.getSharedName());
+    if (globalOpIt == symbolRefMap.end()) continue;
+    auto globalOp = std::get<1>(*globalOpIt);
+
+    builder.setInsertionPoint(read);
+
+    Value load = builder.create<mlir::tosa::VariableReadOp>(
+        read.getLoc(), mlir::tosa::getVariableType(globalOp),
+        llvm::StringRef(globalOp.getName()));
+
+    if (read.getType() != load.getType()) {
+      load = builder
+                 .create<UnrealizedConversionCastOp>(read.getLoc(),
+                                                     read.getType(), load)
+                 .getResult(0);
+    }
+    read.getResult().replaceAllUsesWith(load);
+    read.erase();
+  }
+
+  for (auto handle : handleOps) {
+    if (handle.getResult().use_empty()) {
+      handle.erase();
+    }
+  }
+}
+
+}  // namespace
+
+// Creates an instance of the TensorFlow Lite dialect LegalizeTFLStateful pass.
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeTFLStatefulPass() {
+  return std::make_unique<TosaLegalizeTFLStateful>();
+}
+
+}  // namespace tosa
+}  // namespace mlir

tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h

@@ -0,0 +1,308 @@
+/* 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_TOSA_TRANSFORMS_LEGALIZE_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_UTILS_H_
+
+#include <cfloat>
+#include <climits>
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <iterator>
+#include <numeric>
+#include <optional>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"             // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"            // from @llvm-project
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"             // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"       // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"                // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"                     // from @llvm-project
+#include "mlir/IR/ImplicitLocOpBuilder.h"             // from @llvm-project
+#include "mlir/IR/PatternMatch.h"                     // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"     // from @llvm-project
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"     // from @llvm-project
+#include "mlir/Support/LLVM.h"                        // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h"
+#include "tensorflow/core/framework/kernel_shape_util.h"
+#include "tensorflow/core/kernels/conv_grad_shape_utils.h"
+#include "tensorflow/core/util/padding.h"
+#include "tensorflow/core/util/tensor_format.h"
+
+namespace mlir {
+namespace tosa {
+
+// returns acc_type attribute for Conv ops with specified input/output element
+// types
+mlir::TypeAttr getConvAccTypeAttr(PatternRewriter& rewriter,
+                                  mlir::Type input_etype,
+                                  mlir::Type output_etype);
+
+LogicalResult getDynamicDims(PatternRewriter& rewriter, Value value,
+                             llvm::SmallVector<Value>& dims);
+
+std::optional<Value> buildReshapeWithDynamicDims(PatternRewriter& rewriter,
+                                                 Operation* op,
+                                                 Value input_value,
+                                                 ShapedType output_type,
+                                                 llvm::ArrayRef<Value> dims);
+
+// Create a TOSA rescale op from TFLite scaling multiplier, scaling shift, zero
+// points and rounding mode
+Value buildRescale(PatternRewriter& rewriter, Operation* op,
+                   ShapedType output_type, Value input_val,
+                   int32_t scale_multiplier, int32_t scale_shit,
+                   int64_t input_zp, int64_t output_zp, tosa::RoundingMode rounding_mode,
+                   bool scale32);
+
+// Create a TOSA rescale op from TFLite scaling, zero points and rounding mode
+Value buildRescale(PatternRewriter& rewriter, Operation* op,
+                   ShapedType output_type, Value input_val, double scale,
+                   int64_t input_zp, int64_t output_zp, tosa::RoundingMode rounding_mode,
+                   bool scale32);
+
+// Removes the zero point and cast to int32, no need to handle roundings modes
+Value removeZeroPointAndCastToInt32(PatternRewriter& rewriter, Operation* op,
+                                    Value input_val, int64_t input_zp);
+
+// Creates TOSA rescale op with int32 output
+Value buildRescaleToInt32(PatternRewriter& rewriter, Operation* op,
+                          Value input_val, int32_t input_scale_multiplier,
+                          int32_t input_scale_shift, int64_t input_zp);
+
+// Creates TOSA rescale op with int32 output
+Value buildRescaleToInt32(PatternRewriter& rewriter, Operation* op,
+                          Value input_val, double input_scale,
+                          int64_t input_zp);
+
+// Creates TOSA rescale op with int32 input
+Value buildRescaleFromInt32(PatternRewriter& rewriter, Operation* op,
+                            ShapedType output_type, Value input_val,
+                            double output_scale, int64_t output_zp);
+
+// Creates a TOSA rescale op based on conv2d parameters.
+Value buildRescaleOpConvOutput(PatternRewriter& rewriter, Operation* op,
+                               Value conv_val, ShapedType input_type,
+                               ShapedType weight_type, ShapedType output_type);
+
+// Create a 8-bit TOSA TABLE constant tensor
+Value getTosaConst8bitTable(PatternRewriter& rewriter, Operation* op,
+                            float input_scale, int32_t input_zp,
+                            float output_scale, int32_t output_zp,
+                            std::function<float(float)> func);
+
+// Create a 16-bit TOSA TABLE constant tensor
+// A float should be used by default for FloatT except if a double is required
+// for backward compatibility
+template <typename FloatT>
+Value getTosaConst16bitTable(PatternRewriter& rewriter, Operation* op,
+                             FloatT input_scale, int32_t input_zp,
+                             FloatT output_scale, int32_t output_zp,
+                             std::function<FloatT(FloatT)> func);
+
+// Create a 32-bit TOSA TABLE for Softmax Exp
+void getTosaConst32bitSoftmaxExpTable(PatternRewriter& rewriter, Operation* op,
+                                      double beta, double input_scale,
+                                      Value& first_const, Value& second_const,
+                                      Value& third_const, Value& fourth_const);
+
+// Create 8 bit TOSA TABLE constant tensor for the RSqrt operator
+Value getTosaConstRsqrt8bitTable(PatternRewriter& rewriter, Operation* op,
+                                 float input_scale, int32_t input_zp,
+                                 float output_scale, int32_t output_zp);
+
+// Create an 8-bit TOSA Table constant tensor for the HardSwish operator
+Value getTosaConstHardSwish8bitTable(PatternRewriter& rewriter, Operation* op,
+                                     float input_scale, int32_t input_zp,
+                                     float output_scale, int32_t output_zp);
+
+// Create a 32-bit float constant operator from a float
+Value getTosaConstTensorSingleF32(PatternRewriter& rewriter, Operation* op,
+                                  float val, int rank);
+
+// Create a 32-bit integer constant operator from an int of specified rank
+Value getTosaConstTensorSingleI32(PatternRewriter& rewriter, Operation* op,
+                                  int32_t val, int rank);
+
+// Create an expected bitwidth integer constant operator based on the type
+// parameter, of specified rank
+Value getTosaConstTensorScalarInt(ImplicitLocOpBuilder& builder, Type type,
+                                  int64_t val, int rank);
+
+// Populate a int32_t vector from a val tensor
+// return failure if val is not a constant value
+// return success otherwise
+LogicalResult getVectorFromValue32(Value val, SmallVectorImpl<int32_t>& vec);
+
+// Populate a int64_t vector from a val tensor
+// return failure if val is not a constant value
+// return success otherwise
+LogicalResult getVectorFromValue64(Value val, SmallVectorImpl<int64_t>& vec);
+
+// Calculates the TOSA padding values based on TF operators padded with
+// SAME/VALID.
+bool getPaddingValuesFromPadType(tensorflow::Padding tf_pad,
+                                 tensorflow::TensorFormat data_format_tf,
+                                 uint32_t first_filter_spatial_dim,
+                                 ShapedType input_type, ShapedType filter_type,
+                                 DenseI64ArrayAttr strides,
+                                 DenseI64ArrayAttr dilations,
+                                 PatternRewriter& rewriter,
+                                 DenseI64ArrayAttr& explicit_pad);
+
+// Calculates the TOSA padding values for explicit-padded TF operators.
+DenseI64ArrayAttr getPaddingValuesFromExplicitPadAttr(
+    ArrayAttr explicit_pad, tensorflow::TensorFormat data_format_tf,
+    PatternRewriter& rewriter);
+
+// Calculates the TOSA padding values for transposeConv2d
+bool getTransposeConv2dPaddingValues(
+    tensorflow::Padding tf_pad, tensorflow::TensorFormat data_format_tf,
+    uint32_t first_filter_spatial_dim, ShapedType input_type,
+    ShapedType filter_type, ShapedType output_type, DenseI64ArrayAttr strides,
+    PatternRewriter& rewriter, DenseI64ArrayAttr& explicit_pad);
+
+// Templated function to create a constant op for given type and shape.
+// T: storage C type.
+// Default template creates a constant tensor in T.
+// To create INT48 TOSA constant, need to pass in llvm::APInt instead.
+template <typename T>
+std::optional<Value> getConstTensor(PatternRewriter& rewriter, Operation* op,
+                                    ArrayRef<T> vec, ArrayRef<int64_t> shape);
+
+// For each spatial dimension, return the remainder of the output size
+// calculation: (I - 1 + pad - (K - 1) * dilation) % stride.
+llvm::SmallVector<int64_t> getOutputSpatialSizeRemainder(
+    tensorflow::TensorFormat data_format_tf, ShapedType input_type,
+    DenseI64ArrayAttr kernel_size, DenseI64ArrayAttr pads,
+    DenseI64ArrayAttr strides, DenseI64ArrayAttr dilations);
+
+// The TOSA specification requires the full size of the input to be used during
+// the convolution (the output size remainder calculation must be 0). If input
+// slicing is necessary to satisfy the condition, return a tosa::SliceOp,
+// otherwise return input_val.
+Value getInputSlicedToItsUsedSize(PatternRewriter& rewriter, Operation* op,
+                                  tensorflow::TensorFormat data_format_tf,
+                                  ShapedType input_type, Value input_val,
+                                  DenseI64ArrayAttr kernel_size,
+                                  DenseI64ArrayAttr pads,
+                                  DenseI64ArrayAttr strides,
+                                  DenseI64ArrayAttr dilations);
+
+// Check if scale32 mode is used for given output_element_type
+bool isScale32(mlir::quant::UniformQuantizedType output_element_type);
+
+// Checks if the multi-dimensional indices supplied by a constant tensor
+// are unique. This is a useful check for legalizations to tosa.scatter
+// which requires indices are unique, while in TF/TFLite they may be
+// non-unique.
+bool checkUniqueConstantScatterIndices(ShapedType indices_type,
+                                       ShapedType result_type,
+                                       ElementsAttr const_data);
+
+// Applies a set of patterns greedily to the specified function, then applies
+// a cleanup to guarantee the function contract and constants are valid. This
+// means patterns can performed shape inference while not altering immutable
+// types.
+LogicalResult ApplyPatternsWithShapeResolution(
+    func::FuncOp func, const FrozenRewritePatternSet& patterns);
+
+// Creates a TOSA operation and performs shape inference on the individual
+// op. This allows shape inference during the TFLite to TOSA lowering.
+template <typename TosaOp, typename... Args>
+TosaOp CreateOpAndInfer(ImplicitLocOpBuilder& builder, Type result_ty,
+                        Args&&... args) {
+  return CreateOpAndInferShape<TosaOp>(builder, result_ty, args...);
+}
+
+template <typename TosaOp, typename... Args>
+TosaOp CreateOpAndInfer(PatternRewriter& rewriter, Location loc, Type result_ty,
+                        Args&&... args) {
+  ImplicitLocOpBuilder builder(loc, rewriter);
+  return CreateOpAndInfer<TosaOp>(builder, result_ty, args...);
+}
+
+template <typename TosaOp, typename... Args>
+void CreateReplaceOpAndInfer(PatternRewriter& rewriter, Operation* op,
+                             Type result_ty, Args&&... args) {
+  auto result =
+      CreateOpAndInfer<TosaOp>(rewriter, op->getLoc(), result_ty, args...);
+  rewriter.replaceOp(op, result->getResults());
+}
+
+// Nan propagation mode is only applied to maximum and mininum.
+template <typename TOSA_OP>
+LogicalResult ConvertBinaryOp(Operation* op, PatternRewriter& rewriter,
+                              std::optional<tosa::NanPropagationMode> nan_mode = std::nullopt) {
+  TensorType output_type = dyn_cast<TensorType>(op->getResults()[0].getType());
+  if (!output_type) return failure();
+
+  Value x = op->getOperands()[0];
+  Value y = op->getOperands()[1];
+
+  RankedTensorType x_type = dyn_cast<RankedTensorType>(x.getType());
+  RankedTensorType y_type = dyn_cast<RankedTensorType>(y.getType());
+  if (!x_type || !y_type) return failure();
+
+  if constexpr (std::is_same_v<tosa::ReduceMaxOp, TOSA_OP> ||
+                std::is_same_v<tosa::ReduceMinOp, TOSA_OP>) {
+    if (!nan_mode) {
+      (void)rewriter.notifyMatchFailure(op, "invalid NaN mode: must be either 'PROPAGATE' or 'IGNORE'");
+      return failure();
+    }
+    const auto nan_mode_attr = tosa::NanPropagationModeAttr::get(
+        rewriter.getContext(), *nan_mode);
+    CreateReplaceOpAndInfer<TOSA_OP>(rewriter, op, output_type, x, y, nan_mode);
+  } else
+    CreateReplaceOpAndInfer<TOSA_OP>(rewriter, op, output_type, x, y);
+
+  return success();
+}
+
+// Create TOSA mul ops and infer the shape of the operation. During the
+// creation, fill in the shift value if applied.
+tosa::MulOp CreateMulOpAndInfer(PatternRewriter& rewriter, Operation* op,
+                                Type result_ty, Value input1, Value input2,
+                                int8_t shift = 0);
+
+void TrimQuantizedIntegerRangeMin(mlir::quant::UniformQuantizedType dtype,
+                                  int64_t& val_min);
+
+void TrimQuantizedIntegerRangeMax(mlir::quant::UniformQuantizedType dtype,
+                                  int64_t& val_max);
+
+void TrimQuantizedIntegerRange(mlir::quant::UniformQuantizedType dtype,
+                               int64_t& val_min, int64_t& val_max);
+
+inline bool IsTFLDoubleRoundingMode() {
+#if TFLITE_SINGLE_ROUNDING
+  return false;
+#else
+  return true;
+#endif  // TFLITE_SINGLE_ROUNDING
+}
+
+Value reshapeScalarTo1D(PatternRewriter& rewriter, Location loc, Value value);
+
+LogicalResult broadcastLowRankTensor(PatternRewriter &rewriter, Operation* op,
+                                     Value &input1, Value &input2);
+
+}  // namespace tosa
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_LEGALIZE_UTILS_H_

tensorflow/compiler/mlir/tosa/transforms/lower_complex_types.cc

@@ -0,0 +1,171 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// TOSA has no notion of complex datatypes, it represents a single complex
+// input tensor using two floating point input tensors corresponding to the
+// "real" and "imag" parts of the complex input.
+//
+// To maintain the correct mapping of these tensors during the legalization
+// from TFL to TOSA, a complex tensor of shape [x, ..., y] is converted to
+// a single floating point tensor of shape [x, ..., y, 2] where each resulting
+// pair of values can be used to represent a complex value, which ensures a
+// 1:1 mapping between TFL and TOSA input/output tensors. In legalization,
+// "unrealized_conversion_cast" operations are inserted to express this
+// conversion.
+//
+// This pass removes complex tensors from the graph by rewriting them using
+// the above [x, ..., y, 2] floating point format. Consequently, it removes
+// any remaining "unrealized_conversion_cast" operations and ensures the
+// resulting graph is free of illegal complex tensors.
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-lower-complex-types"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSALOWERCOMPLEXTYPESPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+class LowerComplexTypes
+    : public impl::TosaLowerComplexTypesPassBase<LowerComplexTypes> {
+ public:
+  explicit LowerComplexTypes() = default;
+  void runOnOperation() override;
+};
+
+class ComplexTypeConverter : public TypeConverter {
+ public:
+  static Type convertTensor(RankedTensorType type) {
+    if (auto elementType = dyn_cast<ComplexType>(type.getElementType())) {
+      llvm::SmallVector<int64_t> newShape;
+      for (auto dim : type.getShape()) {
+        newShape.push_back(dim);
+      }
+      newShape.push_back(2);
+      return RankedTensorType::get(newShape, elementType.getElementType());
+    }
+    return type;
+  }
+
+  explicit ComplexTypeConverter() { addConversion(convertTensor); }
+};
+
+// Handles the type conversion component of the TypeConversion. This updates
+// conversion patterns that used the original complex tensor types to be
+// updated to the non-complex variants.
+class GenericTypeConvert : public ConversionPattern {
+ public:
+  GenericTypeConvert(MLIRContext* context, TypeConverter& converter)
+      : ConversionPattern(converter, MatchAnyOpTypeTag(), 0, context) {}
+  LogicalResult matchAndRewrite(
+      Operation* op, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const override {
+    if (isa<func::FuncOp>(op)) {
+      return failure();
+    }
+
+    llvm::SmallVector<Type, 4> newResults;
+    (void)getTypeConverter()->convertTypes(op->getResultTypes(), newResults);
+    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
+                         newResults, op->getAttrs(), op->getSuccessors());
+    for (Region& r : op->getRegions()) {
+      Region* newRegion = state.addRegion();
+      rewriter.inlineRegionBefore(r, *newRegion, newRegion->begin());
+      TypeConverter::SignatureConversion result(newRegion->getNumArguments());
+      (void)getTypeConverter()->convertSignatureArgs(
+          newRegion->getArgumentTypes(), result);
+      rewriter.applySignatureConversion(&newRegion->front(), result);
+    }
+    Operation* newOp = rewriter.create(state);
+    rewriter.replaceOp(op, newOp->getResults());
+    return success();
+  }
+};
+
+static bool isIllegalType(Type type) {
+  if (auto shapedType = dyn_cast<ShapedType>(type)) {
+    return mlir::isa<ComplexType>(shapedType.getElementType());
+  }
+  return false;
+}
+
+void LowerComplexTypes::runOnOperation() {
+  ComplexTypeConverter converter;
+  ConversionTarget target(getContext());
+
+  // Operations are legal if they don't contain any illegal type.
+  target.markUnknownOpDynamicallyLegal([](Operation* op) {
+    if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
+      for (Type type : funcOp.getFunctionType().getInputs()) {
+        if (isIllegalType(type)) return false;
+      }
+      for (Type type : funcOp.getFunctionType().getResults()) {
+        if (isIllegalType(type)) return false;
+      }
+    }
+    for (Type type : op->getResultTypes()) {
+      if (type && isIllegalType(type)) return false;
+    }
+    for (Type type : op->getOperandTypes()) {
+      if (type && isIllegalType(type)) return false;
+    }
+    return true;
+  });
+
+  auto func = getOperation();
+  auto* ctx = &getContext();
+
+  RewritePatternSet patterns(ctx);
+  patterns.add<GenericTypeConvert>(ctx, converter);
+  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                 converter);
+  if (failed(applyFullConversion(func, target, std::move(patterns)))) {
+    signalPassFailure();
+  }
+
+  // We need to run folders post rewrite to cleanup conversion casts.
+  RewritePatternSet emptyRewriters(ctx);
+  if (failed(applyPatternsGreedily(func, std::move(emptyRewriters)))) {
+    signalPassFailure();
+  }
+}
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createLowerComplexTypesPass() {
+  return std::make_unique<LowerComplexTypes>();
+}
+}  // namespace tosa
+}  // namespace mlir

tensorflow/compiler/mlir/tosa/transforms/passes.h

@@ -0,0 +1,94 @@
+/* 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_TOSA_TRANSFORMS_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_PASSES_H_
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_set>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+
+namespace quant {
+class QuantDialect;
+}
+
+namespace quantfork {
+class QuantizationForkDialect;
+}
+
+namespace TFL {
+class TFLDialect;
+}
+
+namespace tosa {
+class TosaDialect;
+
+void populateLegalizeTFPatterns(MLIRContext* ctx, RewritePatternSet& patterns);
+void populateLegalizeTFLPatterns(MLIRContext* ctx, RewritePatternSet& patterns);
+
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeTFPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createFuseBiasTFPass();
+
+// `disabledPatterns` is a set of labels used to filter out input patterns with
+// a debug label or debug name in this set.
+// `enabledPatterns` is a set of labels used to filter out input patterns that
+//  do not have one of the labels in this set.
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeTFLPass(
+    ArrayRef<std::string> disabled_patterns = {},
+    ArrayRef<std::string> enabled_patterns = {});
+
+std::unique_ptr<OperationPass<ModuleOp>> createRetainCallOnceFuncsPass();
+std::unique_ptr<OperationPass<ModuleOp>> createStripModuleMetadataPass();
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertTFLUnsignedIntToSignedPass();
+std::unique_ptr<OperationPass<func::FuncOp>>
+createConvertFunctionMetadataPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createDequantizeTFLSoftmaxPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeTFTFLPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createLowerComplexTypesPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createStripFunctionMetadataPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createStripQuantTypesPass();
+std::unique_ptr<OperationPass<func::FuncOp>> createVerifyFullyConvertedPass();
+std::unique_ptr<OperationPass<ModuleOp>> createLegalizeTFLStatefulPass();
+
+#define GEN_PASS_REGISTRATION
+#define GEN_PASS_DECL_TOSALEGALIZETFPASS
+#define GEN_PASS_DECL_TOSALEGALIZETFLPASS
+#define GEN_PASS_DECL_TOSALEGALIZETFTFLPASS
+#define GEN_PASS_DECL_TOSAFUSEBIASTFPASS
+#define GEN_PASS_DECL_TOSACONVERTTFLUNSIGNEDINTTOSIGNEDPASS
+#define GEN_PASS_DECL_TOSASTRIPQUANTTYPESPASS
+#define GEN_PASS_DECL_TOSALOWERCOMPLEXTYPESPASS
+#define GEN_PASS_DECL_TOSADEQUANTIZETFLSOFTMAXPASS
+#define GEN_PASS_DECL_RETAINCALLONCEFUNCS
+#define GEN_PASS_DECL_STRIPFUNCTIONMETADATA
+#define GEN_PASS_DECL_STRIPMODULEMETADATA
+#define GEN_PASS_DECL_VERIFYFULLYCONVERTED
+#define GEN_PASS_DECL_CONVERTFUNCTIONMETADATA
+#define GEN_PASS_DECL_TOSALEGALIZESTATEFULPASS
+
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+}  // namespace tosa
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOSA_TRANSFORMS_PASSES_H_

tensorflow/compiler/mlir/tosa/transforms/passes.td

@@ -0,0 +1,129 @@
+/* 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 "mlir/Pass/PassBase.td"
+
+def TosaLegalizeTFPass : Pass<"tosa-legalize-tf", "mlir::func::FuncOp"> {
+  let summary = "Legalize from TensorFlow to TOSA";
+  let constructor = "createLegalizeTFPass()";
+  let dependentDialects = ["TosaDialect",
+    "quant::QuantDialect","quantfork::QuantizationForkDialect"
+  ];
+}
+
+def TosaLegalizeTFLPass : Pass<"tosa-legalize-tfl", "mlir::func::FuncOp"> {
+  let summary = "Legalize from TensorFlow Lite to TOSA. Options allow selective disabling/enabling of rewrites";
+  let constructor = "createLegalizeTFLPass()";
+
+  let dependentDialects = ["TosaDialect"];
+  let options = [
+    ListOption<"disabled_patterns_", "disable-patterns", "std::string",
+               "Labels of patterns that should be filtered out during application">,
+    ListOption<"enabled_patterns_", "enable-patterns", "std::string",
+               "Labels of patterns that should be used during application, all other patterns are filtered out">,
+  ];
+}
+
+def TosaLegalizeTFTFLPass : Pass<"tosa-legalize-tf-tfl", "mlir::func::FuncOp"> {
+  let summary = "Legalize from TensorFlow / TensorFlow Lite to TOSA";
+  let constructor = "createLegalizeTFTFLPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaFusebiasTFPass : Pass<"tosa-fuse-bias-tf", "mlir::func::FuncOp"> {
+  let summary = "Fuse tf.Op + tf.BiasAdd and legalized to TOSA";
+  let constructor = "createFuseBiasTFPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaConvertTFLUnsignedIntToSignedPass : Pass<"tosa-convert-tfl-uint-to-int", "mlir::func::FuncOp"> {
+  let summary = "Convert unsigned integer graph to integer graph";
+  let constructor = "createConvertTFLUnsignedIntToSignedPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaStripQuantTypesPass : Pass<"tosa-strip-quant-types", "mlir::func::FuncOp"> {
+  let summary = "Convert all quant types to their storage type.";
+  let constructor = "createStripQuantTypesPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaLowerComplexTypesPass : Pass<"tosa-lower-complex-types", "mlir::func::FuncOp"> {
+  let summary = "Convert all complex types to their storage type.";
+  let description = [{
+    Complex tensors are not supported in TOSA and are instead represented by two floating
+    point tensors corresponding to the "real" and "imag" parts of the complex tensor. This
+    pass replaces all complex tensors of shape [x, ..., y] with floating point tensors of
+    shape [x, ..., y, 2] where each inner pair of values represents a complex value. In
+    doing so, a 1:1 mapping between TFL and TOSA input/output tensors is maintained. The
+    pass will fail if not all complex tensors are converted.
+  }];
+  let constructor = "createLowerComplexTypesPass()";
+  let dependentDialects = ["TosaDialect"];
+}
+
+def TosaDequantizeTFLSoftmaxPass : Pass<"tosa-dequantize-tfl-softmax", "mlir::func::FuncOp"> {
+  let summary = "Dequantize TFLite Softmax ops.";
+  let description = [{
+    This pass rewrites quantized TFLite Softmax ops as: Dequantize, (float) Softmax, Quantize.
+    It is a work around for current performance issues with quantized Softmax codegen.
+    For instance it is a 20% end-to-end speedup on certain Softmax-heavy BERTs.
+    The compromise here is that this breaks bit-for-bit agreement with the quantized
+    Softmax lowering. But as Softmax isn't currently a TOSA op, this isn't a TOSA
+    spec violation - just breaking an invariant in the TFL-to-TOSA import pipeline.
+    See https://github.com/google/iree/issues/8974.
+  }];
+  let constructor = "createDequantizeTFLSoftmaxPass()";
+  let dependentDialects = ["mlir::TFL::TFLDialect", "quantfork::QuantizationForkDialect"];
+}
+
+def RetainCallOnceFuncs :
+    Pass<"tflite-retain-call-once-funcs", "mlir::ModuleOp"> {
+  let summary = "Guarantees that functions used by tfl.call_once are retained.";
+  let constructor = "createRetainCallOnceFuncsPass()";
+}
+
+def StripFunctionMetadata :
+    Pass<"tosa-tflite-strip-function-metadata", "mlir::func::FuncOp"> {
+  let summary = "Strip all unneeded TF/TFLite specific metadata.";
+  let constructor = "createStripFunctionMetadataPass()";
+}
+
+def StripModuleMetadata :
+    Pass<"tosa-tflite-strip-module-metadata", "mlir::ModuleOp"> {
+  let summary = "Strip all unneeded TF/TFLite specific metadata.";
+  let constructor = "createStripModuleMetadataPass()";
+}
+
+def VerifyFullyConverted :
+    Pass<"tosa-tflite-verify-fully-converted", "mlir::func::FuncOp"> {
+  let summary = "Verifies that all TFLite frontend ops were converted and none remain.";
+  let constructor = "createVerifyFullyConvertedPass()";
+}
+
+def ConvertFunctionMetadata :
+    Pass<"tosa-tflite-convert-function-metadata", "mlir::func::FuncOp"> {
+  let summary = "Converts TFLite input attributes to MLProgram arg attributes on functions.";
+  let constructor = "createConvertFunctionMetadataPass()";
+}
+
+def TosaLegalizeTFLStatefulPass : Pass<"tosa-legalize-tfl-stateful-tensors", "mlir::ModuleOp"> {
+  let summary = "Legalize tfl stateful operators to tosa stateful operators";
+  let description = [{
+    This pass is legalizing the tfl.call_once op to tosa stateful operators
+  }];
+  let constructor = "createLegalizeTFLStatefulPass()";
+  let dependentDialects = ["mlir::TFL::TFLDialect"];
+}

tensorflow/compiler/mlir/tosa/transforms/retain_call_once_funcs.cc

@@ -0,0 +1,68 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "retain-call-once-funcs"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir::tosa {
+
+#define GEN_PASS_DEF_RETAINCALLONCEFUNCS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+namespace {
+
+class RetainCallOnceFuncsPass
+    : public impl::RetainCallOnceFuncsBase<RetainCallOnceFuncsPass> {
+ public:
+  void getDependentDialects(DialectRegistry& registry) const override {
+    registry.insert<mlir::TFL::TensorFlowLiteDialect>();
+  }
+
+  void runOnOperation() override {
+    auto moduleOp = getOperation();
+
+    llvm::DenseMap<StringRef, func::FuncOp> funcMap;
+    for (auto func : moduleOp.getOps<mlir::func::FuncOp>()) {
+      funcMap[func.getSymName()] = func;
+    }
+
+    for (auto func : moduleOp.getOps<mlir::func::FuncOp>()) {
+      for (auto callOnce : func.getOps<mlir::TFL::CallOnceOp>()) {
+        auto callFunc = funcMap[callOnce.getSessionInitFunction()];
+        callOnce->setAttr("session_init_function_symbol",
+                          SymbolRefAttr::get(callFunc));
+      }
+    }
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> createRetainCallOnceFuncsPass() {
+  return std::make_unique<RetainCallOnceFuncsPass>();
+}
+
+}  // namespace mlir::tosa

tensorflow/compiler/mlir/tosa/transforms/strip_metadata.cc

@@ -0,0 +1,104 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+
+#include "mlir/Interfaces/FunctionInterfaces.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-strip-metadata"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir::tosa {
+
+#define GEN_PASS_DEF_STRIPFUNCTIONMETADATA
+#define GEN_PASS_DEF_STRIPMODULEMETADATA
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+namespace {
+
+static bool isTFLAttr(NamedAttribute &namedAttr) {
+  // TFLite uses both tf and tfl in attribute annotations.
+  auto name = namedAttr.getName().strref();
+  // Don't trim attributes from tf_saved_model---they carry ABI information.
+  if (name.starts_with("tf_saved_model.")) return false;
+
+  if (name.starts_with("tf.") || name.starts_with("tf_") ||
+      name.starts_with("tfl.") || name.starts_with("tfl_")) {
+    return true;
+  }
+  StringRef attrNamespace = namedAttr.getValue().getDialect().getNamespace();
+  return attrNamespace == "tf" || attrNamespace == "tfl";
+}
+
+class StripModuleMetadataPass
+    : public impl::StripModuleMetadataBase<StripModuleMetadataPass> {
+ public:
+  void runOnOperation() override {
+    auto moduleOp = getOperation();
+    auto stripAttrs = llvm::to_vector<4>(llvm::make_filter_range(
+        moduleOp->getAttrs(),
+        [](NamedAttribute namedAttr) { return isTFLAttr(namedAttr); }));
+    for (auto namedAttr : stripAttrs) {
+      moduleOp->removeAttr(namedAttr.getName());
+    }
+  }
+};
+
+class StripFunctionMetadataPass
+    : public impl::StripFunctionMetadataBase<StripFunctionMetadataPass> {
+ public:
+  void runOnOperation() override {
+    auto funcOp = getOperation();
+    auto stripAttrs = llvm::to_vector<4>(llvm::make_filter_range(
+        funcOp->getAttrs(),
+        [](NamedAttribute namedAttr) { return isTFLAttr(namedAttr); }));
+    for (auto namedAttr : stripAttrs) {
+      funcOp->removeAttr(namedAttr.getName());
+    }
+
+    for (int i = 0, e = funcOp.getNumArguments(); i < e; ++i) {
+      auto stripAttrs = llvm::to_vector<4>(llvm::make_filter_range(
+          mlir::function_interface_impl::getArgAttrs(funcOp, i),
+          [](NamedAttribute namedAttr) { return isTFLAttr(namedAttr); }));
+      for (auto namedAttr : stripAttrs) {
+        funcOp.removeArgAttr(i, namedAttr.getName());
+      }
+    }
+
+    for (int i = 0, e = funcOp.getNumResults(); i < e; ++i) {
+      auto stripAttrs = llvm::to_vector<4>(llvm::make_filter_range(
+          mlir::function_interface_impl::getResultAttrs(funcOp, i),
+          [](NamedAttribute namedAttr) { return isTFLAttr(namedAttr); }));
+      for (auto namedAttr : stripAttrs) {
+        funcOp.removeResultAttr(i, namedAttr.getName());
+      }
+    }
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> createStripModuleMetadataPass() {
+  return std::make_unique<StripModuleMetadataPass>();
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>> createStripFunctionMetadataPass() {
+  return std::make_unique<StripFunctionMetadataPass>();
+}
+
+}  // namespace mlir::tosa

tensorflow/compiler/mlir/tosa/transforms/strip_quant_types.cc

@@ -0,0 +1,170 @@
+/* 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 pass converts a TFLite uint8 graph to the int8 domain, with adaptors at
+// input and output tensors. This is needed because TOSA precision is
+// implemented in the int8 domain. This pass does:
+// 1. match TFL::QConst with uint8, generate TFL::QConst with int8 with value
+// remapped.
+// 2. insert tosa.RESCALE uint8 -> int8 if block argument (placeholder of graph)
+// is uint8 typed.
+// 3. insert tosa.RESCALE int8 -> uint8 if original returned tensor is uint8
+// typed.
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tosa/Utils/QuantUtils.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/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_common.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/legalize_utils.h"
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h"
+
+#define PASS_NAME "tosa-convert-tfl-uint8"
+#define DEBUG_TYPE PASS_NAME
+
+namespace mlir {
+namespace tosa {
+namespace {
+
+#define GEN_PASS_DEF_TOSASTRIPQUANTTYPESPASS
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+class StripQuantTypes
+    : public impl::TosaStripQuantTypesPassBase<StripQuantTypes> {
+ public:
+  explicit StripQuantTypes() = default;
+  void runOnOperation() override;
+};
+
+class QuantTypeConverter : public TypeConverter {
+ public:
+  static Type convertType(Type type) {
+    if (auto qType = dyn_cast<quant::QuantizedType>(type)) {
+      if (qType.isSigned() || qType.getStorageTypeIntegralWidth() != 8) {
+        return IntegerType::get(type.getContext(),
+                                qType.getStorageTypeIntegralWidth());
+      }
+
+      return IntegerType::get(type.getContext(),
+                              qType.getStorageTypeIntegralWidth(),
+                              IntegerType::SignednessSemantics::Unsigned);
+    }
+    return type;
+  }
+  static Type convertTensor(RankedTensorType type) {
+    auto newType = RankedTensorType::get(type.getShape(),
+                                         convertType(type.getElementType()));
+    return newType;
+  }
+  explicit QuantTypeConverter() {
+    addConversion([](Type type) { return convertType(type); });
+    addConversion(convertTensor);
+  }
+};
+
+// Handles the type conversion component of the TypeConversion. This updates
+// conversion patterns that used the original Quant types to be updated to
+// the non-quant variants.
+class GenericTypeConvert : public ConversionPattern {
+ public:
+  GenericTypeConvert(MLIRContext* context, TypeConverter& converter)
+      : ConversionPattern(converter, MatchAnyOpTypeTag(), 0, context) {}
+  LogicalResult matchAndRewrite(
+      Operation* op, ArrayRef<Value> operands,
+      ConversionPatternRewriter& rewriter) const override {
+    llvm::SmallVector<Type, 4> newResults;
+    if (isa<func::FuncOp>(op)) {
+      return failure();
+    }
+
+    (void)getTypeConverter()->convertTypes(op->getResultTypes(), newResults);
+    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
+                         newResults, op->getAttrs(), op->getSuccessors());
+    for (Region& r : op->getRegions()) {
+      Region* newRegion = state.addRegion();
+      rewriter.inlineRegionBefore(r, *newRegion, newRegion->begin());
+      TypeConverter::SignatureConversion result(newRegion->getNumArguments());
+      (void)getTypeConverter()->convertSignatureArgs(
+          newRegion->getArgumentTypes(), result);
+      rewriter.applySignatureConversion(&newRegion->front(), result);
+    }
+    Operation* newOp = rewriter.create(state);
+    rewriter.replaceOp(op, newOp->getResults());
+    return success();
+  }
+};
+
+static bool isIllegalType(Type type) {
+  if (mlir::isa<quant::QuantizedType>(type)) return true;
+  if (auto shapedType = dyn_cast<ShapedType>(type)) {
+    return isIllegalType(shapedType.getElementType());
+  }
+  return false;
+}
+
+void StripQuantTypes::runOnOperation() {
+  QuantTypeConverter converter;
+  ConversionTarget target(getContext());
+
+  target.addIllegalDialect<quantfork::QuantizationForkDialect>();
+  // Operations are legal if they don't contain any illegal type.
+  target.markUnknownOpDynamicallyLegal([](Operation* op) {
+    if (auto funcOp = dyn_cast<func::FuncOp>(op)) {
+      for (Type type : funcOp.getFunctionType().getInputs()) {
+        if (isIllegalType(type)) return false;
+      }
+      for (Type type : funcOp.getFunctionType().getResults()) {
+        if (isIllegalType(type)) return false;
+      }
+    }
+    for (Type type : op->getResultTypes()) {
+      if (type && isIllegalType(type)) return false;
+    }
+    for (Type type : op->getOperandTypes()) {
+      if (type && isIllegalType(type)) return false;
+    }
+    return true;
+  });
+
+  auto* ctx = &getContext();
+  auto func = getOperation();
+
+  RewritePatternSet patterns(&getContext());
+  patterns.add<GenericTypeConvert>(ctx, converter);
+  populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
+                                                                 converter);
+
+  if (failed(applyFullConversion(func, target, std::move(patterns)))) {
+    signalPassFailure();
+  }
+}
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createStripQuantTypesPass() {
+  return std::make_unique<StripQuantTypes>();
+}
+}  // namespace tosa
+}  // namespace mlir

tensorflow/compiler/mlir/tosa/transforms/tf_legalize_patterns.td

@@ -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.
+==============================================================================*/
+
+// TensorFlow legalization patterns
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "mlir/Dialect/Tosa/IR/TosaOps.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+
+// Nullary ops patterns.
+
+def : Pat<(TF_ConstOp ElementsAttr : $value), (Tosa_ConstOp $value)>;
+
+// Unary ops patterns.
+
+def : Pat<(TF_IdentityOp $value), (replaceWithValue $value)>;
+def : Pat<(TF_AbsOp $arg), (Tosa_AbsOp $arg)>;
+def : Pat<(TF_CeilOp $arg), (Tosa_CeilOp $arg)>;
+def : Pat<(TF_FloorOp $arg), (Tosa_FloorOp $arg)>;
+def : Pat<(TF_ExpOp $arg), (Tosa_ExpOp $arg)>;
+def : Pat<(TF_ErfOp $arg), (Tosa_ErfOp $arg)>;
+def : Pat<(TF_LogOp $arg), (Tosa_LogOp $arg)>;
+def : Pat<(TF_ReciprocalOp $arg), (Tosa_ReciprocalOp $arg)>;
+def : Pat<(TF_RsqrtOp $arg), (Tosa_RsqrtOp $arg)>;
+def : Pat<(TF_LogicalNotOp $arg), (Tosa_LogicalNotOp $arg)>;
+def : Pat<(TF_InvertOp $arg1), (Tosa_BitwiseNotOp $arg1)>;
+def : Pat<(TF_CastOp $in, BoolAttr : $truncate), (Tosa_CastOp $in)>;
+
+// Binary ops patterns.

tensorflow/compiler/mlir/tosa/transforms/tfl_legalize_patterns.td

@@ -0,0 +1,44 @@
+/* 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.
+==============================================================================*/
+
+// TFLite legalization patterns
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/PatternBase.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"
+include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.td"
+include "mlir/Dialect/Tosa/IR/TosaOps.td"
+
+//===----------------------------------------------------------------------===//
+// Unary ops patterns.
+//===----------------------------------------------------------------------===//
+
+def ConvertTFLAbsOp : Pat<(TFL_AbsOp $arg), (Tosa_AbsOp $arg)>;
+def ConvertTFLCeilOp : Pat<(TFL_CeilOp $arg), (Tosa_CeilOp $arg)>;
+def ConvertTFLFloorOp : Pat<(TFL_FloorOp $arg), (Tosa_FloorOp $arg)>;
+def ConvertTFLLogicalNotOp : Pat<(TFL_LogicalNotOp $arg), (Tosa_LogicalNotOp $arg)>;
+
+// Removing the quant.stats op for unquantized models.
+def ConvertQuantStatOp : Pat<(quantfork_StatisticsOp $value, $layer_stats, $axis_stats, $axis),
+                         (replaceWithValue $value)>;
+
+//===----------------------------------------------------------------------===//
+// Binary ops patterns.
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// Ternary ops patterns.
+//===----------------------------------------------------------------------===//

tensorflow/compiler/mlir/tosa/transforms/verify_fully_converted.cc

@@ -0,0 +1,90 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"
+
+namespace mlir::tosa {
+
+#define GEN_PASS_DEF_VERIFYFULLYCONVERTED
+#include "tensorflow/compiler/mlir/tosa/transforms/passes.h.inc"
+
+namespace {
+
+static void emitLegalizationErrors(Location loc,
+                                   const DenseSet<Operation *> &illegalOps) {
+  // Print op errors for each of the illegal ops that still remain.
+  llvm::MapVector<StringRef, int> opNameCounts;
+  for (Operation *illegalOp : illegalOps) {
+    StringRef opName = illegalOp->getName().getStringRef();
+    opNameCounts[opName]++;
+    illegalOp->emitOpError() << ": illegal op still exists";
+  }
+
+  std::vector<std::string> errorMessages;
+  errorMessages.reserve(opNameCounts.size());
+  for (const auto &opInfo : opNameCounts) {
+    errorMessages.push_back(
+        llvm::formatv("\t{0} (count: {1})", opInfo.first, opInfo.second));
+  }
+  emitError(loc) << "The following illegal operations still remain: \n"
+                 << llvm::join(errorMessages, "\n") << "\n";
+}
+
+LogicalResult verifyAllOperationsAreLegal(Operation *op,
+                                          const ConversionTarget &target) {
+  DenseSet<Operation *> illegalOps;
+  op->walk([&](Operation *op) {
+    if (!target.isLegal(op)) {
+      illegalOps.insert(op);
+    }
+  });
+  if (illegalOps.empty()) return success();
+  emitLegalizationErrors(op->getLoc(), illegalOps);
+  return failure();
+}
+
+class VerifyFullyConvertedPass
+    : public impl::VerifyFullyConvertedBase<VerifyFullyConvertedPass> {
+ public:
+  // Validates that no TFLite frontends ops are in the function.
+  void runOnOperation() override {
+    // We don't just use applyPartialConversion with no patterns because this
+    // pass shouldn't alter the IR at all (including via folding or
+    // canonicalizations that dialect conversion does automatically).
+    ConversionTarget target(getContext());
+    target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+    target.addIllegalDialect<mlir::TFL::TensorFlowLiteDialect>();
+    target.addIllegalOp<mlir::UnrealizedConversionCastOp>();
+    if (failed(verifyAllOperationsAreLegal(getOperation(), target)))
+      return signalPassFailure();
+  }
+};
+
+}  // anonymous namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> createVerifyFullyConvertedPass() {
+  return std::make_unique<VerifyFullyConvertedPass>();
+}
+
+}  // namespace mlir::tosa