chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/stablehlo/BUILD

@@ -0,0 +1,341 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@xla//third_party/rules_python/python:py_test.bzl", "py_test")
+load("@xla//xla/tsl:tsl.default.bzl", "tsl_pybind_extension")
+load("@xla//xla/tsl/platform:build_config_root.bzl", "if_static")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        ":friends",
+        "//tensorflow/tools/pip_package:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    packages = [
+        "//tensorflow/compiler/mlir/lite/...",
+        "//tensorflow/compiler/mlir/quantization/...",
+        "//tensorflow/compiler/mlir/quantization/tensorflow/...",
+        "//tensorflow/compiler/tests/...",
+    ],
+)
+
+tsl_pybind_extension(
+    name = "stablehlo_extension",
+    srcs = [
+        "stablehlo.cc",
+        "@stablehlo//:stablehlo/integrations/python/StablehloApi.cpp",
+    ],
+    hdrs = [
+        "@stablehlo//:stablehlo/integrations/python/StablehloApi.h",
+    ],
+    copts = [
+        "-fexceptions",
+        "-frtti",
+    ],
+    features = ["-use_header_modules"],
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MLIRBindingsPythonHeadersAndDeps",
+        "@nanobind",
+        "@stablehlo//:stablehlo_capi",
+        "@xla//third_party/python_runtime:headers",
+    ],
+)
+
+pytype_strict_library(
+    name = "stablehlo",
+    srcs = ["stablehlo.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":stablehlo_extension",
+    ],
+)
+
+py_test(
+    name = "stablehlo_test",
+    srcs = ["stablehlo_test.py"],
+    strict_deps = True,
+    deps = [
+        ":stablehlo",
+        #internal proto upb dep
+    ],
+)
+
+gentbl_cc_library(
+    name = "legalize_tf_patterns_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = [
+        (
+            ["-gen-rewriters"],
+            "transforms/generated_legalize_tf.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "transforms/legalize_tf_patterns.td",
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncTdFiles",
+        "@llvm-project//mlir:TensorOpsTdFiles",
+        "@xla//xla/mlir_hlo:hlo_ops_td_files",
+    ],
+)
+
+cc_library(
+    name = "fold_broadcast_pass",
+    srcs = [
+        "transforms/fold_broadcast_pass.cc",
+    ],
+    hdrs = [
+        "transforms/fold_broadcast_pass.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = [
+        "-Ithird_party",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@xla//xla/mlir_hlo",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "legalize_utils",
+    srcs = ["transforms/utils.cc"],
+    hdrs = ["transforms/utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@xla//xla/mlir_hlo",
+    ],
+)
+
+tf_cc_test(
+    name = "legalize_utils_test",
+    srcs = ["transforms/utils_test.cc"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":legalize_utils",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/mlir_hlo",
+    ],
+)
+
+cc_library(
+    name = "legalize_tf",
+    srcs = [
+        "transforms/generated_legalize_tf.inc",
+        "transforms/legalize_tf.cc",
+    ],
+    hdrs = [
+        "transforms/legalize_tf_passes.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":legalize_tf_patterns_inc_gen",
+        ":legalize_utils",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:dynamic_shape_utils",
+        "//tensorflow/compiler/mlir/tensorflow:xla_sharding_util",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/kernels:conv_grad_shape_utils",
+        "@com_google_absl//absl/status",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:Dialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:chlo_ops",
+        "@stablehlo//:stablehlo_ops",
+        "@tsl//tsl/platform:bfloat16",
+        "@tsl//tsl/platform:tensor_float_32_hdr_lib",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/builder:padding",
+        "@xla//xla/hlo/builder:sharding_builder",
+        "@xla//xla/hlo/builder/lib:conv_grad_size_util",
+        "@xla//xla/hlo/translate/hlo_to_mhlo:attribute_importer",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/mlir_hlo:convert_op_folder",
+        "@xla//xla/tsl/platform:status",
+    ] + if_static(["@tsl//tsl/platform:tensor_float_32_utils"]),
+)
+
+cc_library(
+    name = "tf_stablehlo",
+    srcs = [
+        "transforms/tf_stablehlo_pass.cc",
+    ],
+    hdrs = [
+        "transforms/tf_stablehlo_pass.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = [
+        "-Ithird_party",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":legalize_tf",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:lower_tf_lib",
+        "//tensorflow/compiler/mlir/tf2xla/transforms:xla_legalize_tf_with_tf2xla",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@stablehlo//:chlo_ops",
+        "@stablehlo//:register",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/mlir_hlo:hlo_dialect_registration",
+        "@xla//xla/mlir_hlo:mhlo_passes",
+        "@xla//xla/mlir_hlo:type_conversion",
+    ],
+    alwayslink = 1,
+)
+
+# LINT.IfChange(legalize_tf_xla_call_module_to_stablehlo_pass)
+cc_library(
+    name = "legalize_tf_xla_call_module_to_stablehlo_pass",
+    srcs = [
+        "transforms/legalize_tf_xla_call_module_to_stablehlo_pass.cc",
+    ],
+    hdrs = [
+        "transforms/legalize_tf_xla_call_module_to_stablehlo_pass.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = [
+        "-Ithird_party",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@stablehlo//:stablehlo_ops",
+        "@stablehlo//:stablehlo_serialization",
+        "@stablehlo//:vhlo_ops",
+    ],
+    alwayslink = 1,
+)
+# LINT.ThenChange(//tensorflow/compiler/mlir/lite/stablehlo:legalize_tf_xla_call_module_to_stablehlo_pass)
+
+cc_library(
+    name = "fuse_convolution_pass",
+    srcs = [
+        "transforms/mhlo_passes/fuse_convolution_pass.cc",
+    ],
+    hdrs = [
+        "transforms/mhlo_passes/fuse_convolution_pass.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = [
+        "-Ithird_party",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/mlir/quantization/common:attrs_and_constraints",
+        "//tensorflow/compiler/mlir/utils:validators",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Dialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@xla//xla/mlir_hlo",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "unfuse_batch_norm_pass",
+    srcs = [
+        "transforms/mhlo_passes/unfuse_batch_norm_pass.cc",
+    ],
+    hdrs = [
+        "transforms/mhlo_passes/unfuse_batch_norm_pass.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = [
+        "-Ithird_party",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@xla//xla/mlir_hlo",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "rename_entrypoint_to_main",
+    srcs = [
+        "transforms/rename_entrypoint_to_main.cc",
+    ],
+    hdrs = [
+        "transforms/rename_entrypoint_to_main.h",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    copts = [
+        "-Ithird_party",
+    ],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+    alwayslink = 1,
+)

tensorflow/compiler/mlir/stablehlo/stablehlo.cc

@@ -0,0 +1,25 @@
+/* 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 "nanobind/nanobind.h"  // from @nanobind
+#include "stablehlo/integrations/python/StablehloApi.h"  // from @stablehlo
+
+namespace mlir {
+namespace stablehlo {
+
+NB_MODULE(stablehlo_extension, m) { mlir::stablehlo::AddPortableApi(m); }
+
+}  // namespace stablehlo
+}  // namespace mlir

tensorflow/compiler/mlir/stablehlo/stablehlo.py

@@ -0,0 +1,26 @@
+# 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.
+# ==============================================================================
+"""StableHLO Portable Python APIs.
+
+This setup only exports the StableHLO Portable C++ APIs, which have
+signatures that do not rely on MLIR classes.
+
+Exporting all of MLIR Python bindings to TF OSS has high maintenance
+implications, especially given the frequency that TF updates the revision of
+LLVM used.
+"""
+
+# pylint: disable=wildcard-import
+from .stablehlo_extension import *

tensorflow/compiler/mlir/stablehlo/stablehlo_test.py

@@ -0,0 +1,40 @@
+# 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.
+# ==============================================================================
+"""Smoke test of functions in StableHLO Portable APIs."""
+
+from tensorflow.compiler.mlir.stablehlo import stablehlo
+
+
+def smoketest():
+  """Test StableHLO Portable APIs."""
+  assert isinstance(stablehlo.get_api_version(), int)
+  assembly = """
+    module @jit_f_jax.0 {
+      func.func public @main(%arg0: tensor<ui32>) -> tensor<i1> {
+        %0 = stablehlo.constant dense<1> : tensor<ui32>
+        %1 = "stablehlo.compare"(%arg0, %0) {compare_type = #stablehlo<comparison_type UNSIGNED>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<ui32>, tensor<ui32>) -> tensor<i1>
+        return %1 : tensor<i1>
+      }
+    }
+  """
+  target = stablehlo.get_current_version()
+  artifact = stablehlo.serialize_portable_artifact_str(assembly, target)
+  deserialized = stablehlo.deserialize_portable_artifact_str(artifact)
+  rountrip = stablehlo.serialize_portable_artifact_str(deserialized, target)
+  assert artifact == rountrip
+
+
+if __name__ == "__main__":
+  smoketest()

tensorflow/compiler/mlir/stablehlo/transforms/fold_broadcast_pass.cc

@@ -0,0 +1,265 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/stablehlo/transforms/fold_broadcast_pass.h"
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <utility>
+
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/APSInt.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Pass/Pass.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 "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace odml {
+
+static const APFloat &addSign(const APFloat &v, Type) { return v; }
+static APSInt addSign(const APInt &v, Type t) {
+  // Add signedness information to the value, treating signless as signed,
+  // unless it's i1.
+  return APSInt(v, t.isUnsignedInteger() || t.isSignlessInteger(1));
+}
+
+// Helper method that given 'shape' and 'current_index' representing
+// index in broadcasted tensor, get the index in the flat original tensor.
+// 'shape' is computed from the original shape and the broadcast dimensions to
+// match result shape.
+int64_t GetElementIndex(llvm::SmallVectorImpl<int64_t> &shape,
+                        llvm::SmallVectorImpl<int64_t> &current_index) {
+  int64_t ind = 0;
+  int64_t mul = 1;
+  for (int i = shape.size() - 1; i >= 0; --i) {
+    ind += (current_index[i] % shape[i]) * mul;
+    mul *= shape[i];
+  }
+  return ind;
+}
+
+// Helper method that increment index represented in 'current_index_ptr'
+// in the shape of 'result_shape'.
+void IncrementIndex(ArrayRef<int64_t> result_shape,
+                    llvm::SmallVectorImpl<int64_t> &current_index) {
+  for (int i = result_shape.size() - 1; i >= 0; --i) {
+    current_index[i]++;
+    if (current_index[i] == result_shape[i]) {
+      current_index[i] = 0;
+    } else {
+      break;
+    }
+  }
+}
+
+template <class AttrElementT,
+          class ElementValueT = typename AttrElementT::ValueType>
+Attribute ConstFoldBroadcastInDim(ShapedType result_type,
+                                  DenseElementsAttr operand,
+                                  DenseIntElementsAttr bcast_dims) {
+  auto dimensions = llvm::to_vector(bcast_dims.getValues<int64_t>());
+  const auto result_shape = result_type.getShape();
+  // Index for the broadcasted matrix.
+  llvm::SmallVector<int64_t, 16> current_index(result_type.getRank(), 0);
+  // Computes the new operand shape using the original shape and the broadcast
+  // dimensions to match result shape.
+  llvm::SmallVector<int64_t, 16> operand_new_shape(result_type.getRank(), 1);
+  for (int i = 0; i < dimensions.size(); ++i) {
+    operand_new_shape[dimensions[i]] = operand.getType().getDimSize(i);
+  }
+
+  llvm::SmallVector<ElementValueT, 16> new_values;
+  auto num_elements = result_type.getNumElements();
+  new_values.reserve(num_elements);
+  auto operand_values = operand.getValues<ElementValueT>();
+  for (int64_t i = 0; i < num_elements; ++i) {
+    const int64_t operand_index =
+        GetElementIndex(operand_new_shape, current_index);
+    new_values.push_back(*(operand_values.begin() + operand_index));
+    IncrementIndex(result_shape, current_index);
+  }
+  return DenseElementsAttr::get(result_type,
+                                ArrayRef<ElementValueT>(new_values));
+}
+
+template <typename Op, typename ElementType = Type, typename ValType,
+          typename Convert>
+static Attribute BinaryFolder(Op *op) {
+  auto lhs_op = op->getLhs().template getDefiningOp<mhlo::ConstantOp>();
+  auto rhs_op = op->getRhs().template getDefiningOp<mhlo::ConstantOp>();
+  if (!lhs_op || !lhs_op) return {};
+
+  auto lhs = dyn_cast_or_null<DenseElementsAttr>(lhs_op.getValue());
+  auto rhs = dyn_cast_or_null<DenseElementsAttr>(rhs_op.getValue());
+  if (!lhs || !rhs) return {};
+
+  ShapedType type = mlir::cast<ShapedType>(op->getType());
+  if (!type.hasStaticShape()) {
+    return {};
+  }
+
+  Type etype = type.getElementType();
+
+  // Evaluate for element types.
+  if (!mlir::isa<ElementType>(etype)) {
+    return {};
+  }
+
+  // Special case for folding splats no matter how large.
+  // Only covers the case of both attrs being splats; operation-specific cases
+  // like adding a zero or multiplying by one are handled elsewhere.
+  SplatElementsAttr splatLhs = mlir::dyn_cast<SplatElementsAttr>(lhs);
+  SplatElementsAttr splatRhs = mlir::dyn_cast<SplatElementsAttr>(rhs);
+  if (splatLhs && splatRhs) {
+    auto signedLhs = addSign(splatLhs.getSplatValue<ValType>(), etype);
+    auto signedRhs = addSign(splatRhs.getSplatValue<ValType>(), etype);
+    FailureOr<decltype(signedLhs)> result(Convert()(signedLhs, signedRhs));
+    return succeeded(result) ? SplatElementsAttr::get(type, *result)
+                             : Attribute();
+  }
+
+  SmallVector<ValType, 6> values;
+  values.reserve(lhs.getNumElements());
+  for (const auto zip : llvm::zip(lhs.template getValues<ValType>(),
+                                  rhs.template getValues<ValType>())) {
+    auto signedLhs = addSign(std::get<0>(zip), etype);
+    auto signedRhs = addSign(std::get<1>(zip), etype);
+    FailureOr<decltype(signedLhs)> result(Convert()(signedLhs, signedRhs));
+    if (failed(result)) {
+      return {};
+    }
+    values.push_back(std::move(*result));
+  }
+
+  return DenseElementsAttr::get(type, values);
+}
+
+template <typename BinaryOpType>
+class FoldBroadcastInDimBeforeBinaryElementwiseOp
+    : public OpRewritePattern<BinaryOpType> {
+ public:
+  using OpRewritePattern<BinaryOpType>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BinaryOpType binary_op,
+                                PatternRewriter &rewriter) const override {
+    auto lhs = binary_op.getLhs();
+    auto rhs = binary_op.getRhs();
+    auto lhs_bcast_op = lhs.template getDefiningOp<mhlo::BroadcastInDimOp>();
+    auto rhs_bcast_op = rhs.template getDefiningOp<mhlo::BroadcastInDimOp>();
+    if ((lhs_bcast_op && rhs_bcast_op) || (!lhs_bcast_op && !rhs_bcast_op)) {
+      return rewriter.notifyMatchFailure(
+          binary_op, "Operands should have exactly one BroadcastInDim op.");
+    }
+    // When the operand other than the broadcast op is not a const op, we
+    // should not fold broadcast op.
+    auto binary_op_const_operand =
+        (lhs_bcast_op ? rhs : lhs).template getDefiningOp<mhlo::ConstantOp>();
+    if (!binary_op_const_operand) return failure();
+    auto bcast_op = lhs_bcast_op ? lhs_bcast_op : rhs_bcast_op;
+    auto const_op =
+        bcast_op.getOperand().template getDefiningOp<mhlo::ConstantOp>();
+    if (!const_op) return failure();
+    auto const_val = dyn_cast_or_null<DenseElementsAttr>(const_op.getValue());
+    if (!const_val) return failure();
+
+    auto result_type =
+        dyn_cast_or_null<ShapedType>(bcast_op.getResult().getType());
+    if (!result_type || !result_type.hasStaticShape())
+      return rewriter.notifyMatchFailure(binary_op,
+                                         "Result type must have static shape.");
+
+    auto bcast_dims = bcast_op.getBroadcastDimensions();
+    auto elem_type = const_val.getElementType();
+    Attribute result;
+    if (mlir::isa<FloatType>(elem_type)) {
+      result = ConstFoldBroadcastInDim<FloatAttr>(result_type, const_val,
+                                                  bcast_dims);
+    } else if (mlir::isa<IntegerType>(elem_type)) {
+      result = ConstFoldBroadcastInDim<IntegerAttr>(result_type, const_val,
+                                                    bcast_dims);
+    } else {
+      return rewriter.notifyMatchFailure(bcast_op, "Unsupported element type.");
+    }
+    Value new_const_op =
+        mhlo::ConstantOp::create(rewriter, bcast_op.getLoc(), result);
+    rewriter.replaceOp(bcast_op, {new_const_op});
+    return success();
+  }
+};
+
+using FoldBroadcastInDimBeforeMulOp =
+    FoldBroadcastInDimBeforeBinaryElementwiseOp<mhlo::MulOp>;
+
+// Constant folds mhlo.mul, this folder doesn't have an upper limit on how many
+// elements can be folded.
+LogicalResult ConstantFoldMul(mhlo::MulOp op, PatternRewriter &rewriter) {
+  ShapedType type = mlir::dyn_cast<ShapedType>(op.getType());
+  Type etype = type.getElementType();
+  Attribute result = {};
+  if (mlir::isa<FloatType>(etype)) {
+    result =
+        BinaryFolder<mhlo::MulOp, FloatType, APFloat, std::multiplies<APFloat>>(
+            &op);
+  } else if (mlir::isa<IntegerType>(etype)) {
+    result =
+        BinaryFolder<mhlo::MulOp, IntegerType, APInt, std::multiplies<APSInt>>(
+            &op);
+  }
+  if (result == Attribute()) return failure();
+  Value new_const_op = mhlo::ConstantOp::create(rewriter, op.getLoc(), result);
+  rewriter.replaceOp(op, {new_const_op});
+  return success();
+}
+
+class FoldBroadcastPass
+    : public PassWrapper<FoldBroadcastPass, OperationPass<func::FuncOp>> {
+ public:
+  StringRef getArgument() const final { return "constant-fold-broadcast-pass"; }
+  StringRef getDescription() const final {
+    return "Constant folds BroadcastInDimOp before binary elementwise ops";
+  }
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {}
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<FoldBroadcastInDimBeforeMulOp>(&getContext());
+    patterns.add(ConstantFoldMul);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+std::unique_ptr<Pass> createFoldBroadcastPass() {
+  return std::make_unique<FoldBroadcastPass>();
+}
+
+static PassRegistration<FoldBroadcastPass> pass;
+
+}  // namespace odml
+}  // namespace mlir

tensorflow/compiler/mlir/stablehlo/transforms/fold_broadcast_pass.h

@@ -0,0 +1,32 @@
+/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_FOLD_BROADCAST_PASS_H_
+#define TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_FOLD_BROADCAST_PASS_H_
+
+#include <memory>
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace odml {
+
+// Constant folds broadcast_in_dim op conditionally.
+std::unique_ptr<Pass> createFoldBroadcastPass();
+
+}  // namespace odml
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_FOLD_BROADCAST_PASS_H_

tensorflow/compiler/mlir/stablehlo/transforms/legalize_tf_passes.h

@@ -0,0 +1,51 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_LEGALIZE_TF_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_LEGALIZE_TF_PASSES_H_
+
+#include <memory>
+#include <optional>
+
+#include "llvm/ADT/StringRef.h"
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+
+namespace mlir {
+
+namespace func {
+class FuncOp;
+}
+class ModuleOp;
+class Operation;
+template <typename T>
+class OperationPass;
+class Pass;
+
+namespace odml {
+
+/// Adds the TF to TF lowerings and TF to XLA rewrite patterns to the pattern
+/// list.
+void PopulateLegalizeTfPatterns(MLIRContext* context,
+                                RewritePatternSet* patterns);
+
+}  // namespace odml
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_LEGALIZE_TF_PASSES_H_

tensorflow/compiler/mlir/stablehlo/transforms/legalize_tf_patterns.td

@@ -0,0 +1,814 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is the legalization pattern definition file for TF to XLA.
+
+include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Shape/IR/ShapeOps.td"
+include "mlir/Dialect/Func/IR/FuncOps.td"
+include "mlir/Dialect/Tensor/IR/TensorOps.td"
+include "stablehlo/dialect/ChloOps.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
+include "mhlo/IR/hlo_ops.td"
+
+def SignedIntTensor : TensorOf<[I1, I8, I16, I32, I64]>;
+def UnsignedIntTensor : TensorOf<[UI8, UI16, UI32, UI64]>;
+
+// IEEE compliant floating point tensors.
+def IEEEFloatTensor : TensorOf<[F16, F32, F64]>;
+
+//===----------------------------------------------------------------------===//
+// BatchNorm op patterns.
+//===----------------------------------------------------------------------===//
+
+def FalseBoolAttr : AttrConstraint<CPred<"!llvm::cast<BoolAttr>($_self).getValue()">>;
+def TrueBoolAttr : AttrConstraint<CPred<"llvm::cast<BoolAttr>($_self).getValue()">>;
+
+def CastValueToI64: NativeCodeCall<
+  "CastValueToI64($0.getLoc(), $1, &$_builder)">;
+
+def CastValueToElementType: NativeCodeCall<
+  "ConvertOp::create($_builder, $0.getLoc(), $1, "
+  "getElementTypeOrSelf($2.getType()))">;
+
+// Here, $0 is an ElementsAttr with exactly one element of type integer. $1 is
+// the corresponding value of ranked tensor type whose axis is referred in $0.
+def GetHLOAxisFromTFAxis : NativeCodeCall<
+  "GetHLOAxisFromTFAxis("
+  "$0, llvm::cast<RankedTensorType>($1.getType()).getRank(), &$_builder)">;
+
+// Same as the above but with $1 of type operand_range from variadic TensorFlow
+// input.
+def GetHLOAxisFromTFAxisVariadic : NativeCodeCall<
+  "GetHLOAxisFromTFAxis("
+  "$0, llvm::cast<RankedTensorType>((*$1.begin()).getType()).getRank(), "
+  "&$_builder)">;
+
+def CastElementsToI64Elements : NativeCodeCall<
+  "llvm::cast<DenseIntElementsAttr>(hlo::convertElementsAttr("
+    "llvm::cast<ElementsAttr>($0), $_builder.getIntegerType(64)))">;
+
+def EmptyDotAlgorithmAttr : NativeCodeCall<"mlir::mhlo::DotAlgorithmAttr{}">;
+
+def ConstDefaultResultAccuracyAttr :
+  ConstantAttr<MHLO_ResultAccuracyAttr, "::mlir::mhlo::ResultAccuracyMode::DEFAULT">;
+
+//===----------------------------------------------------------------------===//
+// ApproximateEqual op pattern.
+//===----------------------------------------------------------------------===//
+
+class MHLO_ComparisonDirectionValue<string enumStr> :
+  ConstantAttr<MHLO_ComparisonDirectionAttr, "::mlir::mhlo::ComparisonDirection::" # enumStr>;
+
+class CHLO_ComparisonDirectionValue<string enumStr> :
+  ConstantAttr<CHLO_ComparisonDirectionAttr, "::mlir::chlo::ComparisonDirection::" # enumStr>;
+
+// TODO(b/228291745): Assert that $x and $y have the same shape.
+def : Pat<(TF_ApproximateEqualOp:$result $x, $y, $tolerance),
+          (CHLO_BroadcastCompareOp
+           (MHLO_AbsOp:$abs (MHLO_SubtractOp $x, $y)),
+           (CastValueToElementType $result, (MHLO_ConstantOp $tolerance), $abs),
+           (NullDenseI64ArrayAttr),
+           CHLO_ComparisonDirectionValue<"LT">,
+           (CHLO_DEFAULT_COMPARISON_TYPE))>;
+
+//===----------------------------------------------------------------------===//
+// Assert op pattern.
+//===----------------------------------------------------------------------===//
+
+// HLO and XLA doesn't support Assertions.
+def LowerAssert : Pattern<(TF_AssertOp $condition, $data, $summarize), []>;
+
+//===----------------------------------------------------------------------===//
+// Binary op patterns.
+//===----------------------------------------------------------------------===//
+
+// Check that two values can be broadcasted together
+def AreBroadcastCompatible : Constraint<CPred<"AreBroadcastCompatible($0, $1)">,
+    "types must be broadcastable">;
+
+class DirectBinaryPat<Op FromOp, Op ToOp>
+  : Pat<(FromOp AnyTensor:$l, AnyTensor:$r),
+        (ToOp $l, $r, (BinBroadcastDimensions $l, $r))>;
+
+foreach fromToBinPair = [[TF_AddV2Op, CHLO_BroadcastAddOp],
+                         [TF_Atan2Op, CHLO_BroadcastAtan2Op],
+                         [TF_ComplexOp, CHLO_BroadcastComplexOp],
+                         [TF_DivOp, CHLO_BroadcastDivOp],
+                         [TF_LeftShiftOp, CHLO_BroadcastShiftLeftOp],
+                         [TF_MaximumOp, CHLO_BroadcastMaxOp],
+                         [TF_MinimumOp, CHLO_BroadcastMinOp],
+                         [TF_ModOp, CHLO_BroadcastRemOp],
+                         [TF_MulOp, CHLO_BroadcastMulOp],
+                         [TF_NextAfterOp, CHLO_BroadcastNextAfterOp],
+                         [TF_PolygammaOp, CHLO_BroadcastPolygammaOp],
+                         [TF_PowOp, CHLO_BroadcastPowOp],
+                         [TF_RealDivOp, CHLO_BroadcastDivOp],
+                         [TF_SubOp, CHLO_BroadcastSubOp],
+                         [TF_ZetaOp, CHLO_BroadcastZetaOp]] in
+  def : DirectBinaryPat<fromToBinPair[0], fromToBinPair[1]>;
+
+def LowerRightShiftSigned :
+  Pat<(TF_RightShiftOp AnyTensor:$l, AnyTensor:$r),
+      (CHLO_BroadcastShiftRightArithmeticOp $l, $r,
+       (BinBroadcastDimensions $l, $r)),
+      [(SignedIntTensor $r)]>;
+
+def LowerRightShiftUnsigned :
+  Pat<(TF_RightShiftOp AnyTensor:$l, AnyTensor:$r),
+      (CHLO_BroadcastShiftRightLogicalOp $l, $r,
+       (BinBroadcastDimensions $l, $r)),
+      [(UnsignedIntTensor $r)]>;
+
+// Performs a substitution of FloorDiv, pseudo code below:
+//
+//  return floor(div(x, y))
+def : Pat<(TF_FloorDivOp AnyTensor:$l, AnyTensor:$r),
+          (MHLO_FloorOp
+           (CHLO_BroadcastDivOp $l, $r, (BinBroadcastDimensions $l, $r))),
+          [(IEEEFloatTensor $l)]>;
+
+// Performs a substitution of FloorDiv for integer tensors, which required
+// additional correction for a negative numerator / denominator. Equivalent
+// pseudocode is shown below:
+//
+// T z = x / y
+// return (z * y != x && (x < 0) != (y < 0)) ? z - 1 : z
+//
+// BroadcastToDimensions is used to compute the broadcast attr to higher
+// dimensions. This computes the broadcast of 'l' to broadcast('l', 'r')
+// without returning the broadcast of 'r' to broadcast('l', 'r').
+def : Pat<(TF_FloorDivOp AnyTensor:$l, AnyTensor:$r),
+      (MHLO_SelectOp
+       (CHLO_BroadcastAndOp
+        (CHLO_BroadcastCompareOp
+         (CHLO_BroadcastMulOp:$mul
+          (CHLO_BroadcastDivOp:$div $l, $r,
+           (BinBroadcastDimensions $l, $r)),
+          $r, (BinBroadcastDimensions $div, $r)),
+         $l, (BinBroadcastDimensions $mul, $l), CHLO_ComparisonDirectionValue<"NE">,
+         (CHLO_DEFAULT_COMPARISON_TYPE)),
+        (CHLO_BroadcastCompareOp
+         (CHLO_BroadcastCompareOp:$l_cmp $l,
+          (MHLO_ConstantOp:$l_zeros (GetScalarOfType<0> $l)),
+          (NullDenseI64ArrayAttr), CHLO_ComparisonDirectionValue<"LT">,
+          (CHLO_DEFAULT_COMPARISON_TYPE)),
+         (CHLO_BroadcastCompareOp:$r_cmp $r,
+          (MHLO_ConstantOp:$r_zeros (GetScalarOfType<0> $r)),
+          (NullDenseI64ArrayAttr), CHLO_ComparisonDirectionValue<"LT">,
+          (CHLO_DEFAULT_COMPARISON_TYPE)),
+         (BinBroadcastDimensions $l_cmp, $r_cmp), CHLO_ComparisonDirectionValue<"NE">,
+         (CHLO_DEFAULT_COMPARISON_TYPE)),
+        (NullDenseI64ArrayAttr)),
+       (CHLO_BroadcastSubOp $div,
+        (MHLO_ConstantOp:$ones (GetScalarOfType<1> $div)),
+        (NullDenseI64ArrayAttr)), $div),
+      [(SignedIntTensor $l)]>;
+
+// FloorDiv of unsigned is just div.
+def : Pat<(TF_FloorDivOp AnyTensor:$l, AnyTensor:$r),
+          (CHLO_BroadcastDivOp $l, $r, (BinBroadcastDimensions $l, $r)),
+          [(UnsignedIntTensor $l)]>;
+
+// Performs a substitution of FloorMod designed to correct for possibly negative
+// values. Pseudocode shown below:
+//
+//   T trunc_mod = std::fmod(x, y);
+//   return trunc_mod != 0 && (y < 0 != trunc_mod < 0) ? trunc_mod + y
+//                                                     : trunc_mod
+def : Pat<(TF_FloorModOp AnyTensor:$l, AnyTensor:$r),
+      (MHLO_SelectOp
+       (CHLO_BroadcastAndOp
+        (CHLO_BroadcastCompareOp
+         (CHLO_BroadcastRemOp:$rem $l, $r, (BinBroadcastDimensions $l, $r)),
+         (MHLO_ConstantOp:$l_zeros (GetScalarOfType<0> $l)),
+         (NullDenseI64ArrayAttr), CHLO_ComparisonDirectionValue<"NE">,
+         (CHLO_DEFAULT_COMPARISON_TYPE)),
+        (CHLO_BroadcastCompareOp
+         (CHLO_BroadcastCompareOp:$r_cmp $r,
+          (MHLO_ConstantOp:$r_zeros (GetScalarOfType<0> $r)),
+          (NullDenseI64ArrayAttr), CHLO_ComparisonDirectionValue<"LT">,
+          (CHLO_DEFAULT_COMPARISON_TYPE)),
+         (CHLO_BroadcastCompareOp:$rem_cmp $rem, $r_zeros,
+          (BinBroadcastDimensions $rem, $r_zeros), CHLO_ComparisonDirectionValue<"LT">,
+          (CHLO_DEFAULT_COMPARISON_TYPE)),
+         (BinBroadcastDimensions $r_cmp, $rem_cmp), CHLO_ComparisonDirectionValue<"NE">,
+         (CHLO_DEFAULT_COMPARISON_TYPE)),
+        (NullDenseI64ArrayAttr)),
+       (CHLO_BroadcastAddOp $r,
+        $rem, (BinBroadcastDimensions $r, $rem)), $rem),
+      [(TensorOf<[I8, I16, I32, I64, F16, F32, F64]> $l)]>;
+
+// FloorMod of unsigned is just mod.
+def : Pat<(TF_FloorModOp AnyTensor:$l, AnyTensor:$r),
+          (CHLO_BroadcastRemOp $l, $r, (BinBroadcastDimensions $l, $r)),
+          [(UnsignedIntTensor $l)]>;
+
+def Get2DTransposePerm: NativeCodeCall<
+  "Get2DTransposePerm($0, &$_builder)">;
+
+def : Pat<(TF_RiscAddOp $l, $r), (MHLO_AddOp $l, $r)>;
+
+def : Pat<(TF_RiscDotOp $a, $b, $transpose_a, $transpose_b),
+          (MHLO_DotOp
+          (TF_TransposeOp $a, (TF_ConstOp (Get2DTransposePerm $transpose_a))),
+          (TF_TransposeOp $b, (TF_ConstOp (Get2DTransposePerm $transpose_b))),
+          /*precision_config=*/(NullArrayAttr))>;
+
+//===----------------------------------------------------------------------===//
+// Logical & bitwise binary op patterns.
+//===----------------------------------------------------------------------===//
+
+class DirectLogicalBinaryPat<Op FromOp, Op ToOp>
+  : Pat<(FromOp AnyTensor:$l, AnyTensor:$r),
+        (ToOp $l, $r, (BinBroadcastDimensions $l, $r)),
+        [(AnyTypeOf<[SignedIntTensor, UnsignedIntTensor]> $l)]>;
+
+foreach fromToBinPair = [[TF_LogicalAndOp, CHLO_BroadcastAndOp],
+                         [TF_LogicalOrOp, CHLO_BroadcastOrOp],
+                         [TF_BitwiseAndOp, CHLO_BroadcastAndOp],
+                         [TF_BitwiseOrOp, CHLO_BroadcastOrOp],
+                         [TF_BitwiseXorOp, CHLO_BroadcastXorOp]] in
+  def : DirectLogicalBinaryPat<fromToBinPair[0], fromToBinPair[1]>;
+
+//===----------------------------------------------------------------------===//
+// Compare op patterns.
+//===----------------------------------------------------------------------===//
+
+class DirectComparePat<Op FromOp, CHLO_ComparisonDirectionValue direction>
+  : Pat<(FromOp AnyTensor:$l, AnyTensor:$r),
+        (CHLO_BroadcastCompareOp
+           $l, $r, (BinBroadcastDimensions $l, $r), direction,
+           (CHLO_DEFAULT_COMPARISON_TYPE))>;
+
+def : DirectComparePat<TF_GreaterOp, CHLO_ComparisonDirectionValue<"GT">>;
+def : DirectComparePat<TF_GreaterEqualOp, CHLO_ComparisonDirectionValue<"GE">>;
+def : DirectComparePat<TF_LessOp, CHLO_ComparisonDirectionValue<"LT">>;
+def : DirectComparePat<TF_LessEqualOp, CHLO_ComparisonDirectionValue<"LE">>;
+
+class EqualityPat<Op FromOp, CHLO_ComparisonDirectionValue direction>
+    : Pat<(FromOp AnyTensor:$l, AnyTensor:$r,
+           TrueBoolAttr:$incompatible_shape_error),
+        (CHLO_BroadcastCompareOp
+         $l, $r, (BinBroadcastDimensions $l, $r), direction,
+         (CHLO_DEFAULT_COMPARISON_TYPE)),
+        [(MHLO_Tensor $l)]>;
+
+def : EqualityPat<TF_EqualOp, CHLO_ComparisonDirectionValue<"EQ">>;
+def : EqualityPat<TF_NotEqualOp, CHLO_ComparisonDirectionValue<"NE">>;
+
+//===----------------------------------------------------------------------===//
+// Concat op patterns.
+//===----------------------------------------------------------------------===//
+
+def OneElementAttrPred
+  : CPred<"llvm::cast<ElementsAttr>($_self).getShapedType().getNumElements() == 1">;
+
+def OneElementAttr
+  : ElementsAttrBase<And<[ElementsAttr.predicate, OneElementAttrPred]>,
+                     "Scalar ElementsAttr">;
+
+def HasRankedFirstOperand
+  : Constraint<CPred<"llvm::isa<RankedTensorType>((*$0.begin()).getType())">>;
+
+def IsShapedTensor
+  : Constraint<CPred<"llvm::isa<RankedTensorType>($0.getType())">>;
+
+// This pattern converts TensorFlow axis format to HLO axis format which
+// doesn't wrap around like TensorFlow and is always positive. For this
+// conversion, use the first input to get inputs rank. Other inputs need not be
+// ranked.
+// Defining op for `axis` is TensorFlow constant op in the pattern as during
+// the conversion, original Concat op operands still refers to the old ops even
+// if HLO constant op is introduced as an replacement for the TensorFlow
+// Constant op.
+def : Pat<(TF_ConcatV2Op $inputs, (ConstantLikeMatcher OneElementAttr:$axis)),
+          (MHLO_ConcatenateOp $inputs,
+            (GetHLOAxisFromTFAxisVariadic $axis, $inputs)),
+          [(HasRankedFirstOperand $inputs)]>;
+
+//===----------------------------------------------------------------------===//
+// CollectivePermute op patterns.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_CollectivePermuteOp $input, (ConstantLikeMatcher ElementsAttr:$source_target_pairs)),
+          (MHLO_CollectivePermuteOp $input,
+            (CastElementsToI64Elements $source_target_pairs),
+            (NullChannelHandleAttr))>;
+
+//===----------------------------------------------------------------------===//
+// CrossReplicaSum op patterns.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_CrossReplicaSumOp $input, (ConstantLikeMatcher ElementsAttr:$group_assignment)),
+          (MHLO_CrossReplicaSumOp $input,
+            (CastElementsToI64Elements $group_assignment))>;
+
+//===----------------------------------------------------------------------===//
+// All2All op patterns.
+//===----------------------------------------------------------------------===//
+
+def ValueToVariadic: NativeCodeCall<"SmallVector<Value, 1>{$0}">;
+def : Pat<(TF_AllToAllOp AnyRankedTensor:$input, (ConstantLikeMatcher ElementsAttr:$group_assignment), I64Attr:$concat_dimension, $split_dimension, $split_count),
+          (MHLO_AllToAllOp (ValueToVariadic $input), $split_dimension, $concat_dimension, $split_count, (CastElementsToI64Elements $group_assignment), (NullChannelHandleAttr))>;
+
+//===----------------------------------------------------------------------===//
+// FFT op patterns.
+//===----------------------------------------------------------------------===//
+
+class MHLO_FftTypeValue<string enumStr> :
+  ConstantAttr<MHLO_FftTypeAttr, "::mlir::mhlo::FftType::" # enumStr>;
+
+def GetInnerDimFromValue : NativeCodeCall<
+  "GetInnerDimFromValue(llvm::cast<ShapedType>($0.getType()), &$_builder)">;
+
+def CheckInnerDimStatic
+  : Constraint<CPred<"CheckInnerDimStatic(llvm::cast<ShapedType>($0.getType()), &$_builder)">>;
+
+def : Pat<(TF_FFTOp:$res $input),
+          (MHLO_FftOp $input, MHLO_FftTypeValue<"FFT">, (GetInnerDimFromValue $res)),
+          [(CheckInnerDimStatic $input)]>;
+
+def : Pat<(TF_IFFTOp:$res $input),
+          (MHLO_FftOp $input, MHLO_FftTypeValue<"IFFT">, (GetInnerDimFromValue $res)),
+          [(CheckInnerDimStatic $input)]>;
+
+//===----------------------------------------------------------------------===//
+// GatherV2 op patterns.
+//===----------------------------------------------------------------------===//
+
+// Here, $params and $indices needs to be ranked so that $axis and $batch_dims
+// attributes can be converted from TensorFlow axis format supporting negative
+// indexing to the HLO format.
+def LegalizeGatherV2 :
+  Pat<(TF_GatherV2Op AnyRankedTensor:$params, AnyRankedTensor:$indices,
+        (ConstantLikeMatcher ElementsAttr:$axis), $batch_dims),
+      (MHLO_TorchIndexSelectOp $params, $indices,
+        (GetHLOAxisFromTFAxis $axis, $params),
+        (GetHLOAxisFromTFAxis $batch_dims, $indices))>;
+
+//===----------------------------------------------------------------------===//
+// Pad op patterns.
+//===----------------------------------------------------------------------===//
+
+class SliceDenseIntElementsAttrColumn2D<string column> : NativeCodeCall<
+  "SliceDenseIntElementsAttrColumn2D(llvm::cast<ElementsAttr>($0), " # column # " )">;
+
+class SliceDenseIntElementsAttr<string index, string axis> : NativeCodeCall<
+  "SliceDenseIntElementsAttr(llvm::cast<ElementsAttr>($0), " # index # ", " # axis # ")">;
+
+// Interior padding attribute based on the TF padding.
+def GetInteriorPadding : NativeCodeCall <
+  "GetInteriorPadding(llvm::cast<ElementsAttr>($0))">;
+
+def : Pat<(TF_PadV2Op $input, (ConstantLikeMatcher ElementsAttr:$padding), $c),
+          (MHLO_PadOp $input, $c,
+           (SliceDenseIntElementsAttrColumn2D<"0"> $padding),
+           (SliceDenseIntElementsAttrColumn2D<"1"> $padding),
+           (GetInteriorPadding $padding))>;
+
+//===----------------------------------------------------------------------===//
+// Identity op patterns.
+//===----------------------------------------------------------------------===//
+
+foreach src = [TF_IdentityOp, TF_StopGradientOp, TF__EagerConstOp] in
+  def : Pat<(src $op), (replaceWithValue $op)>;
+
+// TODO(b/32223192): Support CheckNumerics in HLO.
+foreach src = [TF_PreventGradientOp, TF_CheckNumericsOp] in
+  def : Pat<(src $op, $msg), (replaceWithValue $op)>;
+
+//===----------------------------------------------------------------------===//
+// MatMul op patterns.
+//===----------------------------------------------------------------------===//
+
+def GetPrecisionConfig: NativeCodeCall<
+  "GetPrecisionConfig(&$_builder)">;
+
+def : Pat<(TF_MatMulOp $a, $b, $transpose_a, $transpose_b, $grad_a, $grad_b),
+          (MHLO_DotOp
+          (TF_TransposeOp $a, (TF_ConstOp (Get2DTransposePerm $transpose_a))),
+          (TF_TransposeOp $b, (TF_ConstOp (Get2DTransposePerm $transpose_b))),
+          /*precision_config=*/(GetPrecisionConfig))>;
+
+//===----------------------------------------------------------------------===//
+// Lower `tf.ZerosLike`
+//===----------------------------------------------------------------------===//
+
+class MHLO_ConstantLike<string value> : NativeCodeCall<
+    "chlo::getConstantLike($_builder, $_loc, " # value # ", $0)">;
+
+def : Pat<(TF_ZerosLikeOp AnyTensor:$arg),
+          (MHLO_ConstantLike<"0"> $arg)>;
+
+//===----------------------------------------------------------------------===//
+// Lower `tf.OnesLike`
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_OnesLikeOp AnyTensor:$arg),
+          (MHLO_ConstantLike<"1"> $arg)>;
+
+//===----------------------------------------------------------------------===//
+// Elu op patterns.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_EluOp AnyTensor:$features),
+          (MHLO_SelectOp
+           (MHLO_CompareOp
+              $features,
+              (MHLO_ConstantLike<"0">:$zero $features),
+              MHLO_ComparisonDirectionValue<"GT">, (MHLO_DEFAULT_COMPARISON_TYPE)),
+           $features,
+           (MHLO_Expm1Op $features, ConstDefaultResultAccuracyAttr))>;
+
+def : Pat<(TF_EluGradOp AnyStaticShapeTensor:$gradients, AnyRankedTensor:$features),
+           (MHLO_SelectOp
+            (CHLO_BroadcastCompareOp
+              $features,
+              (MHLO_ConstantOp:$zero (GetScalarOfType<0> $features)),
+              (BinBroadcastDimensions $zero, $features),
+              CHLO_ComparisonDirectionValue<"GT">, (CHLO_DEFAULT_COMPARISON_TYPE)),
+            $gradients,
+            (MHLO_MulOp
+             $gradients,
+             (CHLO_BroadcastAddOp
+               $features,
+               (MHLO_ConstantOp:$one (GetScalarOfType<1> $features)),
+               (BinBroadcastDimensions $one, $features))))>;
+
+//===----------------------------------------------------------------------===//
+// Relu op patterns.
+//===----------------------------------------------------------------------===//
+
+// TODO(hinsu): Make these patterns to TF to TF lowering. Relu6 lowering will
+// require HLO canonicalization of min and max on a tensor to ClampOp.
+
+// TODO(hinsu): Lower quantized types after supporting them in GetScalarOfType.
+def : Pat<(TF_ReluOp AnyTensor:$input),
+          (CHLO_BroadcastMaxOp
+               (MHLO_ConstantOp:$zero (GetScalarOfType<0> $input)), $input,
+               (BinBroadcastDimensions $zero, $input)),
+          [(TF_IntOrFpTensor $input)]>;
+
+// TODO(hinsu): Lower quantized types after supporting them in GetScalarOfType.
+def : Pat<(TF_Relu6Op AnyRankedTensor:$input),
+          (MHLO_ClampOp (MHLO_ConstantOp (GetScalarOfType<0> $input)), $input,
+                       (MHLO_ConstantOp (GetScalarOfType<6> $input))),
+          [(TF_IntOrFpTensor $input)]>;
+
+// ReluGrad(gradients, features) = gradients * (features > 0)
+// The condition that $gradients and $features need to have the same shape is
+// implicitly enforced: $zero is created to have the same shape as $features,
+// MHLO_SelectOp enforces that $gradients and $zero have the same shape.
+def : Pat<(TF_ReluGradOp AnyTensor:$gradients, AnyTensor:$features),
+          (MHLO_SelectOp
+            (MHLO_CompareOp $features, (MHLO_ConstantLike<"0">:$zero $features),
+              MHLO_ComparisonDirectionValue<"GT">, (MHLO_DEFAULT_COMPARISON_TYPE)),
+            $gradients, $zero)>;
+
+//===----------------------------------------------------------------------===//
+// Softsign op patterns.
+//===----------------------------------------------------------------------===//
+
+/// Converts a TF::SoftsignOp to HLO.
+/// Softsign(features) = features / (1 + abs(features))
+def : Pat<(TF_SoftsignOp AnyTensor:$input),
+          (MHLO_DivOp
+            $input,
+            (MHLO_AddOp (MHLO_ConstantLike<"1"> $input), (MHLO_AbsOp $input))
+          )
+         >;
+
+/// Converts a TF::SoftsignGradOp to HLO.
+/// SoftsignGrad(gradient, features) = gradient / ((1 + abs(features)) ^ 2)
+def : Pattern<
+        (TF_SoftsignGradOp AnyRankedTensor:$gradients, AnyRankedTensor:$features),
+        [(CHLO_BroadcastAddOp:$add
+          (MHLO_ConstantOp:$one (GetScalarOfType<1> $features)), (MHLO_AbsOp $features),
+          (BinBroadcastDimensions $one, $features)
+         ),
+         (CHLO_BroadcastDivOp
+           $gradients,
+           (MHLO_MulOp $add, $add),
+           (BinBroadcastDimensions $gradients, $add)
+         )
+        ]>;
+
+//===----------------------------------------------------------------------===//
+// Slice op patterns.
+//===----------------------------------------------------------------------===//
+
+def UnpackStartingIndices: NativeCodeCall<
+  "UnpackTensorAlongZeroDim($0.getLoc(), $1, &$_builder).getOutput()">;
+
+def CanBeTranslatedToDynamicSlice : Constraint<CPred<
+  "CanBeTranslatedToDynamicSlice($0, $1, llvm::cast<DenseIntElementsAttr>($2))">>;
+
+def TFSliceSizes2HLOSliceSizes : NativeCodeCall<
+    "TFSliceSizes2HLOSliceSizes($0, $1, llvm::cast<DenseIntElementsAttr>($2),"
+    "&$_builder)">;
+
+def : Pat<(TF_SliceOp:$op MHLO_Tensor:$input, MHLO_Tensor:$starting_indices,
+           (ConstantLikeMatcher AnyAttr:$slice_sizes)),
+          (MHLO_DynamicSliceOp $input,
+           (UnpackStartingIndices $op, $starting_indices),
+           (TFSliceSizes2HLOSliceSizes $input, $starting_indices, $slice_sizes)),
+          [(CanBeTranslatedToDynamicSlice $input, $starting_indices,
+            $slice_sizes)]>;
+
+//===----------------------------------------------------------------------===//
+// Select op patterns.
+//===----------------------------------------------------------------------===//
+
+ def : Pat<(TF_SelectV2Op MHLO_Tensor:$pred, MHLO_Tensor:$on_true,
+            MHLO_Tensor:$on_false),
+           (CHLO_BroadcastSelectOp $pred, $on_true, $on_false)>;
+
+//===----------------------------------------------------------------------===//
+// PartitionedCall and LegacyCall op patterns.
+//===----------------------------------------------------------------------===//
+
+def ArgTypesMatchCallee : Constraint<
+    // $0 is a resultset (possibly empty), and $_op isn't assigned. So retrieve
+    // the op using the builder.
+    CPred<"ArgTypesMatchCallee(&*$_builder.getInsertionPoint(), $1, $2)">>;
+
+foreach callOp = [TF_PartitionedCallOp, TF_StatefulPartitionedCallOp] in {
+  def : Pat<(callOp:$op $args, $args_attrs, $res_attrs, FlatSymbolRefAttr:$f,
+             $config, $config_proto, $executor_type),
+            (CallOp $f, $args, $args_attrs, $res_attrs, ConstantAttr<UnitAttr, "false">),
+          [(ArgTypesMatchCallee $op, $args, $f)]>;
+}
+
+// The extra attr on this op is _disable_call_shape_inference, which we ignore
+// in the bridge.
+def : Pat<(TF_LegacyCallOp:$op $args, $args_attrs, $res_attrs, FlatSymbolRefAttr:$f, $attr),
+          (CallOp $f, $args, $args_attrs, $res_attrs, ConstantAttr<UnitAttr, "false">),
+        [(ArgTypesMatchCallee $op, $args, $f)]>;
+
+//===----------------------------------------------------------------------===//
+// Reverse op patterns.
+//===----------------------------------------------------------------------===//
+
+// Handles axis conversion for TF reverse.
+def ConvertAxisAttr : NativeCodeCall<"ConvertAxisAttr($0, llvm::cast<ElementsAttr>($1), &$_builder)">;
+
+def : Pat<(TF_ReverseV2Op AnyRankedTensor:$values, (ConstantLikeMatcher ElementsAttr:$axis)),
+    (MHLO_ReverseOp $values, (ConvertAxisAttr $values, $axis))>;
+
+//===----------------------------------------------------------------------===//
+// Unary op patterns.
+//===----------------------------------------------------------------------===//
+
+foreach Mapping = [
+                   [TF_AbsOp, MHLO_AbsOp],
+                   [TF_CeilOp, MHLO_CeilOp],
+                   [TF_ComplexAbsOp, MHLO_AbsOp],
+                   [TF_ErfOp, MHLO_ErfOp],
+                   [TF_FloorOp, MHLO_FloorOp],
+                   [TF_ImagOp, MHLO_ImagOp],
+                   [TF_InvertOp, MHLO_NotOp],
+                   [TF_IsFiniteOp, MHLO_IsFiniteOp],
+                   [TF_LogicalNotOp, MHLO_NotOp],
+                   [TF_NegOp, MHLO_NegOp],
+                   [TF_RealOp, MHLO_RealOp],
+                  ] in {
+ def : Pat<(Mapping[0] MHLO_Tensor:$input),
+           (Mapping[1] $input)>;
+}
+
+foreach Mapping = [
+                  [TF_CosOp, MHLO_CosineOp],
+                  [TF_ExpOp, MHLO_ExpOp],
+                  [TF_Expm1Op, MHLO_Expm1Op],
+                  [TF_LogOp, MHLO_LogOp],
+                  [TF_Log1pOp, MHLO_Log1pOp],
+                  [TF_RsqrtOp, MHLO_RsqrtOp],
+                  [TF_SigmoidOp, MHLO_LogisticOp],
+                  [TF_SinOp, MHLO_SineOp],
+                  [TF_SqrtOp, MHLO_SqrtOp],
+                  [TF_TanhOp, MHLO_TanhOp],
+                  [TF_TanOp, MHLO_TanOp],
+                  ] in {
+ def : Pat<(Mapping[0] MHLO_Tensor:$input),
+           (Mapping[1] $input, ConstDefaultResultAccuracyAttr)>;
+}
+
+foreach Mapping = [
+                   [TF_AcosOp, CHLO_AcosOp],
+                   [TF_AcoshOp, CHLO_AcoshOp],
+                   [TF_AsinOp, CHLO_AsinOp],
+                   [TF_AsinhOp, CHLO_AsinhOp],
+                   [TF_AtanOp, CHLO_AtanOp],
+                   [TF_AtanhOp, CHLO_AtanhOp],
+                   [TF_CoshOp, CHLO_CoshOp],
+                   [TF_ConjOp, CHLO_ConjOp],
+                   [TF_DigammaOp, CHLO_DigammaOp],
+                   [TF_ErfcOp, CHLO_ErfcOp],
+                   [TF_IsInfOp, CHLO_IsInfOp],
+                   [TF_LgammaOp, CHLO_LgammaOp],
+                   [TF_SinhOp, CHLO_SinhOp],
+                  ] in {
+ def : Pat<(Mapping[0] MHLO_AnyTensor:$input),
+           (Mapping[1] $input)>;
+}
+
+def : Pat<(TF_AngleOp $x), (MHLO_Atan2Op (MHLO_ImagOp $x), (MHLO_RealOp $x))>;
+
+// TODO(bixia): Lower with Truncate=True for floating point value conversions.
+def : Pat<(TF_CastOp $arg, ConstBoolAttrFalse), (MHLO_ConvertOp $arg)>;
+
+def : Pat<(TF_TransposeOp:$res $arg, (ConstantLikeMatcher ElementsAttr:$permutation)),
+          (MHLO_TransposeOp $arg, (CastElementsToI64Elements $permutation))>;
+
+
+// Lowering these ops with static shape to mhlo.reshape
+foreach TfOp = [TF_ExpandDimsOp, TF_ReshapeOp, TF_SqueezeOp, ] in {
+  def : Pat<(TfOp:$res MHLO_Tensor:$arg, $ignored),
+            (MHLO_ReshapeOp $arg), [(AnyStaticShapeTensor $res)], [],
+            (addBenefit 2)>;
+}
+
+// Returns NaN if x is NaN, 0 if x is 0, -1 if x < 0 and 1 if x > 0.
+def : Pat<(TF_SignOp $x), (MHLO_SignOp $x)>;
+
+def BothElementTypesSameWidthIntOrFloat : Constraint<CPred<
+  "getElementTypeOrSelf($0.getType()).isIntOrFloat() && "
+  "getElementTypeOrSelf($1.getType()).isIntOrFloat()">,
+  "element types must be integers or floats">;
+
+// TODO(mgester): Due to restrictions of xla::BitcastConvertType we currently
+// only lower if both input and output types are int or float and have same width
+
+def : Pat<(TF_BitcastOp:$res MHLO_Tensor:$arg),
+          (MHLO_BitcastConvertOp $arg),
+          [(BothElementTypesSameWidthIntOrFloat $res, $arg)]>;
+
+// TODO(jpienaar): Lower constant like to constant to broadcast if dynamic
+// and going to MHLO.
+
+//===----------------------------------------------------------------------===//
+// Random ops.
+//===----------------------------------------------------------------------===//
+// TODO(b/148269299): handle random number generator seeds/states correctly.
+
+class MHLO_RngDistributionValue<string enumStr> :
+  ConstantAttr<MHLO_RngDistributionAttr, "::mlir::mhlo::RngDistribution::" # enumStr>;
+
+def : Pat<(TF_RandomUniformOp:$old $shape, $seed, $seed2),
+          (MHLO_RngOp
+            (MHLO_ConstantOp
+              (NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 0.0)">)),
+            (MHLO_ConstantOp
+              (NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 1.0)">)),
+            (CastValueToI64 $old, $shape),
+            MHLO_RngDistributionValue<"UNIFORM">),
+          [(IsShapedTensor $shape)]>;
+
+def : Pat<(TF_RandomStandardNormalOp:$old $shape, $seed, $seed2),
+          (MHLO_RngOp
+            (MHLO_ConstantOp
+              (NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 0.0)">)),
+            (MHLO_ConstantOp
+              (NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 1.0)">)),
+            (CastValueToI64 $old, $shape),
+            MHLO_RngDistributionValue<"NORMAL">),
+          [(IsShapedTensor $shape)]>;
+
+//===----------------------------------------------------------------------===//
+// Sigmoid grad op.
+//===----------------------------------------------------------------------===//
+
+// TODO(hinsu): Handle unranked inputs by broadcasting constant one to the
+// shape of $l instead of having it as a constant.
+def : Pat<(TF_SigmoidGradOp AnyRankedTensor:$l, AnyRankedTensor:$r),
+          (MHLO_MulOp
+           (MHLO_MulOp $r, $l),
+           (MHLO_SubtractOp (MHLO_ConstantOp (ConstantSplat<"1"> $l)), $l))>;
+
+//===----------------------------------------------------------------------===//
+// Softplus op.
+//===----------------------------------------------------------------------===//
+
+def EpsilonValue : NativeCodeCall<"GetEpsilonValue($0.getType())">;
+
+def : Pattern<(TF_SoftplusOp AnyTensor:$features),
+              [
+                (MHLO_ExpOp:$features_exp $features, ConstDefaultResultAccuracyAttr),
+                (CHLO_BroadcastAddOp:$threshold
+                 (MHLO_LogOp (MHLO_ConstantOp (EpsilonValue $features)), ConstDefaultResultAccuracyAttr),
+                 (MHLO_ConstantOp (GetScalarOfType<2> $features)),
+                 (NullDenseI64ArrayAttr)
+                ),
+                (MHLO_SelectOp:$output
+                 (CHLO_BroadcastCompareOp
+                  $features,
+                  (MHLO_NegOp $threshold),
+                  (NullDenseI64ArrayAttr),
+                  CHLO_ComparisonDirectionValue<"GT">,
+                  (CHLO_DEFAULT_COMPARISON_TYPE)
+                 ),
+                 $features,
+                 (MHLO_SelectOp
+                  (CHLO_BroadcastCompareOp
+                   $features,
+                   $threshold,
+                   (NullDenseI64ArrayAttr),
+                   CHLO_ComparisonDirectionValue<"LT">,
+                   (CHLO_DEFAULT_COMPARISON_TYPE)
+                  ),
+                  $features_exp,
+                  (MHLO_Log1pOp $features_exp, ConstDefaultResultAccuracyAttr)
+                 )
+                ),
+                (replaceWithValue $output)
+              ]>;
+
+//===----------------------------------------------------------------------===//
+// XlaReplicaId op.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_XlaReplicaIdOp),
+          (TF_CastOp (MHLO_ReplicaIdOp), /*truncate=*/ConstBoolAttrFalse)>;
+
+//===----------------------------------------------------------------------===//
+// XlaGather op.
+//===----------------------------------------------------------------------===//
+
+def ToGatherDimNumsAttr : NativeCodeCall<"GetGatherDimNumsAttr($0, &$_builder)">;
+
+def HasValidGatherDims : Constraint<CPred<"HasValidGatherDims($0)">>;
+
+def : Pat<(TF_XlaGatherOp $operand, $start_indices, (ConstantLikeMatcher ElementsAttr:$slice_sizes),
+                          $dimension_numbers, $indices_are_sorted),
+          (MHLO_GatherOp $operand, $start_indices,
+                        (ToGatherDimNumsAttr $dimension_numbers),
+                        (CastElementsToI64Elements $slice_sizes),
+                        $indices_are_sorted),
+          [(HasValidGatherDims $dimension_numbers)]>;
+
+//===----------------------------------------------------------------------===//
+// XlaDotOp op.
+//===----------------------------------------------------------------------===//
+
+def ToDotDimNumsAttr : NativeCodeCall<"GetDotDimNumsAttr($0, &$_builder)">;
+
+def ToPrecisionConfigsAttr : NativeCodeCall<"GetPrecisionConfigAttr($0, &$_builder)">;
+
+def HasValidDotDims : Constraint<CPred<"HasValidDotDims($0)">>;
+
+def HasValidPrecisionConfig : Constraint<CPred<"HasValidPrecisionConfig($0)">>;
+
+def : Pat<(TF_XlaDotOp $lhs, $rhs, $dimension_numbers, $precision_config),
+          (MHLO_DotGeneralOp $lhs, $rhs,
+                            (ToDotDimNumsAttr $dimension_numbers),
+                            (ToPrecisionConfigsAttr $precision_config),
+                            (EmptyDotAlgorithmAttr)),
+          [(HasValidDotDims $dimension_numbers), (HasValidPrecisionConfig $precision_config)]>;
+
+//===----------------------------------------------------------------------===//
+// XlaDotV2Op op.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_XlaDotV2Op $lhs, $rhs, $dimension_numbers, $precision_config),
+          (MHLO_DotGeneralOp $lhs, $rhs,
+                            (ToDotDimNumsAttr $dimension_numbers),
+                            (ToPrecisionConfigsAttr $precision_config),
+                            (EmptyDotAlgorithmAttr)),
+          [(HasValidDotDims $dimension_numbers), (HasValidPrecisionConfig $precision_config)]>;
+
+//===----------------------------------------------------------------------===//
+// XlaDynamicSlice op.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_XlaDynamicSliceOp:$op MHLO_Tensor:$input, MHLO_Tensor:$starting_indices,
+           (ConstantLikeMatcher AnyAttr:$slice_sizes)),
+          (MHLO_DynamicSliceOp $input,
+           (UnpackStartingIndices $op, $starting_indices),
+           (TFSliceSizes2HLOSliceSizes $input, $starting_indices, $slice_sizes))>;
+
+//===----------------------------------------------------------------------===//
+// XlaEisumOp op.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_XlaEinsumOp $lhs, $rhs, $equation),
+          (MHLO_EinsumOp $lhs, $rhs, $equation)>;
+
+//===----------------------------------------------------------------------===//
+// XlaOptimizationBarrierOp op.
+//===----------------------------------------------------------------------===//
+
+def : Pat<(TF_XlaOptimizationBarrierOp $args),
+          (MHLO_OptimizationBarrierOp $args)>;

tensorflow/compiler/mlir/stablehlo/transforms/legalize_tf_xla_call_module_to_stablehlo_pass.cc

@@ -0,0 +1,266 @@
+/* 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.
+==============================================================================*/
+// LINT.IfChange
+#include "tensorflow/compiler/mlir/stablehlo/transforms/legalize_tf_xla_call_module_to_stablehlo_pass.h"
+
+#include <cassert>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.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/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.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 "stablehlo/dialect/Serialization.h"  // from @stablehlo
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+#include "stablehlo/dialect/VhloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace mlir {
+namespace stablehlo {
+
+static constexpr absl::string_view kStablehloModuleDefaultEntryFuncName =
+    "main";
+static constexpr absl::string_view kStablehloFuncNamePrefix = "XlaCallModule";
+static constexpr char kShardingAttr[] = "mhlo.sharding";
+static constexpr char kShardingName[] = "Sharding";
+
+class RemoveCustomCallWithSharding
+    : public OpRewritePattern<stablehlo::CustomCallOp> {
+  using OpRewritePattern<stablehlo::CustomCallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(stablehlo::CustomCallOp op,
+                                PatternRewriter &rewriter) const override {
+    // Removes the custom call with sharding op if the operand type is the
+    // same as the result type.
+    if (op->hasAttr(kShardingAttr) && op.getCallTargetName() == kShardingName &&
+        op.getNumOperands() == 1 && op.getNumResults() == 1 &&
+        op.getOperands().front().getType() ==
+            op.getResults().front().getType()) {
+      rewriter.replaceOp(op, op.getOperands());
+      return success();
+    }
+    return failure();
+  }
+};
+
+namespace {
+
+bool IsShloMainFuncOp(func::FuncOp func_op) {
+  if (func_op == nullptr) {
+    return false;
+  }
+
+  if (!func_op.getSymName().contains(kStablehloModuleDefaultEntryFuncName)) {
+    return false;
+  }
+
+  if (func_op.getSymVisibility() == "nested" ||
+      func_op.getSymVisibility() == "private") {
+    return false;
+  }
+
+  return true;
+}
+
+// Returns true if XlaCallModuleOp has the "platform index argument". The
+// platform index argument is an extra 0-dimensional i32 tensor argument at
+// index 0 when the XlaCallModuleOp contains more than one platform specified at
+// the "platform" attribute.
+//
+// See:
+// https://github.com/tensorflow/tensorflow/blob/eba24f41ba9d661d2f58a515921720cf90708cd4/tensorflow/compiler/tf2xla/ops/xla_ops.cc#L1376-L1385
+bool ContainsPlatformIndexArg(TF::XlaCallModuleOp xla_call_module_op) {
+  return xla_call_module_op.getPlatforms().size() > 1;
+}
+
+}  // namespace
+
+class ConvertTFXlaCallModuleOp : public OpRewritePattern<TF::XlaCallModuleOp> {
+ public:
+  explicit ConvertTFXlaCallModuleOp(MLIRContext *context, ModuleOp module_op)
+      : OpRewritePattern<TF::XlaCallModuleOp>(context), module_op_(module_op) {}
+  using OpRewritePattern<TF::XlaCallModuleOp>::OpRewritePattern;
+
+ private:
+  ModuleOp module_op_;
+  LogicalResult matchAndRewrite(TF::XlaCallModuleOp op,
+                                PatternRewriter &rewriter) const override {
+    OwningOpRef<ModuleOp> stablehlo_module_op =
+        stablehlo::deserializePortableArtifact(op.getModuleAttr(),
+                                               getContext());
+    if (stablehlo_module_op.get() == nullptr) {
+      return failure();
+    }
+    SymbolTable parent_module_symbol_table(module_op_);
+    SymbolTable stablehlo_module_symbol_table(stablehlo_module_op.get());
+    {
+      auto main_func_op = stablehlo_module_symbol_table.lookup<func::FuncOp>(
+          kStablehloModuleDefaultEntryFuncName);
+      // TODO(b/291988976): move enforcement of this variable outside of this
+      // rewrite pattern such that it's only checked once. Currently, this
+      // approach results in duplicate error messages as this pattern executes
+      // more than once.
+      if (!IsShloMainFuncOp(main_func_op)) {
+        auto error_msg =
+            "'main' FuncOp in XlaCallModuleOp missing or has visibility other "
+            "than 'public'";
+        if (main_func_op) {
+          main_func_op->emitError(error_msg);
+        }
+        return rewriter.notifyMatchFailure(op, error_msg);
+      }
+    }
+    Builder stablehlo_builder(stablehlo_module_op.get().getContext());
+    // Rename XlaCallModuleOp's functions to avoid naming conflicts.
+    for (auto func_op : stablehlo_module_op.get().getOps<func::FuncOp>()) {
+      const std::string new_func_name =
+          CreateNewFuncName(func_op.getSymName(), parent_module_symbol_table);
+      if (failed(stablehlo_module_symbol_table.replaceAllSymbolUses(
+              func_op, stablehlo_builder.getStringAttr(new_func_name),
+              stablehlo_module_op.get()))) {
+        return failure();
+      }
+      SymbolTable::setSymbolName(func_op, new_func_name);
+    }
+    // Move all functions from XlaCallModuleOp's stablehlo module, to parent
+    // module. Also marks the stablehlo module entry function as private.
+    func::FuncOp main_fn;
+    for (auto func_op : stablehlo_module_op.get().getOps<func::FuncOp>()) {
+      func::FuncOp cloned_func_op = func_op.clone();
+      if (IsShloMainFuncOp(cloned_func_op)) {
+        main_fn = cloned_func_op;
+      }
+      cloned_func_op.setSymVisibility(
+          stablehlo_builder.getStringAttr("private"));
+      parent_module_symbol_table.insert(cloned_func_op);
+    }
+
+    // When the `XlaCallModuleOp`'s callee accepts a platform index argument,
+    // add a dummy platform index argument in order to match the number of
+    // the arguments of the callee function.
+    //
+    // This is because `XlaCallModuleOp` doesn't explicitly take it as an
+    // operand. See:
+    // https://github.com/tensorflow/tensorflow/blob/eba24f41ba9d661d2f58a515921720cf90708cd4/tensorflow/compiler/tf2xla/ops/xla_ops.cc#L1376-L1385
+
+    SmallVector<Value, 4> call_op_operands(op.getOperands());
+    if (ContainsPlatformIndexArg(op)) {
+      Value dummy_const = TF::ConstOp::create(
+          rewriter, op.getLoc(),
+          DenseIntElementsAttr::get(
+              RankedTensorType::get({}, rewriter.getIntegerType(32)), {0}));
+      call_op_operands.insert(call_op_operands.begin(), dummy_const);
+    }
+
+    // The stablehlo module main function's input tensor types might be
+    // different from the XlaCallModuleOp's input tensor types. For example,
+    // The XlaCallModuleOp's input is tensor<*xf32> while the function's
+    // argument type is tensor<1x2f32>.
+    SmallVector<Value, 4> casted_operands;
+    casted_operands.reserve(main_fn.getNumArguments());
+    assert(call_op_operands.size() == main_fn.getNumArguments());
+    for (const auto &operand_and_type :
+         zip(call_op_operands, main_fn.getFunctionType().getInputs())) {
+      Value operand = std::get<0>(operand_and_type);
+      Type expected_type = std::get<1>(operand_and_type);
+      if (operand.getType() != expected_type) {
+        operand =
+            TF::CastOp::create(rewriter, op.getLoc(), expected_type, operand,
+                               /*Truncate=*/rewriter.getBoolAttr(false));
+      }
+      casted_operands.push_back(operand);
+    }
+
+    auto call =
+        func::CallOp::create(rewriter, op->getLoc(), main_fn.getSymName(),
+                             main_fn.getResultTypes(), casted_operands);
+    rewriter.replaceOp(op, call->getResults());
+
+    return success();
+  }
+
+  // Creates a new function name to avoid collision. The naming scheme is
+  // XlaCallModule_%s_%d where %s is the original function name and %d is the
+  // counter.
+  std::string CreateNewFuncName(const StringRef func_name,
+                                SymbolTable &symbol_table) const {
+    int suffix_id = 0;
+    std::string new_func_name = absl::StrCat(kStablehloFuncNamePrefix, "_",
+                                             func_name.str(), "_", suffix_id);
+    while (symbol_table.lookup(new_func_name)) {
+      suffix_id++;
+      new_func_name = absl::StrCat(kStablehloFuncNamePrefix, "_",
+                                   func_name.str(), "_", suffix_id);
+    }
+    return new_func_name;
+  }
+};
+
+class TFXlaCallModuleOpToStablehloPass
+    : public PassWrapper<TFXlaCallModuleOpToStablehloPass,
+                         OperationPass<ModuleOp>> {
+ public:
+  StringRef getArgument() const final {
+    return "tf-xla-callmodule-op-to-stablehlo-pass";
+  }
+  StringRef getDescription() const final {
+    return "Legalize TF_XlaCallModule Op to stablehlo";
+  }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<stablehlo::StablehloDialect, vhlo::VhloDialect,
+                    quant::QuantDialect, shape::ShapeDialect>();
+  }
+
+  void runOnOperation() override {
+    ModuleOp module_op = getOperation();
+    RewritePatternSet patterns(&getContext());
+    patterns.add<ConvertTFXlaCallModuleOp>(&getContext(), module_op);
+    patterns.add<RemoveCustomCallWithSharding>(&getContext());
+    if (failed(applyPatternsGreedily(module_op, std::move(patterns)))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+std::unique_ptr<OperationPass<ModuleOp>>
+CreateLegalizeTFXlaCallModuleToStablehloPass() {
+  return std::make_unique<TFXlaCallModuleOpToStablehloPass>();
+}
+
+static PassRegistration<TFXlaCallModuleOpToStablehloPass> pass;
+
+}  // namespace stablehlo
+}  // namespace mlir
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_tf_xla_call_module_to_stablehlo_pass.cc)

tensorflow/compiler/mlir/stablehlo/transforms/legalize_tf_xla_call_module_to_stablehlo_pass.h

@@ -0,0 +1,36 @@
+/* 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.
+==============================================================================*/
+// LINT.IfChange
+#ifndef TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_LEGALIZE_TF_XLA_CALL_MODULE_TO_STABLEHLO_PASS_H_
+#define TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_LEGALIZE_TF_XLA_CALL_MODULE_TO_STABLEHLO_PASS_H_
+
+#include <memory>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace stablehlo {
+
+// Adds passes which transform TF_XlaCallModule Op to StableHLO Ops.
+// Note that this pass only supports static shape tensors for now.
+std::unique_ptr<mlir::OperationPass<ModuleOp>>
+CreateLegalizeTFXlaCallModuleToStablehloPass();
+
+}  // namespace stablehlo
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_LEGALIZE_TF_XLA_CALL_MODULE_TO_STABLEHLO_PASS_H_
+// LINT.ThenChange(//tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_tf_xla_call_module_to_stablehlo_pass.h)

tensorflow/compiler/mlir/stablehlo/transforms/mhlo_passes/README.md

@@ -0,0 +1,5 @@
+This temporary directory was created to store MHLO pass .cc and .h files. These
+files have been migrated to StableHLO but are still used by inactive or
+potentially outdated compilation paths. Once all MHLO passes have been migrated
+to StableHLO, revisit this directory. At that point, we can replace the uses of
+MHLO passes from this directory with the StableHLO passes.

tensorflow/compiler/mlir/stablehlo/transforms/mhlo_passes/fuse_convolution_pass.cc

@@ -0,0 +1,201 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/stablehlo/transforms/mhlo_passes/fuse_convolution_pass.h"
+
+#include <iterator>
+#include <memory>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/Dialect/Traits.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.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/quantization/common/attrs_and_constraints.h"
+#include "tensorflow/compiler/mlir/utils/validators.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace odml {
+
+class FuseMhloMulAndConvolutionPattern : public OpRewritePattern<mhlo::MulOp> {
+ public:
+  using OpRewritePattern<mhlo::MulOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::MulOp mul_op,
+                                PatternRewriter &rewriter) const override {
+    // Variables for capturing values and attributes used while creating ops.
+    mhlo::ConvolutionOp conv_op;
+    Operation *bcast_or_const_op;
+    shape::ShapeOfOp shape_of_op;
+    mhlo::ConstantOp filter;
+    mhlo::ConstantOp multiplier;
+    mlir::ElementsAttr filter_value, mul_value;
+    mlir::DenseIntElementsAttr broadcast_dims;
+
+    // Match and capture values/attributes.
+    Value lhs = mul_op.getLhs();
+    Value rhs = mul_op.getRhs();
+    conv_op = lhs.getDefiningOp<mhlo::ConvolutionOp>();
+    if (conv_op == nullptr) {
+      return failure();
+    }
+    filter = conv_op.getRhs().getDefiningOp<mhlo::ConstantOp>();
+    if (filter == nullptr) {
+      return failure();
+    }
+    // Try to match static broadcast or dynamic broadcast.
+    bcast_or_const_op = rhs.getDefiningOp();
+    bool is_dynamic_broadcast =
+        isa<mhlo::DynamicBroadcastInDimOp>(bcast_or_const_op);
+    multiplier = isa<mhlo::ConstantOp>(bcast_or_const_op)
+                     ? dyn_cast_or_null<mhlo::ConstantOp>(bcast_or_const_op)
+                     : bcast_or_const_op->getOperand(0)
+                           .getDefiningOp<mhlo::ConstantOp>();
+    if (multiplier == nullptr) {
+      return failure();
+    }
+
+    auto result_type = OpTrait::util::getBroadcastedType(filter.getType(),
+                                                         multiplier.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(mul_op, [&](::mlir::Diagnostic &diag) {
+        diag << "entities 'filter, multiplier' failed to satisfy constraint: "
+                "non-broadcastable operands";
+      });
+    }
+    filter_value = filter.getValue();
+    mul_value = multiplier.getValue();
+    // In MHLO, Conv filter is in HWIO format, Depthwise conv filter is in HW1O
+    // format and backprop input conv filter is in HWOI format.
+    // Only fuses multiplier if all dimensions other than the out channel
+    // dimension are equal to 1.
+    if (!TF::IsDimensionsDegenerateExceptLastOne(
+            mul_value.getShapedType().getShape())) {
+      return rewriter.notifyMatchFailure(mul_op, [&](::mlir::Diagnostic &diag) {
+        diag << "entities 'mul_value' failed to satisfy constraint: "
+                "unsupported dimensions";
+      });
+    }
+    if (!is_dynamic_broadcast &&
+        !((*conv_op.getODSResults(0).begin()).hasOneUse())) {
+      return rewriter.notifyMatchFailure(mul_op, [&](::mlir::Diagnostic &diag) {
+        diag << "entities 'conv' failed to satisfy constraint: has one use";
+      });
+    }
+    // For dynamic case, the result of conv should be used by shape_of and mul.
+    if (is_dynamic_broadcast) {
+      auto conv_uses = (*conv_op.getODSResults(0).begin()).getUses();
+      if (std::distance(conv_uses.begin(), conv_uses.end()) != 2 ||
+          quant::FindUserOfType<shape::ShapeOfOp>(conv_op) == nullptr ||
+          quant::FindUserOfType<mhlo::MulOp>(conv_op) == nullptr) {
+        return rewriter.notifyMatchFailure(mul_op, [&](::mlir::Diagnostic
+                                                           &diag) {
+          diag << "entities 'conv' failed to satisfy constraint: has two uses "
+                  "for dynamic case";
+        });
+      }
+    }
+
+    // Rewrite
+    // For dynamic case, we use filter's shape to create a static broadcast.
+    broadcast_dims =
+        !isa<mhlo::ConstantOp>(bcast_or_const_op) && !is_dynamic_broadcast
+            ? dyn_cast_or_null<mhlo::BroadcastInDimOp>(bcast_or_const_op)
+                  .getBroadcastDimensions()
+            : nullptr;
+    if (broadcast_dims == nullptr) {
+      const auto filter_rank = filter_value.getShapedType().getRank();
+      auto dimsType = RankedTensorType::get({1}, rewriter.getIntegerType(64));
+      broadcast_dims = DenseIntElementsAttr::get(dimsType, {filter_rank - 1});
+    }
+    Value broadcast_multiplier = mhlo::BroadcastInDimOp::create(
+        rewriter, mul_op.getLoc(), filter.getType(), multiplier,
+        broadcast_dims);
+    Value new_filter =
+        mhlo::MulOp::create(rewriter, mul_op.getLoc(), filter.getType(), filter,
+                            broadcast_multiplier);
+    Value new_conv = mhlo::ConvolutionOp::create(
+        rewriter, mul_op.getLoc(), conv_op.getType(), conv_op.getLhs(),
+        new_filter, conv_op.getWindowStridesAttr(), conv_op.getPaddingAttr(),
+        conv_op.getLhsDilationAttr(), conv_op.getRhsDilationAttr(),
+        conv_op.getWindowReversalAttr(), conv_op.getDimensionNumbers(),
+        conv_op.getFeatureGroupCount(), conv_op.getBatchGroupCount(),
+        conv_op.getPrecisionConfigAttr());
+    // For static case, replace the convolution op now.
+    if (!is_dynamic_broadcast) {
+      rewriter.replaceOp(mul_op, {new_conv});
+    } else {
+      // For dynamic case, create new shape_of op and replace uses.
+      shape_of_op =
+          dyn_cast_or_null<mhlo::DynamicBroadcastInDimOp>(bcast_or_const_op)
+              .getOutputDimensions()
+              .getDefiningOp<shape::ShapeOfOp>();
+      // Check if the shape come from the original conv op.
+      if (!shape_of_op ||
+          shape_of_op.getArg().getDefiningOp<mhlo::ConvolutionOp>() !=
+              conv_op) {
+        return failure();
+      }
+      Value new_shape_of = shape::ShapeOfOp::create(
+          rewriter, mul_op.getLoc(), shape_of_op.getType(), new_conv);
+      shape_of_op.replaceAllUsesWith(new_shape_of);
+      rewriter.replaceOp(mul_op, {new_conv});
+    }
+
+    return success();
+  }
+};
+
+class FuseMhloConvolutionPass
+    : public PassWrapper<FuseMhloConvolutionPass, OperationPass<func::FuncOp>> {
+ public:
+  StringRef getArgument() const final { return "fuse-mhlo-convolution-pass"; }
+  StringRef getDescription() const final {
+    return "Fuses MHLO binary element-wise ops and convolution op";
+  }
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<FuseMhloMulAndConvolutionPattern>(&getContext());
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+std::unique_ptr<Pass> createFuseConvolutionPass() {
+  return std::make_unique<FuseMhloConvolutionPass>();
+}
+
+static PassRegistration<FuseMhloConvolutionPass> pass;
+
+}  // namespace odml
+}  // namespace mlir

tensorflow/compiler/mlir/stablehlo/transforms/mhlo_passes/fuse_convolution_pass.h

@@ -0,0 +1,32 @@
+/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_MHLO_PASSES_FUSE_CONVOLUTION_PASS_H_
+#define TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_MHLO_PASSES_FUSE_CONVOLUTION_PASS_H_
+
+#include <memory>
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace odml {
+
+// Fuses MHLO binary element-wise ops and convolution op.
+std::unique_ptr<Pass> createFuseConvolutionPass();
+
+}  // namespace odml
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_MHLO_PASSES_FUSE_CONVOLUTION_PASS_H_

tensorflow/compiler/mlir/stablehlo/transforms/mhlo_passes/unfuse_batch_norm_pass.cc

@@ -0,0 +1,266 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/stablehlo/transforms/mhlo_passes/unfuse_batch_norm_pass.h"
+
+#include <cassert>
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/ImplicitLocOpBuilder.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Pass/Pass.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/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace odml {
+
+// Broadcasts the 1D value tensor 'value_1d' to the shape of 'result_type'. If
+// 'shape_value' is initialized, creates a dynamic broadcast, otherwise creates
+// a static broadcast.
+Value broadcastToFeatureDim(Location loc, RankedTensorType result_type,
+                            Value value1d, Value shape_value,
+                            int64_t feature_dim, PatternRewriter &rewriter) {
+  auto dims_type =
+      RankedTensorType::get(/*shape=*/{1}, rewriter.getIntegerType(64));
+  auto dims = DenseIntElementsAttr::get(dims_type, {feature_dim});
+  if (shape_value) {
+    return rewriter.createOrFold<mhlo::DynamicBroadcastInDimOp>(
+        loc, result_type, value1d, shape_value, dims);
+  }
+  assert(result_type.hasStaticShape());
+  return mhlo::BroadcastInDimOp::create(rewriter, loc, result_type, value1d,
+                                        dims);
+}
+
+// Gets the shape of operand, assuming it is a dynamic shape with static rank.
+Value getShapeValue(Location loc, Value operand, PatternRewriter &rewriter) {
+  RankedTensorType resultType =
+      mlir::dyn_cast<RankedTensorType>(operand.getType());
+  return shape::ShapeOfOp::create(
+      rewriter, loc,
+      RankedTensorType::get(/*shape=*/{resultType.getRank()},
+                            rewriter.getIndexType()),
+      operand);
+}
+
+Value materializeEpsilon(Operation *op, FloatAttr epsilon_attr,
+                         FloatType fp_type, Value broadcast_to,
+                         RankedTensorType broadcast_to_type,
+                         PatternRewriter &rewriter) {
+  ImplicitLocOpBuilder b(op->getLoc(), rewriter);
+  if (epsilon_attr.getType() != fp_type) {
+    // Need to convert.
+    bool loses_info;
+    APFloat epsilon_float = epsilon_attr.getValue();
+    auto status = epsilon_float.convert(
+        fp_type.getFloatSemantics(), APFloat::rmNearestTiesToEven, &loses_info);
+    if ((status & (~APFloat::opInexact)) != APFloat::opOK) {
+      op->emitWarning() << "Could not convert batch_norm epsilon to target fp "
+                           "type: opStatus = "
+                        << static_cast<int>(status);
+      return nullptr;
+    }
+    if (loses_info) {
+      op->emitWarning("Conversion of epsilon loses precision");
+    }
+    epsilon_attr = b.getFloatAttr(fp_type, epsilon_float);
+  }
+
+  auto scalar_type = RankedTensorType::get(/*shape=*/{}, fp_type);
+  auto epsilon_tensor_attr = DenseElementsAttr::get(
+      scalar_type, {mlir::cast<Attribute>(epsilon_attr)});
+  Value epsilon = mhlo::ConstantOp::create(b, epsilon_tensor_attr);
+  auto dims_type = RankedTensorType::get(/*shape=*/{0}, b.getIntegerType(64));
+  auto dims = DenseIntElementsAttr::get(dims_type, SmallVector<int64_t, 1>{});
+  if (broadcast_to_type.hasStaticShape()) {
+    return mhlo::BroadcastInDimOp::create(b, broadcast_to_type, epsilon, dims);
+  }
+  Value shape_value = getShapeValue(op->getLoc(), broadcast_to, rewriter);
+  return b.createOrFold<mhlo::DynamicBroadcastInDimOp>(
+      broadcast_to_type, epsilon, shape_value, dims);
+}
+
+class UnfuseBatchNormTrainingPattern
+    : public OpRewritePattern<mhlo::BatchNormTrainingOp> {
+ public:
+  using OpRewritePattern<mhlo::BatchNormTrainingOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::BatchNormTrainingOp bn_op,
+                                PatternRewriter &rewriter) const override {
+    auto inputs = bn_op.getOperand();
+    auto input_type = mlir::dyn_cast<RankedTensorType>(inputs.getType());
+    if (!input_type) {
+      return failure();
+    }
+    auto feature_index = bn_op.getFeatureIndex();
+
+    // Compute mean
+    int64_t input_last_dim = input_type.getRank() - 1;
+    auto dims_type = RankedTensorType::get(/*shape=*/{input_last_dim},
+                                           rewriter.getIntegerType(32));
+    ::mlir::SmallVector<int32_t> reduce_dim_axes;
+    for (int i = 0; i < input_type.getRank(); ++i) {
+      if (i != feature_index) {
+        reduce_dim_axes.push_back(i);
+      }
+    }
+    auto mean_dims = DenseIntElementsAttr::get(dims_type, reduce_dim_axes);
+    ::mlir::TF::ConstOp reduce_dim_op =
+        TF::ConstOp::create(rewriter, bn_op.getLoc(), mean_dims);
+    int64_t feature_dim_size = input_type.getDimSize(feature_index);
+    auto mean_var_type = RankedTensorType::get(/*shape=*/{feature_dim_size},
+                                               rewriter.getF32Type());
+    ::mlir::Value mean = TF::MeanOp::create(
+        rewriter, bn_op.getLoc(), mean_var_type, inputs, reduce_dim_op,
+        /*keep_dims=*/rewriter.getBoolAttr(false));
+
+    // Compute variance
+    Value shape_value =
+        getShapeValue(bn_op.getLoc(), bn_op.getOperand(), rewriter);
+    auto broadcast_mean = broadcastToFeatureDim(
+        bn_op.getLoc(), input_type, mean, shape_value, feature_index, rewriter);
+    ::mlir::Value square_diff = TF::SquaredDifferenceOp::create(
+        rewriter, bn_op.getLoc(), inputs, broadcast_mean);
+    ::mlir::Value variance = TF::MeanOp::create(
+        rewriter, bn_op.getLoc(), mean_var_type, square_diff, reduce_dim_op,
+        /*keep_dims=*/rewriter.getBoolAttr(false));
+
+    // Invoke BatchNormInferenceOp
+    ::mlir::FloatAttr epsilon = bn_op.getEpsilonAttr();
+    ::mlir::Value batch_norm = mhlo::BatchNormInferenceOp::create(
+        rewriter, bn_op.getLoc(), inputs, bn_op.getScale(), bn_op.getOffset(),
+        mean, variance, epsilon, rewriter.getI64IntegerAttr(feature_index));
+
+    // Return normalized values, mean, variable.
+    rewriter.replaceOp(bn_op, ::mlir::ValueRange{batch_norm, mean, variance});
+    return success();
+  }
+};
+
+class UnfuseBatchNormInferencePattern
+    : public OpRewritePattern<mhlo::BatchNormInferenceOp> {
+ public:
+  using OpRewritePattern<mhlo::BatchNormInferenceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::BatchNormInferenceOp bn_op,
+                                PatternRewriter &rewriter) const override {
+    // Enforce type invariants.
+    // Note that we deduce the actual element type from the variance,
+    // which should not be subject to quantization at a higher level.
+    auto input_type =
+        mlir::dyn_cast<RankedTensorType>(bn_op.getOperand().getType());
+    auto variance_type =
+        mlir::dyn_cast<RankedTensorType>(bn_op.getVariance().getType());
+    if (!input_type || !variance_type) {
+      return failure();
+    }
+    auto fp_type = mlir::dyn_cast<FloatType>(variance_type.getElementType());
+    if (!fp_type) {
+      return failure();
+    }
+
+    // result = (x - mean) * scale / sqrt(variance + epsilon) + offset
+    // Let multiplier = scale / sqrt(variance + epsilon), to compute
+    // (x - mean) * scale / sqrt(variance + epsilon) + offset,
+    // is then to compute (x * multiplier) + (offset - mean * multiplier).
+
+    auto epsilon = materializeEpsilon(
+        bn_op.getOperation(), bn_op.getEpsilonAttr(), fp_type,
+        bn_op.getVariance(), variance_type, rewriter);
+    if (!epsilon) {
+      return failure();
+    }
+
+    // Compute multiplier = scale / sqrt(variance + epsilon)
+    Value multiplier = mhlo::AddOp::create(rewriter, bn_op.getLoc(),
+                                           bn_op.getVariance(), epsilon);
+    multiplier = mhlo::RsqrtOp::create(rewriter, bn_op.getLoc(), multiplier);
+    multiplier = mhlo::MulOp::create(rewriter, bn_op.getLoc(), multiplier,
+                                     bn_op.getScale());
+
+    // Compute rhs = offset - mean * multiplier
+    Value rhs = mhlo::MulOp::create(rewriter, bn_op.getLoc(), multiplier,
+                                    bn_op.getMean());
+    rhs = mhlo::SubtractOp::create(rewriter, bn_op.getLoc(), bn_op.getOffset(),
+                                   rhs);
+
+    // Broadcast `multiplier` and `rhs`
+    Value shape_value;
+    if (!input_type.hasStaticShape()) {
+      shape_value = getShapeValue(bn_op.getLoc(), bn_op.getOperand(), rewriter);
+    }
+    int64_t feature_dim = bn_op.getFeatureIndex();
+    auto broadcast_multiplier =
+        broadcastToFeatureDim(bn_op.getLoc(), input_type, multiplier,
+                              shape_value, feature_dim, rewriter);
+
+    // Computes x * multiplier + rhs
+    Value lhs = mhlo::MulOp::create(rewriter, bn_op.getLoc(),
+                                    bn_op.getOperand(), broadcast_multiplier);
+    auto broadcast_rhs = broadcastToFeatureDim(
+        bn_op.getLoc(), input_type, rhs, shape_value, feature_dim, rewriter);
+    rewriter.replaceOpWithNewOp<mhlo::AddOp>(bn_op, lhs, broadcast_rhs);
+
+    return success();
+  }
+};
+
+class UnfuseMhloBatchNormPass
+    : public PassWrapper<UnfuseMhloBatchNormPass, OperationPass<func::FuncOp>> {
+ public:
+  StringRef getArgument() const final { return "unfuse-mhlo-batch-norm-pass"; }
+  StringRef getDescription() const final {
+    return "Unfuses MHLO batch norm inference op into arithmetic ops";
+  }
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    // TODO(b/299514833): Remove TensorFlowDialect usage.
+    registry.insert<shape::ShapeDialect, mlir::TF::TensorFlowDialect>();
+  }
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<UnfuseBatchNormTrainingPattern>(&getContext());
+    patterns.add<UnfuseBatchNormInferencePattern>(&getContext());
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+std::unique_ptr<Pass> createUnfuseBatchNormPass() {
+  return std::make_unique<UnfuseMhloBatchNormPass>();
+}
+
+static PassRegistration<UnfuseMhloBatchNormPass> pass;
+
+}  // namespace odml
+}  // namespace mlir

tensorflow/compiler/mlir/stablehlo/transforms/mhlo_passes/unfuse_batch_norm_pass.h

@@ -0,0 +1,32 @@
+/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_MHLO_PASSES_UNFUSE_BATCH_NORM_PASS_H_
+#define TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_MHLO_PASSES_UNFUSE_BATCH_NORM_PASS_H_
+
+#include <memory>
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace odml {
+
+// Unfuses MHLO batch norm inference op into arithmetic ops.
+std::unique_ptr<Pass> createUnfuseBatchNormPass();
+
+}  // namespace odml
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_MHLO_PASSES_UNFUSE_BATCH_NORM_PASS_H_

tensorflow/compiler/mlir/stablehlo/transforms/rename_entrypoint_to_main.cc

@@ -0,0 +1,110 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/stablehlo/transforms/rename_entrypoint_to_main.h"
+
+#include <memory>
+#include <string>
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
+
+namespace mlir {
+namespace odml {
+
+class RenameEntrypointToMainPass
+    : public PassWrapper<RenameEntrypointToMainPass, OperationPass<ModuleOp>> {
+ public:
+  StringRef getArgument() const final { return "rename-entrypoint-to-main"; }
+  StringRef getDescription() const final {
+    return "Renames the entrypoint in SavedModel to `main`";
+  }
+
+  void runOnOperation() override {
+    auto fail = [&](Operation* op, std::string message) {
+      op->emitError(message);
+      signalPassFailure();
+    };
+
+    DenseMap<StringRef, func::FuncOp> entrypoints;
+    auto module = getOperation();
+    module.walk([&](func::FuncOp op) {
+      auto visibility = SymbolTable::getSymbolVisibility(op);
+      if (visibility != SymbolTable::Visibility::Public) return;
+      entrypoints[op.getSymName()] = op;
+    });
+
+    if (auto session_initializer =
+            tf_saved_model::GetSessionInitializerOp(module)) {
+      // clang-format off
+      // Skip session initializer functions which are present in saved models.
+      // For example:
+      //   "tf_saved_model.session_initializer"() {initializers = [@init_all_tables]} : () -> () // NOLINT
+      //   func @init_all_tables() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init_all_tables"]} { // NOLINT
+      //     "tf.NoOp"() {device = ""} : () -> ()
+      //     return
+      //   }
+      // clang-format on
+      for (auto attr : session_initializer.getInitializers()) {
+        auto sym_attr = mlir::dyn_cast<FlatSymbolRefAttr>(attr);
+        if (!sym_attr) break;
+        entrypoints.erase(sym_attr.getValue());
+      }
+    }
+
+    if (entrypoints.empty()) {
+      return fail(module, "No entrypoints found");
+    }
+    if (entrypoints.size() == 1) {
+      auto entrypoint = entrypoints.begin()->second;
+      Builder builder(entrypoint);
+      entrypoint.setName(builder.getStringAttr("main"));
+      return;
+    }
+
+    // In case we have more than 1 entry points, choose the one with
+    // 'tf.entry_function' attribute set.
+    llvm::SmallVector<func::FuncOp, 4> candidate_funcs;
+    for (auto& entrypoint : entrypoints) {
+      if (entrypoint.second->hasAttr("tf.entry_function")) {
+        candidate_funcs.push_back(entrypoint.second);
+      }
+    }
+
+    if (candidate_funcs.empty()) {
+      return fail(module, "No entrypoints found");
+    }
+    if (candidate_funcs.size() > 1) {
+      return fail(module, "Too many entrypoints found");
+    }
+    // Found entrypoint
+    Builder builder(candidate_funcs[0]);
+    candidate_funcs[0].setName(builder.getStringAttr("main"));
+  }
+};
+
+std::unique_ptr<Pass> CreateRenameEntrypointToMainPass() {
+  return std::make_unique<RenameEntrypointToMainPass>();
+}
+
+static PassRegistration<RenameEntrypointToMainPass> pass;
+
+}  // namespace odml
+}  // namespace mlir

tensorflow/compiler/mlir/stablehlo/transforms/rename_entrypoint_to_main.h

@@ -0,0 +1,31 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_RENAME_ENTRYPOINT_TO_MAIN_H_
+#define TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_RENAME_ENTRYPOINT_TO_MAIN_H_
+
+#include <memory>
+
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace mlir {
+namespace odml {
+
+std::unique_ptr<Pass> CreateRenameEntrypointToMainPass();
+
+}  // namespace odml
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_RENAME_ENTRYPOINT_TO_MAIN_H_

tensorflow/compiler/mlir/stablehlo/transforms/tf_stablehlo_pass.cc

@@ -0,0 +1,184 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/stablehlo/transforms/tf_stablehlo_pass.h"
+
+#include <memory>
+#include <utility>
+
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo
+#include "stablehlo/dialect/Register.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/stablehlo/transforms/legalize_tf_passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/lower_tf.h"
+#include "tensorflow/compiler/mlir/tf2xla/transforms/legalize_tf_with_tf2xla_passes.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+#include "xla/mlir_hlo/mhlo/IR/register.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+#include "xla/mlir_hlo/mhlo/transforms/rewriters.h"
+#include "xla/mlir_hlo/mhlo/utils/type_conversion.h"
+
+namespace mlir {
+namespace odml {
+
+class TFToMhloPass
+    : public mlir::PassWrapper<TFToMhloPass,
+                               mlir::OperationPass<mlir::func::FuncOp>> {
+ public:
+  explicit TFToMhloPass(bool skip_quantization_ops = false,
+                        bool skip_resize = false,
+                        bool skip_partitioned_calls = false)
+      : PassWrapper() {
+    skip_quantization_ops_ = skip_quantization_ops;
+    skip_resize_ = skip_resize;
+    skip_partitioned_calls_ = skip_partitioned_calls;
+  }
+
+  TFToMhloPass(const TFToMhloPass &pass) {
+    skip_quantization_ops_ = pass.skip_quantization_ops_;
+    skip_resize_ = pass.skip_resize_;
+    skip_partitioned_calls_ = pass.skip_partitioned_calls_;
+  }
+
+ private:
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    mlir::mhlo::registerAllMhloDialects(registry);
+    mlir::stablehlo::registerAllDialects(registry);
+    registry.insert<mlir::func::FuncDialect, mlir::arith::ArithDialect>();
+    registry.insert<shape::ShapeDialect>();
+  }
+
+ public:
+  StringRef getArgument() const final { return "tf-mhlo"; }
+  StringRef getDescription() const final {
+    return "This pass will legalize TF Ops to MHLO Ops.";
+  }
+
+ protected:
+  Option<bool> skip_quantization_ops_{
+      *this, "skip-quantization-ops",
+      ::llvm::cl::desc("Skip quantization ops")};
+
+  Option<bool> skip_resize_{
+      *this, "skip-resize",
+      ::llvm::cl::desc("Skip tf.ResizeBilinear and tf.ResizeNearestNeighbor")};
+
+  Option<bool> skip_partitioned_calls_{
+      *this, "skip-partitioned-calls",
+      ::llvm::cl::desc(
+          "Skip tf.StatefulPartitionedCall and tf.PartitionedCall")};
+};
+
+void TFToMhloPass::runOnOperation() {
+  auto func = getOperation();
+  MLIRContext *context = func->getContext();
+
+  RewritePatternSet patterns(context);
+  odml::PopulateLegalizeTfPatterns(context, &patterns);
+  TF::PopulateTFLoweringBeforeHLOPatterns(context, &patterns);
+  mhlo::Tf2XlaTypeConverter converter;
+  mhlo::PopulateLegalizeTfWithTf2XlaPatterns(
+      "XLA_CPU_JIT", patterns, context, converter, /*prefer_tf2xla=*/false);
+  stablehlo::StablehloToHloTypeConverter hlo_converter;
+  chlo::populateChloToHloPatterns(context, &hlo_converter, &patterns);
+  chlo::ConstantLikeOp::getCanonicalizationPatterns(patterns, context);
+
+  ConversionTarget target(*context);
+  target.addIllegalDialect<chlo::ChloDialect>();
+  target.addLegalDialect<mhlo::MhloDialect>();
+  target.addLegalDialect<arith::ArithDialect>();
+  target.addLegalDialect<func::FuncDialect>();
+  target.addLegalDialect<tensor::TensorDialect>();
+  target.addLegalDialect<shape::ShapeDialect>();
+  target.addLegalOp<func::CallOp>();
+
+  if (skip_quantization_ops_) {
+    target.addLegalOp<TF::FakeQuantWithMinMaxVarsOp>();
+    target.addLegalOp<TF::FakeQuantWithMinMaxVarsPerChannelOp>();
+    target.addLegalOp<TF::FakeQuantWithMinMaxArgsOp>();
+    target.addLegalOp<TF::QuantizeAndDequantizeV2Op>();
+    target.addLegalOp<TF::QuantizeAndDequantizeV3Op>();
+  }
+  if (skip_resize_) {
+    target.addLegalOp<TF::ResizeBilinearOp>();
+    target.addLegalOp<TF::ResizeNearestNeighborOp>();
+  }
+  if (skip_partitioned_calls_) {
+    target.addLegalOp<TF::PartitionedCallOp>();
+    target.addLegalOp<TF::StatefulPartitionedCallOp>();
+  }
+
+  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
+  if (failed(applyPartialConversion(func, target, frozen_patterns))) {
+    return signalPassFailure();
+  }
+}
+
+struct TFToStablehloOptions : public PassPipelineOptions<TFToStablehloOptions> {
+  Option<bool> skip_quantization_ops{*this, "skip-quantization-ops",
+                                     ::llvm::cl::desc("Skip quantization ops")};
+  Option<bool> skip_resize{
+      *this, "skip-resize",
+      ::llvm::cl::desc("Skip tf.ResizeBilinear and tf.ResizeNearestNeighbor")};
+  Option<bool> skip_partitioned_calls{
+      *this, "skip-partitioned-calls",
+      ::llvm::cl::desc(
+          "Skip tf.StatefulPartitionedCall and tf.PartitionedCall")};
+};
+
+void PopulateLegalizeTFToStablehloPipeline(
+    OpPassManager &pm, const TFToStablehloOptions &options) {
+  // TODO(burmako): Migrate this pass from producing MHLO to producing StableHLO
+  // by aligning with the TF/XLA bridge on the corresponding functionality and
+  // reusing their work, perhaps through `LowerToMlProgramAndHlo`.
+  pm.addNestedPass<func::FuncOp>(std::make_unique<TFToMhloPass>(
+      options.skip_quantization_ops, options.skip_resize,
+      options.skip_partitioned_calls));
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(mhlo::createHloLegalizeToStablehloPass());
+}
+
+static PassPipelineRegistration<TFToStablehloOptions>
+    legalize_tf_to_stablehlo_pipeline("tf-stablehlo",
+                                      "Legalize TF ops to StableHLO ops",
+                                      PopulateLegalizeTFToStablehloPipeline);
+
+void AddLegalizeTFToStablehloPasses(OpPassManager &pm,
+                                    bool skip_quantization_ops,
+                                    bool skip_resize,
+                                    bool skip_partitioned_calls) {
+  TFToStablehloOptions options;
+  options.skip_quantization_ops = skip_quantization_ops;
+  options.skip_resize = skip_resize;
+  options.skip_partitioned_calls = skip_partitioned_calls;
+  PopulateLegalizeTFToStablehloPipeline(pm, options);
+}
+
+}  // namespace odml
+}  // namespace mlir

tensorflow/compiler/mlir/stablehlo/transforms/tf_stablehlo_pass.h

@@ -0,0 +1,33 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_TF_STABLEHLO_PASS_H_
+#define TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_TF_STABLEHLO_PASS_H_
+
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+
+namespace mlir {
+namespace odml {
+
+// Adds passes which transform TF Ops to StableHLO Ops.
+void AddLegalizeTFToStablehloPasses(OpPassManager& pm,
+                                    bool skip_quantization_ops,
+                                    bool skip_resize,
+                                    bool skip_partitioned_calls);
+
+}  // namespace odml
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_TF_STABLEHLO_PASS_H_

tensorflow/compiler/mlir/stablehlo/transforms/utils.cc

@@ -0,0 +1,55 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/stablehlo/transforms/utils.h"
+
+#include <cstdint>
+
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+#include "xla/mlir_hlo/utils/hlo_utils.h"
+
+namespace mlir {
+namespace odml {
+
+mhlo::ConstantOp GetScalarConstOfType(Type ty, Location loc, int64_t raw_value,
+                                      OpBuilder* builder) {
+  return mhlo::ConstantOp::create(*builder, loc,
+                                  hlo::getScalarOfType(ty, raw_value));
+}
+
+mhlo::ConstantOp GetScalarNegZeroOfType(Type ty, Location loc,
+                                        OpBuilder* builder) {
+  return mhlo::ConstantOp::create(*builder, loc,
+                                  hlo::getScalarNegZeroOfType(ty));
+}
+
+DenseIntElementsAttr GetI64ElementsAttr(ArrayAttr attr) {
+  RankedTensorType ty =
+      RankedTensorType::get(static_cast<int64_t>(attr.size()),
+                            IntegerType::get(attr.getContext(), 64));
+  return DenseIntElementsAttr::get(ty, attr.getValue());
+}
+
+DenseIntElementsAttr GetI64ElementsAttr(ArrayRef<int64_t> values,
+                                        Builder* builder) {
+  RankedTensorType ty = RankedTensorType::get(
+      {static_cast<int64_t>(values.size())}, builder->getIntegerType(64));
+  return DenseIntElementsAttr::get(ty, values);
+}
+
+}  // namespace odml
+}  // namespace mlir

tensorflow/compiler/mlir/stablehlo/transforms/utils.h

@@ -0,0 +1,63 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_UTILS_H_
+
+#include <cstdint>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace odml {
+
+// Builds body for reduce op by using the template binary op as the
+// reducer op.
+template <typename Op>
+void BuildReduceBody(Type element_type, Region* body, OpBuilder* builder) {
+  OpBuilder::InsertionGuard guard(*builder);
+  Block* block = builder->createBlock(body);
+
+  // Block arguments are scalars of the given element type.
+  Type type = RankedTensorType::get(/*shape=*/{}, element_type);
+  Location loc = body->getLoc();
+  block->addArguments({type, type}, SmallVector<Location, 2>(2, loc));
+
+  auto reducer =
+      Op::create(*builder, loc, block->getArgument(0), block->getArgument(1));
+  mhlo::ReturnOp::create(*builder, loc, reducer.getResult());
+}
+
+mhlo::ConstantOp GetScalarConstOfType(Type ty, Location loc, int64_t raw_value,
+                                      OpBuilder* builder);
+
+mhlo::ConstantOp GetScalarNegZeroOfType(Type ty, Location loc,
+                                        OpBuilder* builder);
+
+// Converts an ArrayAttr to a 1D 64-bit dense elements attribute.
+DenseIntElementsAttr GetI64ElementsAttr(ArrayAttr attr);
+DenseIntElementsAttr GetI64ElementsAttr(llvm::ArrayRef<int64_t> values,
+                                        Builder* builder);
+
+}  // namespace odml
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_STABLEHLO_TRANSFORMS_UTILS_H_

tensorflow/compiler/mlir/stablehlo/transforms/utils_test.cc

@@ -0,0 +1,82 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/stablehlo/transforms/utils.h"
+
+#include <cstdint>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace odml {
+namespace {
+
+TEST(UtilsTest, GetScalarConstOfType) {
+  MLIRContext context;
+  context.loadDialect<mlir::mhlo::MhloDialect>();
+  OpBuilder builder(&context);
+  Location loc = UnknownLoc::get(&context);
+  Type ty = builder.getI32Type();
+  mhlo::ConstantOp op = GetScalarConstOfType(ty, loc, 123, &builder);
+  EXPECT_EQ(op.getValue().getValues<int32_t>()[0], 123);
+
+  op->destroy();
+}
+
+TEST(UtilsTest, GetScalarNegZeroOfType) {
+  MLIRContext context;
+  context.loadDialect<mlir::mhlo::MhloDialect>();
+  OpBuilder builder(&context);
+  Location loc = UnknownLoc::get(&context);
+  Type ty = builder.getF32Type();
+  mhlo::ConstantOp op = GetScalarNegZeroOfType(ty, loc, &builder);
+  EXPECT_EQ(op.getValue().getValues<float>()[0], -0.f);
+
+  op->destroy();
+}
+
+TEST(UtilsTest, GetI64ElementsAttr) {
+  MLIRContext context;
+  context.loadDialect<mlir::mhlo::MhloDialect>();
+  OpBuilder builder(&context);
+  Location loc = UnknownLoc::get(&context);
+  SmallVector<int64_t> values = {1, 2, 3};
+  auto valuesAttr = builder.getI64ArrayAttr(values);
+  DenseIntElementsAttr attr = GetI64ElementsAttr(valuesAttr);
+  EXPECT_THAT(SmallVector<int64_t>(attr.getValues<int64_t>()),
+              testing::ElementsAreArray(values));
+}
+
+TEST(UtilsTest, GetI64ElementsAttrBuilder) {
+  MLIRContext context;
+  context.loadDialect<mlir::mhlo::MhloDialect>();
+  OpBuilder builder(&context);
+  Location loc = UnknownLoc::get(&context);
+  SmallVector<int64_t> values = {1, 2, 3};
+  DenseIntElementsAttr attr = GetI64ElementsAttr(values, &builder);
+  EXPECT_THAT(SmallVector<int64_t>(attr.getValues<int64_t>()),
+              testing::ElementsAreArray(values));
+}
+
+}  // namespace
+
+}  // namespace odml
+}  // namespace mlir