chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/transforms/attr_lowering_utils.cc

@@ -0,0 +1,192 @@
+/* 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/tfrt/transforms/attr_lowering_utils.h"
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_attributes.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tfrt/core_runtime/opdefs/attributes.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace {
+
+mlir::TypeAttr ConvertTypeAttribute(mlir::TypeAttr type_attr,
+                                    mlir::Builder& builder) {
+  auto type = type_attr.getValue();
+
+  if (IsSupportedTfrtNumericDType(type)) return type_attr;
+
+  // For TF custom types, we convert it to custom corert types.
+  if (mlir::isa<mlir::TF::StringType>(type))
+    return mlir::TypeAttr::get(
+        tfrt::corert::StringType::get(builder.getContext()));
+
+  if (mlir::isa<mlir::TF::ResourceType>(type))
+    return mlir::TypeAttr::get(
+        tfrt::corert::ResourceType::get(builder.getContext()));
+
+  if (mlir::isa<mlir::TF::VariantType>(type))
+    return mlir::TypeAttr::get(
+        tfrt::corert::VariantType::get(builder.getContext()));
+
+  if (mlir::isa<mlir::TF::Quint8Type>(type)) {
+    return mlir::TypeAttr::get(
+        tfrt::corert::Quint8Type::get(builder.getContext()));
+  }
+
+  if (mlir::isa<mlir::TF::Quint16Type>(type)) {
+    return mlir::TypeAttr::get(
+        tfrt::corert::Quint16Type::get(builder.getContext()));
+  }
+
+  if (mlir::isa<mlir::TF::Qint8Type>(type)) {
+    return mlir::TypeAttr::get(
+        tfrt::corert::Qint8Type::get(builder.getContext()));
+  }
+
+  if (mlir::isa<mlir::TF::Qint16Type>(type)) {
+    return mlir::TypeAttr::get(
+        tfrt::corert::Qint16Type::get(builder.getContext()));
+  }
+
+  if (mlir::isa<mlir::TF::Qint32Type>(type)) {
+    return mlir::TypeAttr::get(
+        tfrt::corert::Qint32Type::get(builder.getContext()));
+  }
+
+  // Return invalid results to emit error for unsupported types.
+  return {};
+}
+
+mlir::Attribute ConvertAttribute(mlir::Attribute attr, mlir::Builder& builder) {
+  // The supported attributes here should be kept consistent with
+  // //third_party/tf_runtime/include/tfrt/core_runtime/op_attr_type.h
+  //
+  // Currently, not all tensorflow data types are supported. Unranked shape
+  // attributes are not supported yet.
+
+  // Return directly if the attribute is already supported.
+  if (mlir::isa<mlir::IntegerAttr, mlir::FloatAttr, mlir::BoolAttr,
+                mlir::StringAttr, mlir::DenseIntOrFPElementsAttr>(attr))
+    return attr;
+
+  // For type attributes, we convert non-standard MLIR types to corresponding
+  // corert types.
+  if (auto type_attr = mlir::dyn_cast<mlir::TypeAttr>(attr)) {
+    if (auto shape_type =
+            mlir::dyn_cast<mlir::TensorType>(type_attr.getValue())) {
+      if (!shape_type.hasRank())
+        return tfrt::corert::ShapeAttr::get(builder.getContext());
+
+      return tfrt::corert::ShapeAttr::get(builder.getContext(),
+                                          shape_type.getShape());
+    }
+
+    return ConvertTypeAttribute(type_attr, builder);
+  }
+
+  // Convert the attribute to the corresponding format in TFRT dialect if
+  // needed.
+  if (auto shape_attr = mlir::dyn_cast<mlir::TF::ShapeAttr>(attr)) {
+    if (!shape_attr.hasRank())
+      return tfrt::corert::ShapeAttr::get(builder.getContext());
+    return tfrt::corert::ShapeAttr::get(builder.getContext(),
+                                        shape_attr.getShape());
+  }
+
+  // For arrays, we recursively convert the elements.
+  if (auto array_attr = mlir::dyn_cast<mlir::ArrayAttr>(attr)) {
+    llvm::SmallVector<mlir::Attribute, 8> attrs;
+    attrs.reserve(array_attr.size());
+    for (auto attr : array_attr) {
+      auto converted = ConvertAttribute(attr, builder);
+      if (!converted) return {};
+      attrs.push_back(converted);
+    }
+    return builder.getArrayAttr(attrs);
+  }
+
+  return {};
+}
+
+}  // namespace
+
+bool IsSupportedTfrtNumericDType(mlir::Type type) {
+  // Most of the tensorflow data types (eg. f32, i64) are supported and they
+  // are standard MLIR types that need no conversion here.
+  if (type.isBF16() || type.isF16() || type.isF32() || type.isF64() ||
+      type.isInteger(1) || type.isInteger(8) || type.isInteger(16) ||
+      type.isInteger(32) || type.isInteger(64) || type.isUnsignedInteger(8) ||
+      type.isUnsignedInteger(16) || type.isUnsignedInteger(32) ||
+      type.isUnsignedInteger(64))
+    return true;
+
+  if (auto complex_type = mlir::dyn_cast<mlir::ComplexType>(type)) {
+    auto element_type = complex_type.getElementType();
+    if (element_type.isF32() || element_type.isF64()) return true;
+  }
+
+  return false;
+}
+
+// TODO(chky): attributes "_output_shapes" should be removed by any tool that
+// generates TF MLIR dialect, as they are not used by CoreRuntime. Remove this
+// filtering logic once unused attributes are cleaned up in the upper layer.
+bool IsUnusedTfrtAttribute(llvm::StringRef name) {
+  // NOTE: attributes "f.*" are function attribute related and
+  // are added during importing graph to MLIR TF Executor dialect. These
+  // attributes are not actually used by TF ops with function attributes.
+  // TODO(b/180399811): Re-evaluate the usage of these attributes.
+  static const char* const kUnusedAttributes[] = {
+      "_output_shapes",
+      "resultSegmentSizes",
+      "operandSegmentSizes",
+  };
+
+  for (auto attr : kUnusedAttributes) {
+    if (name == attr) {
+      return true;
+    }
+  }
+
+  return name.contains("f.");
+}
+
+mlir::ArrayAttr CreateTfrtOpAttrs(llvm::ArrayRef<mlir::NamedAttribute> attrs,
+                                  mlir::Builder& builder) {
+  llvm::SmallVector<mlir::Attribute, 4> attr_array;
+  for (auto key_and_value : attrs) {
+    if (!IsUnusedTfrtAttribute(key_and_value.getName())) {
+      auto converted = ConvertAttribute(key_and_value.getValue(), builder);
+      if (!converted) return {};
+
+      mlir::StringAttr key =
+          builder.getStringAttr(key_and_value.getName().strref());
+      attr_array.push_back(builder.getArrayAttr({key, converted}));
+    }
+  }
+  return builder.getArrayAttr(attr_array);
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/attr_lowering_utils.h

@@ -0,0 +1,42 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_ATTR_LOWERING_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_ATTR_LOWERING_UTILS_H_
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#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
+
+namespace tensorflow {
+
+// TODO(chky): attributes "_output_shapes" should be removed by any tool that
+// generates TF MLIR dialect, as they are not used by CoreRuntime. Remove this
+// filtering logic once unused attributes are cleaned up in the upper layer.
+bool IsUnusedTfrtAttribute(llvm::StringRef name);
+
+// Create a single attribute that contains the named attribute lists. It is an
+// array of pairs. The key must be a string attribute, and the value can be
+// any attribute that is supported by CoreRuntime.
+mlir::ArrayAttr CreateTfrtOpAttrs(llvm::ArrayRef<mlir::NamedAttribute> attrs,
+                                  mlir::Builder& builder);
+
+bool IsSupportedTfrtNumericDType(mlir::Type type);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_ATTR_LOWERING_UTILS_H_

tensorflow/compiler/mlir/tfrt/transforms/corert_converter.cc

@@ -0,0 +1,229 @@
+/* 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/tfrt/transforms/corert_converter.h"
+
+#include <optional>
+#include <string>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/Interfaces/DerivedAttributeOpInterface.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/attr_lowering_utils.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/utils.h"
+#include "tensorflow/core/util/device_name_utils.h"
+#include "tfrt/basic_kernels/opdefs/basic_kernels.h"  // from @tf_runtime
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+CoreRTConverter::CoreRTConverter(
+    mlir::MLIRContext *context,
+    const mlir::TF::SideEffectAnalysis::Info *side_effect_analysis)
+    : builder_(context), side_effect_analysis_(*side_effect_analysis) {
+  addConversion([](tfrt::compiler::ChainType type) { return type; });
+  addConversion([](tfrt::corert::OpHandlerType type) { return type; });
+  addConversion([](tfrt::corert::TensorHandleType type) { return type; });
+  addConversion([=](mlir::TensorType type) -> std::optional<mlir::Type> {
+    // Ref types are not supported in both compiler and runtime.
+    if (mlir::isa<mlir::TF::TensorFlowRefType>(type.getElementType()))
+      return std::nullopt;
+    return tensor_handle_type();
+  });
+  addConversion([=](mlir::Type type) -> std::optional<mlir::Type> {
+    if (type == builder_.getI1Type()) return type;
+    return std::nullopt;
+  });
+}
+
+void CoreRTConverter::MaterializeDerivedAttributes(mlir::Operation *op) {
+  if (auto interface = llvm::dyn_cast<mlir::DerivedAttributeOpInterface>(op)) {
+    auto derived_attrs = interface.materializeDerivedAttributes();
+    for (auto named_attr : derived_attrs) {
+      op->setAttr(named_attr.getName(), named_attr.getValue());
+    }
+  }
+}
+
+mlir::ArrayAttr CoreRTConverter::CreateOpFuncAttrs(
+    const mlir::SymbolTable &symbol_table, ArrayRef<NamedAttribute> attrs,
+    llvm::SmallVector<mlir::StringAttr, 4> *func_attr_keys,
+    bool use_mlir_func_name) {
+  llvm::SmallVector<mlir::Attribute, 4> attr_array;
+  for (auto key_and_value : attrs) {
+    auto attr_key = key_and_value.getName();
+    auto attr_value = key_and_value.getValue();
+    if (!IsUnusedTfrtAttribute(attr_key) &&
+        mlir::isa<mlir::FlatSymbolRefAttr, mlir::SymbolRefAttr>(attr_value)) {
+      auto func_attr = mlir::dyn_cast<mlir::FlatSymbolRefAttr>(attr_value);
+      auto converted = CanonicalizeTensorflowFunctionName(
+          symbol_table, func_attr.getValue().str(), use_mlir_func_name);
+      if (!converted) return {};
+
+      mlir::StringAttr key = builder_.getStringAttr(attr_key.strref());
+      attr_array.push_back(
+          builder_.getArrayAttr({key, builder_.getStringAttr(*converted)}));
+
+      // Remove the attribute to avoid being converted again.
+      func_attr_keys->push_back(attr_key);
+    }
+  }
+  return builder_.getArrayAttr(attr_array);
+}
+
+// TODO(chky): Add support for multiple device instances.
+std::optional<ParseDeviceNameResult> CoreRTConverter::ParseDeviceName(
+    llvm::StringRef device_name) const {
+  std::string tf_device_name = device_name.str();
+
+  if (tf_device_name.empty()) {
+    return std::nullopt;
+  }
+
+  ParseDeviceNameResult result;
+  result.device_name = tf_device_name;
+
+  // Parse the device name in format of the current tensorflow.
+  DeviceNameUtils::ParsedName parsed_name;
+  if (!DeviceNameUtils::ParseFullName(result.device_name, &parsed_name)) {
+    return std::nullopt;
+  }
+  if (!parsed_name.has_type) {
+    return std::nullopt;
+  }
+  result.device_type = parsed_name.type;
+
+  result.op_handler_name = tf_device_name;
+
+  return result;
+}
+
+std::optional<ParseDeviceNameResult> CoreRTConverter::ParseDeviceName(
+    mlir::Operation *op) const {
+  auto device_attr = op->getAttr("device");
+  if (!device_attr) {
+    return std::nullopt;
+  }
+
+  auto parsed_device_name =
+      ParseDeviceName(mlir::cast<mlir::StringAttr>(device_attr).getValue());
+  if (!parsed_device_name) op->emitWarning("failed to parse device name.");
+  return parsed_device_name;
+}
+
+mlir::Value CoreRTConverter::ConvertOpHandler(
+    mlir::Operation *op, llvm::StringRef op_handler_name,
+    ConversionPatternRewriter *rewriter) {
+  auto iter = op_handler_by_name_.find(op_handler_name);
+  if (iter != op_handler_by_name_.end()) return iter->second;
+
+  mlir::Block *block = op->getBlock();
+  ConversionPatternRewriter::InsertionGuard insertion_guard(*rewriter);
+  rewriter->setInsertionPointToStart(block);
+
+  func::FuncOp func_op = op->getParentOfType<mlir::func::FuncOp>();
+  mlir::Value in_chain = func_op.getArgument(0);
+  auto get_op_handler_op = tfrt::corert::GetOpHandler::create(
+      *rewriter, block->getParent()->getLoc(), op_handler_type(), in_chain,
+      op_handler_name);
+  op_handler_by_name_[op_handler_name] = get_op_handler_op.getResult();
+  return get_op_handler_op.getResult();
+}
+
+mlir::Value CoreRTConverter::GetDistributedContext(
+    mlir::Operation *op, mlir::ConversionPatternRewriter *rewriter) {
+  mlir::func::FuncOp func_op = op->getParentOfType<mlir::func::FuncOp>();
+  auto iter = distributed_context_by_func_.find(func_op.getOperation());
+  if (iter != distributed_context_by_func_.end()) {
+    return iter->second;
+  }
+  return mlir::Value();
+}
+
+mlir::Value CoreRTConverter::GetRemoteChainManager(
+    mlir::Operation *op, mlir::ConversionPatternRewriter *rewriter) {
+  mlir::func::FuncOp func_op = op->getParentOfType<mlir::func::FuncOp>();
+  auto iter = remote_chain_mgr_by_func_.find(func_op.getOperation());
+  if (iter != remote_chain_mgr_by_func_.end()) {
+    return iter->second;
+  }
+  return mlir::Value();
+}
+
+mlir::Value CoreRTConverter::GetLocalSideEffectChain(
+    mlir::Operation *op, mlir::ConversionPatternRewriter *rewriter) {
+  auto func_op = op->getParentOfType<mlir::func::FuncOp>();
+
+  llvm::SmallVector<mlir::Operation *, 4> predecessors;
+  if (llvm::isa<mlir::func::ReturnOp>(op)) {
+    auto sinks = side_effect_analysis_.ControlSinks();
+    predecessors.assign(sinks.begin(), sinks.end());
+  } else {
+    predecessors = side_effect_analysis_.DirectControlPredecessors(op);
+  }
+
+  llvm::SmallVector<mlir::Value, 2> chains;
+  for (auto *pred : predecessors) {
+    // TODO(chky): ReadVariableOp is removed in the pass and not converted.
+    // Ideally, every side-effecting op should be converted to a
+    // tfrt_fallback.executeop.seq op. The special rewrite logic of
+    // ReadVariableOp should be done in a previous pass.
+    if (auto chain = local_side_effect_chains_.lookup(pred))
+      chains.push_back(chain);
+  }
+
+  // If there is no side-effect predecessor, then the input side-effect chain
+  // is used.
+  if (chains.empty()) return func_op.getArgument(0);
+
+  if (chains.size() == 1) return chains[0];
+
+  // If there are multiple side-effect predecessors, insert a merge_chains
+  // kernel and return the merged chain.
+  ConversionPatternRewriter::InsertionGuard insertion_guard(*rewriter);
+  rewriter->setInsertionPoint(op);
+  return tfrt::compiler::MergeChainsOp::create(*rewriter, op->getLoc(),
+                                               chain_type(), chains);
+}
+
+mlir::Value CoreRTConverter::GetTaskHandle(
+    mlir::Operation *op, StringRef task_name,
+    mlir::ConversionPatternRewriter *rewriter) {
+  mlir::func::FuncOp func_op = op->getParentOfType<mlir::func::FuncOp>();
+  llvm::StringMap<mlir::Value> &task_handle_by_name =
+      task_handles_by_func_[func_op.getOperation()];
+  auto iter = task_handle_by_name.find(task_name);
+  if (iter != task_handle_by_name.end()) {
+    return iter->second;
+  }
+
+  return mlir::Value();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/corert_converter.h

@@ -0,0 +1,175 @@
+/* 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_TFRT_TRANSFORMS_CORERT_CONVERTER_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_CORERT_CONVERTER_H_
+
+#include <array>
+#include <memory>
+#include <optional>
+#include <string>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#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/SymbolTable.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+struct ParseDeviceNameResult {
+  std::string device_type;
+  std::string device_name;
+  std::string op_handler_name;
+};
+
+// A helper class for converting CoreRT types and attributes.
+class CoreRTConverter : public mlir::TypeConverter {
+ public:
+  CoreRTConverter(
+      mlir::MLIRContext *context,
+      const mlir::TF::SideEffectAnalysis::Info *side_effect_analysis);
+  // Materialize all derived attributes. Note that this is only needed by
+  // CoreRT ops and fallback ops.
+  void MaterializeDerivedAttributes(mlir::Operation *op);
+
+  // Similar to CreateOpAttrs, create a single attribute that contains the
+  // named attribute lists, which is an array of pairs, with keys and values
+  // both being string attributes. The values represent function names.
+  // This method also populates a vector of attribute keys to be removed.
+  // If `use_mlir_func_name` is true, the function name given by MLIR will be
+  // used, which could be different from the original function name in the graph
+  // function library. This is used when the original function has been changed
+  // by lowering passes, and hence it needs to be exported to function library
+  // for runtime to use.
+  mlir::ArrayAttr CreateOpFuncAttrs(
+      const mlir::SymbolTable &symbol_table,
+      llvm::ArrayRef<mlir::NamedAttribute> attrs,
+      llvm::SmallVector<mlir::StringAttr, 4> *func_attr_keys,
+      bool use_mlir_func_name = false);
+
+  // Parse the device name of `op` to TFRT's device name. For example, "/CPU:0"
+  // will be parsed as "cpu". Return None if no device is assigned.
+  std::optional<ParseDeviceNameResult> ParseDeviceName(
+      llvm::StringRef device_name) const;
+  std::optional<ParseDeviceNameResult> ParseDeviceName(
+      mlir::Operation *op) const;
+
+  // Convert the device name in a TF op to a op_handler value produced by the
+  // corresponding GetOpHandler in the current block. If there does not exist
+  // one, insert a GetOpHandler to the beginning of the block and return the
+  // device value.
+  mlir::Value ConvertOpHandler(mlir::Operation *op, llvm::StringRef device_name,
+                               mlir::ConversionPatternRewriter *rewriter);
+
+  // Get a DistributedContext value to be used by the given op. The
+  // DistributedContext value should be shared by all operations in the body
+  // of the same FuncOp. If there does not exist one, return a null Value.
+  mlir::Value GetDistributedContext(mlir::Operation *op,
+                                    mlir::ConversionPatternRewriter *rewriter);
+
+  // Get a RemoteChainManager value to be used by the given op. The
+  // RemoteChainManager value should be shared by all operations in the body
+  // of the same FuncOp. If there does not exist one, return a null Value.
+  mlir::Value GetRemoteChainManager(mlir::Operation *op,
+                                    mlir::ConversionPatternRewriter *rewriter);
+
+  // Get a TaskHandle value with the given task name. If the TaskHandle value
+  // has already been created for the given task name within the same FuncOp,
+  // return this TaskHandle value. Otherwise, return a null Value.
+  mlir::Value GetTaskHandle(mlir::Operation *op, StringRef task_name,
+                            mlir::ConversionPatternRewriter *rewriter);
+
+  // Any local operation which uses any result of the `op` should depend on the
+  // given `chain`.
+  void RegisterLocalSideEffectChain(mlir::Operation *op, mlir::Value chain) {
+    local_side_effect_chains_[op] = chain;
+  }
+
+  // Return a local chain for side effects for `op`. If there are multiple
+  // chains, a merge_chains kernel will be inserted and the merged chain will be
+  // returned.
+  mlir::Value GetLocalSideEffectChain(
+      mlir::Operation *op, mlir::ConversionPatternRewriter *rewriter);
+
+  mlir::Type op_handler_type() {
+    return builder_.getType<::tfrt::corert::OpHandlerType>();
+  }
+
+  mlir::Type tensor_handle_type() {
+    return builder_.getType<::tfrt::corert::TensorHandleType>();
+  }
+
+  mlir::Type chain_type() {
+    return builder_.getType<::tfrt::compiler::ChainType>();
+  }
+
+  mlir::Builder &builder() { return builder_; }
+
+ private:
+  // TODO(chky): attributes "_output_shapes" should be removed by any tool that
+  // generates TF MLIR dialect, as they are not used by CoreRuntime. Remove this
+  // filtering logic once unused attributes are cleaned up in the upper layer.
+  bool IsUnusedAttribute(llvm::StringRef name) const {
+    // NOTE: attributes "f.*" are function attribute related and
+    // are added during importing graph to MLIR TF Executor dialect. These
+    // attributes are not actually used by TF ops with function attributes.
+    // TODO(b/180399811): Re-evaluate the usage of these attributes.
+    static const char *const kUnusedAttributes[] = {
+        "_output_shapes",
+        "result_segment_sizes",
+        "operand_segment_sizes",
+    };
+
+    for (auto attr : kUnusedAttributes) {
+      if (name == attr) {
+        return true;
+      }
+    }
+
+    return name.contains("f.");
+  }
+
+  // Returns the converted attribute in TFRT dialect. If the conversion fails,
+  // returns a null attribute instead.
+  mlir::Attribute ConvertAttribute(mlir::Attribute attr);
+
+  mlir::TypeAttr ConvertTypeAttribute(mlir::TypeAttr type_attr);
+
+  mlir::Builder builder_;
+
+  const mlir::TF::SideEffectAnalysis::Info &side_effect_analysis_;
+
+  llvm::DenseMap<mlir::Operation *, mlir::Value> local_side_effect_chains_;
+  llvm::DenseMap<mlir::Operation *, mlir::Value> distributed_context_by_func_;
+  llvm::DenseMap<mlir::Operation *, mlir::Value> remote_chain_mgr_by_func_;
+  llvm::DenseMap<mlir::Operation *, llvm::StringMap<mlir::Value>>
+      task_handles_by_func_;
+  llvm::StringMap<mlir::Value> op_handler_by_name_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_CORERT_CONVERTER_H_

tensorflow/compiler/mlir/tfrt/transforms/cross_device_transfer.cc

@@ -0,0 +1,201 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This pass inserts corert.transfer op to make sure any argument of any op is
+// on the same device of the op itself.
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/Casting.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/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/core/util/device_name_utils.h"
+#include "tfrt/basic_kernels/opdefs/basic_kernels.h"  // from @tf_runtime
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+namespace {
+
+using DeviceNameUtils = ::tensorflow::DeviceNameUtils;
+
+constexpr const char *kDeviceAttr = "device";
+constexpr const char *kTFRTDeviceAttr = "tfrt.device";
+// TODO(b/175480458): Do not assign default device once every op in the TF
+// dialect has the device attribute.
+constexpr const char *kDefaultDevice =
+    "/job:localhost/replica:0/task:0/device:CPU:0";
+
+// This method canonicalizes the device name so that we can use string
+// comparison to see if two devices are the same. It does the following
+// transformations:
+// 1) Set device ID to 0 if device ID is not already specified.
+// 2) Change the device type to uppercase string.
+static std::string CanonicalizeDeviceName(const std::string &device) {
+  if (device.empty()) return kDefaultDevice;
+
+  DeviceNameUtils::ParsedName parsed_name;
+  if (!device.empty() && device.at(0) == '/') {
+    DeviceNameUtils::ParseFullName(device, &parsed_name);
+  } else {
+    DeviceNameUtils::ParseFullName("/device:" + device, &parsed_name);
+  }
+
+  if (!parsed_name.has_id) {
+    parsed_name.has_id = true;
+    parsed_name.id = 0;
+  }
+
+  if (parsed_name.type == "cpu")
+    parsed_name.type = "CPU";
+  else if (parsed_name.type == "gpu")
+    parsed_name.type = "GPU";
+  else if (parsed_name.type == "tpu")
+    parsed_name.type = "TPU";
+  return DeviceNameUtils::ParsedNameToString(parsed_name);
+}
+
+// Return the device of the given operation.
+static std::string GetDevice(Operation *op) {
+  std::string device = "";
+  if (StringAttr device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
+    device = device_attr.getValue().str();
+  } else if (auto execute_op = llvm::dyn_cast<tfrt::corert::ExecuteOp>(op)) {
+    SmallVector<std::pair<StringRef, Attribute>, 4> attrs;
+    execute_op.getOpAttrs(&attrs);
+    for (std::pair<StringRef, Attribute> entry : attrs) {
+      if (entry.first == kDeviceAttr && mlir::isa<StringAttr>(entry.second)) {
+        device = mlir::cast<StringAttr>(entry.second).getValue().str();
+        break;
+      }
+    }
+  }
+
+  return CanonicalizeDeviceName(device);
+}
+
+// Return the device of the given value.
+static std::string GetDevice(mlir::Value value, func::FuncOp parent_func_op) {
+  std::string device = "";
+  if (BlockArgument block_arg = mlir::dyn_cast<BlockArgument>(value)) {
+    if (StringAttr device_attr = parent_func_op.getArgAttrOfType<StringAttr>(
+            block_arg.getArgNumber(), kTFRTDeviceAttr)) {
+      device = device_attr.getValue().str();
+    }
+  } else {
+    device = GetDevice(value.getDefiningOp());
+  }
+
+  return CanonicalizeDeviceName(device);
+}
+
+struct CrossDeviceTransferPass
+    : public PassWrapper<CrossDeviceTransferPass, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CrossDeviceTransferPass)
+
+  void runOnOperation() override;
+
+  llvm::StringRef getArgument() const final {
+    return "tfrt-cross-device-transfer";
+  }
+
+  llvm::StringRef getDescription() const final {
+    return "This pass inserts corert.transfer op to make sure any argument of "
+           "any op is on the same device of the op itself.";
+  }
+};
+
+void CrossDeviceTransferPass::runOnOperation() {
+  func::FuncOp func_op = getOperation();
+  llvm::DenseMap<mlir::Value, llvm::StringMap<mlir::Value>>
+      transferred_value_by_value_and_device;
+
+  func_op.getBody().walk([&](Operation *op) {
+    if (op->hasTrait<OpTrait::IsTerminator>()) return WalkResult::advance();
+    // Do not transfer the argument of corert.transfer op.
+    if (llvm::isa<tfrt::corert::TransferOp>(op)) return WalkResult::advance();
+
+    OpBuilder builder(op);
+    std::string dst_device = GetDevice(op);
+    mlir::Type tensor_type_type =
+        builder.getType<::tfrt::compiler::TensorTypeType>();
+    mlir::Type device_type = builder.getType<::tfrt::compiler::DeviceType>();
+
+    for (mlir::Value arg : op->getOperands()) {
+      // Do not transfer non-TensorHandle values.
+      if (!mlir::isa<tfrt::corert::TensorHandleType>(arg.getType())) continue;
+
+      // Do not transfer the result of corert.transfer op.
+      if (OpResult op_result = mlir::dyn_cast<OpResult>(arg)) {
+        Operation *defining_op = arg.getDefiningOp();
+        if (llvm::isa<tfrt::corert::TransferOp>(defining_op)) continue;
+      }
+
+      std::string src_device = GetDevice(arg, func_op);
+
+      if (DeviceNameUtils::LocalName(src_device) ==
+          DeviceNameUtils::LocalName(dst_device))
+        continue;
+
+      // Re-use the value already transferred to the given device.
+      llvm::StringMap<mlir::Value> &transferred_value_by_device =
+          transferred_value_by_value_and_device[arg];
+      auto iter = transferred_value_by_device.find(dst_device);
+      if (iter != transferred_value_by_device.end()) {
+        op->replaceUsesOfWith(arg, iter->second);
+        continue;
+      }
+
+      mlir::Value chain_in = func_op.getArgument(0);
+      auto get_device_op = tfrt::compiler::GetDeviceOp::create(
+          builder, op->getLoc(), device_type, chain_in, dst_device);
+      auto get_tensor_type_op = tfrt::corert::GetDstTensorTypeOp::create(
+          builder, op->getLoc(), tensor_type_type, arg,
+          get_device_op.getResult());
+      auto transfer_op = tfrt::corert::TransferOp::create(
+          builder, op->getLoc(), arg.getType(), arg, get_device_op.getResult(),
+          get_tensor_type_op.getResult());
+      mlir::Value new_arg = transfer_op.getResult();
+      transferred_value_by_device[dst_device] = new_arg;
+      op->replaceUsesOfWith(arg, new_arg);
+    }
+    return WalkResult::advance();
+  });
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> CreateCrossDeviceTransferPass() {
+  return std::make_unique<CrossDeviceTransferPass>();
+}
+
+static PassRegistration<CrossDeviceTransferPass> pass;
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/deduplicate_batch_function.cc

@@ -0,0 +1,177 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <iterator>
+#include <memory>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/iterator_range.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.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/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+using ::mlir::ArrayRef;
+using ::mlir::ModuleOp;
+using ::mlir::Operation;
+using ::mlir::SymbolTable;
+using ::mlir::SymbolTableCollection;
+using ::mlir::SymbolUserMap;
+
+// This only includes some preliminary checks as this is a short term solution.
+bool AreEquivalent(mlir::func::FuncOp& lhs, mlir::func::FuncOp& rhs) {
+  if (lhs.getFunctionType() != rhs.getFunctionType()) return false;
+
+  for (auto arg_pair : llvm::zip(lhs.getArguments(), rhs.getArguments())) {
+    auto& lhs_arg = std::get<0>(arg_pair);
+    auto& rhs_arg = std::get<1>(arg_pair);
+    if (lhs_arg.getType() != rhs_arg.getType()) return false;
+  }
+
+  auto lhs_ops = lhs.getBody().getOps();
+  auto rhs_ops = rhs.getBody().getOps();
+  if (std::distance(lhs_ops.begin(), lhs_ops.end()) !=
+      std::distance(rhs_ops.begin(), rhs_ops.end()))
+    return false;
+
+  for (auto op_pair : llvm::zip(lhs_ops, rhs_ops)) {
+    auto& lhs_op = std::get<0>(op_pair);
+    auto& rhs_op = std::get<1>(op_pair);
+    if (lhs_op.getName() != rhs_op.getName()) return false;
+    if (lhs_op.getNumRegions() != rhs_op.getNumRegions()) return false;
+    if (lhs_op.getNumSuccessors() != rhs_op.getNumSuccessors()) return false;
+    if (!std::equal(lhs_op.getOperandTypes().begin(),
+                    lhs_op.getOperandTypes().end(),
+                    rhs_op.getOperandTypes().begin()))
+      return false;
+    if (!std::equal(lhs_op.getResultTypes().begin(),
+                    lhs_op.getResultTypes().end(),
+                    rhs_op.getResultTypes().begin()))
+      return false;
+  }
+
+  return true;
+}
+
+// Deduplicate the functions if all users are BatchFunctionOp and have the same
+// shared_name.
+//
+// TODO(b/192463730): this is the short term solution and not needed anymore
+// after the shape inference pass is revamped with ideal solution
+// (b/192463730#comment11).
+class DeduplicateFunctionsInovkedByBatchFunction
+    : public mlir::PassWrapper<DeduplicateFunctionsInovkedByBatchFunction,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      DeduplicateFunctionsInovkedByBatchFunction)
+
+ private:
+  llvm::StringRef getArgument() const final {
+    return "tfrt-deduplicate-functions-invoked-by-batch-function";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Deduplicate the functions invoked by tf.BatchFunction with the "
+           "same shared_name";
+  }
+  void runOnOperation() override {
+    if (failed(Run())) {
+      signalPassFailure();
+    }
+  }
+
+  mlir::LogicalResult Run();
+};
+
+mlir::LogicalResult DeduplicateFunctionsInovkedByBatchFunction::Run() {
+  ModuleOp module = getOperation();
+  SymbolTableCollection symbol_table_collection;
+  SymbolTable& symbol_table = symbol_table_collection.getSymbolTable(module);
+  SymbolUserMap symbol_users(symbol_table_collection, module);
+
+  // Categorize the functions invoked by BatchFunctionOp by its shared_name.
+  llvm::StringMap<llvm::SmallVector<mlir::func::FuncOp, 2>>
+      shared_name_to_func_ops;
+
+  for (auto func :
+       llvm::make_early_inc_range(module.getOps<mlir::func::FuncOp>())) {
+    ArrayRef<Operation*> users = symbol_users.getUsers(func);
+    llvm::StringRef shared_name;
+    // Deduplicate the function only if all users are BatchFunctionOp and have
+    // the same shared_name
+    if (!users.empty() && llvm::all_of(users, [&shared_name](Operation* user) {
+          auto op = llvm::dyn_cast_or_null<mlir::TF::BatchFunctionOp>(user);
+          // User is not a BatchFunctionOp
+          if (!op) return false;
+          if (shared_name.empty()) {
+            shared_name = op.getSharedName();
+            return true;
+          }
+          return shared_name == op.getSharedName();
+        })) {
+      shared_name_to_func_ops[shared_name].push_back(func);
+    }
+  }
+
+  for (auto& it : shared_name_to_func_ops) {
+    auto& func_ops = it.second;
+    mlir::func::FuncOp& func_op_to_keep = func_ops.front();
+    for (mlir::func::FuncOp& func_op_to_remove : llvm::drop_begin(func_ops)) {
+      if (!AreEquivalent(func_op_to_keep, func_op_to_remove)) {
+        return func_op_to_remove.emitError(
+            "func_ops for BatchFunctionOp with the same shared name are "
+            "different");
+      }
+      if (failed(SymbolTable::replaceAllSymbolUses(
+              func_op_to_remove, func_op_to_keep.getSymNameAttr(), module))) {
+        return func_op_to_remove.emitError("unable to replace the symbol use");
+      }
+      symbol_table.erase(func_op_to_remove);
+    }
+  }
+
+  return mlir::success();
+}
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<ModuleOp>>
+CreateDeduplicateFunctionsInovkedByBatchFunctionPass() {
+  return std::make_unique<DeduplicateFunctionsInovkedByBatchFunction>();
+}
+
+static mlir::PassRegistration<DeduplicateFunctionsInovkedByBatchFunction>
+    register_pass;
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/deduplicate_if_result_pass.cc

@@ -0,0 +1,232 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+
+#include "absl/log/check.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/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/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+struct ResultMapping {
+  bool has_duplicate = false;
+  llvm::SmallVector<int> old_to_new;
+  llvm::SmallVector<int> new_to_old;
+};
+
+ResultMapping FindDuplicateResultIndices(mlir::func::FuncOp func) {
+  llvm::SmallDenseMap<mlir::Value, int> results;
+  auto &block = func.getBody().front();
+
+  auto return_op = llvm::cast<mlir::func::ReturnOp>(block.getTerminator());
+
+  ResultMapping mapping;
+  mapping.old_to_new.reserve(return_op.getNumOperands());
+
+  int j = 0;
+  for (int i = 0; i < return_op.getNumOperands(); ++i) {
+    auto value = return_op.getOperand(i);
+    auto iter = results.find(value);
+    if (iter != results.end()) {
+      mapping.has_duplicate = true;
+      mapping.old_to_new.push_back(iter->second);
+    } else {
+      results[value] = j;
+      mapping.old_to_new.push_back(j++);
+    }
+  }
+
+  mapping.new_to_old.resize(j);
+  for (int i = 0; i < mapping.old_to_new.size(); ++i) {
+    mapping.new_to_old[mapping.old_to_new[i]] = i;
+  }
+
+  return mapping;
+}
+
+mlir::func::FuncOp CreateBranchFunctionWithDeduplicatedResults(
+    mlir::OpBuilder &builder, mlir::Location loc, absl::string_view name,
+    mlir::func::FuncOp func, const ResultMapping &mapping) {
+  auto arg_types = func.getFunctionType().getInputs();
+  auto result_types = func.getFunctionType().getResults();
+
+  llvm::SmallVector<mlir::Type> new_result_types;
+  new_result_types.reserve(mapping.new_to_old.size());
+
+  for (int i : mapping.new_to_old) {
+    new_result_types.push_back(result_types[i]);
+  }
+
+  auto new_func_type = mlir::FunctionType::get(builder.getContext(), arg_types,
+                                               new_result_types);
+
+  auto new_func = mlir::func::FuncOp::create(builder, loc, name, new_func_type);
+  new_func.setVisibility(mlir::func::FuncOp::Visibility::Private);
+
+  mlir::OpBuilder::InsertionGuard guard(builder);
+
+  // In the body of newly created function, we insert
+  // tf.PartitionedCall ops to call the original func.
+  auto *block = new_func.addEntryBlock();
+  builder.setInsertionPointToStart(block);
+  auto empty_string_attr = builder.getStringAttr("");
+
+  llvm::SmallVector<mlir::Value, 4> results;
+  results.reserve(new_func_type.getNumResults());
+
+  // Create the call op to the original func. The arguments are simply
+  // the arguments from the wrapper function.
+  auto call_op = mlir::TF::PartitionedCallOp::create(
+      builder, loc, result_types, block->getArguments(), /*args_attrs=*/nullptr,
+      /*res_attrs=*/nullptr,
+      mlir::FlatSymbolRefAttr::get(func.getSymNameAttr()), empty_string_attr,
+      empty_string_attr, empty_string_attr);
+
+  for (int i : mapping.new_to_old) {
+    results.push_back(call_op.getResult(i));
+  }
+
+  mlir::func::ReturnOp::create(builder, loc, results);
+
+  return new_func;
+}
+
+void DeduplicateIfOps(mlir::ModuleOp module) {
+  mlir::SymbolTable symbol_table(module);
+
+  llvm::DenseMap<mlir::func::FuncOp, ResultMapping> visited;
+  llvm::DenseMap<mlir::func::FuncOp, mlir::func::FuncOp> new_funcs;
+  llvm::SmallVector<mlir::TF::IfOp> candidates;
+
+  module.walk([&](mlir::TF::IfOp op) { candidates.push_back(op); });
+
+  llvm::SmallVector<mlir::TF::IfOp> to_erase;
+
+  mlir::OpBuilder builder(module.getRegion());
+  for (auto op : candidates) {
+    auto find_mapping = [&](mlir::func::FuncOp func) {
+      if (auto iter = visited.find(func); iter != visited.end()) {
+        return iter->second;
+      }
+
+      auto mapping = FindDuplicateResultIndices(func);
+      visited[func] = mapping;
+      return mapping;
+    };
+
+    auto then_branch =
+        symbol_table.lookup<mlir::func::FuncOp>(op.getThenBranch());
+    auto else_branch =
+        symbol_table.lookup<mlir::func::FuncOp>(op.getElseBranch());
+
+    auto then_mapping = find_mapping(then_branch);
+    auto else_mapping = find_mapping(else_branch);
+
+    if (then_mapping.has_duplicate &&
+        then_mapping.old_to_new == else_mapping.old_to_new) {
+      auto get_or_create = [&](mlir::func::FuncOp func,
+                               const ResultMapping &mapping) {
+        auto iter = new_funcs.find(func);
+        if (iter != new_funcs.end()) {
+          return iter->second;
+        }
+
+        auto new_func = CreateBranchFunctionWithDeduplicatedResults(
+            builder, op.getLoc(),
+            absl::StrCat(func.getSymName().str(), "/tfrt_dedup_results"), func,
+            mapping);
+        new_funcs[func] = new_func;
+        return new_func;
+      };
+      auto new_then_func = get_or_create(then_branch, then_mapping);
+      auto new_else_func = get_or_create(else_branch, else_mapping);
+
+      mlir::OpBuilder::InsertionGuard guard(builder);
+      builder.setInsertionPoint(op);
+
+      llvm::SmallVector<mlir::Type> new_result_types;
+      for (int i : then_mapping.new_to_old) {
+        new_result_types.push_back(op->getResult(i).getType());
+      }
+
+      auto new_if_op = mlir::TF::IfOp::create(
+          builder, op.getLoc(), new_result_types, op.getCond(), op.getInput(),
+          new_then_func.getSymName(), new_else_func.getSymName(),
+          op.getIsStateless());
+
+      DCHECK_EQ(then_mapping.old_to_new.size(), op.getNumResults());
+      for (int i = 0; i < then_mapping.old_to_new.size(); ++i) {
+        op.getResult(i).replaceAllUsesWith(
+            new_if_op.getResult(then_mapping.old_to_new[i]));
+      }
+      to_erase.push_back(op);
+    }
+  }
+
+  for (auto op : to_erase) {
+    op->erase();
+  }
+}
+
+class DeduplicateIfResultPass
+    : public mlir::PassWrapper<DeduplicateIfResultPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  llvm::StringRef getArgument() const final {
+    return "tfrt-deduplicate-if-result";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Deduplicate the results of tf.If ops";
+  }
+
+  void runOnOperation() override { DeduplicateIfOps(getOperation()); }
+
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(DeduplicateIfResultPass)
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateDeduplicateIfResultPass() {
+  return std::make_unique<DeduplicateIfResultPass>();
+}
+
+static mlir::PassRegistration<DeduplicateIfResultPass> register_pass(
+    CreateDeduplicateIfResultPass);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/fallback_converter.cc

@@ -0,0 +1,162 @@
+/* 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/tfrt/transforms/fallback_converter.h"
+
+#include <optional>
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h"
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+FallbackConverter::FallbackConverter(mlir::MLIRContext *context)
+    : builder_(context) {
+  addConversion([](tfrt::compiler::ChainType type) { return type; });
+  addConversion([](tfrt::fallback::TFTensorType type) { return type; });
+  addConversion([=](mlir::TensorType type) -> std::optional<mlir::Type> {
+    // Ref types are not supported in both compiler and runtime.
+    if (mlir::isa<mlir::TF::TensorFlowRefType>(type.getElementType())) {
+      return std::nullopt;
+    }
+
+    return builder_.getType<tfrt::fallback::TFTensorType>();
+  });
+  addConversion([=](mlir::Type type) -> std::optional<mlir::Type> {
+    if (type == builder_.getI1Type()) return type;
+    return std::nullopt;
+  });
+}
+
+mlir::Value ConvertCoreRTTensorHandleToFallbackTensor(
+    mlir::Location loc, llvm::StringRef device, mlir::Value value,
+    mlir::ConversionPatternRewriter &rewriter) {
+  if (mlir::isa<tfrt::fallback::TFTensorType>(value.getType())) return value;
+
+  if (!mlir::isa<tfrt::corert::TensorHandleType>(value.getType())) return {};
+
+  mlir::OpBuilder::InsertionGuard guard(rewriter);
+
+  if (device.ends_with("CPU:0") && !device.starts_with("/job:")) {
+    // Canonicalize CPU device name. This is needed as corert library only uses
+    // the default CPU device name (i.e.
+    // "/job:localhost/replica:0/task:0/device:CPU:0") and cannot recoganize
+    // other legal variants (e.g. "/device:CPU:0").
+    //
+    // Note that we don't want to make change to the device name if it is
+    // already canonicalized by users.
+    // e.g. "/job:tpu_worker/replica:0/task:x/device:CPU:0".
+    // TODO(tfrt-devs): to make the canonicalization more robust we should
+    // introduce a util to check each component of the TF device name.
+    device = GetDefaultCpuDeviceName();
+  }
+
+  auto *def = value.getDefiningOp();
+  if (def) {
+    rewriter.setInsertionPointAfter(def);
+  } else {
+    rewriter.setInsertionPointToStart(value.getParentBlock());
+  }
+
+  return tfrt::fallback_async::CoreRTTensorHandleToFallbackTensorOp::create(
+             rewriter, loc, rewriter.getType<tfrt::fallback::TFTensorType>(),
+             value, device)
+      .getResult(0);
+}
+
+mlir::Value ConvertFallbackTensorToCoreRTTensorHandle(
+    mlir::Location loc, mlir::Value value,
+    mlir::ConversionPatternRewriter &rewriter) {
+  if (mlir::isa<tfrt::corert::TensorHandleType>(value.getType())) return value;
+
+  if (!mlir::isa<tfrt::fallback::TFTensorType>(value.getType())) return {};
+
+  // Use CPU device by default if no device is specified.
+  llvm::StringRef device = GetDefaultCpuDeviceName();
+  if (auto *def = value.getDefiningOp()) {
+    if (auto device_attr = def->getAttrOfType<mlir::StringAttr>("device")) {
+      // NOTE: The TPU_SYSTEM check is just a short term workaround. The long
+      // term solution should be checking the HostMemory annotation of the
+      // defining op (it should be defined in TF OpKernel). If HostMemory
+      // annotation is set for an output tensor, we should use CPU device here.
+      // TODO(b/200896904): Support HostMemory annotation.
+      if (!device_attr.getValue().ends_with("TPU_SYSTEM:0")) {
+        device = device_attr.getValue();
+      }
+    }
+  }
+
+  return tfrt::fallback_async::FallbackTensorToCoreRTTensorHandleOp::create(
+             rewriter, loc, rewriter.getType<tfrt::corert::TensorHandleType>(),
+             value, device)
+      .getResult(0);
+}
+
+mlir::LogicalResult ConvertCoreRTOperands(
+    mlir::Operation *op, mlir::ValueRange operands,
+    llvm::SmallVectorImpl<mlir::Value> *new_operands,
+    mlir::ConversionPatternRewriter &rewriter) {
+  mlir::OpBuilder::InsertionGuard guard(rewriter);
+  // Insert before the current op.
+  rewriter.setInsertionPoint(op);
+
+  for (auto operand : operands) {
+    auto value = ConvertFallbackTensorToCoreRTTensorHandle(op->getLoc(),
+                                                           operand, rewriter);
+    if (!value) {
+      return op->emitWarning("failed to convert to !corert.tensorhandle")
+             << operand.getType();
+    }
+
+    new_operands->push_back(value);
+  }
+  return success();
+}
+
+mlir::LogicalResult ConvertFallbackOperands(
+    mlir::Operation *op, llvm::StringRef device, mlir::ValueRange operands,
+    llvm::SmallVectorImpl<mlir::Value> *new_operands,
+    mlir::ConversionPatternRewriter &rewriter) {
+  for (auto operand : operands) {
+    if (!mlir::isa<tfrt::fallback::TFTensorType>(operand.getType())) {
+      auto new_operand = ConvertCoreRTTensorHandleToFallbackTensor(
+          op->getLoc(), device, operand, rewriter);
+      if (!new_operand)
+        return op->emitWarning(
+            "failed to convert the operand to fallback tensor.");
+      new_operands->push_back(new_operand);
+    } else {
+      new_operands->push_back(operand);
+    }
+  }
+  return success();
+}
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/fallback_converter.h

@@ -0,0 +1,96 @@
+/* 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_TFRT_TRANSFORMS_FALLBACK_CONVERTER_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_FALLBACK_CONVERTER_H_
+
+#include <cstdint>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+inline llvm::StringRef GetDefaultCpuDeviceName() {
+  static constexpr char kCpuDeviceName[] =
+      "/job:localhost/replica:0/task:0/device:CPU:0";
+  return kCpuDeviceName;
+}
+
+class FallbackConverter : public mlir::TypeConverter {
+ public:
+  explicit FallbackConverter(mlir::MLIRContext *context);
+
+  // Return the next dense key for fallback ops. The key is simply an array
+  // index so that in runtime, the fallback ops can be efficiently retrieved.
+  int64_t GetNextFallbackKey() const { return fallback_ops_.size(); }
+
+  void RegisterFallbackOp(mlir::Operation *op) { fallback_ops_.push_back(op); }
+
+  void ReplaceFallbackOp(int64_t key, mlir::Operation *op) {
+    fallback_ops_[key] = op;
+  }
+
+  llvm::ArrayRef<mlir::Operation *> GetFallbackOps() const {
+    return fallback_ops_;
+  }
+
+ private:
+  mlir::Builder builder_;
+  // Using a vector to keep fallback ops in order, and the key for a fallback op
+  // is its corresponding index here.
+  llvm::SmallVector<mlir::Operation *, 8> fallback_ops_;
+};
+
+// Convert the `value` that is a !corert.tensorhandle to
+// !tfrt_fallback.tf_tensor. If needed, tensor conversion kernels will be added.
+// On error it returns nullptr.
+mlir::Value ConvertCoreRTTensorHandleToFallbackTensor(
+    mlir::Location loc, llvm::StringRef device, mlir::Value value,
+    mlir::ConversionPatternRewriter &rewriter);
+
+// Convert the `value` that is a !tfrt_fallback.tf_tensor to
+// !corert.tensorhandle. If needed, tensor conversion kernels will be added. On
+// error it returns nullptr.
+mlir::Value ConvertFallbackTensorToCoreRTTensorHandle(
+    mlir::Location loc, mlir::Value value,
+    mlir::ConversionPatternRewriter &rewriter);
+
+// Convert operands that might be !tfrt_fallback.tf_tensor for corert operations
+// that take only !corert.tensorhandle.
+mlir::LogicalResult ConvertCoreRTOperands(
+    mlir::Operation *op, mlir::ValueRange operands,
+    llvm::SmallVectorImpl<mlir::Value> *new_operands,
+    mlir::ConversionPatternRewriter &rewriter);
+
+// Convert operands that might be !corert.tensorhandle for fallback operations
+// that take only !tfrt_fallback.tf_tensor.
+mlir::LogicalResult ConvertFallbackOperands(
+    mlir::Operation *op, llvm::StringRef device, mlir::ValueRange operands,
+    llvm::SmallVectorImpl<mlir::Value> *new_operands,
+    mlir::ConversionPatternRewriter &rewriter);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_FALLBACK_CONVERTER_H_

tensorflow/compiler/mlir/tfrt/transforms/fuse_tpu_compile_and_execute_ops.cc

@@ -0,0 +1,282 @@
+/* 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 <cassert>
+#include <memory>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+
+namespace tensorflow {
+namespace {
+
+void RecursivelyMoveOp(mlir::TF::_TPUCompileMlirOp compile_op,
+                       mlir::Operation *op_candidate,
+                       llvm::SmallDenseSet<mlir::Operation *> *ops_to_move) {
+  if (!op_candidate || !ops_to_move->contains(op_candidate)) return;
+  // Move the parent first.
+  for (const auto &operand : op_candidate->getOperands()) {
+    RecursivelyMoveOp(compile_op, operand.getDefiningOp(), ops_to_move);
+  }
+  op_candidate->moveBefore(compile_op);
+  // Erase the op to avoid moving the common ancestor.
+  ops_to_move->erase(op_candidate);
+}
+
+// Move exec_op's args defining ops before the compile op to group compile op
+// and execute op.
+void GroupCompileOpAndExecuteOp(mlir::func::FuncOp func,
+                                mlir::TF::_TPUCompileMlirOp compile_op,
+                                mlir::TF::TPUExecuteOp exec_op) {
+  // Collect the ops between compile op and execute op.
+  llvm::SmallDenseSet<mlir::Operation *> ops_to_move;
+  bool collect = false;
+  func.walk(
+      [&ops_to_move, &collect, &compile_op, &exec_op](mlir::Operation *op) {
+        if (collect) ops_to_move.insert(op);
+        if (op == compile_op.getOperation()) collect = true;
+        return op == exec_op.getOperation() ? mlir::WalkResult::interrupt()
+                                            : mlir::WalkResult::advance();
+      });
+  // Recursively move the defining op of the execute op argument in front of the
+  // compile op such that the compile op and execute op are grouped together.
+  for (const auto &operand : exec_op.getArgs()) {
+    RecursivelyMoveOp(compile_op, operand.getDefiningOp(), &ops_to_move);
+  }
+}
+
+mlir::Value FindFullShapedOperand(mlir::TF::SplitOp split_op) {
+  auto operand = split_op.getOperand(1);
+  auto defining_op = operand.getDefiningOp();
+  if (!defining_op) {
+    return operand;
+  }
+  if (auto parent_split_op = llvm::dyn_cast<mlir::TF::SplitOp>(defining_op))
+    return FindFullShapedOperand(parent_split_op);
+
+  return operand;
+}
+
+bool MaybeFindUsedExecuteOp(
+    const llvm::SmallVector<mlir::TF::TPUExecuteOp, 4> &exec_op_in_group,
+    mlir::TF::TPUExecuteOp &used_exec_op) {
+  // TODO(b/218763089): Here we assume the TPU output is replicated on all
+  // cores. Once the b/218763089 is fixed in MLIR Bridge that it properly
+  // supports output being sharded on TPU cores, then the results from all
+  // of the TPUExecuteOp are used.
+  bool found_used_exec_op = false;
+  for (auto exec_op : exec_op_in_group) {
+    if (!exec_op.use_empty()) {
+      if (found_used_exec_op) {
+        exec_op.emitOpError(
+            "More than 1 TPUExecuteOp has dependencies for the same "
+            "_TPUCompileMlirOp");
+        return false;
+      }
+      used_exec_op = exec_op;
+      found_used_exec_op = true;
+    }
+  }
+  if (!exec_op_in_group.empty() && !found_used_exec_op) {
+    used_exec_op = exec_op_in_group[0];
+    found_used_exec_op = true;
+  }
+  return found_used_exec_op;
+}
+
+void FuseCompileAndExecuteOps(
+    mlir::TF::_TPUCompileMlirOp &compile_op,
+    const llvm::SmallVector<mlir::TF::TPUExecuteOp, 4> &exec_op_in_group,
+    mlir::TF::TPUExecuteOp &used_exec_op,
+    llvm::SmallDenseMap<
+        mlir::TF::TPUExecuteOp,
+        llvm::SmallDenseMap<int, mlir::TF::SetStaticDimensionBoundsOp>>
+        &exec_to_static_shaped_operands_map,
+    mlir::OpBuilder &builder, mlir::MLIRContext *context) {
+  builder.setInsertionPointAfter(compile_op);
+  llvm::SmallVector<mlir::Type, 4> output_types;
+  output_types.push_back(mlir::RankedTensorType::get(
+      {3}, builder.getType<mlir::TF::StringType>()));
+  output_types.insert(output_types.end(), used_exec_op.getResultTypes().begin(),
+                      used_exec_op.getResultTypes().end());
+
+  llvm::SmallVector<int> static_shaped_operand_indices_attr;
+  llvm::SmallVector<mlir::Value> static_shape_tensors;
+  llvm::SmallVector<mlir::Value> exec_op_args;
+  exec_op_args.resize(used_exec_op.getArgs().size());
+
+  auto &static_shaped_operands =
+      exec_to_static_shaped_operands_map[used_exec_op];
+  for (int i = 0; i < used_exec_op.getArgs().size(); ++i) {
+    auto iter = static_shaped_operands.find(i);
+    if (iter != static_shaped_operands.end()) {
+      static_shaped_operand_indices_attr.push_back(iter->first);
+      static_shape_tensors.push_back(iter->second.getStaticShape());
+      exec_op_args[i] = iter->second.getInput();
+      // The first operand is the input tensor, while the second operand is
+      // the static shape tensor, hence the drop_back here.
+      iter->second->replaceAllUsesWith(
+          mlir::ValueRange({iter->second.getInput()}));
+      iter->second->erase();
+    } else {
+      auto split_op = ::llvm::dyn_cast_or_null<mlir::TF::SplitOp>(
+          used_exec_op->getOperand(i).getDefiningOp());
+      exec_op_args[i] = split_op ? FindFullShapedOperand(split_op)
+                                 : used_exec_op->getOperand(i);
+    }
+  }
+
+  auto producer_name =
+      used_exec_op->getAttrOfType<mlir::StringAttr>("_producer_name");
+  if (!producer_name) producer_name = mlir::StringAttr::get(context, "default");
+  auto compile_and_execute_op = mlir::TF::TPUCompileMlirAndExecuteOp::create(
+      builder, used_exec_op.getLoc(), output_types, exec_op_args,
+      static_shape_tensors,
+      builder.getI32ArrayAttr(static_shaped_operand_indices_attr),
+      compile_op.getMlirModule(), compile_op.getMetadata(), producer_name);
+
+  for (auto exec_op : exec_op_in_group) {
+    exec_op.replaceAllUsesWith(compile_and_execute_op.getResults());
+    assert(exec_op.use_empty());
+    exec_op.erase();
+  }
+
+  for (auto program_result : compile_op.getProgram()) {
+    program_result.replaceAllUsesWith(
+        compile_and_execute_op.getRendezvousKeyBase());
+  }
+  assert(compile_op.use_empty());
+  compile_op.erase();
+}
+
+// This pass rewrites tf._TPUCompileMlirOp and tf.TPUExecuteOp into a single
+// tf.TPUCompileMlirAndExecuteOp. Also it removes the unnecessary
+// TPUCompileSucceededAssertOp.
+class FuseTpuCompileAndExecutePass
+    : public mlir::PassWrapper<FuseTpuCompileAndExecutePass,
+                               mlir::OperationPass<mlir::func::FuncOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FuseTpuCompileAndExecutePass)
+
+  llvm::StringRef getArgument() const final {
+    return "tfrt-fuse-tpu-compile-and-execute-ops";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Fuse TPU Ops according to TFRT's requirements.";
+  }
+
+  void runOnOperation() override {
+    auto func = getOperation();
+
+    // remove TPUCompileSucceededAssertOp
+    func.walk([&](mlir::Operation *op) {
+      if (llvm::isa<mlir::TF::TPUCompileSucceededAssertOp>(op)) {
+        op->erase();
+      }
+    });
+
+    // A map from an exec op to a struct containing the static shape tensor from
+    // a SetDynamicDimensionBoundsOp and the operand index.
+    llvm::SmallDenseMap<
+        mlir::TF::TPUExecuteOp,
+        llvm::SmallDenseMap<int, mlir::TF::SetStaticDimensionBoundsOp>>
+        exec_to_static_shaped_operands_map;
+
+    llvm::SmallVector<mlir::TF::TPUExecuteOp, 4> tpu_execute_ops;
+    func.walk([&](mlir::Operation *op) {
+      if (auto exec_op = llvm::dyn_cast<mlir::TF::TPUExecuteOp>(op)) {
+        tpu_execute_ops.push_back(exec_op);
+        // Collect any operands to this tf.Execute op that are defined by a
+        // SetStaticDimensionBoundsOp along with the operand index.
+        for (const auto &operand : llvm::enumerate(exec_op.getOperands())) {
+          if (auto defining_op =
+                  operand.value()
+                      .getDefiningOp<mlir::TF::SetStaticDimensionBoundsOp>()) {
+            exec_to_static_shaped_operands_map[exec_op][operand.index()] =
+                defining_op;
+          }
+        }
+      }
+    });
+
+    mlir::OpBuilder builder(&func.getBody());
+
+    llvm::SmallDenseMap<mlir::TF::_TPUCompileMlirOp,
+                        llvm::SmallVector<mlir::TF::TPUExecuteOp, 4>>
+        compile_and_execute_groups;
+
+    for (auto exec_op : tpu_execute_ops) {
+      auto compile_cache_entry = exec_op.getKey();
+      auto compile_op = ::llvm::dyn_cast<mlir::TF::_TPUCompileMlirOp>(
+          compile_cache_entry.getDefiningOp());
+      if (!compile_op) {
+        exec_op.emitOpError("could not get the _TPUCompileMlirOp");
+        signalPassFailure();
+        return;
+      }
+
+      GroupCompileOpAndExecuteOp(func, compile_op, exec_op);
+      compile_and_execute_groups[compile_op].push_back(exec_op);
+    }
+    for (const auto &kv : compile_and_execute_groups) {
+      auto compile_op = kv.first;
+      const llvm::SmallVector<mlir::TF::TPUExecuteOp, 4> &exec_op_in_group =
+          kv.second;
+      mlir::TF::TPUExecuteOp used_exec_op;
+      if (!MaybeFindUsedExecuteOp(exec_op_in_group, used_exec_op)) {
+        signalPassFailure();
+        return;
+      }
+
+      FuseCompileAndExecuteOps(compile_op, exec_op_in_group, used_exec_op,
+                               exec_to_static_shaped_operands_map, builder,
+                               &getContext());
+    }
+  }
+};
+
+}  // namespace
+
+namespace tfrt_compiler {
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateFuseTpuCompileAndExecutePass() {
+  return std::make_unique<FuseTpuCompileAndExecutePass>();
+}
+
+static mlir::PassRegistration<FuseTpuCompileAndExecutePass>
+    fuse_tpu_compile_and_execute_ops_pass;
+
+}  // namespace tfrt_compiler
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/gpu_passes.cc

@@ -0,0 +1,36 @@
+/* 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/tfrt/transforms/gpu_passes.h"
+
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/gpu_ops.h"
+
+namespace tensorflow {
+
+void RegisterGpuDialects(mlir::DialectRegistry *registry) {
+  registry->insert<tfrt::gpu::GpuRuntimeDialect>();
+}
+
+void AddGpuTargetDialectAndPatterns(mlir::MLIRContext *context,
+                                    mlir::ConversionTarget *target,
+                                    mlir::RewritePatternSet *patterns) {
+  target->addLegalDialect<tfrt::gpu::GpuRuntimeDialect>();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/gpu_passes.h

@@ -0,0 +1,38 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_GPU_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_GPU_PASSES_H_
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassOptions.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+
+namespace tensorflow {
+
+// Registers dialects used in TFRT GPU lowering.
+void RegisterGpuDialects(mlir::DialectRegistry *registry);
+
+// Adds a target dialect and rewrite patterns for TFRT GPU lowering.
+void AddGpuTargetDialectAndPatterns(mlir::MLIRContext *context,
+                                    mlir::ConversionTarget *target,
+                                    mlir::RewritePatternSet *patterns);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_GPU_PASSES_H_

tensorflow/compiler/mlir/tfrt/transforms/ifrt/BUILD

@@ -0,0 +1,338 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load(
+    "@xla//xla/tsl/platform:build_config.bzl",
+    "tf_proto_library",
+)
+load("//tensorflow:tensorflow.bzl", "if_google", "tf_cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    packages = [
+        "//tensorflow/compiler/mlir/tfrt/...",
+        "//tensorflow/core/tfrt/ifrt/...",
+        "//tensorflow/core/tfrt/saved_model/tests/...",
+    ] + if_google([
+        "//learning/brain/tfrt/cpp_tests/...",
+        "//learning/serving/servables/tfrt/...",
+        "//learning/brain/tfrt/ifrt/...",
+        "//learning/brain/tfrt/tfrt_session/...",
+        # Allow visibility from the mlir language server.
+        "//learning/brain/mlir/mlir_lsp_server/...",
+        "//smartass/brain/util/...",
+    ]),
+)
+
+gentbl_cc_library(
+    name = "pass_inc_gen",
+    tbl_outs = {"passes.h.inc": [
+        "-gen-pass-decls",
+        "-name=TfrtIfrtServing",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes.td",
+    deps = [
+        "@llvm-project//mlir:PassBaseTdFiles",
+    ],
+)
+
+cc_library(
+    name = "ifrt_constants",
+    srcs = [],
+    hdrs = ["ifrt_constants.h"],
+    deps = [
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "ifrt_types",
+    srcs = [],
+    hdrs = ["ifrt_types.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+    ],
+)
+
+cc_library(
+    name = "tf_ifrt_passes",
+    srcs = [
+        "lower_to_ifrt_restore_variable.cc",
+        "propagate_static_shapes.cc",
+        "rewrite_cluster_to_ifrt_call.cc",
+        "sink_variable_as_named_array.cc",
+        "tf_device_cleanup.cc",
+        "tf_identity_propagation.cc",
+        "tf_ifrt_passes.cc",
+        "tf_restore_merging.cc",
+        "tf_restore_pruning.cc",
+        "tf_restore_splitting.cc",
+    ],
+    hdrs = [
+        "tf_ifrt_passes.h",
+    ],
+    textual_hdrs = ["passes.h.inc"],
+    deps = [
+        ":ifrt_constants",
+        ":pass_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:bridge_logger",
+        "//tensorflow/compiler/mlir/tensorflow:convert_tensor",
+        "//tensorflow/compiler/mlir/tensorflow:convert_type",
+        "//tensorflow/compiler/mlir/tensorflow:device_util",
+        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/compiler/mlir/tensorflow:tpu_rewrite_device_util",
+        "//tensorflow/compiler/mlir/tensorflow:translate_utils",
+        "//tensorflow/compiler/mlir/tensorflow/ir/host_runtime:tensorflow_tfrt_ops",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:tpu_metadata_utils",
+        "//tensorflow/compiler/mlir/tfrt:tf_to_tfrt",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/ir/types:Dialect",
+        "//tensorflow/core/platform:protobuf",
+        "//tensorflow/core/platform:random",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
+        "//tensorflow/core/tfrt/ifrt:ifrt_config_proto_cc",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@tsl//tsl/platform:protobuf",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/service:computation_placer_hdr",
+    ],
+)
+
+cc_library(
+    name = "tf2hlo",
+    srcs = ["tf2hlo.cc"],
+    hdrs = ["tf2hlo.h"],
+    deps = [
+        ":ifrt_compilation_proto_cc",
+        ":ifrt_constants",
+        ":ifrt_types",
+        "//tensorflow/compiler/jit:xla_cpu_jit",
+        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
+        "//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
+        "//tensorflow/compiler/mlir/tf2xla:compile_mlir_util",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:legalize_tf",
+        "//tensorflow/compiler/tf2xla:layout_util",
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib_headers_for_pybind",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
+        "//tensorflow/core/protobuf/tpu:topology_proto_cc",
+        "//tensorflow/core/tpu/kernels:tpu_compile_op_support",
+        "//tensorflow/core/tpu/kernels/xla:host_compute_ops",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tsl//tsl/platform:fingerprint",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/client:client_library",
+        "@xla//xla/hlo/translate/hlo_to_mhlo:hlo_to_mlir_hlo",
+        "@xla//xla/pjrt:pjrt_compiler",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/service:hlo_proto_cc",
+        "@xla//xla/stream_executor:platform_manager",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "extract_callback",
+    srcs = ["extract_callback.cc"],
+    hdrs = ["extract_callback.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_a_m_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_n_z_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow:visitor",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+tf_cc_test(
+    name = "tf2hlo_test",
+    srcs = [
+        "tf2hlo_test.cc",
+    ],
+    data = [
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_types",
+        ":tf2hlo",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:resource_loader",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:RegisterAllDialects",  # buildcleaner: keep
+        "@tsl//tsl/platform:protobuf",
+        "@xla//xla:debug_options_flags",
+        "@xla//xla:shape_util",
+        "@xla//xla/backends/cpu:target_machine_options",
+        "@xla//xla/pjrt:pjrt_compiler",
+        "@xla//xla/pjrt/plugin/xla_cpu:cpu_topology_description",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:mock",
+        "@xla//xla/python/ifrt:test_util",
+        "@xla//xla/python/pjrt_ifrt",
+        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
+        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "ifrt_backend_compiler",
+    srcs = ["ifrt_backend_compiler.cc"],
+    hdrs = ["ifrt_backend_compiler.h"],
+    deps = [
+        ":tf_ifrt_passes",
+        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:visitor",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:cluster_tf",
+        "//tensorflow/compiler/mlir/tfrt:backend_compiler",
+        "//tensorflow/compiler/mlir/tfrt:tpu_passes",
+        "//tensorflow/core/tfrt/ifrt:ifrt_executable_registry",
+        "//tensorflow/core/tfrt/ifrt:ifrt_model_context",
+        "//tensorflow/core/tfrt/ifrt:ifrt_serving_executable",
+        "//tensorflow/core/tfrt/runtime",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@tsl//tsl/profiler/lib:traceme",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "ifrt_backend_compiler_test",
+    srcs = [
+        "ifrt_backend_compiler_test.cc",
+    ],
+    data = [
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_backend_compiler",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:resource_loader",
+        "//tensorflow/core/tfrt/graph_executor:graph_execution_options",
+        "//tensorflow/core/tfrt/ifrt:ifrt_executable_registry",
+        "//tensorflow/core/tfrt/ifrt:ifrt_model_context",
+        "//tensorflow/core/tfrt/ifrt:ifrt_serving_core_selector",
+        "//tensorflow/core/tfrt/ifrt:sharding_utils",
+        "//tensorflow/core/tfrt/runtime",
+        "//tensorflow/core/tfrt/saved_model:saved_model_testutil",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:RegisterAllDialects",  # buildcleaner: keep
+        "@tf_runtime//:hostcontext",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:test_util",
+        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
+        "@xla//xla/tsl/framework/test_util:mock_serving_device_selector",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_proto_library(
+    name = "ifrt_compilation_proto",
+    srcs = ["ifrt_compilation.proto"],
+    protodeps = [
+        "//tensorflow/compiler/tf2xla:host_compute_metadata_proto",
+        "@xla//xla/service:hlo_proto",
+        "//tensorflow/core:protos_all",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto",
+    ],
+    deps = [
+        "@xla//xla:xla_data_proto",
+        "@xla//xla/python/ifrt:layout_proto",
+    ],
+)
+
+cc_library(
+    name = "pack_inputs_pass",
+    srcs = ["pack_inputs_pass.cc"],
+    hdrs = ["pack_inputs_pass.h"],
+    deps = [
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+    ],
+)

tensorflow/compiler/mlir/tfrt/transforms/ifrt/extract_callback.cc

@@ -0,0 +1,81 @@
+/* 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/tfrt/transforms/ifrt/extract_callback.h"
+
+#include <utility>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h.inc"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h.inc"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/visitor.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ExtractCallbackModule(
+    mlir::ModuleOp module, absl::string_view callback_key) {
+  // Find the entry function name first.
+  mlir::func::FuncOp callback_entry_func;
+  module.walk([&](mlir::func::FuncOp func) {
+    if (func.getSymName().str() == callback_key) {
+      callback_entry_func = func;
+      return mlir::WalkResult::skip();
+    }
+    return mlir::WalkResult::advance();
+  });
+
+  if (!callback_entry_func) {
+    return absl::NotFoundError(
+        absl::StrCat("Callback key ", callback_key, " not found"));
+  }
+
+  mlir::StatusScopedDiagnosticHandler diag_handler(module->getContext());
+  auto entry_function_name = callback_entry_func.getSymName();
+  auto submodule = mlir::TF::CreatePrunedModule(module, entry_function_name);
+  if (mlir::failed(submodule)) {
+    return diag_handler.ConsumeStatus();
+  }
+
+  // Remove the attribute inherited from saved model loading. They impose
+  // additional constraint on public functions that are not necessary.
+  submodule->get()->removeAttr("tf_saved_model.semantics");
+  submodule->get().walk([&](mlir::func::FuncOp func) {
+    if (func.getSymName() == entry_function_name) {
+      func.setPublic();
+    }
+  });
+  return std::move(*submodule);
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/extract_callback.h

@@ -0,0 +1,36 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_EXTRACT_CALLBACK_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_EXTRACT_CALLBACK_H_
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// Extracts a module that consists of a public callback function in name of
+// `callback_key` and all its reachables.
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ExtractCallbackModule(
+    mlir::ModuleOp module, absl::string_view callback_key);
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_EXTRACT_CALLBACK_H_

tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_backend_compiler.cc

@@ -0,0 +1,236 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_backend_compiler.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/APInt.h"
+#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/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Verifier.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/visitor.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/cluster_tf.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_ifrt_passes.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tpu_passes.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_model_context.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+#include "tsl/profiler/lib/traceme.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+absl::StatusOr<std::vector<ServingExecutableRegistry::Handle>>
+CompileAndRegisterIfrtPrograms(absl::string_view model_name,
+                               mlir::ModuleOp module,
+                               IfrtModelContext& ifrt_model_context) {
+  std::vector<ServingExecutableRegistry::Handle> handles;
+
+  // Compile Ifrt programs and register the executables. Outlined Ifrt
+  // programs are marked with `tfrt_ifrt_serving.program_id` attributes.
+  for (auto func : module.getOps<mlir::func::FuncOp>()) {
+    int64_t program_id;
+    if (auto attr = func->getAttrOfType<mlir::IntegerAttr>(
+            "tfrt_ifrt_serving.program_id")) {
+      program_id = attr.getInt();
+    } else {
+      continue;
+    }
+
+    mlir::StatusScopedDiagnosticHandler diag_handler(module->getContext());
+    auto entry_function_name = func.getSymName();
+    auto submodule = mlir::TF::CreatePrunedModule(module, entry_function_name);
+    if (mlir::failed(submodule)) {
+      return diag_handler.ConsumeStatus();
+    }
+
+    // Remove the attribute inherited from saved model loading. They impose
+    // additional constraint on public functions that are not necessary.
+    submodule->get()->removeAttr("tf_saved_model.semantics");
+    submodule->get().walk([&](mlir::func::FuncOp func) {
+      if (func.getSymName() == entry_function_name) {
+        func.setName("main");
+        func.setSymName("main");
+        func.setPublic();
+      }
+    });
+    // Remove the program id attribute from the submodule because they are not
+    // needed and will prevent us generating consistent cache key.
+    // program id is already in ifrt_call op's attribute and that part is not
+    // touched here.
+    submodule->get()->walk([](mlir::func::FuncOp func) {
+      func->removeAttr("tfrt_ifrt_serving.program_id");
+    });
+
+    TF_ASSIGN_OR_RETURN(
+        auto executable,
+        IfrtServingExecutable::Create(
+            program_id, model_name, entry_function_name.str(),
+            *std::move(submodule), ifrt_model_context.GetClient(),
+            &ifrt_model_context.GetThreadPool(),
+            &ifrt_model_context.GetLoadedVariableRegistry(),
+            &ifrt_model_context.GetRestoreTensorRegistry(),
+            ifrt_model_context.checkpoint_loader_queue(),
+            ifrt_model_context.GetDeviceMgr(),
+            ifrt_model_context.GetShapeRepresentationFn(),
+            ifrt_model_context.GetIfrtServingCoreSelector(),
+            ifrt_model_context.GetCompilationEnvOrOverrides(),
+            ifrt_model_context.GetTfToHloCompiler(),
+            ifrt_model_context.GetPersistentCompilationCache(),
+            ifrt_model_context.GetH2DTransferExecutorFactory(),
+            ifrt_model_context.use_output_arena()));
+
+    // Register the Ifrt program to `ServingExecutableRegistry` so that
+    // the client TF program can invoke them via `IfrtCall` op.
+    TF_ASSIGN_OR_RETURN(auto handle, ServingExecutableRegistry::Register(
+                                         program_id, std::move(executable)));
+
+    handles.push_back(std::move(handle));
+  }
+
+  return handles;
+}
+
+absl::Status CompileTensorflowForIfrtServing(
+    absl::string_view model_name, IfrtModelContext& ifrt_model_context,
+    mlir::ModuleOp module, bool enable_async_ifrt) {
+  tsl::profiler::TraceMe trace_me("CompileTensorflowForIfrtServing");
+  mlir::Builder builder(module.getContext());
+
+  TF_RETURN_IF_ERROR(RunClusterToIfrtRuntimeOpsPassPipeline(
+      module, model_name,
+      ifrt_model_context.enable_propagate_static_shapes_pass(),
+      enable_async_ifrt));
+
+  TF_ASSIGN_OR_RETURN(
+      auto handles,
+      CompileAndRegisterIfrtPrograms(model_name, module, ifrt_model_context));
+
+  for (auto& handle : handles) {
+    ifrt_model_context.RegisterHandle(std::move(handle));
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+// Compile ifrt programs in TF dialect into ifrt executables.
+// Remove ifrt programs afterwards.
+absl::Status IfrtBackendCompiler::CompileTensorflow(
+    tensorflow::tfrt_stub::ModelRuntimeContext& model_context,
+    mlir::ModuleOp module) const {
+  auto ifrt_model_context =
+      model_context.resource_context().GetResource<IfrtModelContext>(
+          kIfrtModelContextName);
+  if (!ifrt_model_context.has_value()) {
+    return absl::InternalError(
+        "Failed to find model context for ifrt serving.");
+  }
+
+  if ((*ifrt_model_context)->IsFrozen()) {
+    return absl::FailedPreconditionError(
+        "Cannot compile IFRT programs after the model is frozen. Please make "
+        "sure warmup covers all signatures by following go/tf-model-warmup.");
+  }
+
+  mlir::StatusScopedDiagnosticHandler diag_handler(module->getContext());
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("ifrt_tpu_bct_conversion_before", module);
+  }
+
+  TfrtTpuCompileOptions options;
+  options.disable_set_default_tpu_device_and_device_assignment_attributes =
+      compile_options_
+          .disable_set_default_tpu_device_and_device_assignment_attributes;
+  options.support_multi_dims_sharding = true;
+
+  if (tpu_compiler_ != nullptr) {
+    // Run backward compat pass so that we can use bridge to do clustering.
+    if (mlir::failed(
+            tpu_compiler_->RunTPUBackwardCompatConversion(module, options))) {
+      return diag_handler.Combine(
+          absl::InternalError("Failed to handle legacy TPU Ops"));
+    }
+  }
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("ifrt_tpu_bct_conversion_after", module);
+  }
+
+  // Use bridge for cluster formation.
+  TF_RETURN_IF_ERROR(tensorflow::tf2xla::v2::RunFunctionTf2xlaClusteringBridge(
+      module, /*is_supported_by_replicated_brige*/ true,
+      /*is_in_fallback_enabled_mode=*/false));
+
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("before_ifrt_outlining", module);
+  }
+
+  // Extract TPU program for IFRT call.
+  TF_RETURN_IF_ERROR(CompileTensorflowForIfrtServing(
+      model_context.name(), **ifrt_model_context, module,
+      model_context.graph_execution_options()
+          .compile_options.enable_async_ifrt));
+
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("after_ifrt_outlining", module);
+  }
+
+  // IFRT program is no longer needed.
+  llvm::SmallVector<mlir::func::FuncOp> to_erase;
+  for (auto func : module.getOps<mlir::func::FuncOp>()) {
+    if (func->getAttr("tfrt_ifrt_serving.program_id")) {
+      to_erase.push_back(func);
+    }
+  }
+  for (auto func : to_erase) {
+    func->erase();
+  }
+
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("after_ifrt_program_removal", module);
+  }
+
+  if (mlir::failed(mlir::verify(module))) {
+    return diag_handler.ConsumeStatus();
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_backend_compiler.h

@@ -0,0 +1,72 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_BACKEND_COMPILER_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_BACKEND_COMPILER_H_
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/backend_compiler.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tpu_passes.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// Implements the custom backend compiler for IFRT based serving in TFRT.
+class IfrtBackendCompiler : public tensorflow::BackendCompiler {
+ public:
+  struct Options {
+    // If true, disable running TFRTSetTPUDeviceAttrPass which set the default
+    // `tf.device` and `device_assignment` attributes.
+    // This is a server-level option for now. We can consider to make it a
+    // per-model option in the future.
+    bool disable_set_default_tpu_device_and_device_assignment_attributes = true;
+  };
+
+  explicit IfrtBackendCompiler(
+      std::unique_ptr<TpuCompiler> tpu_compiler = nullptr)
+      : tpu_compiler_(std::move(tpu_compiler)) {}
+
+  explicit IfrtBackendCompiler(
+      const Options& ifrt_backend_compile_options,
+      std::unique_ptr<TpuCompiler> tpu_compiler = nullptr)
+      : tpu_compiler_(std::move(tpu_compiler)),
+        compile_options_(ifrt_backend_compile_options) {}
+
+  void GetDependentDialects(mlir::DialectRegistry& registry) const override {
+    if (tpu_compiler_) {
+      tpu_compiler_->RegisterTPUDialects(&registry);
+    }
+  }
+
+  // Rewrites the tensorflow graph in MLIR for IFRT serving. The methods
+  // extracts regions for IFRT execution on accelerator (e.g. TPU).
+  absl::Status CompileTensorflow(
+      tensorflow::tfrt_stub::ModelRuntimeContext& model_context,
+      mlir::ModuleOp module) const override;
+
+ private:
+  std::unique_ptr<TpuCompiler> tpu_compiler_;
+  Options compile_options_;
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_BACKEND_COMPILER_H_

tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_backend_compiler_test.cc

@@ -0,0 +1,186 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_backend_compiler.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/test_util.h"
+#include "xla/tsl/framework/test_util/mock_serving_device_selector.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/tfrt/graph_executor/graph_execution_options.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_model_context.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+#include "tensorflow/core/tfrt/saved_model/saved_model_testutil.h"
+#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+tsl::thread::ThreadPool& GetThreadPool() {
+  constexpr int kMaxParallelism = 16;
+  static tsl::thread::ThreadPool* thread_pool =
+      new tsl::thread::ThreadPool(tsl::Env::Default(), tsl::ThreadOptions(),
+                                  "IfrtSharding", kMaxParallelism);
+  return *thread_pool;
+}
+
+class IfrtBackendCompilerTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    mlir::registerAllDialects(registry_);
+    mlir::RegisterAllTensorFlowDialects(registry_);
+    context_.appendDialectRegistry(registry_);
+
+    // Create contexts required for the compiler execution.
+    TF_ASSERT_OK_AND_ASSIGN(client_, xla::ifrt::test_util::GetClient());
+
+    core_selector_ = std::make_unique<IfrtServingCoreSelector>(
+        &mock_serving_device_selector_, client_->addressable_device_count());
+    h2d_transfer_executor_factory_ =
+        std::make_unique<H2DTransferExecutorFactory>();
+
+    runtime_context_.resource_context().CreateResource<IfrtModelContext>(
+        "IfrtModelContext", client_, core_selector_.get(), &GetThreadPool(),
+        /*compilation_environment_proto=*/nullptr,
+        /*h2d_transfer_executor_factory=*/h2d_transfer_executor_factory_.get());
+  }
+
+  void verifyModules() {
+    absl::MutexLock l(ServingExecutableRegistry::mu_);
+    for (const auto& [_, executable] :
+         *ServingExecutableRegistry::executables_) {
+      absl::MutexLock l(executable->mutex_);
+      executable->module_->walk([](mlir::func::FuncOp func) {
+        ASSERT_FALSE(func->hasAttr("tfrt_ifrt_serving.program_id"));
+      });
+    }
+  }
+
+  mlir::DialectRegistry registry_;
+  mlir::MLIRContext context_;
+  std::shared_ptr<xla::ifrt::Client> client_;
+
+  std::unique_ptr<tensorflow::tfrt_stub::Runtime> runtime_ =
+      tensorflow::tfrt_stub::DefaultTfrtRuntime(/*num_threads=*/1);
+  tensorflow::tfrt_stub::GraphExecutionOptions graph_execution_options_ =
+      tensorflow::tfrt_stub::GraphExecutionOptions(runtime_.get());
+  tfrt::ResourceContext resource_context_;
+  tensorflow::tfrt_stub::ModelRuntimeContext runtime_context_ =
+      tensorflow::tfrt_stub::ModelRuntimeContext(
+          &graph_execution_options_, /*export_dir=*/"", &resource_context_);
+
+  tsl::test_util::MockServingDeviceSelector mock_serving_device_selector_;
+  std::unique_ptr<IfrtServingCoreSelector> core_selector_;
+  std::unique_ptr<H2DTransferExecutorFactory> h2d_transfer_executor_factory_;
+  IfrtBackendCompiler compiler_;
+};
+
+namespace {
+using ::testing::HasSubstr;
+
+struct IfrtBackendCompilerTestParams {
+  std::string mlir_file_name;
+};
+
+class IfrtBackendCompilerParameterizedTest
+    : public IfrtBackendCompilerTest,
+      public ::testing::WithParamInterface<IfrtBackendCompilerTestParams> {};
+
+TEST_P(IfrtBackendCompilerParameterizedTest, CompilesOk) {
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/", GetParam().mlir_file_name));
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK(
+      compiler_.CompileTensorflow(runtime_context_, mlir_module.get()));
+  verifyModules();
+}
+
+INSTANTIATE_TEST_SUITE_P(IfrtBackendCompilerParameterizedTest,
+                         IfrtBackendCompilerParameterizedTest,
+                         ::testing::ValuesIn<IfrtBackendCompilerTestParams>({
+                             {.mlir_file_name = "ifrt_cluster.mlir"},
+                             {.mlir_file_name = "restore_with_reference.mlir"},
+                         }));
+
+TEST_F(IfrtBackendCompilerTest, CompileShallFailAfterModelIsFrozen) {
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/ifrt_cluster.mlir"));
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK(
+      compiler_.CompileTensorflow(runtime_context_, mlir_module.get()));
+
+  std::optional<IfrtModelContext*> ifrt_model_context =
+      runtime_context_.resource_context().GetResource<IfrtModelContext>(
+          "IfrtModelContext");
+  ASSERT_TRUE(ifrt_model_context.has_value());
+
+  TF_ASSERT_OK((*ifrt_model_context)->Freeze());
+
+  mlir::OwningOpRef<mlir::ModuleOp> another_mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
+
+  EXPECT_THAT(
+      compiler_.CompileTensorflow(runtime_context_, another_mlir_module.get()),
+      absl_testing::StatusIs(
+          absl::StatusCode::kFailedPrecondition,
+          HasSubstr("Cannot compile IFRT programs after the model is frozen")));
+}
+
+}  // namespace
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_compilation.proto

@@ -0,0 +1,32 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.ifrt_serving;
+
+import "tensorflow/compiler/tf2xla/host_compute_metadata.proto";
+import "xla/python/ifrt/layout.proto";
+import "xla/service/hlo.proto";
+import "xla/xla_data.proto";
+import "tensorflow/core/framework/types.proto";
+import "tensorflow/core/protobuf/tpu/compile_metadata.proto";
+
+message Tf2HLOResultProto {
+  xla.HloModuleProto hlo_module_proto = 1;
+  tensorflow.tpu.TPUCompileMetadataProto compile_metadata = 2;
+  tensorflow.tf2xla.HostComputeMetadata host_compute_metadata = 3;
+  repeated xla.ShapeProto xla_input_shapes = 4;
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_constants.h

@@ -0,0 +1,40 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_CONSTANTS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_CONSTANTS_H_
+
+#include "absl/strings/string_view.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// Attribute name of a text TpuCompileMetadataProto. Note that the text proto is
+// not backward compatible and shall not be serialized.
+inline constexpr absl::string_view kMetadataTextAttrName =
+    "__tpu_compile_metadata_text";
+
+// Name of a variable as loaded IFRT array .
+inline constexpr absl::string_view kVariableArrayNameAttr =
+    "__variable_array_name";
+
+// Attribute of a text `VariableDeviceShardingConfigProto`.
+inline constexpr absl::string_view kVariableShardingConfigTextAttr =
+    "__variable_sharding_config_text";
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_CONSTANTS_H_

tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h

@@ -0,0 +1,61 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_TYPES_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_TYPES_H_
+
+#include <optional>
+
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// TODO - b/445480506: Refactor this struct to a class so that accesses are
+// always through accessors.
+struct DtypeAndShape {
+  tensorflow::DataType dtype;
+  tensorflow::TensorShape shape;
+  // If available, use static shape (the upper bound of the actual input shapes)
+  // for compilation and caching.
+  std::optional<tensorflow::TensorShape> static_shape;
+
+  bool operator==(const DtypeAndShape& other) const {
+    return dtype == other.dtype && shape == other.shape &&
+           static_shape == other.static_shape;
+  }
+
+  // Returns the shape to be used for compilation. If static shape is
+  // available, use static shape, so that the same executable can be reused for
+  // all input shapes with the same static shape.
+  tensorflow::TensorShape& GetMutableShapeForCompilation() {
+    // Avoid `static_shape.value_or(shape)`. Since `value_or` returns by value,
+    // using it here would return a reference to a temporary object, leading to
+    // a dangling reference.
+    return static_shape.has_value() ? *static_shape : shape;
+  }
+  const tensorflow::TensorShape& GetShapeForCompilation() const {
+    // Avoid `static_shape.value_or(shape)`. Since `value_or` returns by value,
+    // using it here would return a reference to a temporary object, leading to
+    // a dangling reference.
+    return static_shape.has_value() ? *static_shape : shape;
+  }
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_IFRT_TYPES_H_

tensorflow/compiler/mlir/tfrt/transforms/ifrt/lower_to_ifrt_restore_variable.cc

@@ -0,0 +1,597 @@
+/* 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 <algorithm>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/Support/Casting.h"
+#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/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/ir/types/dialect.h"
+#include "tensorflow/core/platform/protobuf.h"  // IWYU pragma: keep
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_config.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+#define GEN_PASS_DEF_LOWERTOIFRTRESTOREVARIABLEPASS
+#define GEN_PASS_DECL_LOWERTOIFRTRESTOREVARIABLEPASS
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+class LowerToIfrtRestoreVariablePass
+    : public impl::LowerToIfrtRestoreVariablePassBase<
+          LowerToIfrtRestoreVariablePass> {
+ public:
+  void runOnOperation() override {
+    mlir::ModuleOp module = getOperation();
+
+    std::vector<mlir::TF::RestoreV2Op> restore_ops;
+    module.walk([&](mlir::TF::RestoreV2Op restore_op) {
+      restore_ops.push_back(restore_op);
+    });
+
+    for (const auto& restore_op : restore_ops) {
+      if (mlir::failed(RewriteRestore(restore_op))) {
+        return signalPassFailure();
+      }
+    }
+  }
+
+ private:
+  // Holds information about a user of a restored tensor from RestoreV2Op.
+  struct RestoreOpUser {
+    // The chain of ops from RestoreV2Op's user. If var_handle_op is
+    // not null, this chain ends with AssignVariableOp and will be deleted.
+    // Otherwise, restored_tensor_to_return will be set, and this chain
+    // ends with an op that uses the restored tensor, which should be returned
+    // by IfrtRestoreVariableOp.
+    std::vector<mlir::Operation*> op_chain_from_restore_output;
+
+    // The name of the restored tensor.
+    mlir::StringRef tensor_name;
+
+    // The attributes of VarHandleOp associated with the AssignVariableOp.
+    mlir::Type var_handle_type;
+    mlir::StringAttr container;
+    mlir::StringAttr shared_name;
+    bool truncate_in_cast =
+        false;  // value of the truncate attribute in the CastOp.
+                // Default to false if CastOp is not present.
+
+    // The restored tensor that is not assigned to a variable and should be
+    // returned by IfrtRestoreVariableOp.
+    std::optional<mlir::Value> restored_tensor_to_return;
+
+    // The VarHandleOp associated with the AssignVariableOp, if it exists in
+    // the same block or function.
+    mlir::TF::VarHandleOp original_var_handle_op;
+  };
+
+  // Groups parameters being collected for IfrtRestoreVariableOp.
+  struct RestorationParams {
+    llvm::SmallVectorImpl<std::string>& tensor_names;
+    llvm::SmallVectorImpl<std::string>& shape_and_slices;
+    llvm::SmallVectorImpl<mlir::Attribute>& dtypes;
+    llvm::SmallVectorImpl<bool>& truncate_in_cast;
+    std::vector<mlir::Value>& var_handle_values;
+  };
+
+  // Holds information about an intercepted Read-Assign chain.
+  struct InterceptedChain {
+    mlir::TF::AssignVariableOp assign_op;
+    mlir::TF::VarHandleOp target_handle;
+    bool truncate = false;
+    std::vector<mlir::Operation*> op_chain;
+  };
+
+  mlir::LogicalResult ValidateThenUpdateUser(
+      mlir::Value source_tensor, mlir::Operation* user,
+      mlir::StringRef tensor_name,
+      std::vector<RestoreOpUser>& restored_tensor_users) {
+    // Traverses the user chain of a restored tensor, validates it, and collects
+    // user information.
+    RestoreOpUser restored_tensor_user;
+    restored_tensor_user.tensor_name = tensor_name;
+    for (;;) {
+      restored_tensor_user.op_chain_from_restore_output.push_back(user);
+      // The user chain must consist of zero or more `TF::CastOp`s followed by
+      // a `TF::AssignVariableOp`.
+      if (auto cast_op = llvm::dyn_cast<mlir::TF::CastOp>(user)) {
+        // The chain can contain intermediate `CastOp`s.
+        if (!cast_op.getResult().hasOneUse()) {
+          return cast_op->emitOpError()
+                 << " has more than one use in the restore user chain";
+        }
+        restored_tensor_user.truncate_in_cast =
+            restored_tensor_user.truncate_in_cast || cast_op.getTruncate();
+        // Move to the next operation in the user chain.
+        user = *cast_op.getResult().getUsers().begin();
+      } else if (auto assign_variable_op =
+                     llvm::dyn_cast<mlir::TF::AssignVariableOp>(user)) {
+        // The chain must end with an `AssignVariableOp`.
+        if (auto var_handle_op = llvm::dyn_cast<mlir::TF::VarHandleOp>(
+                assign_variable_op.getResource().getDefiningOp())) {
+          // The AssignVariableOp must be associated with a VarHandleOp.
+          restored_tensor_user.var_handle_type = var_handle_op.getType();
+          restored_tensor_user.container = var_handle_op.getContainerAttr();
+          restored_tensor_user.shared_name = var_handle_op.getSharedNameAttr();
+          restored_tensor_user.original_var_handle_op = var_handle_op;
+          break;
+        } else {
+          return assign_variable_op->emitOpError()
+                 << "does not have any associated VarHandle";
+        }
+      } else {
+        // If the restored tensor is not assigned to a variable, it is returned
+        // by the IfrtRestoreVariableOp. A VarHandleOp is created to track the
+        // restored tensor.
+        mlir::OpBuilder builder(user);
+
+        mlir::Type resource_type = mlir::RankedTensorType::get(
+            /*shape=*/{}, /*element_type=*/mlir::tf_type::ResourceType::get(
+                {mlir::cast<mlir::TensorType>(source_tensor.getType())},
+                builder.getContext()));
+        restored_tensor_user.var_handle_type = resource_type;
+        // In TensorFlow, variables are identified by a (container,
+        // shared_name) pair. An empty string for `container` is valid and
+        // indicates that the default container should be used. For restored
+        // tensors that are returned as outputs instead of being assigned to
+        // variables, we create a logical VarHandleOp with an empty container
+        // and use `tensor_name` as `shared_name`.
+        restored_tensor_user.container = builder.getStringAttr("");
+        restored_tensor_user.shared_name = builder.getStringAttr(tensor_name);
+
+        restored_tensor_user.restored_tensor_to_return = source_tensor;
+        break;
+      }
+    }
+
+    // If the user chain is valid, add the collected information.
+    restored_tensor_users.push_back(restored_tensor_user);
+    return mlir::success();
+  }
+
+  // Traces and validates an operation chain starting from a `ReadVariableOp`
+  // user. Returns the chain if it's a valid pattern (optional `Cast`s,
+  // `Slice`s, or `Identity`s followed by an `AssignVariableOp`).
+  std::optional<InterceptedChain> TraceReadChain(mlir::Operation* read_user) {
+    InterceptedChain chain;
+    mlir::Operation* current = read_user;
+    while (true) {
+      if (auto cast_op = llvm::dyn_cast<mlir::TF::CastOp>(current)) {
+        chain.truncate |= cast_op.getTruncate();
+        chain.op_chain.push_back(cast_op);
+        if (!cast_op.getResult().hasOneUse()) return std::nullopt;
+        current = *cast_op.getResult().getUsers().begin();
+      } else if (auto slice_op = llvm::dyn_cast<mlir::TF::SliceOp>(current)) {
+        chain.op_chain.push_back(slice_op);
+        if (!slice_op.getOutput().hasOneUse()) return std::nullopt;
+        current = *slice_op.getOutput().getUsers().begin();
+      } else if (auto identity_op =
+                     llvm::dyn_cast<mlir::TF::IdentityOp>(current)) {
+        chain.op_chain.push_back(identity_op);
+        if (!identity_op.getOutput().hasOneUse()) return std::nullopt;
+        current = *identity_op.getOutput().getUsers().begin();
+      } else {
+        break;
+      }
+    }
+
+    if (auto assign_op = llvm::dyn_cast<mlir::TF::AssignVariableOp>(current)) {
+      if (auto target_handle = llvm::dyn_cast<mlir::TF::VarHandleOp>(
+              assign_op.getResource().getDefiningOp())) {
+        chain.assign_op = assign_op;
+        chain.target_handle = target_handle;
+        chain.op_chain.push_back(assign_op);
+        return chain;
+      }
+    }
+    return std::nullopt;
+  }
+
+  // Processes a validly intercepted chain: updates restoration parameters and
+  // manages `VarHandleOp`s.
+  void ProcessInterceptedChain(InterceptedChain& chain,
+                               mlir::StringRef tensor_name,
+                               mlir::StringRef slice, mlir::Attribute dtype,
+                               RestorationParams& params,
+                               mlir::OpBuilder& handle_builder,
+                               mlir::Operation* restore_op) {
+    params.tensor_names.push_back(tensor_name.str());
+    params.shape_and_slices.push_back(slice.str());
+    params.dtypes.push_back(dtype);
+    params.truncate_in_cast.push_back(chain.truncate);
+
+    mlir::Value handle_result;
+    if (chain.target_handle->getBlock() == restore_op->getBlock()) {
+      handle_result = chain.target_handle.getResult();
+      // Since we are reusing an op that originally appeared later, we need to
+      // move it before the new `IfrtRestoreVariableOp` to avoid dominance
+      // violations.
+      chain.target_handle->moveBefore(restore_op);
+    } else {
+      auto derived_local_handle = handle_builder.create<mlir::TF::VarHandleOp>(
+          restore_op->getLoc(), chain.target_handle.getType(),
+          chain.target_handle.getContainerAttr(),
+          chain.target_handle.getSharedNameAttr());
+      handle_result = derived_local_handle.getResult();
+    }
+    params.var_handle_values.push_back(handle_result);
+  }
+
+  // Identifies and handles cases where a restored variable is read and
+  // subsequently assigned to another variable. For example, if VarA is restored
+  // from a checkpoint, and VarB is initialized with VarA's value (e.g., via
+  // ReadVariable(VarA) -> AssignVariable(VarB)), then VarB should also be
+  // treated as being restored from the checkpoint. This function detects such
+  // patterns, adds VarB to list of tensors to be restored by
+  // IfrtRestoreVariableOp, and marks the intermediate ReadVariableOp and
+  // AssignVariableOp chain for deletion.
+  // Returns intercepted chains if all users of a `ReadVariableOp` lead to
+  // assignments.
+  std::optional<std::vector<InterceptedChain>> GetInterceptedChains(
+      mlir::TF::ReadVariableOp read_op) {
+    std::vector<InterceptedChain> intercepted_chains;
+    // Follow the user chain for each user of the `ReadVariableOp`.
+    for (mlir::Operation* read_user : read_op.getResult().getUsers()) {
+      if (std::optional<InterceptedChain> intercepted_chain =
+              TraceReadChain(read_user)) {
+        intercepted_chains.push_back(std::move(*intercepted_chain));
+      } else {
+        return std::nullopt;
+      }
+    }
+    return intercepted_chains;
+  }
+
+  // Identifies and handles cases where a restored variable is read and
+  // subsequently assigned to another variable. For example, if VarA is restored
+  // from a checkpoint, and VarB is initialized with VarA's value (e.g., via
+  // ReadVariable(VarA) -> AssignVariable(VarB)), then VarB should also be
+  // treated as being restored from the checkpoint. This function detects such
+  // patterns, adds VarB to list of tensors to be restored by
+  // IfrtRestoreVariableOp, and marks the intermediate ReadVariableOp and
+  // AssignVariableOp chain for deletion.
+  void InterceptReadAssignVariable(
+      const llvm::StringMap<std::vector<mlir::TF::ReadVariableOp>>&
+          shared_name_to_reads,
+      mlir::StringRef tensor_name, mlir::StringRef slice, mlir::Attribute dtype,
+      RestoreOpUser& restored_tensor_user, mlir::OpBuilder& handle_builder,
+      mlir::Operation* restore_op, RestorationParams& params,
+      llvm::SmallSetVector<mlir::Operation*, 16>& ops_to_erase) {
+    auto it =
+        shared_name_to_reads.find(restored_tensor_user.shared_name.getValue());
+    // If no `ReadVariableOp` reads the variable associated with
+    // `restored_tensor_user`, then there is no read-assign chain to intercept.
+    if (it == shared_name_to_reads.end()) {
+      return;
+    }
+
+    auto add_to_erase = [&](mlir::Operation* op) { ops_to_erase.insert(op); };
+
+    // Iterate through all `ReadVariableOp`s that read the variable.
+    for (mlir::TF::ReadVariableOp read_op : it->second) {
+      // If all users of this specific `ReadVariableOp` instance were
+      // successfully mapped to `AssignVariableOp` chains, we can erase the
+      // `ReadVariableOp` and all intermediate ops in these chains.
+      if (auto intercepted_chains = GetInterceptedChains(read_op)) {
+        if (intercepted_chains->empty()) continue;
+
+        add_to_erase(read_op);
+        for (auto& chain : *intercepted_chains) {
+          ProcessInterceptedChain(chain, tensor_name, slice, dtype, params,
+                                  handle_builder, restore_op);
+          for (mlir::Operation* op : chain.op_chain) {
+            add_to_erase(op);
+          }
+        }
+      }
+    }
+  }
+
+  // Creates a TF::ConstOp with a 1D tensor of strings from the given values.
+  mlir::Value CreateStringTensorConst(
+      mlir::OpBuilder& builder, mlir::Location loc,
+      const llvm::SmallVector<std::string, 4>& values) {
+    llvm::SmallVector<llvm::StringRef> ref_values;
+    for (const auto& value : values) {
+      ref_values.push_back(value);
+    }
+    auto type = mlir::RankedTensorType::get(
+        {static_cast<int64_t>(values.size())},
+        mlir::tf_type::StringType::get(builder.getContext()));
+
+    return builder.create<mlir::TF::ConstOp>(
+        loc, mlir::DenseStringElementsAttr::get(type, ref_values));
+  }
+
+  mlir::LogicalResult RewriteRestore(mlir::TF::RestoreV2Op restore_op) {
+    // Find and validate all users of the RestoreV2Op's output tensors.
+
+    mlir::DenseStringElementsAttr tensor_names_attr;
+    if (!matchPattern(restore_op.getTensorNames(),
+                      m_Constant(&tensor_names_attr))) {
+      return restore_op.emitOpError("expects tensor_names to be a constant");
+    }
+    auto name_values = tensor_names_attr.getValues<mlir::StringRef>();
+    llvm::SmallVector<mlir::StringRef> tensor_names_vec(name_values.begin(),
+                                                        name_values.end());
+
+    if (restore_op.getTensors().size() != tensor_names_vec.size()) {
+      return restore_op.emitOpError(absl::StrCat(
+          "expects the number of tensors and tensor names to be the same, got ",
+          restore_op.getTensors().size(), " tensors and ",
+          tensor_names_vec.size(), " tensor names."));
+    }
+
+    // The RestoreV2 op contains shape and slice specifications for each tensor
+    // to be restored. These are expected to be provided as a constant string
+    // tensor.
+    mlir::DenseStringElementsAttr shape_and_slices_attr;
+    if (!matchPattern(restore_op.getShapeAndSlices(),
+                      m_Constant(&shape_and_slices_attr))) {
+      return restore_op.emitOpError(
+          "expects shape_and_slices to be a constant");
+    }
+    // Extract shape and slice strings into a vector.
+    // This vector is used later to associate each restored tensor with its
+    // corresponding shape and slice specification when creating the
+    // IfrtRestoreVariableOp.
+    auto slice_values = shape_and_slices_attr.getValues<mlir::StringRef>();
+    llvm::SmallVector<mlir::StringRef> shape_and_slices_vec(
+        slice_values.begin(), slice_values.end());
+
+    std::vector<RestoreOpUser> restored_tensor_users;
+    for (int i = 0; i < restore_op.getTensors().size(); ++i) {
+      mlir::Value out_tensor = restore_op.getTensors()[i];
+      mlir::StringRef tensor_name = tensor_names_vec[i];
+      for (mlir::Operation* user : out_tensor.getUsers()) {
+        if (mlir::failed(ValidateThenUpdateUser(out_tensor, user, tensor_name,
+                                                restored_tensor_users))) {
+          return mlir::failure();
+        }
+      }
+    }
+
+    // Check that each tensor has exactly one valid user chain.
+    if (restored_tensor_users.size() != restore_op.getTensors().size()) {
+      return restore_op->emitOpError()
+             << "expects " << restore_op.getTensors().size()
+             << " valid users, but got " << restored_tensor_users.size();
+    }
+
+    // The transformation needs to identify cases where a restored variable is
+    // read and then assigned to another variable (e.g., RestoreV2 ->
+    // AssignVariable(var1), ReadVariable(var1) -> AssignVariable(var2)).
+    // In such cases, `var2` should also be treated as a restored variable
+    // initialized from the checkpoint. To achieve this, we first map all
+    // variable shared names to their VarHandleOps and all ReadVariableOps
+    // that read them. This allows us to trace reads of restored variables and
+    // check if they are part of an assignment to another variable.
+    llvm::StringMap<mlir::TF::VarHandleOp> shared_name_to_handle;
+    llvm::StringMap<std::vector<mlir::TF::ReadVariableOp>> shared_name_to_reads;
+    restore_op->getParentOfType<mlir::ModuleOp>().walk([&](mlir::Operation*
+                                                               op) {
+      if (auto var_handle = llvm::dyn_cast<mlir::TF::VarHandleOp>(op)) {
+        shared_name_to_handle[var_handle.getSharedName()] = var_handle;
+      } else if (auto read_op = llvm::dyn_cast<mlir::TF::ReadVariableOp>(op)) {
+        if (auto var_handle =
+                read_op.getResource().getDefiningOp<mlir::TF::VarHandleOp>()) {
+          shared_name_to_reads[var_handle.getSharedName()].push_back(read_op);
+        }
+      }
+    });
+
+    // Collect tensor dtypes for the new op.
+    std::vector<mlir::Attribute> dtypes;
+    for (const auto& dtype : restore_op.getDtypes()) {
+      dtypes.push_back(mlir::TypeAttr::get(dtype));
+    }
+
+    // The following vectors collect information for each tensor that will be
+    // restored by the new IfrtRestoreVariableOp. This includes tensors
+    // directly restored from RestoreV2Op, as well as variables initialized
+    // by reading other restored variables (e.g., ReadVariable(restored_var) ->
+    // AssignVariable(new_var)). This information is used to construct the
+    // arguments for the IfrtRestoreVariableOp that replaces RestoreV2Op and
+    // its user chains.
+    llvm::SmallVector<std::string, 4> new_tensor_names;
+    llvm::SmallVector<std::string, 4> new_shape_and_slices;
+    llvm::SmallVector<mlir::Attribute, 4> new_dtypes;
+    llvm::SmallVector<bool, 4> new_truncate_in_cast;
+
+    std::vector<mlir::Value> var_handle_values;
+
+    RestorationParams params{new_tensor_names, new_shape_and_slices, new_dtypes,
+                             new_truncate_in_cast, var_handle_values};
+
+    llvm::SmallVector<mlir::Type> result_types;
+    llvm::DenseMap<mlir::Value, int> returned_tensors_indices;
+    llvm::SmallVector<mlir::StringRef, 4> returned_tensor_names;
+
+    // The following data structures are used to manage operations that need to
+    // be erased after the IfrtRestoreVariableOp is created. This includes
+    // chains of operations from RestoreV2Op to AssignVariableOp, as well as
+    // ReadVariableOp -> AssignVariableOp chains that are optimized away by
+    // InterceptReadAssignVariable.
+    llvm::SmallSetVector<mlir::Operation*, 16> ops_to_erase;
+    // Helper function to add an operation to ops_to_erase, ensuring no
+    // duplicates are added.
+    auto add_to_erase = [&](mlir::Operation* op) { ops_to_erase.insert(op); };
+    // OpBuilder for creating new VarHandleOps before the RestoreV2Op.
+    mlir::OpBuilder handle_builder(restore_op);
+
+    for (int i = 0; i < restored_tensor_users.size(); ++i) {
+      RestoreOpUser& restored_tensor_user = restored_tensor_users[i];
+      mlir::StringRef tensor_name = tensor_names_vec[i];
+      mlir::StringRef slice = shape_and_slices_vec[i];
+      mlir::Attribute dtype = dtypes[i];
+
+      params.tensor_names.push_back(tensor_name.str());
+      params.shape_and_slices.push_back(slice.str());
+      params.dtypes.push_back(dtype);
+      params.truncate_in_cast.push_back(restored_tensor_user.truncate_in_cast);
+
+      // Reuse the original VarHandleOp if it exists and is in the same block as
+      // the RestoreV2Op. This avoids redundant VarHandleOps in the IR and
+      // makes the output cleaner.
+      mlir::Value handle_result;
+      if (restored_tensor_user.original_var_handle_op &&
+          restored_tensor_user.original_var_handle_op->getBlock() ==
+              restore_op->getBlock()) {
+        handle_result = restored_tensor_user.original_var_handle_op.getResult();
+        // Since we are reusing an op that originally appeared later, we need to
+        // move it before the new IfrtRestoreVariableOp to avoid dominance
+        // violations.
+        restored_tensor_user.original_var_handle_op->moveBefore(restore_op);
+      } else {
+        mlir::TF::VarHandleOp local_handle =
+            handle_builder.create<mlir::TF::VarHandleOp>(
+                restore_op->getLoc(), restored_tensor_user.var_handle_type,
+                restored_tensor_user.container,
+                restored_tensor_user.shared_name);
+        handle_result = local_handle.getResult();
+      }
+      params.var_handle_values.push_back(handle_result);
+
+      if (!restored_tensor_user.restored_tensor_to_return.has_value()) {
+        // This tensor is assigned to a variable.
+        // Delete the path from the RestoreV2Op to the AssignVariableOp.
+        for (mlir::Operation* op :
+             restored_tensor_user.op_chain_from_restore_output) {
+          add_to_erase(op);
+        }
+
+        // Intercept usages where the restored variable is assigned to a var
+        // handle 1, and then read from varhandle1 and assigned to var handle 2.
+        InterceptReadAssignVariable(shared_name_to_reads, tensor_name, slice,
+                                    dtype, restored_tensor_user, handle_builder,
+                                    restore_op, params, ops_to_erase);
+      } else {
+        // This tensor is not assigned to a variable and should be returned by
+        // IfrtRestoreVariableOp.
+        mlir::Value tensor = *restored_tensor_user.restored_tensor_to_return;
+        if (returned_tensors_indices.find(tensor) ==
+            returned_tensors_indices.end()) {
+          returned_tensors_indices.insert({tensor, result_types.size()});
+          result_types.push_back(tensor.getType());
+          returned_tensor_names.push_back(restored_tensor_user.tensor_name);
+        }
+      }
+    }
+
+    for (mlir::Operation* op : llvm::reverse(ops_to_erase)) {
+      if (!op->use_empty()) {
+        return op->emitOpError() << "is expected to be erased but has uses";
+      }
+      op->erase();
+    }
+
+    mlir::OpBuilder builder(restore_op);
+
+    // If tensor_names or shape_and_slices have been modified (e.g., due to
+    // intercepted read-assign variables), create new constant tensors.
+    // Otherwise, reuse the original tensors from RestoreV2Op to avoid
+    // redundancy.
+    mlir::Value new_tensor_names_op;
+    if (new_tensor_names.size() == tensor_names_vec.size() &&
+        std::equal(new_tensor_names.begin(), new_tensor_names.end(),
+                   tensor_names_vec.begin())) {
+      new_tensor_names_op = restore_op.getTensorNames();
+    } else {
+      new_tensor_names_op = CreateStringTensorConst(
+          builder, restore_op->getLoc(), new_tensor_names);
+    }
+    mlir::Value new_shape_and_slices_op;
+    if (new_shape_and_slices.size() == shape_and_slices_vec.size() &&
+        std::equal(new_shape_and_slices.begin(), new_shape_and_slices.end(),
+                   shape_and_slices_vec.begin())) {
+      new_shape_and_slices_op = restore_op.getShapeAndSlices();
+    } else {
+      new_shape_and_slices_op = CreateStringTensorConst(
+          builder, restore_op->getLoc(), new_shape_and_slices);
+    }
+
+    auto ifrt_restore_variable_op = mlir::TF::IfrtRestoreVariableOp::create(
+        builder, restore_op->getLoc(), result_types, restore_op.getPrefix(),
+        new_tensor_names_op, new_shape_and_slices_op, var_handle_values,
+        builder.getArrayAttr(new_dtypes),
+        builder.getDenseBoolArrayAttr(new_truncate_in_cast),
+        builder.getStrArrayAttr(returned_tensor_names));
+
+    // Replace the original restored tensors with the results of the new
+    // IfrtRestoreVariableOp.
+    llvm::SmallPtrSet<mlir::Value, 4> replaced_tensors;
+    for (auto& restored_tensor_user : restored_tensor_users) {
+      if (restored_tensor_user.restored_tensor_to_return.has_value()) {
+        mlir::Value tensor = *restored_tensor_user.restored_tensor_to_return;
+        // Replace all uses of the original tensor with the new op's result.
+        // The `replaced_tensors` set ensures that we only do this once per
+        // tensor.
+        if (replaced_tensors.insert(tensor).second) {
+          tensor.replaceAllUsesWith(ifrt_restore_variable_op.getResult(
+              returned_tensors_indices[tensor]));
+        }
+      }
+    }
+
+    // Finally, erase the original RestoreV2Op.
+    if (!restore_op->use_empty()) {
+      return restore_op->emitOpError() << "failed to identify all users"
+                                          "associated with this RestoreV2Op.";
+    } else {
+      restore_op.erase();
+    }
+
+    return mlir::success();
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateLowerToIfrtRestoreVariablePass() {
+  return std::make_unique<LowerToIfrtRestoreVariablePass>();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/pack_inputs_pass.cc

@@ -0,0 +1,238 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/pack_inputs_pass.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Block.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/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "stablehlo/dialect/StablehloOps.h"  // from @stablehlo
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+PackInputsPass::PackInputsPass(llvm::ArrayRef<SliceInfo> slices)
+    : slices_(slices.begin(), slices.end()) {}
+
+// TODO(b/445201291): This pass currently has a caveat of losing the original
+// user inputs' layout. We should look into this if it turns out to be a
+// problem.
+void PackInputsPass::runOnOperation() {
+  mlir::ModuleOp module = getOperation();
+  auto main_func = module.lookupSymbol<mlir::func::FuncOp>("main");
+  if (!main_func || main_func.isPrivate() || main_func.isExternal()) {
+    module.emitError("Expected a public function 'main'.");
+    signalPassFailure();
+    return;
+  }
+
+  // If command line option is provided, use it to populate slices_
+  if (!slices_flat_list_.empty()) {
+    if (slices_flat_list_.size() % 3 != 0) {
+      module.emitError(
+          "The 'slices' option must be a flat list of integers "
+          "with a length multiple of 3.");
+      signalPassFailure();
+      return;
+    }
+    slices_.clear();
+    for (size_t i = 0; i < slices_flat_list_.size(); i += 3) {
+      slices_.push_back({static_cast<unsigned>(slices_flat_list_[i]),
+                         slices_flat_list_[i + 1], slices_flat_list_[i + 2]});
+    }
+  }
+
+  if (slices_.empty()) {
+    return;
+  }
+
+  llvm::SmallVector<unsigned> small_arg_indices;
+  int64_t total_small_size = 0;
+
+  for (const auto& slice : slices_) {
+    if (slice.arg_index >= main_func.getNumArguments()) {
+      module->emitError(absl::StrCat("Slice arg_index ", slice.arg_index,
+                                     " is out of range."));
+      return signalPassFailure();
+    }
+
+    if (llvm::is_contained(small_arg_indices, slice.arg_index)) {
+      module->emitError(
+          absl::StrCat("Duplicate slice arg_index: ", slice.arg_index));
+      return signalPassFailure();
+    }
+
+    mlir::BlockArgument arg = main_func.getArgument(slice.arg_index);
+    auto tensor_type = llvm::dyn_cast<mlir::RankedTensorType>(arg.getType());
+    if (!tensor_type) {
+      module->emitError(absl::StrCat("Argument ", slice.arg_index,
+                                     " is not a ranked tensor type."));
+      return signalPassFailure();
+    }
+    if (!tensor_type.hasStaticShape()) {
+      module->emitError(absl::StrCat("Argument ", slice.arg_index,
+                                     " does not have static shape."));
+      return signalPassFailure();
+    }
+
+    int64_t bitwidth = tensor_type.getElementType().getIntOrFloatBitWidth();
+    if (bitwidth < 8 || bitwidth % 8 != 0) {
+      module->emitError(
+          absl::StrCat("Argument ", slice.arg_index,
+                       " has invalid element bit width: ", bitwidth));
+      return signalPassFailure();
+    }
+
+    int64_t expected_size = tensor_type.getNumElements() * (bitwidth / 8);
+    if (slice.size != expected_size) {
+      module->emitError(absl::StrCat(
+          "Slice size ", slice.size, " for argument ", slice.arg_index,
+          " does not match expected byte size ", expected_size));
+      return signalPassFailure();
+    }
+
+    small_arg_indices.push_back(slice.arg_index);
+    total_small_size = std::max(total_small_size, slice.start + slice.size);
+  }
+
+  // Check for overlapping slices
+  for (size_t i = 0; i < slices_.size(); ++i) {
+    for (size_t j = i + 1; j < slices_.size(); ++j) {
+      int64_t start_i = slices_[i].start;
+      int64_t end_i = start_i + slices_[i].size;
+      int64_t start_j = slices_[j].start;
+      int64_t end_j = start_j + slices_[j].size;
+      if (std::max(start_i, start_j) < std::min(end_i, end_j)) {
+        module->emitError(absl::StrCat("Slices for argument ",
+                                       slices_[i].arg_index, " and ",
+                                       slices_[j].arg_index, " overlap."));
+        return signalPassFailure();
+      }
+    }
+  }
+
+  mlir::OpBuilder builder(&getContext());
+  mlir::RankedTensorType combined_arg_type =
+      mlir::RankedTensorType::get({total_small_size}, builder.getI8Type());
+
+  mlir::Block& entry_block = main_func.front();
+  mlir::BlockArgument combined_arg =
+      entry_block.addArgument(combined_arg_type, main_func.getLoc());
+
+  builder.setInsertionPointToStart(&entry_block);
+  for (const auto& slice : slices_) {
+    mlir::BlockArgument arg = entry_block.getArgument(slice.arg_index);
+    llvm::SmallVector<int64_t> start_indices = {slice.start};
+    llvm::SmallVector<int64_t> limit_indices = {slice.start + slice.size};
+    llvm::SmallVector<int64_t> strides = {1};
+    auto slice_op = builder.create<mlir::stablehlo::SliceOp>(
+        arg.getLoc(), combined_arg, builder.getDenseI64ArrayAttr(start_indices),
+        builder.getDenseI64ArrayAttr(limit_indices),
+        builder.getDenseI64ArrayAttr(strides));
+
+    auto tensor_type = llvm::dyn_cast<mlir::RankedTensorType>(arg.getType());
+    int64_t targetEltBitWidth =
+        tensor_type.getElementType().getIntOrFloatBitWidth();
+
+    // If the original type is one byte, we can just do a reshape. If it is
+    // larger than one byte, we need to do a reshape and a bitcast to convert
+    // to the original type. For example, if the original type is F32 of shape
+    // (x, y), we need to reshape to {x, y, 4} as F32 is 4 bytes, and then
+    // bitcast to F32 with shape (x, y).
+    if (targetEltBitWidth == 8) {
+      auto reshape_op = builder.create<mlir::stablehlo::ReshapeOp>(
+          arg.getLoc(), arg.getType(), slice_op.getResult());
+      arg.replaceAllUsesWith(reshape_op.getResult());
+    } else {
+      llvm::SmallVector<int64_t> intermediate_shape(
+          tensor_type.getShape().begin(), tensor_type.getShape().end());
+      intermediate_shape.push_back(targetEltBitWidth / 8);
+      auto intermediate_type =
+          mlir::RankedTensorType::get(intermediate_shape, builder.getI8Type());
+
+      auto reshape_op = builder.create<mlir::stablehlo::ReshapeOp>(
+          arg.getLoc(), intermediate_type, slice_op.getResult());
+
+      auto bitcast_op = builder.create<mlir::stablehlo::BitcastConvertOp>(
+          arg.getLoc(), arg.getType(), reshape_op.getResult());
+
+      arg.replaceAllUsesWith(bitcast_op.getResult());
+    }
+  }
+
+  // Erase arguments
+  llvm::BitVector small_args_to_erase_bv(entry_block.getNumArguments());
+  for (unsigned index : small_arg_indices) {
+    small_args_to_erase_bv.set(index);
+  }
+  entry_block.eraseArguments(small_args_to_erase_bv);
+
+  // Set new function type and attributes
+  mlir::FunctionType old_func_type = main_func.getFunctionType();
+  llvm::SmallVector<mlir::Type> new_input_types;
+  llvm::SmallVector<mlir::DictionaryAttr> new_arg_attrs;
+  for (unsigned i = 0; i < old_func_type.getNumInputs(); ++i) {
+    if (!llvm::is_contained(small_arg_indices, i)) {
+      new_input_types.push_back(old_func_type.getInput(i));
+      new_arg_attrs.push_back(main_func.getArgAttrDict(i));
+    }
+  }
+  new_input_types.push_back(combined_arg_type);
+  new_arg_attrs.push_back(nullptr);
+
+  auto new_func_type = mlir::FunctionType::get(&getContext(), new_input_types,
+                                               old_func_type.getResults());
+  main_func.setFunctionType(new_func_type);
+  main_func.setAllArgAttrs(new_arg_attrs);
+}
+
+void PackInputsPass::getDependentDialects(
+    mlir::DialectRegistry& registry) const {
+  registry.insert<mlir::stablehlo::StablehloDialect>();
+}
+
+mlir::StringRef PackInputsPass::getArgument() const { return "pack-inputs"; }
+
+mlir::StringRef PackInputsPass::getDescription() const {
+  return "Pack specified tensor inputs of main function into a single i8 "
+         "tensor.";
+}
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreatePackInputsPass(
+    llvm::ArrayRef<SliceInfo> slices) {
+  return std::make_unique<PackInputsPass>(slices);
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/pack_inputs_pass.h

@@ -0,0 +1,80 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_PACK_INPUTS_PASS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_PACK_INPUTS_PASS_H_
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/CommandLine.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+
+namespace mlir {
+class DialectRegistry;
+}  // namespace mlir
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+struct SliceInfo {
+  unsigned arg_index;
+  int64_t start;
+  int64_t size;
+};
+
+// This pass packs specified tensor inputs of main function into a single i8
+// tensor buffer using explicit slice coordinates.
+class PackInputsPass
+    : public mlir::PassWrapper<PackInputsPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PackInputsPass)
+
+  PackInputsPass() = default;
+  PackInputsPass(const PackInputsPass& other)
+      : PassWrapper(other), slices_(other.slices_) {}
+  explicit PackInputsPass(llvm::ArrayRef<SliceInfo> slices);
+
+  void runOnOperation() final;
+
+  void getDependentDialects(mlir::DialectRegistry& registry) const override;
+
+  mlir::StringRef getArgument() const override;
+  mlir::StringRef getDescription() const override;
+
+ private:
+  std::vector<SliceInfo> slices_;
+  ListOption<int64_t> slices_flat_list_{
+      *this, "slices",
+      llvm::cl::desc(
+          "Flat list of integers specifying {arg_index, start, size} "
+          "for each slice to be packed.")};
+};
+
+// Creates an instance of the PackInputsPass.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreatePackInputsPass(
+    llvm::ArrayRef<SliceInfo> slices = {});
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_PACK_INPUTS_PASS_H_

tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.td

@@ -0,0 +1,158 @@
+/* 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 "mlir/Pass/PassBase.td"
+
+
+def RewriteClusterToIfrtCallPass: Pass<"rewrite-cluster-to-ifrt-call", "mlir::ModuleOp"> {
+  let summary = "Lowers tf_device.cluster_func to tf.ifrt_call";
+
+  let description = [{
+    This pass converts the clustered tpu program into an IFRT program and rewrites 
+    tf_device.cluster_fun to tf.ifrt_call.
+  }];
+
+  let dependentDialects = ["mlir::tf_device::TensorFlowDeviceDialect"];
+
+  let constructor = "CreateRewriteClusterToIfrtCallPass()";
+  let options = [
+    Option<"enable_async_ifrt", "enable-async-ifrt", "bool", "false",
+           "Enable async IFRT execution.">
+  ];
+ }
+
+def PropagateStaticShapesPass : Pass<"propagate-static-shapes", "mlir::ModuleOp"> {
+  let summary = "Propagates static shapes from `tf.SetStaticDimensionBounds` to `tf.IfrtCall`";
+
+  let description = [{
+    This pass propagates static shape information from `tf.SetStaticDimensionBounds`
+    to `tf.IfrtCall` operands and the arguments of its callee function.
+    `tf.SetStaticDimensionBounds` is then replaced with its input.
+  }];
+
+  let constructor = "CreatePropagateStaticShapesPass()";
+}
+
+def SinkVariableAsNamedArrayPass: Pass<"sink-variable-as-named-array", "mlir::ModuleOp"> {
+  let summary = "Sink variable tensor for tpu device as named IFRT array for tf.IfrtCall";
+
+  let description = [{
+    This pass sinks variable tensor argument to `tf.IfrtCall` as variable_arg_indices
+    and variable_names attributes and also lowers `tf.ReadVariableOp` to
+    `tf.IfrtLoadVariableOp`.
+
+    The runtime ensures that `tf.IfrtCall` kernel can bind the IFRT array by
+    its name as input to the TPU program.
+
+ }];
+
+  let constructor = "CreateSinkVariableAsNamedArrayPass()";
+}
+
+def LowerToIfrtRestoreVariablePass: Pass<"lower-to-ifrt-restore-variable", "mlir::ModuleOp"> {
+  let summary = "Lowers tf.RestoreV2 to tf.IfrtRestoreVariable";
+
+  let description = [{
+    This pass rewrites `RestoreV2Op` and `AssignVariableOp` to `IfrtRestoreVariableOp`.
+
+    This pass is ran so that the runtime may run restore operation asynchronously
+    such that the expensive resotre operation can be parallerized with other 
+    model compilation operation.
+  }];
+
+  let constructor = "CreateLowerToIfrtRestoreVariablePass()";
+ }
+
+def TfRestorePruningPass
+    : Pass<"tf-restore-pruning", "mlir::func::FuncOp"> {
+  let summary = "Prune unused`tf.RestoreV2` ops";
+
+  let description = [{
+    This pass prune unused `tf.RestoreV2` op. A typical use case is to combine
+    `TfRestoreSplittingPass`, this pass and `TfRestoreMergingPass` in sequence
+    so that the un-used restored tensors are not read into host memory.
+  }];
+
+  let constructor = "CreateTfRestorePruningPass()";
+}
+
+
+def TfRestoreSplittingPass
+    : Pass<"tf-restore-splitting", "mlir::func::FuncOp"> {
+  let summary = "Splits `tf.RestoreV2` ops";
+
+  let description = [{
+    This pass splits each `tf.RestoreV2` op so that one restore op handles one
+    variable only. This pass can split restore ops only if the tensor names and
+    the shape/slices arguments are constants, which is usually the case.
+
+    Splitting monolithic restore ops into per-tensor restore ops makes it easier
+    to shard SavedModel initialization across multiple clusters.
+  }];
+
+  let constructor = "CreateTfRestoreSplittingPass()";
+}
+
+def TfRestoreMergingPass : Pass<"tf-restore-merging", "mlir::func::FuncOp"> {
+  let summary = "Merges `tf.RestoreV2` ops";
+
+  let description = [{
+    This pass merges multiple `tf.RestoreV2` ops into one `tf.RestoreV2` op
+    using variadic results. The current implementation merges restore ops only
+    if they have the same `prefix` and have constant tensor names and
+    shape/slice arguments.
+
+    This pass is run in order to undo `tf-restore-splitting` after cluster
+    formation and reduce the op dispatch overhead.
+  }];
+
+  let constructor = "CreateTfRestoreMergingPass()";
+}
+
+def TfIdentityPropagationPass
+    : Pass<"tf-identity-propagation", "mlir::func::FuncOp"> {
+  let summary = "Propagates inputs of no-op identity ops to their outputs";
+
+  let description = [{
+    This pass finds identity ops that are no-op and propagates their inputs
+    directly to outputs so that identity ops can be skipped.
+
+    One example of identity ops that are not no-op is identity ops with XLA
+    sharding annotation. Since some models use identity ops with `_XlaSharding`
+    attributes to change output sharding, this pass doesn't propagate the inputs
+    of such identity ops in order to preserve the sharding changes.
+
+    This pass is useful to make sure that ineffective identity ops don't affect
+    the graph partitioning. For example, in a pipelined model, if there is a CPU
+    identity op between two TPU computation stages (which sometimes happens
+    because TensorFlow inserts it), this will unnecessarily route the
+    intermediate tensors through the CPU device. By forwarding the inputs of the
+    identity op directly to its outputs, we can avoid such inefficiency.
+  }];
+
+  let constructor = "CreateTfIdentityPropagationPass()";
+}
+
+def TfDeviceCleanupPass : Pass<"tf-device-cleanup", "mlir::func::FuncOp"> {
+  let summary = "Cleans up device attributes from all ops";
+
+  let description = [{
+    This pass removes `device` attributes from all TF ops. Some Serving
+    doesn't rely on `device` attributes from SavedModel.
+  }];
+
+  let constructor = "CreateTfDeviceCleanupPass()";
+}
+

tensorflow/compiler/mlir/tfrt/transforms/ifrt/propagate_static_shapes.cc

@@ -0,0 +1,332 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "absl/log/log.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.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/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/ValueRange.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_constants.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_ifrt_passes.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+using mlir::TF::AsyncIfrtCallOp;
+using mlir::TF::IfrtCallOp;
+using mlir::TF::SetStaticDimensionBoundsOp;
+
+// Identifies operands of `IfrtCallOp` or `AsyncIfrtCallOp` that are defined by
+// `SetStaticDimensionBoundsOp` and returns a map from operand index to the
+// defining op.
+template <typename OpT>
+llvm::SmallDenseMap<size_t, SetStaticDimensionBoundsOp>
+GetArgIdxToStaticShapeOpMap(OpT ifrt_call) {
+  llvm::SmallDenseMap<size_t, SetStaticDimensionBoundsOp>
+      arg_idx_to_static_shape_op;
+  for (const auto& [i, arg] : llvm::enumerate(ifrt_call.getArgs())) {
+    mlir::Value cur_arg = arg;
+    // We only check `IdentityOp` (not `IdentityNOp`) because
+    // `SetStaticDimensionBoundsOp` returns a single tensor.
+    while (auto identity = cur_arg.getDefiningOp<mlir::TF::IdentityOp>()) {
+      cur_arg = identity.getInput();
+    }
+    if (auto static_shape_op =
+            cur_arg.getDefiningOp<SetStaticDimensionBoundsOp>()) {
+      VLOG(2) << "Found SetStaticDimensionBoundsOp for arg " << i;
+      arg_idx_to_static_shape_op[i] = static_shape_op;
+    }
+  }
+  return arg_idx_to_static_shape_op;
+}
+
+// Replaces `old_op` with a new op with `updated_args` and
+// `static_shape_args`. Returns the new op.
+template <typename OpT>
+OpT UpdateIfrtCallOp(OpT old_op,
+                     const llvm::SmallVector<mlir::Value, 4>& updated_args,
+                     const llvm::SmallVector<mlir::Value, 4>& static_shape_args,
+                     mlir::OpBuilder& builder) {
+  // Clone to a new op.
+  builder.setInsertionPoint(old_op);
+  auto new_op = mlir::cast<OpT>(builder.clone(*old_op));
+  // Wire up inputs.
+  llvm::SmallVector<mlir::Value, 4> new_operands;
+  new_operands.append(updated_args.begin(), updated_args.end());
+  new_operands.append(static_shape_args.begin(), static_shape_args.end());
+  new_op->setOperands(new_operands);
+  // This attr is to delimit the boundary between original arguments and static
+  // shape arguments.
+  new_op->setAttr("operandSegmentSizes",
+                  builder.getDenseI32ArrayAttr(
+                      {static_cast<int32_t>(updated_args.size()),
+                       static_cast<int32_t>(static_shape_args.size())}));
+  // Wire up outputs.
+  old_op.replaceAllUsesWith(new_op.getResults());
+  // Get rid of the old op.
+  old_op.erase();
+  return new_op;
+}
+
+// Updates TPUCompileMetadataProto in function attributes to include new
+// arguments for static shapes. Each static shape is added as a new parameter
+// argument with replicated sharding.
+mlir::LogicalResult UpdateTpuCompileMetadata(
+    mlir::func::FuncOp func_op,
+    const llvm::SmallVector<mlir::Value, 4>& static_shape_args) {
+  // Read the old metadata from the func op.
+  auto metadata_attr =
+      func_op->getAttrOfType<mlir::StringAttr>(kMetadataTextAttrName);
+  if (!metadata_attr) return mlir::success();
+  tensorflow::tpu::TPUCompileMetadataProto metadata;
+  if (!tsl::protobuf::TextFormat::ParseFromString(
+          metadata_attr.getValue().str(), &metadata)) {
+    return func_op.emitError() << "Failed to parse " << kMetadataTextAttrName;
+  }
+  // Update metadata proto to include static shape args.
+  for (const auto& static_shape_arg : static_shape_args) {
+    auto* new_compile_arg = metadata.add_args();
+    // NOTE: tf2hlo.cc strictly requires PARAMETER kind for all args.
+    new_compile_arg->set_kind(
+        tensorflow::tpu::TPUCompileMetadataProto::Arg::PARAMETER);
+    tensorflow::DataType dtype;
+    if (const auto status =
+            tensorflow::ConvertToDataType(static_shape_arg.getType(), &dtype);
+        !status.ok()) {
+      return func_op.emitError() << "Failed to convert static shape type to "
+                                    "TensorFlow DataType: "
+                                 << status.message();
+    }
+    new_compile_arg->set_dtype(dtype);
+    tensorflow::ConvertTypeToTensorShape(static_shape_arg.getType())
+        .AsProto(new_compile_arg->mutable_shape());
+    new_compile_arg->mutable_sharding()->set_type(xla::OpSharding::REPLICATED);
+  }
+  // Write the new metadata to the func op.
+  std::string new_metadata_str;
+  if (!tsl::protobuf::TextFormat::PrintToString(metadata, &new_metadata_str)) {
+    return func_op.emitError()
+           << "Failed to serialize " << kMetadataTextAttrName;
+  }
+  func_op->setAttr(
+      kMetadataTextAttrName,
+      mlir::StringAttr::get(func_op.getContext(), new_metadata_str));
+  return mlir::success();
+}
+
+// Updates the `FuncOp` identified by `program_id` to incorporate
+// `static_shape_args`. This includes setting `tf._static_shape_arg` attributes
+// on original arguments, updating the function signature, adding block
+// arguments, and updating the `TPUCompileMetadataProto`.
+template <typename OpT>
+mlir::LogicalResult UpdateCalleeFuncOp(
+    mlir::ModuleOp module, OpT ifrt_call,
+    const llvm::SmallVector<mlir::Value, 4>& static_shape_args,
+    const llvm::SmallDenseMap<size_t, size_t>& arg_idx_to_static_shape_idx) {
+  // Find the callee `FuncOp` by program ID.
+  mlir::func::FuncOp func_op;
+  // NOTE: IFRT program functions are generated by
+  // `RewriteClusterToIfrtCallPass` and inserted into the module's symbol table,
+  // so they are always top-level operations.
+  for (auto f : module.getOps<mlir::func::FuncOp>()) {
+    auto id_attr =
+        f->getAttrOfType<mlir::IntegerAttr>("tfrt_ifrt_serving.program_id");
+    // NOTE: Cast is used because this ID is uint64_t in `IfrtCallOp` but
+    // signless (aka int64_t in MLIR's sense) in `FuncOp`.
+    if (id_attr &&
+        static_cast<uint64_t>(id_attr.getInt()) == ifrt_call.getProgramId()) {
+      func_op = f;
+      break;
+    }
+  }
+  if (!func_op) {
+    return ifrt_call.emitOpError() << "callee func with program_id "
+                                   << ifrt_call.getProgramId() << " not found";
+  }
+  // Point original args to new static shape args.
+  //
+  // The `tf._static_shape_arg_idx` attribute on an original argument indicates
+  // the index of the corresponding static shape argument in the updated
+  // function signature. For example, if the original arguments are (arg0, arg1,
+  // arg2) and only arg0 and arg2 have static shapes (ss0, ss2), the new
+  // signature is (arg0, arg1, arg2, ss0, ss2). Then arg0 will have
+  // `tf._static_shape_arg_idx = 3` and arg2 will have
+  // `tf._static_shape_arg_idx = 4`.
+  mlir::OpBuilder builder(func_op);
+  for (const auto& [arg_idx, static_shape_idx] : arg_idx_to_static_shape_idx) {
+    func_op.setArgAttr(
+        arg_idx, "tf._static_shape_arg_idx",
+        builder.getI32IntegerAttr(static_cast<int32_t>(static_shape_idx)));
+  }
+  // Update function type to include static shape args.
+  const auto& old_function_type = func_op.getFunctionType();
+  llvm::SmallVector<mlir::Type, 4> new_input_types(
+      old_function_type.getInputs().begin(),
+      old_function_type.getInputs().end());
+  for (const auto& arg : static_shape_args) {
+    new_input_types.push_back(arg.getType());
+  }
+  func_op.setType(mlir::FunctionType::get(module.getContext(), new_input_types,
+                                          old_function_type.getResults()));
+  // Append static shape args to the entry block.
+  for (const auto& arg : static_shape_args) {
+    func_op.getBody().front().addArgument(arg.getType(), func_op.getLoc());
+  }
+  return UpdateTpuCompileMetadata(func_op, static_shape_args);
+}
+
+template <typename OpT>
+mlir::LogicalResult ProcessIfrtCall(
+    OpT ifrt_call, mlir::ModuleOp module, mlir::OpBuilder& op_builder,
+    llvm::DenseSet<uint64_t>& processed_programs) {
+  llvm::SmallDenseMap<size_t, SetStaticDimensionBoundsOp>
+      arg_idx_to_static_shape_op = GetArgIdxToStaticShapeOpMap(ifrt_call);
+  if (arg_idx_to_static_shape_op.empty()) return mlir::success();
+
+  VLOG(2) << "Found `IfrtCallOp` with " << arg_idx_to_static_shape_op.size()
+          << " static shaped args. IfrtCallOp: " << ifrt_call;
+
+  mlir::OperandRange old_args = ifrt_call.getArgs();
+  llvm::SmallVector<mlir::Value, 4> updated_args;
+  updated_args.resize(old_args.size());
+  llvm::SmallVector<mlir::Value, 4> static_shape_args;
+  llvm::SmallDenseMap<size_t, size_t> arg_idx_to_static_shape_idx;
+  for (const auto& [i, arg] : llvm::enumerate(old_args)) {
+    auto iter = arg_idx_to_static_shape_op.find(i);
+    if (iter != arg_idx_to_static_shape_op.end()) {
+      SetStaticDimensionBoundsOp static_shape_op = iter->second;
+      arg_idx_to_static_shape_idx[i] =
+          old_args.size() + static_shape_args.size();
+      static_shape_args.push_back(static_shape_op.getStaticShape());
+      updated_args[i] = static_shape_op.getInput();
+    } else {
+      updated_args[i] = arg;
+    }
+  }
+  uint64_t program_id = ifrt_call.getProgramId();
+  auto new_ifrt_call =
+      UpdateIfrtCallOp(ifrt_call, updated_args, static_shape_args, op_builder);
+
+  // Defensive coding behavior: Ensure we only update the callee FuncOp once
+  // per program_id. In theory, it is highly unlikely that the same program
+  // is called by both IfrtCallOp and AsyncIfrtCallOp (or multiple times with
+  // different static shape bounds) in the same module, because the choice of
+  // op type is usually controlled by a session-wide flag (e.g.,
+  // enable_async_ifrt). However, protecting against it avoids duplicate
+  // argument promotion in the callee signature.
+  if (!processed_programs.contains(program_id)) {
+    if (mlir::failed(UpdateCalleeFuncOp(module, new_ifrt_call,
+                                        static_shape_args,
+                                        arg_idx_to_static_shape_idx))) {
+      return mlir::failure();
+    }
+    processed_programs.insert(program_id);
+  }
+
+  // Erase all `SetStaticDimensionBounds` ops.
+  for (auto [_, op] : arg_idx_to_static_shape_op) {
+    op.replaceAllUsesWith(op.getOperand(0));
+    op.erase();
+  }
+  return mlir::success();
+}
+
+// Pass that propagates static shapes from `tf.SetStaticDimensionBoundsOp` to
+// `IfrtCallOp` and callee `FuncOp`.
+struct PropagateStaticShapesPass : public mlir::OperationPass<mlir::ModuleOp> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PropagateStaticShapesPass)
+
+  PropagateStaticShapesPass()
+      : mlir::OperationPass<mlir::ModuleOp>(
+            mlir::TypeID::get<PropagateStaticShapesPass>()) {}
+  PropagateStaticShapesPass(const PropagateStaticShapesPass& other) = default;
+
+  llvm::StringRef getArgument() const override {
+    return "propagate-static-shapes";
+  }
+  llvm::StringRef getDescription() const override {
+    return "Propagates static shapes from tf.SetStaticDimensionBoundsOp to "
+           "IfrtCallOp and callee FuncOp.";
+  }
+  llvm::StringRef getName() const override {
+    return "PropagateStaticShapesPass";
+  }
+
+  std::unique_ptr<mlir::Pass> clonePass() const override {
+    return std::make_unique<PropagateStaticShapesPass>(*this);
+  }
+
+  void runOnOperation() override {
+    mlir::ModuleOp module = getOperation();
+    mlir::OpBuilder op_builder(module.getContext());
+
+    llvm::DenseSet<uint64_t> processed_programs;
+
+    // Collect all calls in the module, before any modifications.
+    llvm::SmallVector<IfrtCallOp, 4> ifrt_calls;
+    llvm::SmallVector<AsyncIfrtCallOp, 4> async_ifrt_calls;
+    module.walk([&ifrt_calls](IfrtCallOp op) { ifrt_calls.push_back(op); });
+    module.walk([&async_ifrt_calls](AsyncIfrtCallOp op) {
+      async_ifrt_calls.push_back(op);
+    });
+
+    for (auto ifrt_call : ifrt_calls) {
+      if (mlir::failed(ProcessIfrtCall(ifrt_call, module, op_builder,
+                                       processed_programs))) {
+        return signalPassFailure();
+      }
+    }
+    for (auto ifrt_call : async_ifrt_calls) {
+      if (mlir::failed(ProcessIfrtCall(ifrt_call, module, op_builder,
+                                       processed_programs))) {
+        return signalPassFailure();
+      }
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreatePropagateStaticShapesPass() {
+  return std::make_unique<PropagateStaticShapesPass>();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/rewrite_cluster_to_ifrt_call.cc

@@ -0,0 +1,387 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/LogicalResult.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/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/WalkResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_structs.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/host_runtime/tpu_metadata_utils.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/device_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/tpu_rewrite_device_util.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_constants.h"
+#include "tensorflow/core/platform/random.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+#define GEN_PASS_DECL_REWRITECLUSTERTOIFRTCALLPASS
+#define GEN_PASS_DEF_REWRITECLUSTERTOIFRTCALLPASS
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+// A pass that inserts tf.ifrt_call and create its callee as a Ifrt
+// Program.
+class RewriteClusterToIfrtCallPass
+    : public impl::RewriteClusterToIfrtCallPassBase<
+          RewriteClusterToIfrtCallPass> {
+ public:
+  explicit RewriteClusterToIfrtCallPass(bool arg_enable_async_ifrt = false) {
+    this->enable_async_ifrt = arg_enable_async_ifrt;
+  }
+
+  void runOnOperation() override {
+    mlir::ModuleOp module = getOperation();
+    mlir::SymbolTable symbol_table(module);
+    mlir::SymbolTableCollection symbol_table_collection;
+
+    mlir::TF::RuntimeDevices devices;
+    if (failed(tensorflow::GetDevicesFromOp(module.getOperation(), &devices)))
+      return signalPassFailure();
+
+    // key: original callee function in tf_device.cluster_func. value: ifrt
+    // program.
+    llvm::DenseMap<mlir::func::FuncOp, mlir::func::FuncOp>
+        cluster_to_ifrt_program;
+
+    std::vector<mlir::tf_device::ClusterFuncOp> cluster_func_ops;
+    module.walk([&](mlir::tf_device::ClusterFuncOp cluster_func) {
+      cluster_func_ops.push_back(cluster_func);
+    });
+    for (auto cluster_func : cluster_func_ops) {
+      Rewrite(symbol_table, symbol_table_collection, cluster_to_ifrt_program,
+              cluster_func, devices);
+    }
+
+    // TODO(b/304839793): Move this to a separate pass. The old remove
+    // compilation result pass rely on TPUPartitionedCall
+    llvm::SmallVector<mlir::TF::TPUCompilationResultOp> compilation_result_ops;
+    module.walk([&](mlir::TF::TPUCompilationResultOp op) {
+      compilation_result_ops.push_back(op);
+    });
+    for (auto op : compilation_result_ops) {
+      if (!op.use_empty()) {
+        module->emitError("TPUCompilationResultOp is under use");
+        return signalPassFailure();
+      }
+      op.erase();
+    }
+  }
+
+ private:
+  // Returns a new unique program id.
+  static int64_t NewProgramId() {
+    const uint64_t id = static_cast<int64_t>(tensorflow::random::New64());
+    // We use a signed int for program ids since TensorFlow doesn't
+    // support uint64_t attributes.
+    return absl::bit_cast<int64_t>(id);
+  }
+
+  mlir::LogicalResult GetTpuCompileMetadata(
+      mlir::tf_device::ClusterFuncOp cluster_func,
+      mlir::TF::RuntimeDevices &devices,
+      tensorflow::tpu::TPUCompileMetadataProto *metadata) {
+    // Collect `num_replicas` and `num_cores_per_replica` attributes.
+    int num_replicas = 1;
+    mlir::tf_device::ReplicateOp replicate =
+        cluster_func->getParentOfType<mlir::tf_device::ReplicateOp>();
+    if (replicate) num_replicas = replicate.getN();
+
+    auto num_cores_per_replica_attr =
+        cluster_func->getAttrOfType<mlir::IntegerAttr>(
+            tensorflow::kNumCoresPerReplicaAttr);
+    if (!num_cores_per_replica_attr)
+      return cluster_func.emitOpError()
+             << "Attribute" << tensorflow::kNumCoresPerReplicaAttr
+             << " is missing";
+    int num_cores_per_replica = num_cores_per_replica_attr.getInt();
+
+    return mlir::TFTPU::SetMetadataProtoFromClusterFuncOp(
+        cluster_func, num_replicas, num_cores_per_replica,
+        /*xla_device_assignment=*/std::nullopt, metadata);
+  }
+
+  // Returns true if the function or any of its callees recursively contains
+  // tf.PwStreamResultsOp.
+  //
+  // PwStreamResultsOp immediately streams results back to the serving
+  // controller. We must disable async execution (falling back to synchronous
+  // IfrtCallOp) if this op is present. Otherwise, the async call might finish
+  // without properly synchronizing or awaiting the streaming results, causing
+  // race conditions or broken serving behavior (see b/500390771).
+  //
+  // Complexity Note: This method must recursively traverse control flow ops
+  // (e.g., tf.If, tf.While) and partitioned calls. Since these ops frequently
+  // reference callee functions via an ArrayAttr of symbols rather than a simple
+  // SymbolRefAttr, we explicitly iterate through ArrayAttr elements to avoid
+  // missing streaming ops hidden inside branches.
+  bool ContainsStreamingOp(mlir::func::FuncOp func,
+                           mlir::SymbolTableCollection& symbol_table_collection,
+                           llvm::DenseSet<mlir::func::FuncOp>& visited) {
+    // Track visited functions to safely handle recursive call graphs and
+    // prevent infinite loops.
+    if (!func || !visited.insert(func).second) return false;
+
+    VLOG(1) << "Checking function for streaming ops: " << func.getName().str();
+
+    // Walk all operations within the current function's body.
+    auto walk_result = func.walk([&](mlir::Operation* op) {
+      // Base case - If we find the streaming op directly, interrupt the
+      // walk and signal discovery.
+      if (mlir::isa<mlir::TF::PwStreamResultsOp>(op)) {
+        VLOG(1) << "Found tf.PwStreamResults in " << func.getName().str();
+        return mlir::WalkResult::interrupt();
+      }
+
+      // Recursive case - Many ops (like tf.If, tf.While, tf.PartitionedCall)
+      // don't hold operations directly but instead reference callee functions
+      // via attributes. We need to inspect these attributes to traverse the
+      // call graph.
+      for (const mlir::NamedAttribute& attr : op->getAttrs()) {
+        if (auto sym_ref =
+                mlir::dyn_cast<mlir::SymbolRefAttr>(attr.getValue())) {
+          // Handle single function references (e.g., tf.PartitionedCall).
+          // Resolve the symbol to the actual function and recursively check it.
+          VLOG(1) << "Looking up symbol: " << sym_ref.getLeafReference().str()
+                  << " from op: " << op->getName().getStringRef().str();
+          if (auto callee = symbol_table_collection
+                                .lookupNearestSymbolFrom<mlir::func::FuncOp>(
+                                    op, sym_ref)) {
+            if (ContainsStreamingOp(callee, symbol_table_collection, visited)) {
+              return mlir::WalkResult::interrupt();
+            }
+          }
+        } else if (auto arr_attr =
+                       mlir::dyn_cast<mlir::ArrayAttr>(attr.getValue())) {
+          // Handle arrays of function references (e.g., branches in tf.If).
+          // We must check every referenced function in the array.
+          for (auto elt : arr_attr.getValue()) {
+            if (auto sym_ref = mlir::dyn_cast<mlir::SymbolRefAttr>(elt)) {
+              VLOG(1) << "Looking up symbol from array: "
+                      << sym_ref.getLeafReference().str()
+                      << " from op: " << op->getName().getStringRef().str();
+              if (auto callee =
+                      symbol_table_collection
+                          .lookupNearestSymbolFrom<mlir::func::FuncOp>(
+                              op, sym_ref)) {
+                if (ContainsStreamingOp(callee, symbol_table_collection,
+                                        visited)) {
+                  return mlir::WalkResult::interrupt();
+                }
+              }
+            }
+          }
+        }
+      }
+      return mlir::WalkResult::advance();
+    });
+
+    // If the walk was interrupted, it means we found the streaming op
+    // somewhere in this function or its callees.
+    return walk_result.wasInterrupted();
+  }
+
+  void Rewrite(mlir::SymbolTable& symbol_table,
+               mlir::SymbolTableCollection& symbol_table_collection,
+               llvm::DenseMap<mlir::func::FuncOp, mlir::func::FuncOp>&
+                   cluster_to_ifrt_program,
+               mlir::tf_device::ClusterFuncOp cluster_func,
+               mlir::TF::RuntimeDevices& devices) {
+    mlir::OpBuilder builder(cluster_func);
+    mlir::FlatSymbolRefAttr callee_symbol = cluster_func.getFuncAttr();
+    mlir::func::FuncOp callee_func =
+        symbol_table.lookup<mlir::func::FuncOp>(callee_symbol.getValue());
+
+    // TODO - b/500390771: Ideally we need to make the await op the preceding
+    // op of the PwStreamResultsOp, which can be achieved in the tf-to-mlrt
+    // lowering pass.
+    bool use_async = enable_async_ifrt;
+    if (use_async) {
+      llvm::DenseSet<mlir::func::FuncOp> visited;
+      if (ContainsStreamingOp(callee_func, symbol_table_collection, visited)) {
+        cluster_func->emitRemark(
+            "Disabling AsyncIfrtCallOp because tf.PwStreamResults is present "
+            "in the cluster.");
+        use_async = false;
+      }
+    }
+
+    auto ifrt_program_name =
+        absl::StrCat("_ifrt_program_", callee_func.getSymName().str());
+    if (mlir::func::FuncOp ifrt_program =
+            cluster_to_ifrt_program[callee_func]) {
+      // ifrt program already exists
+      builder.setInsertionPoint(cluster_func);
+
+      mlir::Operation* ifrt_call_op;
+      if (use_async) {
+        ifrt_call_op = mlir::TF::AsyncIfrtCallOp::create(
+            builder, cluster_func->getLoc(), cluster_func.getResultTypes(),
+            cluster_func->getOperands());
+      } else {
+        ifrt_call_op = mlir::TF::IfrtCallOp::create(
+            builder, cluster_func->getLoc(), cluster_func.getResultTypes(),
+            cluster_func->getOperands());
+      }
+
+      ifrt_call_op->setAttr(
+          "operandSegmentSizes",
+          builder.getDenseI32ArrayAttr(
+              {static_cast<int32_t>(cluster_func.getNumOperands()), 0}));
+
+      int64_t program_id;
+      if (auto attr = ifrt_program->getAttrOfType<mlir::IntegerAttr>(
+              "tfrt_ifrt_serving.program_id")) {
+        program_id = attr.getInt();
+      } else {
+        return signalPassFailure();
+      }
+
+      auto metadata_attr =
+          ifrt_program->getAttrOfType<mlir::StringAttr>(kMetadataTextAttrName);
+      auto device_assignment_attr =
+          ifrt_program->getAttrOfType<mlir::ArrayAttr>(kDeviceAssignmentAttr);
+      if (!metadata_attr || !device_assignment_attr) {
+        return signalPassFailure();
+      }
+
+      // For better debuggability, attach attributes such as
+      // tpu_compile_metadata and device_assignment to IfrtCallOp.
+      ifrt_call_op->setAttr(kMetadataTextAttrName, metadata_attr);
+      ifrt_call_op->setAttr(kDeviceAssignmentAttr, device_assignment_attr);
+
+      // TODO(b/304839793): populate variable names after adding a variable
+      // hoisting pass.
+      ifrt_call_op->setAttr("variable_arg_indices",
+                            builder.getI32ArrayAttr({}));
+      ifrt_call_op->setAttr("program_id",
+                            builder.getI64IntegerAttr(program_id));
+
+      cluster_func->replaceAllUsesWith(ifrt_call_op->getResults());
+      cluster_func->erase();
+
+      return;
+    }
+    mlir::OpBuilder::InsertionGuard insertion_guard(builder);
+    builder.setInsertionPoint(callee_func);
+
+    mlir::func::FuncOp cloned_ifrt_program = mlir::func::FuncOp::create(
+        builder, callee_func->getLoc(), ifrt_program_name,
+        callee_func.getFunctionType());
+    mlir::IRMapping mapper;
+    callee_func.cloneInto(cloned_ifrt_program, mapper);
+
+    tensorflow::tpu::TPUCompileMetadataProto metadata;
+    if (mlir::failed(GetTpuCompileMetadata(cluster_func, devices, &metadata))) {
+      return signalPassFailure();
+    }
+    std::string serialized_metadata;
+    tsl::protobuf::TextFormat::Printer printer;
+    printer.SetSingleLineMode(true);
+    printer.PrintToString(metadata, &serialized_metadata);
+
+    cloned_ifrt_program->setAttr(kMetadataTextAttrName,
+                                 builder.getStringAttr(serialized_metadata));
+
+    auto device_assignment_attr =
+        cluster_func->getAttrOfType<mlir::ArrayAttr>(kDeviceAssignmentAttr);
+    if (!device_assignment_attr) {
+      device_assignment_attr = builder.getArrayAttr({});
+    }
+    cloned_ifrt_program->setAttr(kDeviceAssignmentAttr, device_assignment_attr);
+
+    cloned_ifrt_program.setName(ifrt_program_name);
+
+    int64_t program_id = NewProgramId();
+    cloned_ifrt_program->setAttr("tfrt_ifrt_serving.program_id",
+                                 builder.getI64IntegerAttr(program_id));
+
+    // Make cloned ifrt program public so that it does not get dropped by
+    // inliner.
+    cloned_ifrt_program.setPublic();
+
+    builder.setInsertionPoint(cluster_func);
+
+    mlir::Operation* ifrt_call_op;
+    if (use_async) {
+      ifrt_call_op = mlir::TF::AsyncIfrtCallOp::create(
+          builder, cluster_func->getLoc(), cluster_func.getResultTypes(),
+          cluster_func->getOperands());
+    } else {
+      ifrt_call_op = mlir::TF::IfrtCallOp::create(
+          builder, cluster_func->getLoc(), cluster_func.getResultTypes(),
+          cluster_func->getOperands());
+    }
+
+    ifrt_call_op->setAttr(
+        "operandSegmentSizes",
+        builder.getDenseI32ArrayAttr(
+            {static_cast<int32_t>(cluster_func.getNumOperands()), 0}));
+
+    // TODO(b/304839793): populate variable names after adding a variable
+    // hoisting pass.
+    ifrt_call_op->setAttr("variable_arg_indices", builder.getI32ArrayAttr({}));
+    ifrt_call_op->setAttr("program_id", builder.getI64IntegerAttr(program_id));
+    // For better debuggability, attach attributes such as tpu_compile_metadata
+    // and device_assignment to IfrtCallOp.
+    ifrt_call_op->setAttr(kMetadataTextAttrName,
+                          builder.getStringAttr(serialized_metadata));
+    ifrt_call_op->setAttr(kDeviceAssignmentAttr, device_assignment_attr);
+
+    cluster_func->replaceAllUsesWith(ifrt_call_op->getResults());
+    cluster_func->erase();
+
+    symbol_table.insert(cloned_ifrt_program);
+    cluster_to_ifrt_program[callee_func] = cloned_ifrt_program;
+  }
+};
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateRewriteClusterToIfrtCallPass(bool enable_async_ifrt) {
+  return std::make_unique<RewriteClusterToIfrtCallPass>(enable_async_ifrt);
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/sink_variable_as_named_array.cc

@@ -0,0 +1,255 @@
+/* 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 <memory>
+#include <numeric>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#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/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_constants.h"
+#include "xla/service/computation_placer.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/platform/protobuf.h"  // IWYU pragma: keep
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_config.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+#define GEN_PASS_DEF_SINKVARIABLEASNAMEDARRAYPASS
+#define GEN_PASS_DECL_SINKVARIABLEASNAMEDARRAYPASS
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+class SinkVariableAsNamedArrayPass
+    : public impl::SinkVariableAsNamedArrayPassBase<
+          SinkVariableAsNamedArrayPass> {
+ public:
+  void runOnOperation() override {
+    mlir::ModuleOp module = getOperation();
+    mlir::OpBuilder builder(&getContext());
+
+    // Rewrite ReadVariableOp with IfrtLoadVariableOp
+    llvm::SmallDenseMap<mlir::TF::ReadVariableOp, mlir::TF::IfrtLoadVariableOp>
+        read_to_load;
+
+    mlir::WalkResult walk_result =
+        module.walk([&](mlir::TF::ReadVariableOp read_variable_op) {
+          // TODO(b/319045348): consider use resource alias analysis for
+          // this.
+          auto var_handle = GetDefiningOp<mlir::TF::VarHandleOp>(
+              read_variable_op.getResource());
+
+          if (!var_handle) {
+            read_variable_op->emitError(
+                "ReadVariableOp has no defining VarHandleOp.");
+            return mlir::WalkResult::interrupt();
+          }
+
+          // Avoid lowering ReadVariableOp to IfrtLoadVariableOp if the
+          // assignment AssignVariableOp happens at the same module because
+          // IfrtLoadVariableOp assumes asynchronous assignment of the variable.
+          for (auto var_handle_user : var_handle->getUsers()) {
+            if (llvm::isa<mlir::TF::AssignVariableOp>(var_handle_user)) {
+              return mlir::WalkResult::advance();
+            }
+          }
+          std::vector<mlir::Type> result_types;
+          result_types.push_back(mlir::RankedTensorType::get(
+              {}, builder.getType<mlir::TF::StringType>()));
+          result_types.push_back(read_variable_op.getResult().getType());
+
+          builder.setInsertionPointAfter(read_variable_op);
+          auto load_variable_op = mlir::TF::IfrtLoadVariableOp::create(
+              builder, read_variable_op->getLoc(), result_types,
+              var_handle.getResult());
+          read_to_load[read_variable_op] = load_variable_op;
+
+          return mlir::WalkResult::advance();
+        });
+
+    if (walk_result.wasInterrupted()) {
+      return signalPassFailure();
+    }
+
+    // Rewrite ifrt call: variable tensors are sunk as attribute.
+    // The runtime guarantees the binding of corresponding loaded ifrt array
+    // based on attributes.
+    auto process_call = [&](auto call) -> mlir::WalkResult {
+      IfrtArgConfigList ifrt_call_argument_configs;
+
+      if (mlir::failed(BuildIfrtCallArgumentConfig(
+              call.getOperation(), ifrt_call_argument_configs))) {
+        return mlir::WalkResult::interrupt();
+      }
+
+      if (!call.getVariableArgIndicesAttr().empty()) {
+        call->emitError() << "Expect empty "
+                          << call.getVariableArgIndicesAttrName().str()
+                          << " attributes, but got "
+                          << call.getVariableArgIndicesAttr().size()
+                          << " elements";
+        return mlir::WalkResult::interrupt();
+      }
+      if (call->getOpOperands().size() != ifrt_call_argument_configs.size()) {
+        call->emitError() << "IfrtCallOp got " << call->getOpOperands().size()
+                          << " operands, but expects "
+                          << ifrt_call_argument_configs.size();
+        return mlir::WalkResult::interrupt();
+      }
+
+      llvm::SmallVector<int> variable_arg_indices;
+      llvm::SmallVector<mlir::Value> updated_args;
+
+      for (const auto& [arg_idx, arg] :
+           llvm::enumerate(ifrt_call_argument_configs)) {
+        if (arg.is_variable) {
+          variable_arg_indices.push_back(arg_idx);
+          // Variable use the key from IfrtLoadVariable.
+          updated_args.push_back(
+              read_to_load[arg.read_variable_op].getArrayKey());
+        } else {
+          // non variable
+          updated_args.push_back(call->getOperand(arg_idx));
+        }
+      }
+
+      builder.setInsertionPointAfter(call);
+      mlir::Operation* updated_ifrt_call;
+      if (llvm::isa<mlir::TF::AsyncIfrtCallOp>(call)) {
+        updated_ifrt_call = mlir::TF::AsyncIfrtCallOp::create(
+            builder, call->getLoc(), call.getResultTypes(), updated_args);
+      } else {
+        updated_ifrt_call = mlir::TF::IfrtCallOp::create(
+            builder, call->getLoc(), call.getResultTypes(), updated_args);
+      }
+
+      updated_ifrt_call->setAttrs(call->getAttrs());
+      // Update variable_arg_indices attribute.
+      updated_ifrt_call->setAttr("variable_arg_indices",
+                                 builder.getI32ArrayAttr(variable_arg_indices));
+
+      call.replaceAllUsesWith(updated_ifrt_call);
+      call.erase();
+      return mlir::WalkResult::advance();
+    };
+
+    if (module
+            .walk([&](mlir::TF::IfrtCallOp call) { return process_call(call); })
+            .wasInterrupted()) {
+      return signalPassFailure();
+    }
+
+    if (module
+            .walk([&](mlir::TF::AsyncIfrtCallOp call) {
+              return process_call(call);
+            })
+            .wasInterrupted()) {
+      return signalPassFailure();
+    }
+
+    // Remove all ReadVariableOp after replacing the CPU usage of
+    // ReadVariableOp.
+    for (auto& [read_variable_op, load_variable_op] : read_to_load) {
+      if (!read_variable_op->use_empty()) {
+        // This variable tensor is used by CPU host.
+        load_variable_op.setUsedByHost(true);
+
+        // Replace CPU use of ReadVariableOp
+        read_variable_op.replaceAllUsesWith(load_variable_op.getTensorFuture());
+      }
+      read_variable_op.erase();
+    }
+  }
+
+ private:
+  struct IfrtArgConfig {
+    bool is_variable;
+    mlir::TF::ReadVariableOp read_variable_op;
+  };
+  using IfrtArgConfigList = llvm::SmallVector<IfrtArgConfig>;
+
+  // Build argument configuration map of a IfrtCallOp.
+  mlir::LogicalResult BuildIfrtCallArgumentConfig(mlir::Operation* call,
+                                                  IfrtArgConfigList& args) {
+    for (const auto& [arg_idx, input] : llvm::enumerate(call->getOperands())) {
+      // Assuming the nested function calls are inlined.
+      if (auto read_variable_op =
+              GetDefiningOp<mlir::TF::ReadVariableOp>(input)) {
+        args.push_back(
+            {.is_variable = true, .read_variable_op = read_variable_op});
+      } else {
+        args.push_back({.is_variable = false});
+      }
+    }
+
+    return mlir::success();
+  }
+
+  template <typename OpT>
+  OpT GetDefiningOp(const mlir::Value& value) {
+    mlir::Operation* op = value.getDefiningOp();
+
+    while (op && !llvm::isa<OpT>(op)) {
+      if (llvm::isa<mlir::TF::IdentityOp>(op)) {
+        op = op->getOperand(0).getDefiningOp();
+      } else {
+        return nullptr;
+      }
+    }
+
+    if (op != nullptr) {
+      return llvm::dyn_cast<OpT>(op);
+    } else {
+      return nullptr;
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateSinkVariableAsNamedArrayPass() {
+  return std::make_unique<SinkVariableAsNamedArrayPass>();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/BUILD

@@ -0,0 +1,15 @@
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//learning/brain/tfrt/ifrt:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+filegroup(
+    name = "testdata",
+    srcs = glob(
+        ["*"],
+    ),
+)

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/ifrt_cluster.mlir

@@ -0,0 +1,24 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:TPU:0"]} {
+  func.func @main() {
+    "tf_device.cluster_func"() {_replication_info = "cluster0", func = @empty_func, num_cores_per_replica = 1, step_marker_location = "", input_sharding_configuration = [], output_sharding_configuration = [], use_spmd_for_xla_partitioning = false} : () -> ()
+    func.return
+  }
+  func.func @empty_func() {
+    func.return
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/restore_with_reference.mlir

@@ -0,0 +1,29 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}, tf_saved_model.semantics} {
+  "tf_saved_model.session_initializer"() <{initializers = [@"save/restore_all_1"]}> : () -> ()
+  func.func @"save/restore_all_1"() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_save/restore_all_1"], tf_saved_model.initializer_type = "restore_op"} {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<["", ""]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<["y", "z"]> : tensor<2x!tf_type.string>}> : () -> tensor<2x!tf_type.string>
+    %0:2 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<i64>, tensor<1x3xf32>)
+    %1 = "tf.VariableV2"() <{container = "x", shape = #tf_type.shape<>, shared_name = "y"}> : () -> tensor<!tf_type.int64ref>
+    %dummy = "tf.Assign"(%1, %0#0) : (tensor<!tf_type.int64ref>, tensor<i64>) ->  tensor<!tf_type.int64ref>
+    %2 = "tf.VarHandleOp"() <{container = "x", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+    "tf.AssignVariableOp"(%2, %0#1) : (tensor<!tf_type.resource<tensor<1x3xf32>>>, tensor<1x3xf32>) -> ()
+    return
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/tf2hlo_empty.mlir

@@ -0,0 +1,20 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+  func.func @main() -> () attributes {__tpu_compile_metadata_text = "num_replicas: 1 num_cores_per_replica: 1"} {
+    func.return
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/tf2hlo_gpu.mlir

@@ -0,0 +1,40 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+// MLIR which is legalize only with the right device type.
+// The mlir is generated by running
+// tensorflow/compiler/mlir/tf-opt -tf-xla-call-module-serialization
+//
+// module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32,
+// producer = 268 : i32}} {
+//   func.func private @_jit_sin(%arg0: tensor<f32>) -> tensor<f32> {
+//       %0 = stablehlo.sine %arg0 : tensor<f32>
+//       return %0 : tensor<f32>
+//     }
+//   func.func @main(%arg0: tensor<f32>) -> tensor<*xf32> {
+//       %0 = "tf.XlaCallModule"(%arg0) {Sout = [#tf_type.shape<*>], device =
+//            "", dim_args_spec = [], _entry_function = @_jit_sin, module = "",
+//            platforms = ["CUDA"], version = 6 : i64} : (tensor<f32>) ->
+//            tensor<*xf32>
+//       func.return %0 : tensor<*xf32>
+//    }
+// }
+//
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+  func.func @main(%arg0: tensor<f32>) -> tensor<*xf32> attributes {__tpu_compile_metadata_text = "args { dtype: DT_FLOAT kind: PARAMETER } retvals { } num_replicas: 1 num_cores_per_replica: 1"} {
+   %0 = "tf.XlaCallModule"(%arg0) {Sout = [#tf_type.shape<*>], device = "", dim_args_spec = [], module = "ML\EFR\03MLIRxxx-trunk\00\01\17\05\01\05\01\03\05\03\07\07\t\0B\03K5\07\01\1B\07\0B\13\0B3\0B\0B\0B\0B\0F\0B\13\0B\03\1B\0F\1B\0B\0B\0B\0B\0B\0F\13\0B\0B\0B\0B\03\07\0F\17\07\02\A7\1F\05\0D\03\03\03\07\05\0F\03\0B\0B\1B\0D'\0F)\031\113\05\11\05\13\05\15\05\17\1D\15\17\05\19\17\19\EF\01\05\1B\03\03\1D\0D\05\1F!#%\1D\1D\1D\1F\1D!\1D##\03\03\03+\0D\03-/\1D%\1D'\1D)\1D+)\01\05\11\03\01\03\01\t\04A\05\01\11\01\05\07\03\01\05\03\11\01\t\05\03\05\0B\03\01\01\05\06\13\03\01\03\01\07\04\01\03\03\06\03\01\05\01\00\9A\04-\0F\0B\03!\1B\1D\05\1B\83/\1F\15\1D\15\11\13\15\11\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00sine_v1\00return_v1\00sym_name\00jit_sin\00arg_attrs\00function_type\00res_attrs\00sym_visibility\00jit(sin)/jit(main)/sin\00third_party/py/jax/experimental/jax2tf/tests/back_compat_test.py\00jax.arg_info\00x\00mhlo.sharding\00{replicated}\00jax.result_info\00\00main\00public\00", platforms = ["CUDA"], version = 6 : i64} : (tensor<f32>) -> tensor<*xf32>
+    func.return %0 : tensor<*xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/tf2hlo_spmd_no_device_assignment.mlir

@@ -0,0 +1,27 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+func.func @main(%arg0: tensor<4x64xf32> {mhlo.sharding = "{devices=[2,1]0,1}"}, %arg1: tensor<64x10xf32> {mhlo.sharding = "{devices=[2,1]0,1}"}, %arg2: tensor<*xf32> {mhlo.sharding = ""}) -> (tensor<?x10xf32> {mhlo.sharding = ""}) attributes {tfrt_ifrt_serving.program_id = -5304490761103885474 : i64, __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { unknown_rank: true } kind: PARAMETER sharding { type: OTHER tile_assignment_dimensions: 2 tile_assignment_dimensions: 1 tile_assignment_devices: 0 tile_assignment_devices: 1 } is_bounded_dynamic_dim: false } args { dtype: DT_FLOAT shape { unknown_rank: true } kind: PARAMETER sharding { type: OTHER tile_assignment_dimensions: 2 tile_assignment_dimensions: 1 tile_assignment_devices: 0 tile_assignment_devices: 1 } is_bounded_dynamic_dim: false } args { dtype: DT_FLOAT shape { unknown_rank: true } kind: PARAMETER is_bounded_dynamic_dim: false } retvals { sharding { } } num_replicas: 1 num_cores_per_replica: 2 use_spmd_for_xla_partitioning: true compile_options { }"} {
+    %cst = "tf.Const"() <{value = dense<[-1, 4]> : tensor<2xi32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<2xi32>
+    %0 = "tf.XlaSharding"(%arg0) <{_XlaSharding = "{devices=[2,1]0,1}", sharding = "{devices=[2,1]0,1}"}> : (tensor<4x64xf32>) -> tensor<4x64xf32>
+    %1 = "tf.XlaSharding"(%arg1) <{_XlaSharding = "{devices=[2,1]0,1}", sharding = "{devices=[2,1]0,1}"}> : (tensor<64x10xf32>) -> tensor<64x10xf32>
+    %2 = "tf.Reshape"(%arg2, %cst) : (tensor<*xf32>, tensor<2xi32>) -> tensor<?x4xf32>
+    %3 = "tf.MatMul"(%2, %0) <{transpose_a = false, transpose_b = false}> : (tensor<?x4xf32>, tensor<4x64xf32>) -> tensor<?x64xf32>
+    %4 = "tf.Relu"(%3) : (tensor<?x64xf32>) -> tensor<?x64xf32>
+    %5 = "tf.MatMul"(%4, %1) <{transpose_a = false, transpose_b = false}> : (tensor<?x64xf32>, tensor<64x10xf32>) -> tensor<?x10xf32>
+    return %5 : tensor<?x10xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/tf2hlo_spmd_shardy.mlir

@@ -0,0 +1,21 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+func.func @main(%arg0: tensor<4x64xf32> {mhlo.sharding = "{devices=[2,1]<=[2]}"}) -> (tensor<4x64xf32> {mhlo.sharding = ""}) attributes {tfrt_ifrt_serving.program_id = -5304490761103885474 : i64, __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { dim { size: 4 } dim { size: 64 } } kind: PARAMETER sharding { type: OTHER tile_assignment_dimensions: 2 tile_assignment_dimensions: 1 iota_reshape_dims: 2 iota_transpose_perm: 0 } is_bounded_dynamic_dim: false } retvals { sharding { } } num_replicas: 1 num_cores_per_replica: 2 device_assignment { replica_count: 1 computation_count: 2 computation_devices { replica_device_ids: 0 } computation_devices { replica_device_ids: 1 } } use_spmd_for_xla_partitioning: true use_shardy_partitioner: true compile_options { }"} {
+    %0 = "tf.Relu"(%arg0) : (tensor<4x64xf32>) -> tensor<4x64xf32>
+    return %0 : tensor<4x64xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/tf2hlo_spmd_with_device_assignment.mlir

@@ -0,0 +1,21 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+func.func @main(%arg0: tensor<4x64xf32> {mhlo.sharding = "{devices=[2,1]0,1}"}) -> (tensor<4x64xf32> {mhlo.sharding = ""}) attributes {tfrt_ifrt_serving.program_id = -5304490761103885474 : i64, __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { dim { size: 4 } dim { size: 64 } } kind: PARAMETER sharding { type: OTHER tile_assignment_dimensions: 2 tile_assignment_dimensions: 1 tile_assignment_devices: 0 tile_assignment_devices: 1 } is_bounded_dynamic_dim: false } retvals { sharding { } } num_replicas: 1 num_cores_per_replica: 2 device_assignment { replica_count: 1 computation_count: 2 computation_devices { replica_device_ids: 0 } computation_devices { replica_device_ids: 1 } } use_spmd_for_xla_partitioning: true compile_options { }"} {
+    %0 = "tf.Relu"(%arg0) : (tensor<4x64xf32>) -> tensor<4x64xf32>
+    return %0 : tensor<4x64xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/tf2hlo_tuple.mlir

@@ -0,0 +1,21 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
+  func.func @main(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> (tensor<1x1xf32>, tensor<1x3xf32>) attributes {__tpu_compile_metadata_text = "args { dtype: DT_FLOAT kind: PARAMETER } args { dtype: DT_FLOAT kind: PARAMETER } retvals { } retvals { } num_replicas: 1 num_cores_per_replica: 1"} {
+    %0 = "tf.MatMul"(%arg0, %arg1): (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+    func.return %0, %arg0: tensor<1x1xf32>, tensor<1x3xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata/xla_call_host_callback.mlir

@@ -0,0 +1,38 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1758 : i32}} {
+
+  func.func private @callee(%arg0: tensor<?xi32>, %arg1: tensor<*xi32>) {
+    "tf.XlaHostCompute"(%arg0, %arg1) <{ancestors = [], key = "@test_callee", recv_key = "", send_key = "", shapes = []}> {_xla_original_oc_node_name = "hcb0", _xla_token_input_nodes = ["_xla_token_arg_node"]} : (tensor<?xi32>, tensor<*xi32>) -> ()
+    return
+  }
+
+  // The mlir module in XlaCallModule is serialized from:
+  //
+  // func.func private @_stablehlo_main_0(%arg0: tensor<?xi32>, %arg1: tensor<*xi32>) -> () attributes {_from_xla_call_module} {
+  //   stablehlo.custom_call @tf.call_tf_function(%arg0, %arg1) {api_version = 2 : i32, has_side_effect = true, tf.backend_config = {called_func = @callee}} : (tensor<?xi32>, tensor<*xi32>) -> ()
+  //   return
+  // }
+  //
+  // func.func @main(%arg0: tensor<?xi32>, %arg1: tensor<*xi32>) -> () {
+  //   "tf.XlaCallModule"(%arg0, %arg1) {Sout = [#tf_type.shape<?>], dim_args_spec = [], _entry_function = @_stablehlo_main_0, _stablehlo_version = "0.17.6", _stablehlo_module_attrs = { mhlo.num_partitions = 1 }, module = "", platforms = [], version = 5 : i64} : (tensor<?xi32>, tensor<*xi32>) -> ()
+  //   func.return
+  // }
+  func.func @main(%arg0: tensor<?xi32>, %arg1: tensor<*xi32>) attributes {tfrt_ifrt_serving.program_id = -2372940092539171444 : i64, __tpu_compile_metadata_text = "args { dtype: DT_INT32 kind: PARAMETER sharding { } }  args { dtype: DT_INT32 kind: PARAMETER sharding { } } num_replicas: 1 num_cores_per_replica: 1 use_spmd_for_xla_partitioning: true compile_options { }"} {
+    "tf.XlaCallModule"(%arg0, %arg1) <{Sout = [], dim_args_spec = [], function_list = [@callee], module = "ML\EFR\0DStableHLO_v0.17.6\00\01\19\05\01\05\09\01\03\0B\03\07\0F\13\17\03M-\0D\01\19\0B\13\0B\0F\13\13\13\13\13\0B\13\13\03\15\0B\0B\0B\0B\13\0B\0F\0B\0B\0B\01\03\0F\03\0B3\07\0B\17\07\02\B1\05\0D\03\03\05\07\05\0F\11\01\05\17\01A\0B\17\01!\07\17\01!Q\17\01!}\03\03\13!\05\11\17\01#\0B\17\01%\0B\03\01\1D\13#\09\1D\15\0D\03#%\1D\17\13\0B\01\0B\05\1D\19\05\03\01\02\04)\03\00\FF\FF\FF\FF\FF\FF\FF\FF\05\1B3\05\11\05\03\07\01\1D\04O\05\01Q\09\03\01\07\04=\03\01\05\03P\0B\03\07\04)\03\05\0B\05\07\0D\0F\0F\00\05E\15\11\05\05\01\03\07\00\17\06\03\01\05\01\00j\03\1B)\1B\0B\03%)\95\15\1F\11\0F\0B\11builtin\00vhlo\00module\00func_v1\00custom_call_v1\00return_v1\00experimental/users/deqiangc/mira/testdata/xla_call_module_serialized.mlir\00mhlo.num_partitions\00tf.backend_config\00\00main\00called_index\00tf.call_tf_function\00\08'\07\05\01\01\0B\19\1D\19\1F\1B\11'\1B)\19+\19\19\19", platforms = [], version = 5 : i64}> : (tensor<?xi32>, tensor<*xi32>) -> ()
+    return
+  }
+}

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.cc

@@ -0,0 +1,323 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/raw_ostream.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/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v1/compile_mlir_util.h"
+#include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_compilation.pb.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_constants.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "tensorflow/compiler/tf2xla/layout_util.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/client/client_library.h"
+#include "xla/hlo/translate/hlo_to_mhlo/hlo_to_mlir_hlo.h"
+#include "xla/pjrt/pjrt_compiler.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/layout.h"
+#include "xla/service/computation_placer.h"
+#include "xla/shape.h"
+#include "xla/stream_executor/platform_manager.h"
+#include "xla/tsl/lib/strings/proto_serialization.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tensorflow/core/protobuf/tpu/topology.pb.h"
+#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
+#include "tsl/platform/fingerprint.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+static constexpr absl::string_view kEntryFuncName = "main";
+uint64_t MlirModuleFingerprint(mlir::ModuleOp module) {
+  std::string s;
+  llvm::raw_string_ostream os(s);
+  mlir::OpPrintingFlags flags;
+  flags.enableDebugInfo(false);
+  module.print(os, flags);
+  return tsl::Fingerprint64(os.str());
+}
+}  // namespace
+
+absl::StatusOr<uint64_t> Tf2HloArg::Fingerprint() const {
+  uint64_t fingerprint = tsl::Fingerprint64(platform_name);
+  if (topology) {
+    TF_ASSIGN_OR_RETURN(fingerprint, topology->Fingerprint());
+  }
+  if (platform_name != xla::CudaName() && !topology) {
+    return absl::FailedPreconditionError(
+        "Topology is required for non-GPU compilation.");
+  }
+  fingerprint =
+      tsl::FingerprintCat64(fingerprint, MlirModuleFingerprint(module));
+  for (const auto& dtype_and_shape : input_dtypes_and_shapes) {
+    fingerprint = tsl::FingerprintCat64(
+        fingerprint,
+        tsl::Fingerprint64(tensorflow::DataType_Name(dtype_and_shape.dtype)));
+
+    std::string serialized_shape;
+    if (!tsl::SerializeToStringDeterministic(
+            dtype_and_shape.GetShapeForCompilation().AsProto(),
+            &serialized_shape)) {
+      return absl::InternalError("Failed to serialize shape");
+    }
+
+    fingerprint = tsl::FingerprintCat64(fingerprint,
+                                        tsl::Fingerprint64(serialized_shape));
+  }
+  fingerprint = tsl::FingerprintCat64(fingerprint,
+                                      tsl::Fingerprint64(entry_function_name));
+  std::string serialized_compile_metadata;
+  if (!tsl::SerializeToStringDeterministic(compile_metadata,
+                                           &serialized_compile_metadata)) {
+    return absl::InternalError("Failed to serialize compile metadata");
+  }
+  fingerprint = tsl::FingerprintCat64(
+      fingerprint, tsl::Fingerprint64(serialized_compile_metadata));
+  fingerprint = tsl::FingerprintCat64(
+      fingerprint,
+      tsl::Fingerprint64(populate_layout_in_xla_input_shapes ? "1" : "0"));
+  return fingerprint;
+}
+
+absl::StatusOr<Tf2HLOResultProto> Tf2HloResult::ToProto() const {
+  Tf2HLOResultProto proto;
+  *proto.mutable_hlo_module_proto() = hlo_module_proto;
+  *proto.mutable_compile_metadata() = compile_metadata;
+  *proto.mutable_host_compute_metadata() = host_compute_metadata;
+  for (const auto& shape : xla_input_shapes) {
+    *proto.add_xla_input_shapes() = shape.ToProto();
+  }
+  return proto;
+}
+
+/*static*/ absl::StatusOr<Tf2HloResult> Tf2HloResult::FromProto(
+    const Tf2HLOResultProto& proto) {
+  Tf2HloResult result;
+  result.hlo_module_proto = proto.hlo_module_proto();
+  result.compile_metadata = proto.compile_metadata();
+  result.host_compute_metadata = proto.host_compute_metadata();
+  for (const auto& shape : proto.xla_input_shapes()) {
+    TF_ASSIGN_OR_RETURN(xla::Shape xla_shape, xla::Shape::FromProto(shape));
+    result.xla_input_shapes.push_back(std::move(xla_shape));
+  }
+  return result;
+}
+
+absl::Status UpdateCompileMetadata(
+    tensorflow::tpu::TPUCompileMetadataProto& metadata,
+    absl::Span<const DtypeAndShape> inputs) {
+  VLOG(3) << "TpuCompileMetadata before shape is populated " << metadata;
+  if (metadata.num_replicas() < 1 || metadata.num_cores_per_replica() < 1) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Number of replicas ", metadata.num_replicas(),
+                     " and number of cores per replica ",
+                     metadata.num_cores_per_replica(), " must be >= 1"));
+  }
+  if (metadata.args_size() != inputs.size()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Number of inputs mismatched! Expected ",
+                     metadata.args_size(), " got ", inputs.size()));
+  }
+
+  for (int i = 0; i < metadata.args_size(); ++i) {
+    if (metadata.args(i).kind() !=
+        tensorflow::tpu::TPUCompileMetadataProto::Arg::PARAMETER) {
+      return absl::FailedPreconditionError(absl::StrCat(
+          "Only support PARAMETER, but got ", metadata.args(i).kind()));
+    }
+
+    if (metadata.args(i).dtype() != inputs[i].dtype) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Dtype mismatched! Expected ", metadata.args(i).dtype(),
+                       " got ", inputs[i].dtype));
+    }
+
+    // Update shape.
+    *metadata.mutable_args(i)->mutable_shape() =
+        inputs[i].GetShapeForCompilation().AsProto();
+  }
+  return absl::OkStatus();
+}
+
+absl::StatusOr<tensorflow::tpu::TPUCompileMetadataProto> GetCompileMetadata(
+    mlir::ModuleOp module, const xla::ifrt::Client& ifrt_client) {
+  tensorflow::tpu::TPUCompileMetadataProto metadata;
+
+  auto op = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
+  if (!op) {
+    return absl::InternalError("Could not find entry function in MLIR Module.");
+  }
+
+  auto metadata_text_attr =
+      op->getAttrOfType<mlir::StringAttr>(kMetadataTextAttrName);
+
+  if (metadata_text_attr && !metadata_text_attr.getValue().empty()) {
+    // Try __tpu_compile_metadata_text attribute. This only for debugging
+    // purpose.
+    VLOG(1) << "Parsing from attribute " << kMetadataTextAttrName
+            << metadata_text_attr.getValue().str();
+    if (!tsl::protobuf::TextFormat::ParseFromString(
+            metadata_text_attr.getValue().str(), &metadata)) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Attribute ", kMetadataTextAttrName, ":",
+          metadata_text_attr.getValue().str(), " cannot be parsed"));
+    }
+  } else {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Missing ", kMetadataTextAttrName));
+  }
+
+  // Create a default device assignment if one is not given by the model.
+  if (!metadata.has_device_assignment()) {
+    TF_ASSIGN_OR_RETURN(
+        auto device_assignment,
+        ifrt_client.GetDefaultDeviceAssignment(
+            metadata.num_replicas(), metadata.num_cores_per_replica()));
+
+    xla::DeviceAssignmentProto device_assignment_proto;
+    device_assignment.Serialize(&device_assignment_proto);
+
+    *metadata.mutable_device_assignment() = device_assignment_proto;
+  }
+
+  return metadata;
+}
+
+absl::StatusOr<Tf2HloResult> CompileTfToHlo(const Tf2HloArg& arg) {
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("ifrt_before_bridge_phase2", arg.module);
+  }
+
+  // Device_type is a string of
+  // tensorflow/compiler/mlir/tf2xla/api/v2/device_type.proto:DeviceType
+  std::string device_type = "XLA_TPU_JIT";
+  if (arg.platform_name == xla::CudaName()) {
+    device_type = "XLA_GPU_JIT";
+  }
+  VLOG(1) << "device_type: " << device_type;
+
+  TF_ASSIGN_OR_RETURN(
+      auto* platform,
+      stream_executor::PlatformManager::PlatformWithName("Host"));
+  TF_ASSIGN_OR_RETURN(
+      auto* client, xla::ClientLibrary::GetOrCreateCompileOnlyClient(platform));
+
+
+  std::vector<TensorShape> arg_shapes;
+  arg_shapes.reserve(arg.input_dtypes_and_shapes.size());
+  for (const auto& input : arg.input_dtypes_and_shapes) {
+    arg_shapes.push_back(input.GetShapeForCompilation());
+  }
+
+  llvm::SmallVector<tensorflow::TensorOrResourceShape>
+      arg_tensor_or_resource_shapes;
+  for (const auto& shape : arg_shapes) {
+    arg_tensor_or_resource_shapes.push_back({shape});
+  }
+
+  TF_RETURN_IF_ERROR(
+      tensorflow::RefineShapes(arg_tensor_or_resource_shapes, arg.module));
+  tpu::MlirToHloArgs mlir_to_hlo_args;
+  std::string module_str = tensorflow::SerializeMlirModule(arg.module);
+  mlir_to_hlo_args.mlir_module = module_str;
+  // Use fallback bridge as other modes may get deprecated.
+  mlir_to_hlo_args.rollout_state =
+      ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_DISABLED;
+
+  bool use_tuple_args = false;
+  std::vector<tpu::ShardingAndIndex> arg_core_mapping;
+  std::vector<std::vector<xla::Shape>> per_core_arg_shapes;
+  std::vector<std::unique_ptr<mlir::Pass>> custom_legalization_passes;
+
+  TF_ASSIGN_OR_RETURN(
+      tensorflow::XlaCompiler::CompilationResult compilation_result,
+      tensorflow::tf2xla::v2::LegalizeMlirToHlo(
+          mlir_to_hlo_args, arg.compile_metadata, use_tuple_args, device_type,
+          custom_legalization_passes,
+          /*shape_determination_fns=*/
+          tensorflow::XlaShapeLayoutHelpers::ShapeDeterminationFns(
+              tensorflow::UseNoPreferenceLayoutFn(),
+              arg.shape_representation_fn),
+          arg_shapes, &arg_core_mapping, &per_core_arg_shapes, client));
+
+  for (auto arg_shapes_iter = per_core_arg_shapes.begin() + 1;
+       arg_shapes_iter != per_core_arg_shapes.end(); ++arg_shapes_iter) {
+    if (per_core_arg_shapes.front() != *arg_shapes_iter) {
+      return absl::UnimplementedError(
+          "Only support even sharding SPMD, but get "
+          "different shapes across cores");
+    }
+  }
+
+  Tf2HloResult result;
+  if (arg.populate_layout_in_xla_input_shapes) {
+    result.xla_input_shapes = std::move(compilation_result.xla_input_shapes);
+  }
+
+  result.hlo_module_proto = std::move(compilation_result.computation->proto());
+  result.compile_metadata = arg.compile_metadata;
+  result.host_compute_metadata =
+      std::move(compilation_result.host_compute_metadata);
+
+  return result;
+}
+
+absl::StatusOr<Tf2HloResult> TfToHloCompiler::CompileTfToHlo(Tf2HloArg& arg) {
+  return tensorflow::ifrt_serving::CompileTfToHlo(arg);
+}
+
+absl::StatusOr<std::string> TfToHloCompiler::Key(const Tf2HloArg& arg) {
+  TF_ASSIGN_OR_RETURN(uint64_t fingerprint, arg.Fingerprint());
+  return absl::StrCat(absl::Hex(fingerprint));
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.h

@@ -0,0 +1,98 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_TF2HLO_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_TF2HLO_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_compilation.pb.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/layout.h"
+#include "xla/python/ifrt/topology.h"
+#include "xla/service/hlo.pb.h"
+#include "xla/shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+struct Tf2HloArg {
+  mlir::ModuleOp module;
+  // `input_dtypes_and_shapes` can be mutable during Tf2HLO compilation.
+  std::vector<DtypeAndShape> input_dtypes_and_shapes;
+  absl::Span<const int> variable_arg_indices;
+  absl::string_view entry_function_name;
+  // `compile_metadata` can be mutable during Tf2HLO compilation.
+  tensorflow::tpu::TPUCompileMetadataProto compile_metadata;
+  tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn;
+  std::shared_ptr<xla::ifrt::Topology> topology;
+  absl::string_view platform_name;
+  bool populate_layout_in_xla_input_shapes = false;
+
+  absl::StatusOr<uint64_t> Fingerprint() const;
+};
+
+struct Tf2HloResult {
+  xla::HloModuleProto hlo_module_proto;
+  tensorflow::tpu::TPUCompileMetadataProto compile_metadata;
+  tf2xla::HostComputeMetadata host_compute_metadata;
+  // `xla_input_shapes[i]` corresponds to the input shape of the i-th argument
+  // in the original `Tf2HloArg`. It will be empty if
+  // `populate_layout_in_xla_input_shapes` is false or there is no input in
+  // `module`.
+  std::vector<xla::Shape> xla_input_shapes;
+
+  absl::StatusOr<Tf2HLOResultProto> ToProto() const;
+
+  static absl::StatusOr<Tf2HloResult> FromProto(const Tf2HLOResultProto& proto);
+};
+
+absl::Status UpdateCompileMetadata(
+    tensorflow::tpu::TPUCompileMetadataProto& metadata,
+    absl::Span<const DtypeAndShape> inputs);
+
+absl::StatusOr<tensorflow::tpu::TPUCompileMetadataProto> GetCompileMetadata(
+    mlir::ModuleOp module, const xla::ifrt::Client& ifrt_client);
+
+class TfToHloCompiler {
+ public:
+  TfToHloCompiler() = default;
+  virtual ~TfToHloCompiler() = default;
+
+  // Returns a cache key that can be used to identify the result of
+  // CompileTfToHlo.
+  virtual absl::StatusOr<std::string> Key(const Tf2HloArg& arg);
+
+  virtual absl::StatusOr<Tf2HloResult> CompileTfToHlo(Tf2HloArg& arg);
+
+  virtual bool IsXlaCompilationDisabled() const { return false; }
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_TF2HLO_H_

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo_test.cc

@@ -0,0 +1,766 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.h"
+
+#include <memory>
+#include <ostream>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "mlir/IR/AsmState.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/backends/cpu/target_machine_options.h"
+#include "xla/debug_options_flags.h"
+#include "xla/layout.h"
+#include "xla/pjrt/pjrt_compiler.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/mock.h"
+#include "xla/python/ifrt/test_util.h"
+#include "xla/python/pjrt_ifrt/pjrt_topology.h"
+#include "xla/service/computation_placer.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/test.h"
+#include "tsl/platform/protobuf.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+using ::testing::Eq;
+using ::testing::HasSubstr;
+using ::testing::Ne;
+
+// TODO(b/229726259): Make EqualsProto available in OSS
+class ProtoStringMatcher {
+ public:
+  explicit ProtoStringMatcher(const tsl::protobuf::Message& expected)
+      : expected_(expected.SerializeAsString()) {}
+
+  template <typename Message>
+  bool MatchAndExplain(const Message& p,
+                       ::testing::MatchResultListener*) const {
+    return p.SerializeAsString() == expected_;
+  }
+
+  void DescribeTo(::std::ostream* os) const { *os << expected_; }
+  void DescribeNegationTo(::std::ostream* os) const {
+    *os << "not equal to expected message: " << expected_;
+  }
+
+ private:
+  const std::string expected_;
+};
+
+inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
+    const tsl::protobuf::Message& x) {
+  return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
+}
+
+class Tf2HloTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    mlir::DialectRegistry registry;
+    mlir::registerAllDialects(registry);
+    mlir::RegisterAllTensorFlowDialects(registry);
+    context_ = std::make_unique<mlir::MLIRContext>(registry);
+  }
+  TfToHloCompiler tf_to_hlo_compiler_;
+  std::unique_ptr<mlir::MLIRContext> context_;
+};
+
+TEST_F(Tf2HloTest, Empty) {
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/tf2hlo_empty.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, context_.get());
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), *client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, {}));
+
+  const xla::CpuTopologyDescription cpu_topology(
+      xla::CpuId(), xla::CpuName(), /*platform_version=*/"",
+      xla::CpuTopology(
+          {}, xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
+  std::shared_ptr<xla::CpuTopologyDescription> cpu_topology_ptr =
+      std::make_shared<xla::CpuTopologyDescription>(cpu_topology);
+
+  std::vector<int> variable_arg_indices;
+  Tf2HloArg arg{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = {},
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .platform_name = xla::CpuName(),
+  };
+  auto result = tf_to_hlo_compiler_.CompileTfToHlo(arg);
+
+  TF_ASSERT_OK(result.status());
+}
+
+// Multiple input and multiple out.
+TEST_F(Tf2HloTest, Tuple) {
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/tf2hlo_tuple.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, context_.get());
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+
+  std::vector<DtypeAndShape> dtype_and_shapes;
+  dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {1, 3}});
+  dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {3, 1}});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), *client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
+
+  const xla::CpuTopologyDescription cpu_topology(
+      xla::CpuId(), xla::CpuName(), /*platform_version=*/"",
+      xla::CpuTopology(
+          {}, xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
+  std::shared_ptr<xla::CpuTopologyDescription> cpu_topology_ptr =
+      std::make_shared<xla::CpuTopologyDescription>(cpu_topology);
+
+  std::vector<int> variable_arg_indices;
+
+  Tf2HloArg arg{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .platform_name = xla::CpuName(),
+  };
+
+  auto result = tf_to_hlo_compiler_.CompileTfToHlo(arg);
+
+  TF_ASSERT_OK(result.status());
+}
+
+struct SpmdTestCase {
+  std::string test_name;
+  std::string mlir_file;
+  std::string expected_metadata;
+};
+
+// Spmd and device assignment is given
+class Tf2HloSpmdTest : public Tf2HloTest,
+                       public ::testing::WithParamInterface<SpmdTestCase> {};
+
+TEST_P(Tf2HloSpmdTest, SpmdTest) {
+  const SpmdTestCase& test_case = GetParam();
+
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/", test_case.mlir_file));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, context_.get());
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+
+  std::vector<DtypeAndShape> dtype_and_shapes;
+  dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {4, 64}});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), *client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
+
+  const xla::CpuTopologyDescription cpu_topology(
+      xla::CpuId(), xla::CpuName(), /*platform_version=*/"",
+      xla::CpuTopology(
+          {}, xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
+  std::shared_ptr<xla::CpuTopologyDescription> cpu_topology_ptr =
+      std::make_shared<xla::CpuTopologyDescription>(cpu_topology);
+
+  std::vector<int> variable_arg_indices;
+  Tf2HloArg arg{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .platform_name = xla::CpuName(),
+  };
+
+  auto result = tf_to_hlo_compiler_.CompileTfToHlo(arg);
+
+  LOG(INFO) << result->compile_metadata;
+  TF_ASSERT_OK(result.status());
+
+  tensorflow::tpu::TPUCompileMetadataProto expected_compile_metadata;
+  ASSERT_TRUE(tsl::protobuf::TextFormat::ParseFromString(
+      test_case.expected_metadata, &expected_compile_metadata));
+
+  EXPECT_THAT(result->compile_metadata, EqualsProto(expected_compile_metadata));
+}
+INSTANTIATE_TEST_SUITE_P(
+    SpmdTests, Tf2HloSpmdTest,
+    ::testing::ValuesIn<SpmdTestCase>({
+        {"SpmdWithDeviceAssignment", "tf2hlo_spmd_with_device_assignment.mlir",
+         R"pb(
+           args {
+             dtype: DT_FLOAT
+             shape {
+               dim { size: 4 }
+               dim { size: 64 }
+             }
+             kind: PARAMETER
+             sharding {
+               type: OTHER
+               tile_assignment_dimensions: 2
+               tile_assignment_dimensions: 1
+               tile_assignment_devices: 0
+               tile_assignment_devices: 1
+             }
+             is_bounded_dynamic_dim: false
+           }
+           retvals { sharding {} }
+           num_replicas: 1
+           num_cores_per_replica: 2
+           device_assignment {
+             replica_count: 1
+             computation_count: 2
+             computation_devices { replica_device_ids: 0 }
+             computation_devices { replica_device_ids: 1 }
+           }
+           use_spmd_for_xla_partitioning: true
+           compile_options {}
+         )pb"},
+        {"SpmdShardy", "tf2hlo_spmd_shardy.mlir",
+         R"pb(
+           args {
+             dtype: DT_FLOAT
+             shape {
+               dim { size: 4 }
+               dim { size: 64 }
+             }
+             kind: PARAMETER
+             sharding {
+               type: OTHER
+               tile_assignment_dimensions: 2
+               tile_assignment_dimensions: 1
+               iota_reshape_dims: 2
+               iota_transpose_perm: 0
+             }
+             is_bounded_dynamic_dim: false
+           }
+           retvals { sharding {} }
+           num_replicas: 1
+           num_cores_per_replica: 2
+           device_assignment {
+             replica_count: 1
+             computation_count: 2
+             computation_devices { replica_device_ids: 0 }
+             computation_devices { replica_device_ids: 1 }
+           }
+           use_spmd_for_xla_partitioning: true
+           use_shardy_partitioner: true
+           compile_options {}
+         )pb"},
+    }),
+    [](const ::testing::TestParamInfo<Tf2HloSpmdTest::ParamType>& info) {
+      return info.param.test_name;
+    });
+
+// Spmd and use default device assignment b/c no device assignment is given
+TEST_F(Tf2HloTest, UsingDefaultDeviceAssignment) {
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/tf2hlo_spmd_no_device_assignment.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, context_.get());
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+
+  std::vector<DtypeAndShape> dtype_and_shapes;
+  dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {4, 64}});
+  dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {64, 10}});
+  dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {1, 4}});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), *client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
+
+  const xla::CpuTopologyDescription cpu_topology(
+      xla::CpuId(), xla::CpuName(), /*platform_version=*/"",
+      xla::CpuTopology(
+          {}, xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
+  std::shared_ptr<xla::CpuTopologyDescription> cpu_topology_ptr =
+      std::make_shared<xla::CpuTopologyDescription>(cpu_topology);
+
+  std::vector<int> variable_arg_indices;
+  Tf2HloArg arg{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .platform_name = xla::CpuName(),
+  };
+
+  auto result = tf_to_hlo_compiler_.CompileTfToHlo(arg);
+
+  LOG(INFO) << result->compile_metadata;
+  TF_ASSERT_OK(result.status());
+
+  tensorflow::tpu::TPUCompileMetadataProto expected_compile_metadata;
+  ASSERT_TRUE(tsl::protobuf::TextFormat::ParseFromString(
+      R"pb(
+        args {
+          dtype: DT_FLOAT
+          shape {
+            dim { size: 4 }
+            dim { size: 64 }
+          }
+          kind: PARAMETER
+          sharding {
+            type: OTHER
+            tile_assignment_dimensions: 2
+            tile_assignment_dimensions: 1
+            tile_assignment_devices: 0
+            tile_assignment_devices: 1
+          }
+          is_bounded_dynamic_dim: false
+        }
+        args {
+          dtype: DT_FLOAT
+          shape {
+            dim { size: 64 }
+            dim { size: 10 }
+          }
+          kind: PARAMETER
+          sharding {
+            type: OTHER
+            tile_assignment_dimensions: 2
+            tile_assignment_dimensions: 1
+            tile_assignment_devices: 0
+            tile_assignment_devices: 1
+          }
+          is_bounded_dynamic_dim: false
+        }
+        args {
+          dtype: DT_FLOAT
+          shape {
+            dim { size: 1 }
+            dim { size: 4 }
+          }
+          kind: PARAMETER
+          is_bounded_dynamic_dim: false
+        }
+        retvals { sharding {} }
+        num_replicas: 1
+        num_cores_per_replica: 2
+        device_assignment {
+          replica_count: 1
+          computation_count: 2
+          computation_devices { replica_device_ids: 0 }
+          computation_devices { replica_device_ids: 1 }
+        }
+        use_spmd_for_xla_partitioning: true
+        compile_options {}
+      )pb",
+      &expected_compile_metadata));
+
+  EXPECT_THAT(result->compile_metadata, EqualsProto(expected_compile_metadata));
+}
+
+// Multiple input and multiple out.
+TEST_F(Tf2HloTest, XlaCallHostCallback) {
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/xla_call_host_callback.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path,
+                                            mlir::ParserConfig(context_.get()));
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+
+  std::vector<DtypeAndShape> dtype_and_shapes;
+  dtype_and_shapes.push_back(DtypeAndShape{DT_INT32, {1}});
+  dtype_and_shapes.push_back(DtypeAndShape{DT_INT32, {1}});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), *client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
+
+  const xla::CpuTopologyDescription cpu_topology(
+      xla::CpuId(), xla::CpuName(), /*platform_version=*/"",
+      xla::CpuTopology(
+          {}, xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
+  std::shared_ptr<xla::CpuTopologyDescription> cpu_topology_ptr =
+      std::make_shared<xla::CpuTopologyDescription>(cpu_topology);
+
+  std::vector<int> variable_arg_indices;
+  Tf2HloArg arg{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .platform_name = xla::CpuName(),
+  };
+
+  auto result = tf_to_hlo_compiler_.CompileTfToHlo(arg);
+
+  TF_ASSERT_OK(result.status());
+
+  ASSERT_EQ((*result).host_compute_metadata.device_to_host().size(), 1);
+  ASSERT_EQ(
+      (*result).host_compute_metadata.device_to_host().begin()->metadata_size(),
+      2);
+  ASSERT_EQ((*result).host_compute_metadata.host_to_device().size(), 0);
+}
+
+// On GPU enabled build, the compilation should pass. On a GPU disabled build,
+// the compilation should fail with a correct error message.
+TEST_F(Tf2HloTest, GpuCompile) {
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/tf2hlo_gpu.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, context_.get());
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  xla::ifrt::MockClient mock_client;
+  ON_CALL(mock_client, GetDefaultDeviceAssignment)
+      .WillByDefault([]() -> absl::StatusOr<xla::DeviceAssignment> {
+        return xla::DeviceAssignment(1, 1);
+      });
+
+  std::vector<DtypeAndShape> dtype_and_shapes;
+  dtype_and_shapes.push_back(DtypeAndShape{DT_FLOAT, {}});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), mock_client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
+
+  std::vector<int> variable_arg_indices;
+  Tf2HloArg arg{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = nullptr,
+      .platform_name = xla::CudaName(),
+  };
+
+  auto result = tf_to_hlo_compiler_.CompileTfToHlo(arg);
+#if defined(GOOGLE_CUDA)
+  LOG(INFO) << "GPU compile success";
+  EXPECT_OK(result);
+#else
+  LOG(INFO) << "Non-GPU compile failure";
+  EXPECT_THAT(result,
+              absl_testing::StatusIs(absl::StatusCode::kUnimplemented,
+                                     HasSubstr("CUDA or ROCM build required")));
+#endif
+}
+
+TEST_F(Tf2HloTest, SameArgProduceSameKeyFingerprint) {
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/xla_call_host_callback.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path,
+                                            mlir::ParserConfig(context_.get()));
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+
+  std::vector<DtypeAndShape> dtype_and_shapes;
+  dtype_and_shapes.push_back(DtypeAndShape{DT_INT32, {1}});
+  dtype_and_shapes.push_back(DtypeAndShape{DT_INT32, {1}});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), *client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
+
+  const xla::CpuTopologyDescription cpu_topology(
+      xla::CpuId(), xla::CpuName(), /*platform_version=*/"",
+      xla::CpuTopology(
+          {}, xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
+  std::shared_ptr<xla::CpuTopologyDescription> cpu_topology_ptr =
+      std::make_shared<xla::CpuTopologyDescription>(cpu_topology);
+
+  std::vector<int> variable_arg_indices;
+  Tf2HloArg arg0{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .populate_layout_in_xla_input_shapes = true,
+  };
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module_clone =
+      mlir::OwningOpRef<mlir::ModuleOp>(mlir_module->clone());
+  Tf2HloArg arg1{
+      .module = mlir_module_clone.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .populate_layout_in_xla_input_shapes = true,
+  };
+
+  TfToHloCompiler tf_to_hlo_compiler;
+  TF_ASSERT_OK_AND_ASSIGN(std::string key0, tf_to_hlo_compiler.Key(arg0));
+  TF_ASSERT_OK_AND_ASSIGN(std::string key1, tf_to_hlo_compiler.Key(arg1));
+
+  EXPECT_THAT(key0, Eq(key1));
+}
+
+TEST_F(Tf2HloTest, DifferentCompileMetadataProduceDifferentKeyFingerprint) {
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/xla_call_host_callback.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path,
+                                            mlir::ParserConfig(context_.get()));
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+
+  std::vector<DtypeAndShape> dtype_and_shapes;
+  dtype_and_shapes.push_back(DtypeAndShape{DT_INT32, {1}});
+  dtype_and_shapes.push_back(DtypeAndShape{DT_INT32, {1}});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), *client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
+
+  const xla::CpuTopologyDescription cpu_topology(
+      xla::CpuId(), xla::CpuName(), /*platform_version=*/"",
+      xla::CpuTopology(
+          {}, xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
+  std::shared_ptr<xla::CpuTopologyDescription> cpu_topology_ptr =
+      std::make_shared<xla::CpuTopologyDescription>(cpu_topology);
+
+  std::vector<int> variable_arg_indices;
+  Tf2HloArg arg0{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+  };
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module_clone =
+      mlir::OwningOpRef<mlir::ModuleOp>(mlir_module->clone());
+  compile_metadata.set_num_replicas(11111);
+  Tf2HloArg arg1{
+      .module = mlir_module_clone.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+  };
+
+  TfToHloCompiler tf_to_hlo_compiler;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::string key0, tf_to_hlo_compiler.Key(arg0));
+  TF_ASSERT_OK_AND_ASSIGN(std::string key1, tf_to_hlo_compiler.Key(arg1));
+  EXPECT_THAT(key0, Ne(key1));
+}
+
+TEST_F(Tf2HloTest,
+       DifferentPopulateLayoutInXlaInputShapesProduceDifferentKeyFingerprint) {
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/compiler/mlir/tfrt/transforms/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/xla_call_host_callback.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path,
+                                            mlir::ParserConfig(context_.get()));
+
+  ASSERT_TRUE(mlir_module);
+  ASSERT_TRUE(mlir_module.get() != nullptr);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+
+  std::vector<DtypeAndShape> dtype_and_shapes;
+  dtype_and_shapes.push_back(DtypeAndShape{DT_INT32, {1}});
+  dtype_and_shapes.push_back(DtypeAndShape{DT_INT32, {1}});
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      tensorflow::tpu::TPUCompileMetadataProto compile_metadata,
+      GetCompileMetadata(mlir_module.get(), *client));
+  TF_ASSERT_OK(UpdateCompileMetadata(compile_metadata, dtype_and_shapes));
+
+  const xla::CpuTopologyDescription cpu_topology(
+      xla::CpuId(), xla::CpuName(), /*platform_version=*/"",
+      xla::CpuTopology(
+          {}, xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
+  std::shared_ptr<xla::CpuTopologyDescription> cpu_topology_ptr =
+      std::make_shared<xla::CpuTopologyDescription>(cpu_topology);
+
+  std::vector<int> variable_arg_indices;
+  Tf2HloArg arg0{
+      .module = mlir_module.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .populate_layout_in_xla_input_shapes = true,
+  };
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module_clone =
+      mlir::OwningOpRef<mlir::ModuleOp>(mlir_module->clone());
+  Tf2HloArg arg1{
+      .module = mlir_module_clone.get(),
+      .input_dtypes_and_shapes = dtype_and_shapes,
+      .variable_arg_indices = variable_arg_indices,
+      .entry_function_name = "main",
+      .compile_metadata = compile_metadata,
+      .shape_representation_fn = tensorflow::IdentityShapeRepresentationFn(),
+      .topology = std::make_shared<xla::ifrt::PjRtTopology>(cpu_topology_ptr),
+      .populate_layout_in_xla_input_shapes = false,
+  };
+
+  TfToHloCompiler tf_to_hlo_compiler;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::string key0, tf_to_hlo_compiler.Key(arg0));
+  TF_ASSERT_OK_AND_ASSIGN(std::string key1, tf_to_hlo_compiler.Key(arg1));
+  EXPECT_THAT(key0, Ne(key1));
+}
+
+TEST_F(Tf2HloTest, ToProtoAndFromProto) {
+  Tf2HloResult result;
+  result.hlo_module_proto.set_name("test_module");
+  result.compile_metadata.set_num_replicas(1);
+
+  xla::Shape shape0 = xla::ShapeUtil::MakeShape(xla::F32, {10, 20});
+  xla::Shape shape1 = xla::ShapeUtil::MakeShape(xla::F32, {30, 40});
+  result.xla_input_shapes = {shape0, shape1};
+
+  TF_ASSERT_OK_AND_ASSIGN(auto proto, result.ToProto());
+  TF_ASSERT_OK_AND_ASSIGN(auto result_from_proto,
+                          Tf2HloResult::FromProto(proto));
+
+  EXPECT_EQ(result_from_proto.hlo_module_proto.name(), "test_module");
+  EXPECT_EQ(result_from_proto.compile_metadata.num_replicas(), 1);
+  ASSERT_EQ(result_from_proto.xla_input_shapes.size(), 2);
+  EXPECT_EQ(result_from_proto.xla_input_shapes[0], shape0);
+  EXPECT_EQ(result_from_proto.xla_input_shapes[1], shape1);
+}
+
+}  // namespace
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_device_cleanup.cc

@@ -0,0 +1,51 @@
+/* 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 <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+#define GEN_PASS_DEF_TFDEVICECLEANUPPASS
+#define GEN_PASS_DECL_TFDEVICECLEANUPPASS
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+class TfDeviceCleanupPass
+    : public impl::TfDeviceCleanupPassBase<TfDeviceCleanupPass> {
+ public:
+  void runOnOperation() override {
+    mlir::func::FuncOp func = getOperation();
+    func.walk([](mlir::Operation* op) {
+      if (llvm::isa<mlir::TF::TensorFlowDialect>(op->getDialect())) {
+        op->removeAttr("device");
+      }
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfDeviceCleanupPass() {
+  return std::make_unique<TfDeviceCleanupPass>();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_identity_propagation.cc

@@ -0,0 +1,88 @@
+/* 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 <memory>
+
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+#define GEN_PASS_DEF_TFIDENTITYPROPAGATIONPASS
+#define GEN_PASS_DECL_TFIDENTITYPROPAGATIONPASS
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+constexpr absl::string_view kXlaShardingAttr = "_XlaSharding";
+
+bool IsTerminator(mlir::Operation* op) {
+  return op->hasTrait<mlir::OpTrait::IsTerminator>();
+}
+
+class TfIdentityPropagationPass
+    : public impl::TfIdentityPropagationPassBase<TfIdentityPropagationPass> {
+ public:
+  void runOnOperation() override {
+    mlir::func::FuncOp func = getOperation();
+
+    func.walk([](mlir::TF::IdentityOp identity) {
+      // Don't propagate inputs of identity ops with sharding annotation since
+      // identity ops are sometimes used to change output sharding.
+      if (identity->hasAttr(kXlaShardingAttr)) {
+        return;
+      }
+      // Identity outputs to terminator ops (e.g., `func.return`) cannot be
+      // replaced unless input/output types are exactly the same. Doing so may
+      // cause mismatch between the enclosing region's return type and the
+      // terminator's arg type.
+      const bool same_type =
+          identity.getInput().getType() == identity.getOutput().getType();
+      identity.getOutput().replaceUsesWithIf(
+          identity.getInput(), [&](mlir::OpOperand& operand) {
+            return same_type || !IsTerminator(operand.getOwner());
+          });
+    });
+
+    func.walk([](mlir::TF::IdentityNOp identity_n) {
+      if (identity_n->hasAttr(kXlaShardingAttr)) {
+        return;
+      }
+      for (auto [input, output] :
+           llvm::zip(identity_n.getInput(), identity_n.getOutput())) {
+        const bool same_type = input.getType() == output.getType();
+        output.replaceUsesWithIf(input, [&](mlir::OpOperand& operand) {
+          return same_type || !IsTerminator(operand.getOwner());
+        });
+      }
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfIdentityPropagationPass() {
+  return std::make_unique<TfIdentityPropagationPass>();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_ifrt_passes.cc

@@ -0,0 +1,167 @@
+
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_ifrt_passes.h"
+
+#include <memory>
+#include <string>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/core/util/debug_data_dumper.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+using mlir::LogicalResult;
+using mlir::OpPassManager;
+using mlir::PassManager;
+using mlir::func::FuncOp;
+
+void AddClusterToIfrtRuntimeOpsPassPipeline(
+    OpPassManager& pm, llvm::StringRef module_name,
+    bool enable_propagate_static_shapes_pass, bool enable_async_ifrt) {
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::CreateExecutorDialectToFunctionalConversionPass());
+
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::TF::CreateCanonicalizeCompileAndReplicateAttributesPass());
+
+  pm.addNestedPass<mlir::func::FuncOp>(CreateTfIdentityPropagationPass());
+
+  pm.addNestedPass<mlir::func::FuncOp>(CreateTfRestoreSplittingPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+  pm.addNestedPass<mlir::func::FuncOp>(CreateTfRestorePruningPass());
+  pm.addNestedPass<mlir::func::FuncOp>(CreateTfRestoreMergingPass());
+  // Convert reference variable to resource variable since
+  // LowerToIfrtRestoreVariablePass does not support reference variable.
+  pm.addPass(CreateConvertReferenceVariableToResourceVariablePass());
+  pm.addPass(CreateLowerToIfrtRestoreVariablePass());
+
+  pm.addPass(CreateRewriteClusterToIfrtCallPass(enable_async_ifrt));
+
+  if (enable_propagate_static_shapes_pass) {
+    pm.addPass(CreatePropagateStaticShapesPass());
+  }
+
+  // After device program is extracted, we can clean up device attributes from
+  // all ops.
+  pm.addNestedPass<mlir::func::FuncOp>(CreateTfDeviceCleanupPass());
+
+  // Sink VarHandle with ReadVariableOp: subsequent SinkVariableAsNamedArrayPass
+  // rely on the co-existence of VarHandle and ReadVariable in the same
+  // function.
+  // First, we inline all the function calls. This will sink VarHandle
+  // with ReadVariable in most cases. Then SinkInvariantOpsPass will sink
+  // VarHandle to a few special Ops that inliner does not handle.
+  // TODO(b/319045348): the bridge before this pipeline already does some
+  // inlining. Consider removing this inliner.
+  pm.addPass(mlir::createInlinerPass());
+  // Convert the region control flow to functional control flow so that the
+  // invariant sinking pass can work correctly.
+  pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
+  pm.addPass(::tensorflow::CreateSinkInInvariantOpsPass());
+  pm.addPass(mlir::TF::CreateTFFunctionalControlFlowToRegions());
+  pm.addPass(mlir::createInlinerPass());
+
+  // Decompose resource ops as resource variables are loaded by ReadVariableOp
+  // and can be lowered to IfrtLoadVariableOp in the subsequent
+  // SinkVariableAsNamedArrayPass.
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::TFDevice::CreateDecomposeResourceOpsPass());
+
+  // Sink variable tensor as named array in IFRT.
+  pm.addPass(CreateSinkVariableAsNamedArrayPass());
+}
+
+}  // namespace
+
+// Setup the input pass manager to enable IR dumping after each pass.
+// Note a side effect of this method is that multi threading will be disabled.
+void EnablePassIRPrinting(PassManager& pm, const std::string& dump_group_name,
+                          llvm::StringRef module_name) {
+  // Print the whole module after each pass, which requires disabling
+  // multi-threading as well.
+  pm.getContext()->disableMultithreading();
+  pm.enableIRPrinting(std::make_unique<::tensorflow::DataDumperLoggerConfig>(
+      [module_name, dump_group_name](const std::string& pass_tag_name,
+                                     mlir::Operation* op) {
+        return DEBUG_DATA_DUMPER()->GetDumpFilename(
+            module_name.str(), dump_group_name, pass_tag_name);
+      },
+      /*pass_prefix=*/"",
+      /*print_module_scope=*/true));
+  pm.enableTiming();
+}
+
+absl::Status RunClusterToIfrtRuntimeOpsPassPipeline(
+    mlir::ModuleOp module, llvm::StringRef module_name,
+    bool enable_propagate_static_shapes_pass, bool enable_async_ifrt) {
+  mlir::StatusScopedDiagnosticHandler diag_handler(
+      module.getContext(), /*propagate=*/false,
+      /*filter_stack=*/!VLOG_IS_ON(1));
+
+  PassManager runtime_lowering(module.getContext());
+  ::tensorflow::applyTensorflowAndCLOptions(runtime_lowering);
+
+  AddClusterToIfrtRuntimeOpsPassPipeline(runtime_lowering, module_name,
+                                         enable_propagate_static_shapes_pass,
+                                         enable_async_ifrt);
+
+  if (VLOG_IS_ON(1)) {
+    ::tensorflow::DumpMlirOpToFile(
+        DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain,
+                                             "ifrt_runtime_lowering_before"),
+        module, llvm::StringRef(), &runtime_lowering);
+  }
+
+  if (VLOG_IS_ON(2)) {
+    EnablePassIRPrinting(runtime_lowering, kDebugGroupRuntimeLowering,
+                         module_name);
+  }
+
+  // Ignore the result since diag_handler consumes it
+  LogicalResult result = runtime_lowering.run(module);
+  (void)result;
+
+  if (VLOG_IS_ON(1)) {
+    ::tensorflow::DumpMlirOpToFile(
+        DEBUG_DATA_DUMPER()->GetDumpFilename(module_name.str(), kDebugGroupMain,
+                                             "ifrt_runtime_lowering_after"),
+        module, llvm::StringRef(), &runtime_lowering);
+  }
+
+  return diag_handler.ConsumeStatus();
+}
+
+// Register all IfrtPass
+void RegisterTfIfrtPasses() { registerTfrtIfrtServingPasses(); }
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_ifrt_passes.h

@@ -0,0 +1,94 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_TF_IFRT_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_TF_IFRT_PASSES_H_
+
+#include <memory>
+#include <string>
+
+#include "absl/status/status.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// Create a pass to convert tf_device.cluster_func to tf.ifrt_program_call.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateRewriteClusterToIfrtCallPass(bool enable_async_ifrt = false);
+
+// Creates a pass that propagates static shapes from
+// `tf.SetStaticDimensionBoundsOp` to `IfrtCallOp` and callee `FuncOp`
+// arguments.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreatePropagateStaticShapesPass();
+
+// Creates a pass that sinks variable tensor argument to `tf.IfrtCall` as named
+// arrays and lowers `tf.ReadVariableOp` to `tf.IfrtLoadVariableOp`.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateSinkVariableAsNamedArrayPass();
+
+// Creates a pass that splits `tf.RestoreV2` ops.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfRestoreSplittingPass();
+
+// Creates a pass that merges `tf.RestoreV2` ops.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfRestoreMergingPass();
+
+// Creates a pass that propagates inputs of no-op identity ops to their outputs.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfIdentityPropagationPass();
+
+// Creates a pass that prunes unused `tf.RestoreV2` ops.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfRestorePruningPass();
+
+// Creates a pass that lower `tf.RestoreVariableOp` to
+// `tf.IfrtRestoreVariableOp`.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateLowerToIfrtRestoreVariablePass();
+
+// Creates a pass that cleans up device attributes from all ops.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfDeviceCleanupPass();
+
+#define GEN_PASS_REGISTRATION
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+// Register all passes.
+void RegisterTfIfrtPasses();
+
+// Setup the input pass manager to enable IR dumping after each pass.
+// Note a side effect of this method is that multi threading will be disabled.
+void EnablePassIRPrinting(mlir::PassManager& pm,
+                          const std::string& dump_group_name,
+                          llvm::StringRef module_name);
+
+// Convert tf_device.cluster_func to tf.ifrt_program_call.
+// The callee function is converted to a ifrt_program.
+absl::Status RunClusterToIfrtRuntimeOpsPassPipeline(
+    mlir::ModuleOp module, llvm::StringRef module_name = llvm::StringRef(),
+    bool enable_propagate_static_shapes_pass = true,
+    bool enable_async_ifrt = false);
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_IFRT_TF_IFRT_PASSES_H_

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_restore_merging.cc

@@ -0,0 +1,164 @@
+/* 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 <stdint.h>
+
+#include <memory>
+#include <vector>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+#define GEN_PASS_DEF_TFRESTOREMERGINGPASS
+#define GEN_PASS_DECL_TFRESTOREMERGINGPASS
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+class TfRestoreMergingPass
+    : public impl::TfRestoreMergingPassBase<TfRestoreMergingPass> {
+ public:
+  void runOnOperation() override {
+    mlir::func::FuncOp func = getOperation();
+    for (mlir::Block& block : func) {
+      // Group `tf.RestoreV2` ops by prefixes and merge each group.
+      llvm::SmallDenseMap<mlir::Value, std::vector<mlir::TF::RestoreV2Op>>
+          restore_groups;
+      for (auto restore : block.getOps<mlir::TF::RestoreV2Op>()) {
+        restore_groups[restore.getPrefix()].push_back(restore);
+      }
+      for (const auto& restores : llvm::make_second_range(restore_groups)) {
+        if (mlir::failed(MergeRestores(restores))) {
+          return signalPassFailure();
+        }
+      }
+    }
+  }
+
+ private:
+  mlir::DenseStringElementsAttr GetStringTensorAttr(
+      llvm::ArrayRef<llvm::StringRef> values) {
+    const int size = values.size();
+    const auto type = mlir::RankedTensorType::get(
+        {size}, mlir::TF::StringType::get(&getContext()));
+    return mlir::DenseStringElementsAttr::get(type, values);
+  }
+
+  // Merges `tf.RestoreV2` ops with the same prefix. Ignores restore ops with
+  // non-constant `tensor_names` and/or `shape_and_slices`.
+  mlir::LogicalResult MergeRestores(
+      llvm::ArrayRef<mlir::TF::RestoreV2Op> restores) {
+    if (restores.size() <= 1) {
+      return mlir::success();
+    }
+
+    // All restore ops must have the same prefix.
+    const mlir::Value prefix =
+        mlir::TF::RestoreV2Op(restores.front()).getPrefix();
+
+    std::vector<mlir::TF::RestoreV2Op> restores_to_merge;
+    std::vector<mlir::Value> values_to_replace;
+    std::vector<llvm::StringRef> merged_tensor_names;
+    std::vector<llvm::StringRef> merged_shape_and_slices;
+
+    std::vector<mlir::Location> restore_locs;
+    std::vector<mlir::Location> tensor_names_locs;
+    std::vector<mlir::Location> shape_and_slices_locs;
+
+    for (mlir::TF::RestoreV2Op restore : restores) {
+      mlir::DenseStringElementsAttr tensor_names;
+      mlir::DenseStringElementsAttr shape_and_slices;
+      if (!mlir::matchPattern(restore,
+                              mlir::m_Op<mlir::TF::RestoreV2Op>(
+                                  mlir::matchers::m_Val(prefix),
+                                  mlir::m_Constant(&tensor_names),
+                                  mlir::m_Constant(&shape_and_slices)))) {
+        continue;
+      }
+      if (tensor_names.size() != restore.getNumResults() ||
+          shape_and_slices.size() != restore.getNumResults()) {
+        return restore.emitOpError()
+               << "returns an inconsistent number of results";
+      }
+
+      restores_to_merge.push_back(restore);
+      llvm::append_range(values_to_replace, restore.getTensors());
+      llvm::append_range(merged_tensor_names,
+                         tensor_names.getValues<llvm::StringRef>());
+      llvm::append_range(merged_shape_and_slices,
+                         shape_and_slices.getValues<llvm::StringRef>());
+
+      restore_locs.push_back(restore.getLoc());
+      tensor_names_locs.push_back(restore.getTensorNames().getLoc());
+      shape_and_slices_locs.push_back(restore.getShapeAndSlices().getLoc());
+    }
+    if (restores_to_merge.size() <= 1) {
+      return mlir::success();
+    }
+
+    // Insert the merged restore op right before the first restore op to be
+    // merged in order to keep the dominance property.
+    mlir::OpBuilder builder(restores_to_merge.front());
+
+    auto new_tensor_names = mlir::TF::ConstOp::create(
+        builder, builder.getFusedLoc(tensor_names_locs),
+        GetStringTensorAttr(merged_tensor_names));
+    auto new_shape_and_slices = mlir::TF::ConstOp::create(
+        builder, builder.getFusedLoc(shape_and_slices_locs),
+        GetStringTensorAttr(merged_shape_and_slices));
+
+    auto new_restore = mlir::TF::RestoreV2Op::create(
+        builder, builder.getFusedLoc(restore_locs),
+        mlir::TypeRange(mlir::ValueRange(values_to_replace)), prefix,
+        new_tensor_names, new_shape_and_slices);
+    for (auto [old_value, new_value] :
+         llvm::zip(values_to_replace, new_restore.getTensors())) {
+      old_value.replaceAllUsesWith(new_value);
+    }
+
+    for (mlir::TF::RestoreV2Op restore : restores_to_merge) {
+      restore.erase();
+    }
+    return mlir::success();
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfRestoreMergingPass() {
+  return std::make_unique<TfRestoreMergingPass>();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_restore_pruning.cc

@@ -0,0 +1,52 @@
+/* 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 <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+#define GEN_PASS_DEF_TFRESTOREPRUNINGPASS
+#define GEN_PASS_DECL_TFRESTOREPRUNINGPASS
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+// Prune unused RestoreV2 Op.
+class TfRestorePruningPass
+    : public impl::TfRestorePruningPassBase<TfRestorePruningPass> {
+ public:
+  void runOnOperation() override {
+    mlir::func::FuncOp func = getOperation();
+    func.walk([&](mlir::TF::RestoreV2Op restore) {
+      if (restore.use_empty()) {
+        restore.erase();
+      }
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfRestorePruningPass() {
+  return std::make_unique<TfRestorePruningPass>();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf_restore_splitting.cc

@@ -0,0 +1,122 @@
+/* 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 <stdint.h>
+
+#include <memory>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+#define GEN_PASS_DEF_TFRESTORESPLITTINGPASS
+#define GEN_PASS_DECL_TFRESTORESPLITTINGPASS
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/passes.h.inc"  // IWYU pragma: keep
+
+class TfRestoreSplittingPass
+    : public impl::TfRestoreSplittingPassBase<TfRestoreSplittingPass> {
+ public:
+  void runOnOperation() override {
+    mlir::func::FuncOp func = getOperation();
+    const mlir::WalkResult result =
+        func.walk([&](mlir::TF::RestoreV2Op restore) {
+          if (mlir::failed(SplitRestore(restore))) {
+            return mlir::WalkResult::interrupt();
+          }
+          return mlir::WalkResult::advance();
+        });
+    if (result.wasInterrupted()) {
+      return signalPassFailure();
+    }
+  }
+
+ private:
+  mlir::DenseStringElementsAttr GetStringTensorAttr(
+      llvm::ArrayRef<llvm::StringRef> values) {
+    const int size = values.size();
+    const auto type = mlir::RankedTensorType::get(
+        {size}, mlir::TF::StringType::get(&getContext()));
+    return mlir::DenseStringElementsAttr::get(type, values);
+  }
+
+  // Splits the `tf.RestoreV2` op into per-variable restore ops if its
+  // `tensor_name` and `shape_and_slices` are constant.
+  mlir::LogicalResult SplitRestore(mlir::TF::RestoreV2Op restore) {
+    mlir::DenseStringElementsAttr tensor_names;
+    mlir::DenseStringElementsAttr shape_and_slices;
+    if (!mlir::matchPattern(restore,
+                            mlir::m_Op<mlir::TF::RestoreV2Op>(
+                                /*prefix=*/mlir::matchers::m_Any(),
+                                mlir::m_Constant(&tensor_names),
+                                mlir::m_Constant(&shape_and_slices)))) {
+      return mlir::success();
+    }
+    if (tensor_names.size() != restore.getNumResults() ||
+        shape_and_slices.size() != restore.getNumResults()) {
+      return restore.emitOpError()
+             << "returns an inconsistent number of results";
+    }
+
+    mlir::OpBuilder builder(restore);
+    for (auto [tensor_name, shape_and_slice, result] :
+         llvm::zip(tensor_names.getValues<llvm::StringRef>(),
+                   shape_and_slices.getValues<llvm::StringRef>(),
+                   restore.getTensors())) {
+      auto new_tensor_names =
+          mlir::TF::ConstOp::create(builder, restore.getTensorNames().getLoc(),
+                                    GetStringTensorAttr({tensor_name}));
+
+      auto new_shape_and_slices = mlir::TF::ConstOp::create(
+          builder, restore.getShapeAndSlices().getLoc(),
+          GetStringTensorAttr({shape_and_slice}));
+
+      auto new_restore = mlir::TF::RestoreV2Op::create(
+          builder, restore.getLoc(), mlir::TypeRange({result.getType()}),
+          restore.getPrefix(), new_tensor_names, new_shape_and_slices);
+      result.replaceAllUsesWith(new_restore.getTensors()[0]);
+    }
+
+    restore.erase();
+    return mlir::success();
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateTfRestoreSplittingPass() {
+  return std::make_unique<TfRestoreSplittingPass>();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/insert_tensor_copy.cc

@@ -0,0 +1,187 @@
+/* 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 <cassert>
+#include <memory>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/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/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tfrt/basic_kernels/opdefs/basic_kernels.h"  // from @tf_runtime
+#include "tfrt/compiler/stream_analysis.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+// This pass inserts copy kernels for fallback tensors when they are passed to
+// multiple threads, to avoid atomic contention on their refcounts.
+class InsertFallbackTensorCopy
+    : public mlir::PassWrapper<InsertFallbackTensorCopy,
+                               mlir::OperationPass<mlir::func::FuncOp>> {
+  void getDependentDialects(mlir::DialectRegistry& registry) const override {
+    registry.insert<tfrt::fallback_async::FallbackAsyncDialect>();
+  }
+
+  llvm::StringRef getArgument() const final {
+    return "tfrt-insert-fallback-tensor-copy";
+  }
+
+  llvm::StringRef getDescription() const final {
+    return "Inserts copy kernels for fallback tensors when they are passed to "
+           "multiple threads, to avoid atomic contention on refcounts.";
+  }
+
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertFallbackTensorCopy)
+
+  void runOnOperation() override {
+    mlir::func::FuncOp func_op = getOperation();
+
+    // Use stream analysis to know whether a value is passed to different
+    // threads.
+    tfrt::compiler::StreamAnalysis stream_analysis(func_op);
+
+    auto builder = mlir::OpBuilder::atBlockBegin(&func_op.front());
+
+    // Process function arguments first.
+    for (auto arg : func_op.getArguments()) {
+      if (!mlir::isa<tfrt::fallback::TFTensorType>(arg.getType())) continue;
+      InsertFallbackTensorCopyForValue(arg, func_op->getLoc(), builder,
+                                       stream_analysis);
+    }
+
+    // Then process each operations in the block.
+    for (mlir::Operation& op : llvm::make_early_inc_range(func_op.front())) {
+      if (llvm::isa<tfrt::fallback_async::ExecuteOp,
+                    tfrt::fallback_async::ExecuteOpSeq>(&op)) {
+        InsertFallbackTensorCopyForFallbackOp(&op, builder, stream_analysis);
+      }
+    }
+  }
+
+ private:
+  void InsertFallbackTensorCopyForFallbackOp(
+      mlir::Operation* op, mlir::OpBuilder& builder,
+      const tfrt::compiler::StreamAnalysis& stream_analysis) {
+    builder.setInsertionPointAfter(op);
+
+    // Process each result value.
+    for (auto result : op->getResults()) {
+      if (!mlir::isa<tfrt::fallback::TFTensorType>(result.getType())) continue;
+      InsertFallbackTensorCopyForValue(result, op->getLoc(), builder,
+                                       stream_analysis);
+    }
+  }
+
+  // Insert copy kernels to copy the result, and allocate new atomic refcount
+  // if the value is going to be used by different streams/threads, in order to
+  // avoid contention on the atomic counter.
+  void InsertFallbackTensorCopyForValue(
+      mlir::Value value, mlir::Location loc, mlir::OpBuilder& builder,
+      const tfrt::compiler::StreamAnalysis& stream_analysis) {
+    llvm::DenseMap<int, llvm::SmallVector<mlir::OpOperand*, 4>> stream_map;
+
+    // Find out streams that use this value and the corresponding uses.
+    for (mlir::OpOperand& use : value.getUses()) {
+      // Skip return op as there should not be atomic contention on the return
+      // op.
+      if (llvm::isa<tfrt::compiler::ReturnOp>(use.getOwner())) continue;
+
+      int stream_id = stream_analysis.GetStream(use.getOwner()).id();
+      stream_map[stream_id].push_back(&use);
+    }
+
+    // Organize these uses into groups. If a stream has many uses of this value,
+    // put these uses into one stream. Otherwise, streams with small number
+    // of uses are grouped with each other to form groups with enough uses.
+    constexpr int kCopyGroupThreshold = 16;
+    llvm::SmallVector<llvm::SmallVector<mlir::OpOperand*, 4>, 4> small_copies;
+    llvm::SmallVector<llvm::SmallVector<mlir::OpOperand*, 4>, 4> copies;
+    for (const auto& iter : stream_map) {
+      if (iter.second.size() >= kCopyGroupThreshold) {
+        copies.push_back(iter.second);
+      } else {
+        if (small_copies.empty() ||
+            small_copies.back().size() >= kCopyGroupThreshold) {
+          small_copies.push_back(iter.second);
+        } else {
+          small_copies.back().append(iter.second.begin(), iter.second.end());
+        }
+      }
+    }
+
+    if (!small_copies.empty())
+      copies.append(small_copies.begin(), small_copies.end());
+
+    // If it is only used by one group, then we don't need to copy.
+    if (copies.size() <= 1) return;
+
+    // Remove one group from the candidates, as we can just use the original
+    // value for this group.
+    copies.pop_back();
+
+    // For each stream, we will create one new value that replaces the uses in
+    // that stream.
+
+    assert(mlir::isa<tfrt::fallback::TFTensorType>(value.getType()));
+
+    // The number of results is the number candidate streams.
+    llvm::SmallVector<mlir::Type, 4> result_types(copies.size(),
+                                                  value.getType());
+    assert(!result_types.empty());
+
+    // Create the tfrt_fallback_async.copy_if_small kernel.
+    auto copy_op = tfrt::fallback_async::CopyIfSmallOp::create(
+        builder, loc, result_types, value);
+
+    // Finally, replaces all uses with the new value.
+    for (int i = 0; i < copies.size(); ++i) {
+      const auto& uses = copies[i];
+      auto new_value = copy_op.getResult(i);
+      for (auto* use : uses) {
+        use->set(new_value);
+      }
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateInsertFallbackTensorCopyPass() {
+  return std::make_unique<InsertFallbackTensorCopy>();
+}
+
+static mlir::PassRegistration<InsertFallbackTensorCopy> register_pass(
+    CreateInsertFallbackTensorCopyPass);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/lower_saved_model.cc

@@ -0,0 +1,772 @@
+/* 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 <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <deque>
+#include <memory>
+#include <tuple>
+#include <utility>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Block.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/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.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/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.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/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_interfaces.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/utils.h"
+
+namespace tensorflow {
+namespace {
+
+using ::mlir::tf_saved_model::kTfSavedModelExportedNamesAttr;
+using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
+
+constexpr char kCpuDeviceName[] =
+    "/job:localhost/replica:0/task:0/device:CPU:0";
+
+mlir::TF::ResourceHandle GetResourceHandle(mlir::Operation *op) {
+  llvm::StringRef device;
+  if (auto attr = op->getAttrOfType<mlir::StringAttr>("device")) {
+    device = attr.getValue();
+  }
+
+  llvm::StringRef container;
+  if (auto attr = op->getAttrOfType<mlir::StringAttr>("container")) {
+    container = attr.getValue();
+  }
+
+  llvm::StringRef shared_name;
+  if (auto attr = op->getAttrOfType<mlir::StringAttr>("shared_name")) {
+    shared_name = attr.getValue();
+  }
+
+  return {container, shared_name, device, /*op=*/nullptr};
+}
+
+struct HoistInfo {
+  // All hoisted ops in the topological order.
+  llvm::SmallVector<mlir::Operation *, 4> hoists_in_topological_order;
+
+  // Mapping from the old values produced by hoisted ops before hoisting to the
+  // new values after hoisting.
+  mlir::IRMapping value_mapping;
+
+  // `hoisted_values` is to keep all values that are produced by hoisted ops
+  // but used by non-hoisted ops. These values will be replaced by results of
+  // tf._TfrtGetResource op. The index of each value in this array will be the
+  // index used in tf._TfrtGetResource and tf._TfrtSetResource op. This also
+  // stores the ResourceHandle which has the shared_name and container
+  // attributes used by later resource alias analysis and side effect analysis
+  // passes.
+  llvm::SmallVector<std::pair<mlir::Value, mlir::TF::ResourceHandle>, 4>
+      hoisted_values;
+};
+
+bool OnlyHasReadOrNoEffect(mlir::Operation *op) {
+  auto interface = llvm::dyn_cast<mlir::MemoryEffectOpInterface>(op);
+  if (!interface) return false;
+  return interface.onlyHasEffect<mlir::MemoryEffects::Read>() ||
+         interface.hasNoEffect();
+}
+
+bool CanHoist(const llvm::DenseSet<mlir::TF::ResourceHandle> &read_only_vars,
+              mlir::Operation *op) {
+  // return ops should not be hoisted.
+  if (op->mightHaveTrait<mlir::OpTrait::IsTerminator>()) return false;
+
+  // Fixes a corner case where hoisting the tf.BatchFunction leads to
+  // a compilation error; such a case may occur in unit tests.
+  if (llvm::isa<mlir::TF::BatchFunctionOp>(op)) return false;
+
+  // Non-side-effecting ops can be hoisted.
+  if (mlir::isMemoryEffectFree(op)) return true;
+
+  // ResourceHandle ops can be hoisted.
+  if (llvm::isa<mlir::TF::VarHandleOp, mlir::TF::HashTableV2Op>(op))
+    return true;
+
+  // If it is ReadVariableOp and the variable is readonly, it can be hoisted.
+  if (auto read_var_op = llvm::dyn_cast<mlir::TF::ReadVariableOp>(op)) {
+    if (auto var_handle_op = llvm::dyn_cast_or_null<mlir::TF::VarHandleOp>(
+            read_var_op.getResource().getDefiningOp())) {
+      if (read_only_vars.count(GetResourceHandle(var_handle_op)) > 0)
+        return true;
+    }
+  }
+
+  // If it is LookupTableSizeOp, it can be hoisted as the size of the hash table
+  // cannot be changed after initialization.
+  if (auto lookup_table_size_op =
+          llvm::dyn_cast<mlir::TF::LookupTableSizeV2Op>(op)) {
+    if (auto hash_table_op = llvm::dyn_cast_or_null<mlir::TF::HashTableV2Op>(
+            lookup_table_size_op.getTableHandle().getDefiningOp())) {
+      if (read_only_vars.count(GetResourceHandle(hash_table_op)) > 0)
+        return true;
+    }
+  }
+
+  // TODO(chky): Allow more readonly ops.
+
+  return false;
+}
+
+void HoistInvariantOpsInFunction(
+    mlir::func::FuncOp func,
+    const llvm::DenseSet<mlir::TF::ResourceHandle> &read_only_vars,
+    const mlir::TF::SideEffectAnalysis::Info &side_effect_analysis,
+    mlir::OpBuilder &builder, HoistInfo &module_hoist_info) {
+  // Keep the hoisted ops in this function.
+  llvm::DenseSet<mlir::Operation *> hoists;
+
+  auto all_operands_in_hoists = [&module_hoist_info](mlir::Operation *op) {
+    for (mlir::Value operand : op->getOperands()) {
+      if (module_hoist_info.value_mapping.lookupOrNull(operand) == nullptr)
+        return false;
+    }
+    return true;
+  };
+
+  auto all_control_predeccessors_in_hoists = [&hoists, &side_effect_analysis](
+                                                 mlir::Operation *op) {
+    auto preds = side_effect_analysis.DirectControlPredecessors(op);
+    return std::all_of(
+        preds.begin(), preds.end(),
+        [&hoists](mlir::Operation *pred) { return hoists.count(pred) > 0; });
+  };
+
+  std::deque<mlir::Operation *> work_list;
+
+  // Start with ops with tf.VarHandleOp ops and tf.Const ops.
+  //
+  // TODO(chky): Consider allowing other ops including custom ops to be hoisted.
+  func.walk([&work_list](mlir::Operation *op) {
+    if (llvm::isa<mlir::TF::VarHandleOp, mlir::TF::HashTableV2Op,
+                  mlir::TF::ConstOp>(op))
+      work_list.push_back(op);
+  });
+
+  while (!work_list.empty()) {
+    auto *op = work_list.front();
+    work_list.pop_front();
+
+    // Skip if it is already hoisted.
+    if (hoists.count(op) > 0) continue;
+
+    // If the op can be hoisted, and all of its data dependencies and control
+    // dependencies are hoisted, then we hoist it. Otherwise, skip.
+    if (!(CanHoist(read_only_vars, op) && all_operands_in_hoists(op) &&
+          all_control_predeccessors_in_hoists(op)))
+      continue;
+
+    // Record the hoisted operation.
+    hoists.insert(op);
+    module_hoist_info.hoists_in_topological_order.push_back(op);
+
+    // Create a copy in the init function.
+    builder.clone(*op, module_hoist_info.value_mapping);
+
+    for (mlir::Operation *user : op->getUsers()) {
+      work_list.push_back(user);
+    }
+  }
+
+  // Find out the values that are produced by hoisted ops but used by
+  // non-hoisted ops. These values need to be replaced.
+  for (auto *op : hoists) {
+    for (auto result : op->getResults()) {
+      if (std::any_of(result.getUsers().begin(), result.getUsers().end(),
+                      [&hoists](mlir::Operation *user) {
+                        return hoists.count(user) == 0;
+                      })) {
+        module_hoist_info.hoisted_values.push_back(
+            {result, GetResourceHandle(op)});
+      }
+    }
+  }
+}
+
+void FindCalleesRecursiveForOp(const mlir::SymbolTable &symbol_table,
+                               mlir::Operation *op,
+                               llvm::StringSet<> &callees) {
+  for (const auto &named_attr : op->getAttrs()) {
+    if (auto symbol_attr =
+            mlir::dyn_cast<mlir::FlatSymbolRefAttr>(named_attr.getValue())) {
+      auto symbol = symbol_attr.getValue();
+      if (!callees.contains(symbol)) {
+        callees.insert(symbol);
+
+        auto func = symbol_table.lookup<mlir::func::FuncOp>(symbol);
+        if (!func) continue;
+
+        func.walk([&](mlir::Operation *op) {
+          FindCalleesRecursiveForOp(symbol_table, op, callees);
+        });
+      }
+    }
+  }
+}
+
+void FindCalleesRecursive(const mlir::SymbolTable &symbol_table,
+                          mlir::func::FuncOp func, llvm::StringSet<> &callees) {
+  assert(func);
+  func.walk([&](mlir::Operation *op) {
+    FindCalleesRecursiveForOp(symbol_table, op, callees);
+  });
+}
+
+// This pass rewrites tf_saved_model dialect's ops according to TFRT's
+// requirements:
+//
+// 1) Remove all tf_saved_model's attributes and ops.
+// 2) Create a function for every exported names of the original function.
+// 3) Hoist invariant ops (ie. guaranteed to return the same value on every
+// invocation) for every non-init function.
+//
+class LowerTFSavedModelPass
+    : public mlir::PassWrapper<LowerTFSavedModelPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  void getDependentDialects(mlir::DialectRegistry &registry) const override {
+    registry.insert<mlir::func::FuncDialect>();
+  }
+
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerTFSavedModelPass)
+
+  explicit LowerTFSavedModelPass(bool hoist_invariant_ops,
+                                 bool fuse_get_resource_ops) {
+    hoist_invariant_ops_ = hoist_invariant_ops;
+    fuse_get_resource_ops_ = fuse_get_resource_ops;
+  }
+  LowerTFSavedModelPass() = default;
+  LowerTFSavedModelPass(const LowerTFSavedModelPass &) {}
+
+  llvm::StringRef getArgument() const final {
+    return "tfrt-lower-tf-savedmodel";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Lower tf-saved-model ops according to TFRT's requirements.";
+  }
+
+  void runOnOperation() override {
+    auto module = getOperation();
+
+    // TODO(b/185928201): Create a standalone pass for hoisting invariant ops so
+    // that it can be reusable and configurable in other contexts than saved
+    // models.
+    if (hoist_invariant_ops_) HoistInvariantOps(module);
+
+    // Skip non-savedmodel MLIR module.
+    if (!mlir::tf_saved_model::HasTfSavedModelSemantics(module)) return;
+
+    mlir::SymbolTable symbol_table(module);
+
+    module->removeAttr("tf_saved_model.semantics");
+
+    mlir::OpBuilder builder(&getContext());
+    auto resource_id = builder.getStringAttr("tf.resource_name");
+    auto bound_id = builder.getStringAttr("tf_saved_model.bound_input");
+    auto path_id = builder.getStringAttr(kTfSavedModelIndexPathAttr);
+
+    module.walk([resource_id, bound_id, path_id,
+                 &builder](mlir::Operation *op) mutable {
+      if (auto func_op = llvm::dyn_cast<mlir::func::FuncOp>(op)) {
+        // Remove tf_saved_model specific function arg attributes.
+        for (unsigned i = 0, e = func_op.getNumArguments(); i != e; ++i) {
+          if (auto sym = func_op.getArgAttrOfType<mlir::FlatSymbolRefAttr>(
+                  i, bound_id)) {
+            func_op.removeArgAttr(i, bound_id);
+            func_op.setArgAttr(i, resource_id,
+                               builder.getStringAttr(sym.getValue()));
+          }
+          func_op.removeArgAttr(i, path_id);
+        }
+        for (unsigned i = 0, e = func_op.getNumResults(); i != e; ++i) {
+          func_op.removeResultAttr(i, bound_id);
+          func_op.removeResultAttr(i, path_id);
+        }
+        if (auto exported_names = func_op->getAttrOfType<mlir::ArrayAttr>(
+                kTfSavedModelExportedNamesAttr)) {
+          bool is_session_initializer = IsSessionInitializer(func_op);
+
+          // Create a function for each exported name.
+          //
+          // TODO(b/148477882): TFRT dialect should have similar concepts of
+          // exported names so that a function can be referenced by multiple
+          // exported names.
+          func_op->removeAttr(kTfSavedModelExportedNamesAttr);
+          for (auto exported_name : exported_names) {
+            auto exported_func_op = func_op.clone();
+            exported_func_op.setName(
+                mlir::cast<mlir::StringAttr>(exported_name));
+
+            // If it is a session initializer, we want to maximize parallelism
+            // and do not perform any stream merge, to minimize latency.
+            //
+            // TODO(b/183219530): This is a workaround as the cost model used
+            // currently is not very accurate, and leads to performance
+            // regression on IO ops that are common in initialization functions.
+            if (is_session_initializer) {
+              exported_func_op->setAttr("tfrt.cost_threshold",
+                                        builder.getI64IntegerAttr(1));
+            }
+
+            builder.setInsertionPoint(func_op);
+            builder.insert(exported_func_op);
+          }
+          func_op.erase();
+        }
+      }
+    });
+
+    module.walk([](mlir::Operation *op) {
+      if (llvm::isa<mlir::tf_saved_model::TensorFlowSavedModelDialect>(
+              op->getDialect())) {
+        // Remove all tf_saved_model ops.
+        op->erase();
+      }
+    });
+  }
+
+ private:
+  void HoistInvariantOps(mlir::ModuleOp module);
+  void ReplaceHoistedValues(
+      llvm::ArrayRef<std::pair<mlir::Value, mlir::TF::ResourceHandle>>
+          hoisted_values,
+      mlir::OpBuilder &builder);
+
+  Option<bool> hoist_invariant_ops_{*this, "hoist-invariant-ops",
+                                    llvm::cl::desc("hoist-invariant-ops"),
+                                    llvm::cl::init(false)};
+  Option<bool> fuse_get_resource_ops_{*this, "fuse-get-resource-ops",
+                                      llvm::cl::desc("fuse get resource ops"),
+                                      llvm::cl::init(true)};
+};
+
+void LowerTFSavedModelPass::HoistInvariantOps(mlir::ModuleOp module) {
+  mlir::SymbolTable symbol_table(module);
+
+  // Find all resources used in non-init functions.
+  llvm::DenseMap<mlir::TF::ResourceHandle,
+                 llvm::SmallVector<mlir::Operation *, 4>>
+      resources;
+
+  // Find all callees referenced in the initialization functions.
+  llvm::StringSet<> init_callees;
+
+  // Recursively find all callees referenced in the tf.XlaLaunch op.
+  // At and after the point of calling this pass, the MLIR xla function is no
+  // longer used. So there is no point to do hoisting for xla functions.
+  llvm::StringSet<> xla_launch_callees;
+
+  module.walk([&](mlir::Operation *op) {
+    if (llvm::isa<mlir::TF::VarHandleOp, mlir::TF::HashTableV2Op>(op)) {
+      auto func = op->getParentOfType<mlir::func::FuncOp>();
+      if (IsSessionInitializer(func)) return;
+      resources[GetResourceHandle(op)].push_back(op);
+    } else if (auto func = llvm::dyn_cast<mlir::func::FuncOp>(op)) {
+      if (!IsSessionInitializer(func)) return;
+      FindCalleesRecursive(symbol_table, func, init_callees);
+    } else if (llvm::isa<mlir::TF::XlaLaunchOp>(op)) {
+      // TODO(b/275095412): Clean up MLIR XLA functions after they are written
+      // back to function library, so that we don't need to do special handling
+      // for those functions here.
+      FindCalleesRecursiveForOp(symbol_table, op, xla_launch_callees);
+    }
+  });
+
+  llvm::DenseSet<mlir::TF::ResourceHandle> read_only_vars;
+  for (const auto &iter : resources) {
+    const auto &key = iter.first;
+    const auto &vars = iter.second;
+    if (std::all_of(vars.begin(), vars.end(), [](mlir::Operation *op) {
+          for (auto *user : op->getUsers()) {
+            if (!OnlyHasReadOrNoEffect(user)) return false;
+          }
+          return true;
+        })) {
+      read_only_vars.insert(key);
+    }
+  }
+
+  mlir::TF::SideEffectAnalysis side_effect_analysis(module);
+
+  mlir::OpBuilder builder(&module.getBodyRegion());
+  // "_tfrt_resource_init" is the special function that executes all invariant
+  // ops (eg. read-only variables) used in the model. This function should be
+  // executed after user-specified initialization.
+  auto init_func_op = mlir::func::FuncOp::create(
+      builder, module.getLoc(), "_tfrt_resource_init",
+      mlir::FunctionType::get(module.getContext(), /*inputs=*/{},
+                              /*results=*/{}));
+  auto *block = init_func_op.addEntryBlock();
+  builder.setInsertionPointToStart(block);
+
+  HoistInfo module_hoist_info;
+
+  for (auto func : module.getOps<mlir::func::FuncOp>()) {
+    // Skips hoisting if this function is an init function or any callees,
+    // including recursive ones, of an init functions, because otherwise the
+    // hoisted values won't be initialized when this function is called.
+    if (IsSessionInitializer(func) ||
+        init_callees.contains(func.getSymName()) || func == init_func_op ||
+        xla_launch_callees.contains(func.getSymName()))
+      continue;
+
+    // Skips hoisting if this function runs on TPU. This is will happen when
+    // fallback to TPUPartitionedCallOp is enabled for SPMD.
+    // TODO(b/214039254): remove this once tfrt support native SPMD.
+    bool has_tpu_op = false;
+    func.walk([&has_tpu_op](mlir::Operation *op) {
+      if (op->hasAttr("_tpu_replicate")) has_tpu_op = true;
+    });
+    if (has_tpu_op) continue;
+
+    HoistInvariantOpsInFunction(func, read_only_vars,
+                                side_effect_analysis.GetAnalysisForFunc(func),
+                                builder, module_hoist_info);
+  }
+
+  // Create tf._TfrtSetResource ops in the init function.
+  for (auto iter : llvm::enumerate(module_hoist_info.hoisted_values)) {
+    mlir::Value value = iter.value().first;
+    int64_t index = iter.index();
+
+    auto new_value = module_hoist_info.value_mapping.lookup(value);
+    auto *new_op = new_value.getDefiningOp();
+    assert(new_op);
+    builder.setInsertionPointAfter(new_op);
+    auto set_resource_op = mlir::TF::_TfrtSetResourceOp::create(
+        builder, new_op->getLoc(), new_value, index);
+
+    // Preserve the device attribute.
+    llvm::StringRef device = kCpuDeviceName;
+    if (auto device_attr = new_op->getAttrOfType<mlir::StringAttr>("device")) {
+      if (!device_attr.getValue().empty()) device = device_attr.getValue();
+    }
+    set_resource_op->setAttr("device", builder.getStringAttr(device));
+  }
+
+  builder.setInsertionPointToEnd(block);
+  // Finish building the init function by inserting an return op.
+  mlir::func::ReturnOp::create(builder, init_func_op.getLoc());
+
+  // Now that we have the index for each value that will be replaced, we can
+  // create the tf._TfrtGetResource op in each function using these indices.
+  ReplaceHoistedValues(module_hoist_info.hoisted_values, builder);
+
+  // Lastly, erase the hoisted ops in reverse topological order.
+  for (auto *op :
+       llvm::reverse(module_hoist_info.hoists_in_topological_order)) {
+    assert(op->use_empty());
+    op->erase();
+  }
+}
+
+void LowerTFSavedModelPass::ReplaceHoistedValues(
+    llvm::ArrayRef<std::pair<mlir::Value, mlir::TF::ResourceHandle>>
+        hoisted_values,
+    mlir::OpBuilder &builder) {
+  struct HoistedValueInfo {
+    llvm::SmallVector<mlir::Value, 4> hoisted_values;
+    llvm::SmallVector<int64_t, 4> indices;
+    llvm::SmallVector<llvm::StringRef, 4> shared_names;
+    llvm::SmallVector<llvm::StringRef, 4> containers;
+  };
+  // Rearrange the hoisted values by each function and each device.
+  llvm::DenseMap<mlir::Block *, llvm::StringMap<HoistedValueInfo>>
+      hoisted_values_by_block_device;
+
+  // Find a block where to place tf._TfrtGetResource operation. We do not place
+  // get resource operations inside the `tf_device.cluster` operations, because
+  // these blocks are intended for later on-device compilation. Insert resource
+  // reads to the closest block outside of the `tf_device.cluster` operation.
+  auto hoist_into_block = [](mlir::Value value) -> mlir::Block * {
+    mlir::Operation *cluster_op =
+        value.getDefiningOp()->getParentOfType<mlir::tf_device::ClusterOp>();
+    return cluster_op ? cluster_op->getBlock() : value.getParentBlock();
+  };
+
+  for (auto iter : llvm::enumerate(hoisted_values)) {
+    auto value = iter.value().first;
+    auto index = iter.index();
+    auto &device_map = hoisted_values_by_block_device[hoist_into_block(value)];
+
+    assert(value.getDefiningOp() && "hoisted values must not be arguments.");
+    llvm::StringRef device = kCpuDeviceName;
+    if (auto device_attr =
+            value.getDefiningOp()->getAttrOfType<mlir::StringAttr>("device")) {
+      if (!device_attr.getValue().empty()) device = device_attr.getValue();
+    }
+
+    auto &item = device_map[device];
+
+    item.hoisted_values.push_back(value);
+    item.indices.push_back(index);
+    item.shared_names.push_back(iter.value().second.name);
+    item.containers.push_back(iter.value().second.container);
+  }
+
+  // Create tf._TfrtGetResource op for each function and device.
+  for (const auto &block_iter : hoisted_values_by_block_device) {
+    auto *block = block_iter.first;
+    const auto &device_map = block_iter.second;
+
+    builder.setInsertionPointToStart(block);
+    for (const auto &device_iter : device_map) {
+      llvm::StringRef device = device_iter.getKey();
+      mlir::ValueRange old_values = device_iter.getValue().hoisted_values;
+      const auto &indices = device_iter.getValue().indices;
+      const auto &shared_name_arr = device_iter.getValue().shared_names;
+      const auto &container_arr = device_iter.getValue().containers;
+
+      llvm::SmallVector<mlir::Value> new_values;
+
+      if (fuse_get_resource_ops_) {
+        auto get_resource_op = mlir::TF::_TfrtGetResourceOp::create(
+            builder, block->getParentOp()->getLoc(), old_values.getTypes(),
+            builder.getI64ArrayAttr(indices),
+            builder.getStrArrayAttr(shared_name_arr),
+            builder.getStrArrayAttr(container_arr));
+        get_resource_op->setAttr("device", builder.getStringAttr(device));
+        new_values = get_resource_op.getResults();
+      } else {
+        for (int i = 0; i < old_values.size(); ++i) {
+          auto get_resource_op = mlir::TF::_TfrtGetResourceOp::create(
+              builder, block->getParentOp()->getLoc(),
+              mlir::TypeRange(old_values[i].getType()),
+              builder.getI64ArrayAttr(indices[i]),
+              builder.getStrArrayAttr(shared_name_arr[i]),
+              builder.getStrArrayAttr(container_arr[i]));
+          get_resource_op->setAttr("device", builder.getStringAttr(device));
+          new_values.append(get_resource_op->result_begin(),
+                            get_resource_op->result_end());
+        }
+      }
+
+      for (auto iter : llvm::zip(old_values, new_values)) {
+        auto old_value = std::get<0>(iter);
+        auto new_value = std::get<1>(iter);
+        old_value.replaceAllUsesWith(new_value);
+      }
+    }
+  }
+}
+
+static llvm::SmallVector<unsigned, 4> CompareTypes(mlir::TypeRange x,
+                                                   mlir::TypeRange y) {
+  llvm::SmallVector<unsigned, 4> results;
+  assert(x.size() == y.size());
+  for (int i = 0, e = x.size(); i < e; ++i) {
+    if (x[i] != y[i]) results.push_back(i);
+  }
+  return results;
+}
+
+// Converts ref variables to resource variables in a few cases.
+//
+// If the users of one variable in the entire module satisfies the following
+// condition, it will be converted to resource variable:
+//
+// 1) tf.Identity op
+// 2) tf.Assign op
+// 3) side-effect-free ops: This is also the TF1 behavior that the TF executor
+//    will automatically convert ref tensors to non-ref tensors if the user is
+//    not expecting a ref tensor. Refer to
+//    http://cs?q=tensorflow/core/common_runtime/executor.cc:932%20at_cl:356873227
+class ConvertReferenceVariableToResourceVariablePass
+    : public mlir::PassWrapper<ConvertReferenceVariableToResourceVariablePass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  llvm::StringRef getArgument() const final {
+    return "tfrt-convert-ref-variables";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Convert reference variable to resource variables.";
+  }
+  void runOnOperation() override;
+
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      ConvertReferenceVariableToResourceVariablePass)
+};
+
+mlir::LogicalResult ConvertReferenceVariableToResourceVariable(
+    mlir::TF::VariableV2Op var_op) {
+  auto tensor_type = mlir::cast<mlir::TensorType>(
+      mlir::TF::DropRefType(var_op.getRef().getType()));
+
+  llvm::SmallVector<mlir::TF::IdentityOp, 4> identity_ops;
+  llvm::SmallVector<mlir::TF::AssignOp, 4> assign_ops;
+  llvm::SmallVector<std::pair<mlir::Operation *, unsigned>, 4>
+      side_effect_free_ops;
+
+  for (mlir::OpOperand &use : var_op.getRef().getUses()) {
+    mlir::Operation *user = use.getOwner();
+
+    if (auto identity = llvm::dyn_cast<mlir::TF::IdentityOp>(user)) {
+      identity_ops.push_back(identity);
+      continue;
+    } else if (auto assign = llvm::dyn_cast<mlir::TF::AssignOp>(user)) {
+      // Conservatively we only allow the case that the output of this tf.Assign
+      // is not consumed by any other ops.
+      if (assign.getOutputRef().use_empty()) {
+        assign_ops.push_back(assign);
+        continue;
+      }
+    } else if (mlir::isMemoryEffectFree(user)) {
+      side_effect_free_ops.push_back({user, use.getOperandNumber()});
+      continue;
+    }
+
+    return var_op.emitOpError()
+           << "failed to convert reference variables with unexpected users. "
+           << *user;
+  }
+
+  mlir::OpBuilder builder(var_op);
+
+  auto var_handle_op = mlir::TF::VarHandleOp::create(
+      builder, var_op.getLoc(),
+      mlir::RankedTensorType::get(
+          {}, mlir::TF::ResourceType::get(
+                  llvm::ArrayRef<mlir::TensorType>{tensor_type},
+                  builder.getContext())),
+      var_op.getContainer(), var_op.getSharedName());
+
+  for (auto op : identity_ops) {
+    // Set insertion point to this identity_op so that the side-effect
+    // visibility is preserved.
+    builder.setInsertionPoint(op);
+    auto read_var_op = mlir::TF::ReadVariableOp::create(
+        builder, op.getLoc(), op.getType(), var_handle_op);
+    op.replaceAllUsesWith(read_var_op.getValue());
+    op.erase();
+  }
+
+  for (auto op : assign_ops) {
+    // Set the insertion point after the assign op so that all operands are
+    // dominating the newly created op.
+    builder.setInsertionPoint(op);
+    mlir::TF::AssignVariableOp::create(builder, op.getLoc(), var_handle_op,
+                                       op.getValue());
+    op.erase();
+  }
+
+  for (auto pair : side_effect_free_ops) {
+    mlir::Operation *op = pair.first;
+    unsigned idx = pair.second;
+    // Set the insertion point after the op so that all operands are dominating
+    // the newly created op.
+    builder.setInsertionPoint(op);
+    // Create a new read variable op, so that the side-effects are preserved.
+    auto read_var_op = mlir::TF::ReadVariableOp::create(
+        builder, op->getLoc(), tensor_type, var_handle_op);
+    op->setOperand(idx, read_var_op.getValue());
+  }
+
+  return mlir::success();
+}
+
+void ConvertReferenceVariableToResourceVariablePass::runOnOperation() {
+  auto module = getOperation();
+
+  // The key here is a tuple of device, container and shared_name to uniquely
+  // identify a variable.
+  llvm::DenseMap<std::tuple<llvm::StringRef, llvm::StringRef, llvm::StringRef>,
+                 llvm::SmallVector<mlir::TF::VariableV2Op, 4>>
+      ref_vars;
+
+  // First, we collect all variables' corresponding tf.VariableV2 ops.
+  module.walk([&ref_vars](mlir::TF::VariableV2Op op) {
+    if (op.getSharedName().empty()) {
+      op.emitOpError()
+          << "unable to convert reference variables with empty shared_names.";
+      return mlir::WalkResult::interrupt();
+    }
+
+    llvm::StringRef device;
+    if (auto device_attr = op->getAttrOfType<mlir::StringAttr>("device")) {
+      device = device_attr.getValue();
+    }
+
+    ref_vars[{device, op.getContainer(), op.getSharedName()}].push_back(op);
+
+    return mlir::WalkResult::advance();
+  });
+
+  // Then we perform rewrite for each variable if possible.
+  for (const auto &iter : ref_vars) {
+    const auto &var_ops = iter.second;
+
+    for (auto var_op : var_ops) {
+      if (mlir::succeeded(ConvertReferenceVariableToResourceVariable(var_op)))
+        var_op.erase();
+    }
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateLowerTFSavedModelPass(bool hoist_invariant_ops,
+                            bool fuse_get_resource_ops) {
+  return std::make_unique<LowerTFSavedModelPass>(hoist_invariant_ops,
+                                                 fuse_get_resource_ops);
+}
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateConvertReferenceVariableToResourceVariablePass() {
+  return std::make_unique<ConvertReferenceVariableToResourceVariablePass>();
+}
+
+static mlir::PassRegistration<LowerTFSavedModelPass> saved_model_pass;
+
+static mlir::PassRegistration<ConvertReferenceVariableToResourceVariablePass>
+    ref_var_pass;
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/merge_tf_if_ops.cc

@@ -0,0 +1,296 @@
+/* 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 <algorithm>
+#include <cassert>
+#include <memory>
+#include <string>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+// A special DenseMapInfo that hashes only operands of a operation, and treats
+// two operations equivalent if their operands are the same.
+struct OpWithSameArgsInfo : llvm::DenseMapInfo<mlir::Operation *> {
+  static unsigned getHashValue(const mlir::Operation *const_op) {
+    auto *op = const_cast<mlir::Operation *>(const_op);
+    return llvm::hash_combine(
+        llvm::hash_combine_range(op->operand_begin(), op->operand_end()));
+  }
+
+  static bool isEqual(const mlir::Operation *const_lhs,
+                      const mlir::Operation *const_rhs) {
+    auto *lhs = const_cast<mlir::Operation *>(const_lhs);
+    auto *rhs = const_cast<mlir::Operation *>(const_rhs);
+    if (lhs == rhs) return true;
+
+    return std::equal(lhs->operand_begin(), lhs->operand_end(),
+                      rhs->operand_begin(), rhs->operand_end());
+  }
+};
+
+// This pass merges non-side-effecting tf.If ops if their operands are the same.
+// For example,
+//  %r0 = tf.If(%cond, %x) {else = @else_0, then = @then_0}
+//  %r1, %r2 = tf.If(%cond, %x) {else = @else_1, then = @then_1}
+//
+// will be converted to:
+//  func private @merge_else(%arg) {
+//    %r0 = tf.PartitionedCall(%arg) {f = @else_0}
+//    %r1, %r2 = tf.PartitionedCall(%arg) {f = @else_1}
+//    return %r0, %r1, %r2
+//  }
+//  func private @merge_then(%arg) {
+//    %r0 = tf.PartitionedCall(%arg) {f = @then_0}
+//    %r1, %r2 = tf.PartitionedCall(%arg) {f = @then_1}
+//    return %r0, %r1, %r2
+//  }
+//
+//  %r0, %r1, %r2 = tf.If(%cond, %arg) {else = @merge_else, then = @merge_then}
+//
+// Then the inliner pass is run on the module, so the bodies of else_0 and
+// else_1 are inlined into the body of merge_else, and the bodies of then_0 and
+// then_1 are inlined into the body of merge_then.
+//
+// Note that the results will be concatenated.
+class MergeTfIfOpsPass
+    : public mlir::PassWrapper<MergeTfIfOpsPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  llvm::StringRef getArgument() const final { return "tfrt-merge-tf-if-ops"; }
+  llvm::StringRef getDescription() const final {
+    return "Merge tf.If ops with the same arguments.";
+  }
+
+  void runOnOperation() override {
+    constexpr int kMaxIter = 20;
+    auto module = getOperation();
+
+    bool changed = true;
+
+    for (int i = 0; i < kMaxIter && changed; ++i) {
+      changed = false;
+      for (auto func_op :
+           llvm::make_early_inc_range(module.getOps<mlir::func::FuncOp>())) {
+        changed |= ProcessFunction(func_op, i);
+      }
+
+      mlir::OperationFingerPrint fingerprint(module);
+
+      // Run optimization passes to expose more merge opportunities among the
+      // then-branch functions and the else-branch functions that are now
+      // respectively merged, for the next iteration.
+      mlir::OpPassManager pm(module.getOperationName());
+      pm.addPass(mlir::createInlinerPass());
+      pm.addPass(mlir::createCSEPass());
+      pm.addPass(CreateDeduplicateIfResultPass());
+      pm.addPass(mlir::createInlinerPass());
+      pm.addPass(mlir::createCSEPass());
+      pm.addNestedPass<mlir::func::FuncOp>(
+          tfrt_compiler::CreateOptimizeTfForTfrtPass());
+      if (mlir::failed(runPipeline(pm, module))) {
+        module.emitWarning(
+            absl::StrCat("could not run inliner pass within the "
+                         "tfrt-merge-tf-if-ops pass iteration ",
+                         i));
+        break;
+      }
+
+      changed |= fingerprint != mlir::OperationFingerPrint(module);
+    }
+  }
+
+  bool ProcessFunction(mlir::func::FuncOp op, int iteration) {
+    // Use a hash map to group tf.If ops with the same operands.
+    llvm::SmallDenseMap<mlir::Operation *, llvm::SmallVector<mlir::TF::IfOp, 2>,
+                        2, OpWithSameArgsInfo>
+        if_ops_to_merge;
+
+    for (mlir::Operation &op : op.front()) {
+      auto if_op = llvm::dyn_cast<mlir::TF::IfOp>(&op);
+
+      // Skip non tf.If ops.
+      if (!if_op) continue;
+
+      if_ops_to_merge[if_op].push_back(if_op);
+    }
+
+    int id = 0;
+
+    // Set the insertion point to the current function, as we will insert new
+    // functions here.
+    mlir::OpBuilder builder(op);
+
+    // Track the tf.If ops that should be removed as they are merged.
+    llvm::SmallVector<mlir::TF::IfOp, 4> if_ops_to_remove;
+
+    bool changed = false;
+    for (auto &iter : if_ops_to_merge) {
+      if (iter.second.size() <= 1) continue;
+
+      // Merge tf.If ops that have the same operands. The merged branches will
+      // be given unique names.
+      MergeIfOpsWithSameArgs(builder, iter.first->getLoc(),
+                             /*branch_prefix=*/
+                             absl::StrCat(op.getSymName().str(), "_merged_if_",
+                                          iteration, "_", id++),
+                             iter.second);
+
+      if_ops_to_remove.append(iter.second.begin(), iter.second.end());
+      changed = true;
+    }
+
+    // Now that we are no longer using `if_ops_to_merge` or any other data
+    // structures that uses the operations that will be removed, we can now
+    // erase these if ops.
+    for (auto op : if_ops_to_remove) op->erase();
+
+    return changed;
+  }
+
+  void MergeIfOpsWithSameArgs(mlir::OpBuilder &builder, mlir::Location loc,
+                              absl::string_view branch_prefix,
+                              llvm::MutableArrayRef<mlir::TF::IfOp> if_ops) {
+    assert(if_ops.size() > 1);
+
+    // The results of the merged tf.If op are the concatenation of results of
+    // the original tf.If ops.
+    llvm::SmallVector<mlir::Type, 4> new_result_types;
+    for (auto if_op : if_ops) {
+      new_result_types.append(if_op->result_type_begin(),
+                              if_op->result_type_end());
+    }
+
+    auto branch_function_type = builder.getFunctionType(
+        if_ops.front().getInput().getTypes(), new_result_types);
+
+    // Create new branches for the merged tf.If op.
+    auto then_branch_name = CreateBranchFunction(
+        builder, loc, branch_prefix,
+        /*branch_suffix=*/"_then", branch_function_type, if_ops,
+        [](mlir::TF::IfOp op) { return op.getThenBranchAttr(); });
+
+    auto else_branch_name = CreateBranchFunction(
+        builder, loc, branch_prefix,
+        /*branch_suffix=*/"_else", branch_function_type, if_ops,
+        [](mlir::TF::IfOp op) { return op.getElseBranchAttr(); });
+
+    mlir::OpBuilder::InsertionGuard guard(builder);
+    builder.setInsertionPoint(if_ops.front());
+
+    bool is_stateless =
+        std::all_of(if_ops.begin(), if_ops.end(),
+                    [](mlir::TF::IfOp op) { return op.getIsStateless(); });
+
+    // Create the merged tf.If op using the new branches.
+    auto new_if_op = mlir::TF::IfOp::create(
+        builder, loc, new_result_types, if_ops.front().getCond(),
+        if_ops.front().getInput(), then_branch_name, else_branch_name,
+        is_stateless);
+
+    // Replace the uses of results of the original tf.If ops with the results of
+    // the merged tf.If op.
+    auto new_result_iter = new_if_op.getOutput().begin();
+    for (auto if_op : if_ops) {
+      for (auto result : if_op.getOutput()) {
+        assert(new_result_iter != new_if_op.getOutput().end());
+        result.replaceAllUsesWith(*new_result_iter);
+        ++new_result_iter;
+      }
+    }
+  }
+
+  llvm::StringRef CreateBranchFunction(
+      mlir::OpBuilder &builder, mlir::Location loc,
+      absl::string_view branch_prefix, absl::string_view branch_suffix,
+      mlir::FunctionType branch_function_type,
+      llvm::ArrayRef<mlir::TF::IfOp> if_ops,
+      llvm::function_ref<mlir::FlatSymbolRefAttr(mlir::TF::IfOp)> get_branch) {
+    std::string branch_name = absl::StrCat(branch_prefix, branch_suffix);
+    auto branch = mlir::func::FuncOp::create(builder, loc, branch_name,
+                                             branch_function_type);
+    branch.setVisibility(mlir::func::FuncOp::Visibility::Private);
+
+    mlir::OpBuilder::InsertionGuard guard(builder);
+
+    // In the body of newly created branch function, we insert
+    // tf.PartitionedCall ops to call the original branches.
+    auto *block = branch.addEntryBlock();
+    builder.setInsertionPointToStart(block);
+    auto empty_string_attr = builder.getStringAttr("");
+
+    llvm::SmallVector<mlir::Value, 4> results;
+    results.reserve(branch_function_type.getNumResults());
+
+    for (auto if_op : if_ops) {
+      // Create the call op to the original branch. The arguments are simply
+      // the arguments from the wrapper function.
+      auto call_op = mlir::TF::PartitionedCallOp::create(
+          builder, if_op.getLoc(), if_op.getResultTypes(),
+          block->getArguments(),
+          /*args_attrs=*/nullptr, /*res_attrs=*/nullptr, get_branch(if_op),
+          empty_string_attr, empty_string_attr, empty_string_attr);
+
+      // The results are the concatenation of the original branches.
+      results.append(call_op.getOutput().begin(), call_op.getOutput().end());
+    }
+
+    mlir::func::ReturnOp::create(builder, loc, results);
+
+    return branch.getSymName();
+  }
+
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeTfIfOpsPass)
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateMergeTfIfOpsPass() {
+  return std::make_unique<MergeTfIfOpsPass>();
+}
+
+static mlir::PassRegistration<MergeTfIfOpsPass> register_pass(
+    CreateMergeTfIfOpsPass);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/BUILD

@@ -0,0 +1,289 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/tfrt:__subpackages__",
+    ],
+)
+
+cc_library(
+    name = "parallelization",
+    srcs = ["parallelization.cc"],
+    hdrs = ["parallelization.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
+        "//tensorflow/compiler/mlir/tfrt:constants",
+        "//tensorflow/compiler/mlir/tfrt:cost_analysis",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "//tensorflow/core/tfrt/fallback:cost_recorder",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@tf_runtime//:stream_analysis",
+    ],
+)
+
+cc_library(
+    name = "assign_op_key",
+    srcs = ["assign_op_key.cc"],
+    hdrs = ["assign_op_key.h"],
+    deps = [
+        ":util",
+        "//tensorflow/compiler/mlir/tfrt:constants",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+    ],
+)
+
+cc_library(
+    name = "mlrt_device_constants",
+    hdrs = ["mlrt_device_constants.h"],
+)
+
+cc_library(
+    name = "util",
+    srcs = ["util.cc"],
+    hdrs = ["util.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_a_m_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_n_z_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow/ir/host_runtime:tensorflow_tfrt_ops_inc_gen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "tf_to_mlrt",
+    srcs = ["tf_to_mlrt.cc"],
+    hdrs = ["tf_to_mlrt.h"],
+    deps = [
+        ":execute_op_registry",
+        ":mlrt_device_constants",
+        ":tpu_conversion_patterns",
+        ":util",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:convert_tensor",
+        "//tensorflow/compiler/mlir/tensorflow:translate_utils",
+        "//tensorflow/compiler/mlir/tensorflow/ir/host_runtime:tensorflow_tfrt_ops_inc_gen",
+        "//tensorflow/compiler/mlir/tfrt:constants",
+        "//tensorflow/compiler/mlir/tfrt:tfrt_pipeline_options",
+        "//tensorflow/compiler/mlir/tfrt:transform_utils",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_tpu_ops",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/tfrt/fallback:fallback_state",
+        "//tensorflow/core/tfrt/fallback:op_kernel_runner_cache",
+        "@com_google_absl//absl/log",
+        "@com_google_protobuf//:protobuf_headers",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncTransforms",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@xla//xla/tsl/platform:status",
+    ],
+)
+
+cc_library(
+    name = "passes",
+    srcs = ["passes.cc"],
+    hdrs = ["passes.h"],
+    deps = [
+        ":assign_op_key",
+        ":async_while",
+        ":fuse_mlrt_ops",
+        ":ifrt_set_tpu_host_allocator",
+        ":parallelization",
+        ":rewrite_ifrt_load_variable",
+        ":tf_to_mlrt",
+        ":while_to_map_fn",
+        "//tensorflow/compiler/mlir/tfrt:tfrt_pipeline_options",
+        "//tensorflow/core/tfrt/fallback:cost_recorder",
+        "//tensorflow/core/tfrt/fallback:fallback_state",
+        "@com_google_absl//absl/log:check",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+cc_library(
+    name = "execute_op_registry",
+    hdrs = ["execute_op_registry.h"],
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "tpu_conversion_patterns",
+    srcs = ["tpu_conversion_patterns.cc"],
+    hdrs = ["tpu_conversion_patterns.h"],
+    deps = [
+        ":execute_op_registry",
+        ":mlrt_device_constants",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tfrt:tfrt_pipeline_options",
+        "//tensorflow/compiler/mlir/tfrt:transform_utils",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_tpu_ops",
+        "@com_google_absl//absl/log:check",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+cc_library(
+    name = "fuse_mlrt_ops",
+    srcs = ["fuse_mlrt_ops.cc"],
+    hdrs = ["fuse_mlrt_ops.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:Pass",
+    ],
+)
+
+cc_library(
+    name = "import_model",
+    srcs = ["import_model.cc"],
+    hdrs = ["import_model.h"],
+    deps = [
+        ":assign_op_key",
+        ":passes",
+        ":while_to_map_fn",
+        "//tensorflow/compiler/mlir/tensorflow:bridge_logger",
+        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tfrt:export",
+        "//tensorflow/compiler/mlir/tfrt:import_model",
+        "//tensorflow/compiler/mlir/tfrt:tf_to_tfrt",
+        "//tensorflow/compiler/mlir/tfrt:tfrt_compile_options",
+        "//tensorflow/compiler/mlir/tfrt:tfrt_pipeline_options",
+        "//tensorflow/compiler/mlir/tfrt/translate/mlrt:mlir_to_bytecode",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:statusor",
+        "//tensorflow/core/tfrt/fallback:cost_recorder",
+        "//tensorflow/core/tfrt/fallback:fallback_state",
+        "//tensorflow/core/tfrt/mlrt/attribute",
+        "//tensorflow/core/tfrt/mlrt/bytecode",
+        "//tensorflow/core/tfrt/runtime",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "while_to_map_fn",
+    srcs = ["while_to_map_fn.cc"],
+    hdrs = ["while_to_map_fn.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow/ir/host_runtime:tensorflow_tfrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+cc_library(
+    name = "async_while",
+    srcs = ["async_while.cc"],
+    hdrs = ["async_while.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+cc_library(
+    name = "rewrite_ifrt_load_variable",
+    srcs = ["rewrite_ifrt_load_variable.cc"],
+    hdrs = ["rewrite_ifrt_load_variable.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow/ir/host_runtime:tensorflow_tfrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+cc_library(
+    name = "ifrt_set_tpu_host_allocator",
+    srcs = ["ifrt_set_tpu_host_allocator.cc"],
+    hdrs = ["ifrt_set_tpu_host_allocator.h"],
+    deps = [
+        ":mlrt_device_constants",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow/ir/host_runtime:tensorflow_tfrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)

tensorflow/compiler/mlir/tfrt/transforms/mlrt/assign_op_key.cc

@@ -0,0 +1,71 @@
+/* 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/tfrt/transforms/mlrt/assign_op_key.h"
+
+#include <stdint.h>
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/constants.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/util.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+namespace {
+
+class AssignOpKeyPass
+    : public mlir::PassWrapper<AssignOpKeyPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  AssignOpKeyPass() = default;
+  AssignOpKeyPass& operator=(const AssignOpKeyPass&) = delete;
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AssignOpKeyPass)
+
+ private:
+  llvm::StringRef getArgument() const final { return "tf-mlrt-assign-op-key"; }
+  llvm::StringRef getDescription() const final {
+    return "tf-mlrt-assign-op-key";
+  }
+
+  void runOnOperation() override;
+};
+
+void AssignOpKeyPass::runOnOperation() {
+  auto module = getOperation();
+  mlir::OpBuilder builder(module);
+
+  int32_t op_key = 0;
+  module.walk([&builder, &op_key](mlir::Operation* op) mutable {
+    if (UseFallback(op)) {
+      op->setAttr(tensorflow::tfrt_compiler::kOpKeyAttrName,
+                  builder.getI32IntegerAttr(op_key));
+      op_key++;
+    }
+  });
+}
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateAssignOpKeyPass() {
+  return std::make_unique<AssignOpKeyPass>();
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/assign_op_key.h

@@ -0,0 +1,32 @@
+/* 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_TFRT_TRANSFORMS_MLRT_ASSIGN_OP_KEY_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_ASSIGN_OP_KEY_H_
+#include <memory>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+// Create a pass that assigns an op_key to every fallback OP. The op_key
+// provides a uniform key to look up online cost for a specific op.
+// This pass is expected to run before parallerization.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateAssignOpKeyPass();
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_ASSIGN_OP_KEY_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/async_while.cc

@@ -0,0 +1,745 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/async_while.h"
+
+#include <linux/limits.h>
+
+#include <algorithm>
+#include <deque>
+#include <list>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/check.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Block.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/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+namespace {
+
+// Move await right before its tensor is used.
+void MoveTFAwaitOpToFirstUse(mlir::Block &block) {
+  llvm::SmallVector<tf_mlrt::TFAwaitOp> await_ops;
+  for (auto &op : block) {
+    if (auto await_op = llvm::dyn_cast<tf_mlrt::TFAwaitOp>(&op)) {
+      await_ops.push_back(await_op);
+    }
+  }
+
+  for (auto op : await_ops) {
+    auto result = op.getResult();
+    if (result.use_empty()) continue;
+
+    mlir::Operation *first_user = *result.user_begin();
+    for (auto *user : result.getUsers()) {
+      if (user->isBeforeInBlock(first_user)) {
+        first_user = user;
+      }
+    }
+
+    op->moveBefore(first_user);
+  }
+}
+
+// Move promise right after the tensor becomes available.
+void MoveTFPromiseOpRightAfterAvailable(mlir::Block &block) {
+  llvm::SmallVector<tf_mlrt::TFPromiseOp> promise_ops;
+  for (auto &op : block) {
+    if (auto promise_op = llvm::dyn_cast<tf_mlrt::TFPromiseOp>(&op)) {
+      promise_ops.push_back(promise_op);
+    }
+  }
+
+  for (auto promise_op : promise_ops) {
+    auto tensor = promise_op.getOperand(1);
+    // Promise is on variant tensor, which must have a defining op inside the
+    // block.
+    DCHECK(tensor.getDefiningOp() != nullptr);
+    promise_op->moveAfter(tensor.getDefiningOp());
+  }
+}
+
+// Converts applicable tf.While to tf_mlrt.AsyncWhile.
+//
+// In pseudo code, the following
+//
+//  %res0, %res1 = tf.While(%arg0, %arg1)(body = @while_body, cond =
+//  @while_predicate)
+//
+// is converted to
+//
+//  %0 = tf.PartitionedCall(%res0) (f =
+//  @while_predicate/TfMlrtAsyncWhilePredicate) %1, %2, %3 =
+//  tf_mlrt.async_while(%0, %arg0, %arg1) (body =
+//  @while_body/TfMlrtAsyncWhileBody, immutable_size = 1) %res0 =
+//  tf_mlrt.await(%2) %res1 = tf_mlrt.await(%3)
+//
+class AsyncWhilePass
+    : public mlir::PassWrapper<AsyncWhilePass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  AsyncWhilePass() = default;
+  AsyncWhilePass &operator=(const AsyncWhilePass &) = delete;
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AsyncWhilePass)
+
+ private:
+  // The input to while op will be re-order and supplied as input to async_while
+  // op.
+  struct ArgumentRemap {
+    std::vector<int> new_to_old_remap;
+    std::vector<int> old_to_new_remap;
+  };
+
+  // The input to while op will grouped into immutables(invariants) between
+  // iterations and mutables(variants).
+  struct ArgumentGroups {
+    std::vector<int> immutables;
+    std::vector<int> mutables;
+  };
+
+  // First argument of the new async_while is a predicate.
+  static constexpr int kNonPredicateStartIndex = 1;
+
+  void getDependentDialects(mlir::DialectRegistry &registry) const override {
+    registry.insert<tensorflow::tf_mlrt::TensorflowMlrtDialect>();
+    registry.insert<mlrt::compiler::MlrtDialect>();
+  }
+
+  llvm::StringRef getArgument() const final { return "tf-mlrt-async-while"; }
+
+  llvm::StringRef getDescription() const final {
+    return "Convert tf.while to tf_mlrt.async_while when possible for parallel "
+           "execution.";
+  }
+
+  void runOnOperation() override {
+    mlir::ModuleOp module = getOperation();
+    mlir::SymbolTable symbol_table(module);
+    mlir::TF::SideEffectAnalysis side_effect_analysis(module);
+
+    // Keep a record of predicate and body functions that we already used to
+    // create new predicate and body functions.
+    // key is the original function name and value is the newly created
+    // function.
+    // We keep our own record instead of relying on SymbolTable to know whether
+    // a function is already converted is because SymbolTable might create new
+    // symbol name for inserted function to avoid conflict.
+    llvm::SmallDenseMap<llvm::StringRef, mlir::func::FuncOp>
+        processed_functions;
+    // Use make_early_inc_range because the processing might insert new node
+    // into the list
+    for (auto func_op :
+         llvm::make_early_inc_range(module.getOps<mlir::func::FuncOp>())) {
+      MayConvertWhileToAsyncWhile(
+          func_op, symbol_table, processed_functions,
+          side_effect_analysis.GetAnalysisForFunc(func_op));
+    }
+  }
+
+  void MayConvertWhileToAsyncWhile(
+      mlir::func::FuncOp op, mlir::SymbolTable &symbol_table,
+      llvm::SmallDenseMap<llvm::StringRef, mlir::func::FuncOp>
+          &processed_functions,
+      const mlir::TF::SideEffectAnalysis::Info &side_effect_analysis) {
+    mlir::OpBuilder builder(op);
+    for (mlir::Operation &op : llvm::make_early_inc_range(op.front())) {
+      auto while_op = llvm::dyn_cast<mlir::TF::WhileOp>(&op);
+      if (!while_op) continue;
+
+      // Fill in bodies
+      mlir::func::FuncOp body_fn =
+          symbol_table.lookup<mlir::func::FuncOp>(while_op.getBody());
+
+      mlir::func::FuncOp predicate_fn =
+          symbol_table.lookup<mlir::func::FuncOp>(while_op.getCond());
+
+      if (ShouldConvertWhileToAsyncWhile(body_fn, predicate_fn)) {
+        std::vector<int> predicate_determinat =
+            GetPredicateDeterminat(predicate_fn);
+
+        ArgumentGroups argument_groups = GetArgumentGroups(body_fn);
+        ArgumentRemap argument_remap =
+            CreateArgumentRemap(body_fn, argument_groups.immutables);
+
+        mlir::func::FuncOp new_predicate_fn =
+            CreatePredicate(builder, symbol_table, processed_functions,
+                            predicate_fn, predicate_determinat);
+
+        mlir::func::FuncOp new_body =
+            CreateBody(builder, symbol_table, processed_functions, body_fn,
+                       new_predicate_fn, predicate_determinat, argument_groups,
+                       argument_remap, side_effect_analysis);
+
+        mlir::OpBuilder::InsertionGuard insertion_guard(builder);
+        builder.setInsertionPoint(while_op);
+
+        // Call the predicate function first. async_while requires the first
+        // predicate to be given.
+        llvm::SmallVector<mlir::Value> predicate_input;
+        for (int i : predicate_determinat) {
+          predicate_input.push_back(while_op->getOperand(i));
+        }
+
+        auto empty_string_attr = builder.getStringAttr("");
+        auto init_predicate = mlir::TF::PartitionedCallOp::create(
+            builder, while_op->getLoc(), predicate_fn.getResultTypes()[0],
+            predicate_input, /*args_attrs=*/nullptr, /*res_attrs=*/nullptr,
+            mlir::FlatSymbolRefAttr::get(new_predicate_fn.getSymNameAttr()),
+            empty_string_attr, empty_string_attr, empty_string_attr);
+
+        // These arguments exclude predicate.
+        llvm::SmallVector<mlir::Value> async_while_input;
+        async_while_input.reserve(while_op->getNumOperands());
+        for (auto old_arg_idx_iter =
+                 argument_remap.new_to_old_remap.begin() + 1;
+             old_arg_idx_iter != argument_remap.new_to_old_remap.end();
+             ++old_arg_idx_iter) {
+          async_while_input.push_back(while_op.getOperand(*old_arg_idx_iter));
+        }
+
+        llvm::SmallVector<mlir::Type> async_while_output_types;
+        async_while_output_types.resize(1 + while_op->getNumResults());
+        std::fill(async_while_output_types.begin(),
+                  async_while_output_types.end(),
+                  builder.getType<mlrt::compiler::FutureType>());
+
+        auto async_while_op = tf_mlrt::TFAsyncWhileOp::create(
+            builder, while_op.getLoc(), async_while_output_types,
+            init_predicate->getResult(0), async_while_input,
+            new_body.getSymName(), argument_groups.immutables.size());
+
+        for (int i = 0; i < while_op.getResults().size(); ++i) {
+          if (!while_op.getResult(i).use_empty()) {
+            auto async_while_result = tf_mlrt::TFAwaitOp::create(
+                builder, while_op.getLoc(), while_op->getResultTypes()[i],
+                async_while_op.getResult(argument_remap.old_to_new_remap[i]));
+            while_op.getResult(i).replaceAllUsesWith(
+                async_while_result.getResult());
+          }
+        }
+
+        while_op.erase();
+      }
+    }
+
+    // TFAwaitOp may have been inserted, move them to first use.
+    MoveTFAwaitOpToFirstUse(op.getBody().front());
+  }
+
+  // Creates a new predicate function that returns a boolean tensor. This new
+  // predicate function is converted from the original predicate function after
+  // removing the unused input arguments. We want the inputs to the new
+  // predicate function to be the minimal set because the new body function
+  // wants to call the predicate as early as possible to maximize the chances of
+  // parallerization. Also since the new body function reorders the input
+  // arguments, the new predicate function follows that new order to be
+  // consistent (just like the original predicate function and body function
+  // have the same order of arguments).
+  mlir::func::FuncOp CreatePredicate(
+      mlir::OpBuilder &builder, mlir::SymbolTable &symbol_table,
+      llvm::SmallDenseMap<llvm::StringRef, mlir::func::FuncOp>
+          &processed_functions,
+      mlir::func::FuncOp original_predicate_fn,
+      const std::vector<int> &predicate_determinat) {
+    auto &processed_function =
+        processed_functions[original_predicate_fn.getName()];
+    if (processed_function) {
+      return processed_function;
+    }
+    mlir::OpBuilder::InsertionGuard insertion_guard(builder);
+    builder.setInsertionPointAfter(original_predicate_fn);
+
+    // Input to the new predicate only contains arguments listed in
+    // predicate_determinat.
+    llvm::SmallVector<mlir::Type> remapped_input_types;
+    for (int i : predicate_determinat) {
+      remapped_input_types.push_back(
+          original_predicate_fn.getArgumentTypes()[i]);
+    }
+
+    std::string new_predicate_fn_name = absl::StrCat(
+        original_predicate_fn.getName().str(), "/TfMlrtAsyncWhilePredicate");
+    mlir::func::FuncOp new_predicate_fn = mlir::func::FuncOp::create(
+        builder, original_predicate_fn->getLoc(), new_predicate_fn_name,
+        mlir::FunctionType::get(original_predicate_fn.getContext(),
+                                remapped_input_types,
+                                original_predicate_fn.getResultTypes()));
+
+    processed_function = new_predicate_fn;
+
+    new_predicate_fn.setPrivate();
+
+    builder.setInsertionPointToEnd(new_predicate_fn.addEntryBlock());
+
+    // Inputs will be reordered according to the body functions' reorder.
+    mlir::IRMapping remap;
+    int j = 0;
+    for (int i : predicate_determinat) {
+      remap.map(original_predicate_fn.getArgument(i),
+                new_predicate_fn.getArgument(j));
+      j++;
+    }
+
+    for (auto &op_it : original_predicate_fn.getBody().front()) {
+      builder.clone(op_it, remap);
+    }
+
+    symbol_table.insert(new_predicate_fn);
+    return new_predicate_fn;
+  }
+
+  // Creates the new body function for the async while.
+  // The new body function has the input signature of
+  // (predicate_promise, mutable0_future, mutable0_promise, mutable1_future,
+  // mutable1_promise, ..., immutable0, immutable1,...) The new body is mostly a
+  // clone from the old body function, but the input arguments are reordered
+  // such that the immutables(invariants) between iterations are always at the
+  // bottom of the argument list. Additionaly, the new body function calls the
+  // new predicate function and puts the boolean predicate for the next
+  // iteration into the predicate_promise.
+  mlir::func::FuncOp CreateBody(
+      mlir::OpBuilder &builder, mlir::SymbolTable &symbol_table,
+      llvm::SmallDenseMap<llvm::StringRef, mlir::func::FuncOp>
+          &processed_functions,
+      mlir::func::FuncOp original_body_fn, mlir::func::FuncOp new_predicate_fn,
+      const std::vector<int> &predicate_determinat,
+      const ArgumentGroups &argument_groups,
+      const ArgumentRemap &argument_remap,
+      const mlir::TF::SideEffectAnalysis::Info &side_effect_analysis) {
+    auto &processed_function = processed_functions[original_body_fn.getName()];
+    if (processed_function) {
+      return processed_function;
+    }
+
+    // Safe to do topology sort here b/c there are no promise/future involved
+    // here.
+    SortBodyFnByFirstReadyFirstRun(original_body_fn, predicate_determinat,
+                                   side_effect_analysis);
+
+    mlir::OpBuilder::InsertionGuard insertion_guard(builder);
+    builder.setInsertionPointAfter(original_body_fn);
+
+    // Create the new body function signature.
+    std::vector<mlir::Type> remapped_input_types;
+    std::vector<mlir::Type> remapped_output_types;
+
+    // predicate promise is the first argument.
+    remapped_input_types.push_back(
+        builder.getType<mlrt::compiler::PromiseType>());
+
+    // future and promise pairs for the mutables follows.
+    for (int i = 0; i < argument_groups.mutables.size(); ++i) {
+      remapped_input_types.push_back(
+          builder.getType<mlrt::compiler::FutureType>());
+      remapped_input_types.push_back(
+          builder.getType<mlrt::compiler::PromiseType>());
+    }
+
+    // immutables are always at the bottom of the argument list.
+    for (int i : argument_groups.immutables) {
+      remapped_input_types.push_back(original_body_fn.getArgumentTypes()[i]);
+    }
+
+    std::string new_body_fn_name =
+        absl::StrCat(original_body_fn.getName().str(), "/TfMlrtAsyncWhileBody");
+    auto body_fn = mlir::func::FuncOp::create(
+        builder, original_body_fn->getLoc(), new_body_fn_name,
+        mlir::FunctionType::get(original_body_fn.getContext(),
+                                remapped_input_types, remapped_output_types));
+    processed_function = body_fn;
+
+    body_fn.setPrivate();
+
+    // Inserts await for futures so that all mutable tensors are ready.
+    // We will move these awaits to first use location later on.
+    builder.setInsertionPointToEnd(body_fn.addEntryBlock());
+    // First future starts at index 1 after predicate promise.
+    int future_index = 1;
+    mlir::IRMapping mapping;
+    int body_arg_idx = kNonPredicateStartIndex;
+    for (int i : argument_groups.mutables) {
+      auto future_value = tf_mlrt::TFAwaitOp::create(
+          builder, body_fn->getLoc(), original_body_fn.getArgumentTypes()[i],
+          body_fn.getArgument(future_index));
+      mapping.map(original_body_fn.getArgument(i), future_value);
+      // Move by 2 b/c each variant correspond to one future and one promise
+      body_arg_idx += 2;
+      future_index += 2;
+    }
+
+    for (int i : argument_groups.immutables) {
+      mapping.map(original_body_fn.getArgument(i),
+                  body_fn.getArgument(body_arg_idx));
+      body_arg_idx++;
+    }
+
+    // All future values are ready; so we can clone from the original body
+    // function.
+    for (auto &op : original_body_fn.getBody().front()) {
+      builder.clone(op, mapping);
+    }
+
+    auto return_op = body_fn.getBody().front().getTerminator();
+
+    // Find the earliest location that the new predicate function can execute.
+    auto *body_block = &body_fn.getBody().front();
+    auto *earliest_op = &body_block->front();
+    for (int i = 0; i < return_op->getOperands().size(); ++i) {
+      if (std::find(predicate_determinat.begin(), predicate_determinat.end(),
+                    i) != predicate_determinat.end()) {
+        auto latest_op = return_op->getOperands()[i].getDefiningOp();
+        if (latest_op != nullptr && earliest_op->isBeforeInBlock(latest_op)) {
+          earliest_op = latest_op;
+        }
+      }
+    }
+
+    llvm::SmallVector<mlir::Value> predicate_inputs;
+    for (int i : predicate_determinat) {
+      predicate_inputs.push_back(return_op->getOperands()[i]);
+    }
+    // Call new predicate when all dependencies are ready.
+    builder.setInsertionPointAfter(earliest_op);
+    llvm::SmallVector<mlir::Type> predicate_returned_types;
+    auto empty_string_attr = builder.getStringAttr("");
+    auto predicate_op = mlir::TF::PartitionedCallOp::create(
+        builder, earliest_op->getLoc(), new_predicate_fn.getResultTypes(),
+        predicate_inputs, /*args_attrs=*/nullptr, /*res_attrs=*/nullptr,
+        mlir::FlatSymbolRefAttr::get(new_predicate_fn.getSymNameAttr()),
+        empty_string_attr, empty_string_attr, empty_string_attr);
+
+    // Insert promise right before return
+    builder.setInsertionPoint(return_op);
+    // builder.create<mlrt::compiler::PromiseOp>(return_op->getLoc());
+    auto predicate_tensor = predicate_op->getResult(0);
+
+    const int kPredicatePromiseIndex = 0;
+    tf_mlrt::TFPromiseOp::create(builder, return_op->getLoc(),
+                                 body_fn.getArgument(kPredicatePromiseIndex),
+                                 predicate_tensor);
+    // First variant promise start at 2 as [predicate_promise, arg0_future,
+    // arg0_promise, ...]
+    int promise_index = 2;
+    for (int i : argument_groups.mutables) {
+      tf_mlrt::TFPromiseOp::create(builder, return_op->getLoc(),
+                                   body_fn.getArgument(promise_index),
+                                   return_op->getOperand(i));
+      promise_index += 2;
+    }
+    mlir::func::ReturnOp::create(builder, return_op->getLoc());
+    return_op->erase();
+
+    // Reorder await and promise to reduce blocking waits.
+    MoveTFAwaitOpToFirstUse(body_fn.getBody().front());
+    MoveTFPromiseOpRightAfterAvailable(body_fn.getBody().front());
+
+    symbol_table.insert(body_fn);
+    return body_fn;
+  }
+
+  // Decides whether we should convert a while to async_while.
+  // This determination is largely by heuristic.
+  bool ShouldConvertWhileToAsyncWhile(mlir::func::FuncOp body_fn,
+                                      mlir::func::FuncOp predicate_fn) {
+    // If the predicate are NOT determined by simple ops such as (less)
+    // comparison and logic and, do not apply conversion.
+    for (auto &op : predicate_fn.getBody().front()) {
+      if (!llvm::isa<mlir::TF::ConstOp, mlir::TF::LessOp, mlir::TF::IdentityOp,
+                     mlir::TF::LogicalAndOp, mlir::TF::ToBoolOp,
+                     mlir::func::ReturnOp>(&op))
+        return false;
+    }
+
+    // We then apply some heuristics:
+    // 1. If all variants are needed to determine the predicate, do not convert
+    // to async while because the next iteration has a low chance of overlapping
+    // with the current iteration.
+    // 2. The predicate shall be updated before any other variants is updated.
+    std::vector<int> predicate_determinat =
+        GetPredicateDeterminat(predicate_fn);
+    ArgumentGroups argument_groups = GetArgumentGroups(body_fn);
+
+    // Find out defining ops for predicate and non-predicate inside the body.
+    llvm::SmallSetVector<mlir::Operation *, 4> predicate_determinant_ops;
+    llvm::SmallSetVector<mlir::Operation *, 4> worker_ops;
+    for (int i : argument_groups.mutables) {
+      auto *defining_op = body_fn.getBody()
+                              .front()
+                              .getTerminator()
+                              ->getOperand(i)
+                              .getDefiningOp();
+      if (std::find(predicate_determinat.begin(), predicate_determinat.end(),
+                    i) != predicate_determinat.end()) {
+        predicate_determinant_ops.insert(defining_op);
+      } else {
+        worker_ops.insert(defining_op);
+      }
+    }
+
+    // simple heuristic to reject pipelining: predicate are determined by all
+    // variants.
+    if (worker_ops.empty()) return false;
+
+    // Validate that the determinant of predicate must be updated before any
+    // other variants and variants are determined in a cluster.
+    int predicate_determinant_cnt = 0;
+    for (auto &body_op : body_fn.getBlocks().front()) {
+      if (predicate_determinant_cnt < predicate_determinant_ops.size()) {
+        if (worker_ops.contains(&body_op)) {
+          return false;
+        } else if (predicate_determinant_ops.contains(&body_op)) {
+          predicate_determinant_cnt++;
+          if (predicate_determinant_cnt >= predicate_determinant_ops.size()) {
+            break;
+          }
+        }
+      }
+    }
+
+    return true;
+  }
+
+  // Returns the list of minimal argument indices that determines the predicate.
+  std::vector<int> GetPredicateDeterminat(mlir::func::FuncOp predicate_fn) {
+    std::vector<int> determinat;
+    for (auto &op : predicate_fn.getArguments()) {
+      if (!op.use_empty()) {
+        determinat.push_back(op.getArgNumber());
+      }
+    }
+    return determinat;
+  }
+
+  // Groups the arguments of the original body function into two groups:
+  // mutables between iterations and immutables across iterations.
+  ArgumentGroups GetArgumentGroups(mlir::func::FuncOp body_fn) {
+    ArgumentGroups groups;
+    auto *return_op = body_fn.getBlocks().front().getTerminator();
+
+    int i = 0;
+    for (auto return_op_operand : return_op->getOperands()) {
+      if (std::find(body_fn.getArguments().begin(),
+                    body_fn.getArguments().end(),
+                    return_op_operand) != body_fn.getArguments().end()) {
+        groups.immutables.push_back(i);
+      } else {
+        groups.mutables.push_back(i);
+      }
+      i++;
+    }
+    return groups;
+  }
+
+  // Re-orders the argument such that all the immutables are at the bottom of
+  // the argument list. This function returns a mapping for such re-order.
+  ArgumentRemap CreateArgumentRemap(mlir::func::FuncOp original_body_fn,
+                                    std::vector<int> invariant) {
+    absl::flat_hash_set<int> invariant_set(invariant.begin(), invariant.end());
+    ArgumentRemap argument_remap;  // new to old
+    argument_remap.old_to_new_remap.resize(
+        original_body_fn.getArguments().size());
+    argument_remap.new_to_old_remap.resize(
+        original_body_fn.getArguments().size() + 1);
+
+    const int invariant_start_index =
+        original_body_fn.getArguments().size() - invariant_set.size();
+    int variant_index = 0 + kNonPredicateStartIndex;
+    int invariant_index = invariant_start_index + kNonPredicateStartIndex;
+    for (int i = 0; i < original_body_fn.getArguments().size(); ++i) {
+      if (invariant_set.contains(i)) {
+        argument_remap.old_to_new_remap[i] = invariant_index;
+        argument_remap.new_to_old_remap[invariant_index] = i;
+        invariant_index++;
+      } else {
+        argument_remap.old_to_new_remap[i] = variant_index;
+        argument_remap.new_to_old_remap[variant_index] = i;
+        variant_index++;
+      }
+    }
+    DCHECK_EQ(variant_index, invariant_start_index + kNonPredicateStartIndex);
+    DCHECK_EQ(invariant_index,
+              original_body_fn.getArguments().size() + kNonPredicateStartIndex);
+    return argument_remap;
+  }
+
+  // Sorts the operations in the block to an order of first ready first run to
+  // increase the parallelism between iterations. The order of ops before the
+  // predicate can be computed remains unchanged because we prefer to have the
+  // next iteration starts as soon as possible. For ops that are
+  // ready at the same time, ops producing the final results will be executed
+  // first.
+  void SortBodyFnByFirstReadyFirstRun(
+      mlir::func::FuncOp body_fn, const std::vector<int> &predicate_determinat,
+      const mlir::TF::SideEffectAnalysis::Info &side_effect_analysis) {
+    mlir::Block &body_block = body_fn.getBody().front();
+    mlir::Operation *return_op = body_block.getTerminator();
+    if (!return_op) return;
+
+    // Ops in final execution order.
+    std::vector<mlir::Operation *> sorted_ops;
+    // Number of outstanding (not executed yet) dependency for an op. When that
+    // number hits zero, this op is ready to be executed.
+    llvm::DenseMap<mlir::Operation *, int> dependency_cnt;
+
+    // Build the graph and mark ops that depends on input arguments in the
+    // `new_ready_ops`. We will prioritize final result producing ops later on
+    // before finalizing their order.
+    std::list<mlir::Operation *> new_ready_ops;
+    for (mlir::Operation &op : body_block.without_terminator()) {
+      // Existing side effecting ops are dependencies on all subsequent ops.
+      dependency_cnt[&op] =
+          op.getNumOperands() +
+          side_effect_analysis.DirectControlPredecessors(&op).size();
+
+      // Block argument is viewed as ready right away
+      for (const mlir::Value &operand : op.getOperands()) {
+        if (!operand.getDefiningOp()) {
+          dependency_cnt[&op]--;
+        }
+      }
+      if (!dependency_cnt[&op]) new_ready_ops.push_back(&op);
+    }
+
+    // Ops that produces final returned results have priority among ops that are
+    // ready at the same time.
+    llvm::DenseSet<mlir::Operation *> final_result_producing_ops;
+    for (const auto &operand : return_op->getOperands()) {
+      if (auto *op = operand.getDefiningOp()) {
+        final_result_producing_ops.insert(op);
+      }
+    }
+
+    // The order of all ops before the predicate for the next iteration
+    // can be computed remains unchanged because we want to kick off next
+    // iteration preparation asap by making the predicate ready asap.
+    llvm::SmallDenseSet<mlir::Operation *> predicate_determined_ops;
+    for (int i : predicate_determinat) {
+      if (auto *defining_op = return_op->getOperand(i).getDefiningOp()) {
+        predicate_determined_ops.insert(defining_op);
+      }
+    }
+    auto iter = body_block.without_terminator().begin();
+    // Make those ops before predicate ready and remember them so that they
+    // don't get resorted.
+    llvm::SmallDenseSet<mlir::Operation *> pre_scheduled_ops;
+    while (!predicate_determined_ops.empty()) {
+      if (predicate_determined_ops.contains(&*iter)) {
+        predicate_determined_ops.erase(&*iter);
+      }
+      sorted_ops.push_back(&*iter);
+      pre_scheduled_ops.insert(&*iter);
+      iter++;
+    }
+
+    // Topologically sorts the ops in a breadth first order.
+    std::deque<mlir::Operation *> ready_ops;
+    MoveNewReadyOps(ready_ops, std::move(new_ready_ops),
+                    final_result_producing_ops);
+    while (!ready_ops.empty()) {
+      mlir::Operation *active_op = ready_ops.front();
+      if (!pre_scheduled_ops.contains(active_op))
+        sorted_ops.push_back(active_op);
+      ready_ops.pop_front();
+      // Do not directly push new ready ops into ready_ops yet.
+      // We will re-arrange the new ready ops by moving those producing the
+      // final results to the top so that they get executed first.
+      std::list<mlir::Operation *> new_ready_ops;
+      for (const auto &result : active_op->getOpResults()) {
+        // Uses instead of users in case users can dedup.
+        for (const auto &use : result.getUses()) {
+          int cnt = --dependency_cnt[use.getOwner()];
+          if (!cnt) new_ready_ops.push_back(use.getOwner());
+        }
+      }
+      for (auto *control_dependent :
+           side_effect_analysis.DirectControlSuccessors(active_op)) {
+        int cnt = --dependency_cnt[control_dependent];
+        if (!cnt) new_ready_ops.push_back(control_dependent);
+      }
+
+      MoveNewReadyOps(ready_ops, std::move(new_ready_ops),
+                      final_result_producing_ops);
+    }
+
+    // Update the block with the new order.
+    sorted_ops.push_back(return_op);
+    for (mlir::Operation *op : sorted_ops) {
+      op->remove();
+      body_block.push_back(op);
+    }
+  }
+
+  // Add ops in the `new_ready_ops` to `ready_ops` in the following order:
+  // ops that produces final returned result are added first and then the rest
+  // ops.
+  void MoveNewReadyOps(
+      std::deque<mlir::Operation *> &ready_ops,
+      std::list<mlir::Operation *> new_ready_ops,
+      const llvm::DenseSet<mlir::Operation *> &final_result_producing_ops) {
+    auto iter = new_ready_ops.begin();
+    while (iter != new_ready_ops.end()) {
+      if (final_result_producing_ops.contains(*iter)) {
+        ready_ops.push_back(*iter);
+        iter = new_ready_ops.erase(iter);
+      } else {
+        iter++;
+      }
+    }
+    for (mlir::Operation *op : new_ready_ops) {
+      ready_ops.push_back(op);
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateAsyncWhilePass() {
+  return std::make_unique<AsyncWhilePass>();
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/async_while.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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_ASYNC_WHILE_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_ASYNC_WHILE_H_
+
+#include <memory>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+// Creates a pass that converts applicable tf.While to tf_mlrt.AsyncWhile.
+// tf_mlrt.AsyncWhile dispatch iterations asynchronously, thus allowing
+// pipelining between iterations to reduce latency. This is intended for
+// tf.While that is not converted from tf.MapFn, but still can benefit from
+// asynchronous execution of iterations to reduce latency.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateAsyncWhilePass();
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_ASYNC_WHILE_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/execute_op_registry.h

@@ -0,0 +1,60 @@
+/* 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_TFRT_TRANSFORMS_MLRT_EXECUTE_OP_REGISTRY_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_EXECUTE_OP_REGISTRY_H_
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+class ExecuteOpRegistry {
+ public:
+  mlir::LogicalResult RegisterExecuteOp(mlir::Operation* op, uint32_t op_key) {
+    if (op_key >= execute_ops_.size()) {
+      execute_ops_.resize(op_key + 1);
+    }
+    if (auto* register_op = execute_ops_[op_key]) {
+      if (register_op->getName() != op->getName() ||
+          register_op->getAttrs() != op->getAttrs()) {
+        return op->emitError() << "Key " << op_key << " already registered.";
+      }
+      return mlir::success();
+    }
+    execute_ops_[op_key] = op;
+    return mlir::success();
+  }
+
+  void ReplaceExecuteOp(int64_t key, mlir::Operation* op) {
+    execute_ops_[key] = op;
+  }
+
+  llvm::ArrayRef<mlir::Operation*> GetExecuteOps() const {
+    return execute_ops_;
+  }
+
+ private:
+  // Using a vector to keep fallback ops in order, and the key for a fallback op
+  // is its corresponding index here.
+  llvm::SmallVector<mlir::Operation*, 8> execute_ops_;
+};
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_EXECUTE_OP_REGISTRY_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/fuse_mlrt_ops.cc

@@ -0,0 +1,175 @@
+/* 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/tfrt/transforms/mlrt/fuse_mlrt_ops.h"
+
+#include <memory>
+
+#include "llvm/ADT/SmallVector.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+namespace {
+
+class FuseMlrtOpPass
+    : public mlir::PassWrapper<FuseMlrtOpPass,
+                               mlir::OperationPass<mlir::func::FuncOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FuseMlrtOpPass)
+
+ private:
+  llvm::StringRef getArgument() const final { return "tf-mlrt-fuse"; }
+
+  llvm::StringRef getDescription() const final {
+    return "Fuse consecutive mlrt ops of the same kind into one.";
+  }
+
+  void runOnOperation() override;
+};
+
+void FuseGetResourceOps(mlir::OpBuilder& builder, mlir::Block& block) {
+  llvm::SmallVector<tf_mlrt::GetResourceOp> get_resource_ops;
+  for (auto& op : llvm::make_early_inc_range(block)) {
+    if (auto get_resource_op = llvm::dyn_cast<tf_mlrt::GetResourceOp>(&op)) {
+      get_resource_ops.push_back(get_resource_op);
+    }
+  }
+
+  if (get_resource_ops.empty()) return;
+
+  // The last op is always a return op, so it is guaranteed to process all
+  // groups of the candidate ops.
+  auto first_get = get_resource_ops.front();
+
+  builder.setInsertionPointAfter(first_get);
+
+  llvm::SmallVector<mlir::Attribute> indices;
+  llvm::SmallVector<mlir::Type> result_types;
+  llvm::SmallVector<mlir::Value> old_values;
+
+  indices.reserve(get_resource_ops.size());
+  result_types.reserve(get_resource_ops.size());
+  old_values.reserve(get_resource_ops.size());
+
+  for (auto op : get_resource_ops) {
+    auto indices_attr = op.getIndices();
+    indices.append(indices_attr.begin(), indices_attr.end());
+    result_types.append(op.result_type_begin(), op.result_type_end());
+    old_values.append(op.result_begin(), op.result_end());
+  }
+
+  auto new_op = tf_mlrt::GetResourceOp::create(
+      builder, first_get.getLoc(), result_types, builder.getArrayAttr(indices));
+
+  for (auto [old_value, new_value] :
+       llvm::zip(old_values, new_op.getResults())) {
+    old_value.replaceAllUsesWith(new_value);
+  }
+
+  for (auto get_resource_op : get_resource_ops) {
+    get_resource_op->erase();
+  }
+}
+
+template <typename AwaitOpType, typename AwaitAllOpType,
+          typename ValueType = void>
+void FuseAwaitOps(mlir::OpBuilder& builder, mlir::Block& block) {
+  llvm::SmallVector<AwaitOpType> await_ops;
+  for (auto& op : llvm::make_early_inc_range(block)) {
+    if (auto await_op = llvm::dyn_cast<AwaitOpType>(&op)) {
+      await_ops.push_back(await_op);
+      continue;
+    }
+
+    // The last op is always a return op, so it is guaranteed to process all
+    // groups of the candidate ops.
+    if (await_ops.size() > 1) {
+      auto last_await = await_ops.back();
+
+      builder.setInsertionPointAfter(last_await);
+
+      llvm::SmallVector<mlir::Value> futures;
+      futures.reserve(await_ops.size());
+      for (auto op : await_ops) {
+        futures.push_back(op.getOperand());
+      }
+
+      llvm::SmallVector<mlir::Type> result_types;
+      if constexpr (!std::is_same_v<ValueType, void>) {
+        result_types.assign(futures.size(), builder.getType<ValueType>());
+      }
+
+      auto await_all =
+          AwaitAllOpType::create(builder, op.getLoc(), result_types, futures);
+
+      if constexpr (!std::is_same_v<ValueType, void>) {
+        for (auto [await_op, new_value] :
+             llvm::zip(await_ops, await_all.getResults())) {
+          await_op.getResult().replaceAllUsesWith(new_value);
+        }
+      }
+
+      for (auto await_op : await_ops) {
+        await_op->erase();
+      }
+    }
+
+    await_ops.clear();
+  }
+}
+
+void FusePromiseReturn(mlir::OpBuilder& builder, mlir::Block& block) {
+  auto* terminator = block.getTerminator();
+  auto return_op = llvm::dyn_cast<mlir::func::ReturnOp>(terminator);
+  if (!return_op || return_op->getNumOperands() > 0) return;
+
+  auto promise_op =
+      llvm::dyn_cast_or_null<tf_mlrt::PromiseOp>(return_op->getPrevNode());
+  if (!promise_op) return;
+
+  builder.setInsertionPointAfter(return_op);
+  tf_mlrt::PromiseReturnOp::create(builder, return_op->getLoc(),
+                                   promise_op->getResultTypes(),
+                                   promise_op->getOperands());
+  return_op->erase();
+  promise_op->erase();
+}
+
+void FuseMlrtOpPass::runOnOperation() {
+  auto func = getOperation();
+
+  mlir::OpBuilder builder(func);
+
+  FuseAwaitOps<tf_mlrt::AwaitOp, tf_mlrt::AwaitAllOp, tf_mlrt::TFTensorType>(
+      builder, func.front());
+  FuseAwaitOps<mlrt::compiler::AwaitHandleOp, mlrt::compiler::AwaitAllHandleOp>(
+      builder, func.front());
+  FuseAwaitOps<mlrt::compiler::AwaitControlOp,
+               mlrt::compiler::AwaitAllControlOp>(builder, func.front());
+  FuseGetResourceOps(builder, func.front());
+  FusePromiseReturn(builder, func.front());
+}
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateFuseMlrtOpPass() {
+  return std::make_unique<FuseMlrtOpPass>();
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/fuse_mlrt_ops.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_TFRT_TRANSFORMS_MLRT_FUSE_MLRT_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_FUSE_MLRT_OPS_H_
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>> CreateFuseMlrtOpPass();
+
+}
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_FUSE_MLRT_OPS_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/ifrt_set_tpu_host_allocator.cc

@@ -0,0 +1,129 @@
+/* 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/tfrt/transforms/mlrt/ifrt_set_tpu_host_allocator.h"
+
+#include <memory>
+#include <vector>
+
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/mlrt_device_constants.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+namespace {
+
+class IfrtSetTpuHostAllocatorPass
+    : public mlir::PassWrapper<IfrtSetTpuHostAllocatorPass,
+                               mlir::OperationPass<mlir::func::FuncOp>> {
+ public:
+  IfrtSetTpuHostAllocatorPass() = default;
+  IfrtSetTpuHostAllocatorPass &operator=(const IfrtSetTpuHostAllocatorPass &) =
+      delete;
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(IfrtSetTpuHostAllocatorPass)
+
+ private:
+  void getDependentDialects(mlir::DialectRegistry &registry) const override {
+    registry.insert<tensorflow::tf_mlrt::TensorflowMlrtDialect>();
+    registry.insert<mlrt::compiler::MlrtDialect>();
+  }
+
+  llvm::StringRef getArgument() const final {
+    return "tf-mlrt-ifrt-set-tpu-host-allocator";
+  }
+
+  llvm::StringRef getDescription() const final {
+    return "Set input producer to IfrtCall to use Tpu Host Allocator";
+  }
+  void runOnOperation() override {
+    mlir::func::FuncOp func = getOperation();
+    mlir::OpBuilder builder(&getContext());
+
+    llvm::SmallDenseSet<mlir::Operation *> producers;
+
+    auto process_call = [&](auto call) {
+      std::vector<int> variable_arg_indices;
+      variable_arg_indices.reserve(call.getVariableArgIndices().size());
+      for (auto variable_index_attr : call.getVariableArgIndices()) {
+        auto variable_index =
+            llvm::dyn_cast_or_null<mlir::IntegerAttr>(variable_index_attr);
+        if (!variable_index) {
+          call->emitError()
+              << "Expect variable_arg_indices to be integer, but get "
+              << call.getVariableArgIndices();
+          return mlir::WalkResult::interrupt();
+        }
+        variable_arg_indices.push_back(variable_index.getInt());
+      }
+
+      int variable_index = 0;
+      for (int i = 0; i < call.getOperands().size(); ++i) {
+        if (variable_index < variable_arg_indices.size() &&
+            i == variable_arg_indices[variable_index]) {
+          variable_index++;
+          continue;
+        }
+        producers.insert(call.getOperands()[i].getDefiningOp());
+      }
+      return mlir::WalkResult::advance();
+    };
+
+    mlir::WalkResult walk_result = func.walk(
+        [&](mlir::TF::IfrtCallOp call) { return process_call(call); });
+    if (walk_result.wasInterrupted()) {
+      return signalPassFailure();
+    }
+
+    walk_result = func.walk(
+        [&](mlir::TF::AsyncIfrtCallOp call) { return process_call(call); });
+    if (walk_result.wasInterrupted()) {
+      return signalPassFailure();
+    }
+
+    for (auto *def : producers) {
+      if (def && llvm::isa<mlir::TF::TensorFlowDialect>(def->getDialect())) {
+        def->setAttr(kTfMlrtCustomDevice,
+                     builder.getStringAttr(kTpuHostDevice));
+      }
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateIfrtSetTpuHostAllocatorPass() {
+  return std::make_unique<IfrtSetTpuHostAllocatorPass>();
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/ifrt_set_tpu_host_allocator.h

@@ -0,0 +1,34 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_IFRT_SET_TPU_HOST_ALLOCATOR_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_IFRT_SET_TPU_HOST_ALLOCATOR_H_
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+// Creates a pass that set tpu input producers to use tpu host allocators.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateIfrtSetTpuHostAllocatorPass();
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_IFRT_SET_TPU_HOST_ALLOCATOR_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.cc

@@ -0,0 +1,206 @@
+/* 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/tfrt/transforms/mlrt/import_model.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/error_util.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/assign_op_key.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/passes.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/while_to_map_fn.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/compiler/mlir/tfrt/translate/import_model.h"
+#include "tensorflow/compiler/mlir/tfrt/translate/mlrt/mlir_to_bytecode.h"
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+#include "tensorflow/compiler/mlir/tfrt/utils/export.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tensorflow/core/tfrt/mlrt/attribute/attribute.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+#include "tensorflow/core/util/debug_data_dumper.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+namespace {
+// Setup the input pass manager to enable IR dumping after each pass.
+// Note a side effect of this method is that multi threading will be disabled.
+void EnablePassIRPrinting(mlir::PassManager& pm,
+                          const std::string& dump_group_name,
+                          llvm::StringRef module_name) {
+  // Print the whole module after each pass, which requires disabling
+  // multi-threading as well.
+  pm.getContext()->disableMultithreading();
+  pm.enableIRPrinting(std::make_unique<::tensorflow::DataDumperLoggerConfig>(
+      [module_name, dump_group_name](const std::string& pass_tag_name,
+                                     mlir::Operation* op) {
+        return DEBUG_DATA_DUMPER()->GetDumpFilename(
+            module_name.str(), dump_group_name, pass_tag_name);
+      },
+      /*pass_prefix=*/"",
+      /*print_module_scope=*/true));
+  pm.enableTiming();
+}
+}  // namespace
+
+absl::StatusOr<mlrt::bc::Buffer> ConvertTfMlirToBytecode(
+    const TfrtCompileOptions& options, tfrt_stub::FallbackState& fallback_state,
+    mlir::ModuleOp module, tfrt_stub::ModelRuntimeContext& model_context,
+    mlir::OwningOpRef<mlir::ModuleOp>* module_with_op_keys,
+    std::vector<std::string>* added_xla_function_names) {
+  mlrt::bc::Buffer bytecode_buffer;
+  TF_RETURN_IF_ERROR(ConvertTfMlirToRuntimeExecutable(
+      options, module,
+      [&bytecode_buffer, &fallback_state, &model_context,
+       backend_compiler = options.backend_compiler,
+       module_with_op_keys](mlir::PassManager& pm, mlir::ModuleOp module,
+                            const TfrtPipelineOptions& options) {
+        if (backend_compiler) {
+          if (auto* flib_def = model_context.function_library_definition()) {
+            // Copy the module before exporting as exporting to graph will
+            // transform the MLIR to TFG dialect.
+            mlir::OwningOpRef<mlir::ModuleOp> copy(module.clone());
+            TF_RETURN_IF_ERROR(
+                ExportFunctionDefs(*copy, [flib_def](FunctionDef function_def) {
+                  VLOG(1) << absl::StrCat(
+                      "Exporting MLIR function as function_def: ",
+                      // NOLINTNEXTLINE
+                      function_def.DebugString());
+
+                  // The TF MLIR compiler may change the function name. Then we
+                  // need to retrieve the original name from the
+                  // _original_func_name attribute.
+                  auto iter = function_def.attr().find("_original_func_name");
+                  if (iter != function_def.attr().end()) {
+                    function_def.mutable_signature()->set_name(
+                        iter->second.s());
+                  }
+
+                  const auto& name = function_def.signature().name();
+                  if (flib_def->Contains(name)) {
+                    TF_RETURN_IF_ERROR(flib_def->RemoveFunction(name));
+                  }
+
+                  return flib_def->AddFunctionDef(function_def);
+                }));
+          }
+        }
+
+        mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());
+
+        tensorflow::CreateTFInvariantOptimizationPipelineHelper(pm, options);
+
+        pm.addPass(mlrt_compiler::CreateWhileToMapFnPass());
+        // Remove unreachable private functions after map_fn conversion.
+        pm.addPass(mlir::createSymbolDCEPass());
+
+        // TODO(b/283481729): Add test to cover unused constants that do not
+        // cause op_key discontinuity
+        pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+        pm.addPass(mlrt_compiler::CreateAssignOpKeyPass());
+
+        // Run passes until (including) AssignOpKeyPass.
+        if (VLOG_IS_ON(4)) {
+          EnablePassIRPrinting(pm, "mlrt_runtime_lowering_tf",
+                               "mlrt_runtime_lowering_tf");
+        }
+
+        if (mlir::failed(pm.run(module))) {
+          return diag_handler.Combine(absl::InternalError(
+              "failed to finish passes before (including) assign op keys."));
+        }
+        if (VLOG_IS_ON(1)) {
+          tensorflow::DumpMlirOpToFile("tf_dialect_after_assign_op_key",
+                                       module);
+        }
+        // Save the module.
+        if (module_with_op_keys != nullptr) {
+          *module_with_op_keys = module.clone();
+        }
+        // Clear passes already run.
+        pm.clear();
+        // Create the remaining pipeline and run.
+        CreateTfToMlrtPipeline(pm, options, &fallback_state);
+        if (mlir::failed(pm.run(module))) {
+          return diag_handler.Combine(absl::InternalError(
+              "failed to lower TF Dialect to MLRT dialect."));
+        }
+        // Generate bytecode.
+        mlrt::AttributeEncoderRegistry registry;
+        registry.Register("tf_mlrt",
+                          &tensorflow::tf_mlrt::EncodeTensorflowAttribute);
+        auto statusor = mlrt::EmitExecutable(registry, module);
+        if (!statusor.ok()) return statusor.status();
+        bytecode_buffer = std::move(*statusor);
+        return absl::OkStatus();
+      },
+      model_context, &fallback_state, added_xla_function_names));
+  return bytecode_buffer;
+}
+
+absl::StatusOr<mlrt::bc::Buffer> ConvertTfMlirWithOpKeysToBytecode(
+    const TfrtCompileOptions& options,
+    const tfrt_stub::FallbackState& fallback_state,
+    mlir::ModuleOp module_with_op_keys,
+    const tfrt_stub::CostRecorder& cost_recorder) {
+  mlir::StatusScopedDiagnosticHandler diag_handler(
+      module_with_op_keys.getContext());
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("tf_dialect_with_op_keys",
+                                 module_with_op_keys);
+  }
+  // Create the reconversion pipeline and run.
+  mlir::PassManager pm(module_with_op_keys.getContext());
+  const auto pipeline_options = GetTfrtPipelineOptions(options);
+  CreateTfToMlrtPipeline(pm, *pipeline_options, &fallback_state,
+                         &cost_recorder);
+  if (mlir::failed(pm.run(module_with_op_keys))) {
+    return diag_handler.Combine(
+        absl::InternalError("failed to lower TF Dialect to MLRT dialect."));
+  }
+  // Generate bytecode.
+  mlrt::AttributeEncoderRegistry registry;
+  registry.Register("tf_mlrt", &tensorflow::tf_mlrt::EncodeTensorflowAttribute);
+  auto statusor = mlrt::EmitExecutable(registry, module_with_op_keys);
+  if (!statusor.ok()) return statusor.status();
+  if (VLOG_IS_ON(1)) {
+    tensorflow::DumpMlirOpToFile("tfrt_dialect_from_tf_dialect_with_op_keys",
+                                 module_with_op_keys);
+  }
+  return std::move(*statusor);
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/import_model.h

@@ -0,0 +1,54 @@
+/* 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_TFRT_TRANSFORMS_MLRT_IMPORT_MODEL_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_IMPORT_MODEL_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+// Converts an MLIR `module` in TF dialect to MLRT's bytecode format. If
+// `module_with_op_keys` is non-null, the intermediate module on which passes
+// until (including) AssignOpKeyPass have run will be cloned to it.
+//
+// This is for initial conversion.
+absl::StatusOr<mlrt::bc::Buffer> ConvertTfMlirToBytecode(
+    const TfrtCompileOptions& options, tfrt_stub::FallbackState& fallback_state,
+    mlir::ModuleOp module, tfrt_stub::ModelRuntimeContext& model_context,
+    mlir::OwningOpRef<mlir::ModuleOp>* module_with_op_keys = nullptr,
+    std::vector<std::string>* added_xla_function_names = nullptr);
+
+// Converts an MLIR `module_with_op_keys` in TF dialect to MLRT's bytecode
+// format, with op costs from `cost_recorder`.
+//
+// This is for re-conversion.
+absl::StatusOr<mlrt::bc::Buffer> ConvertTfMlirWithOpKeysToBytecode(
+    const TfrtCompileOptions& options,
+    const tfrt_stub::FallbackState& fallback_state,
+    mlir::ModuleOp module_with_op_keys,
+    const tfrt_stub::CostRecorder& cost_recorder);
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_IMPORT_MODEL_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/mlrt_device_constants.h

@@ -0,0 +1,27 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_MLRT_DEVICE_CONSTANTS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_MLRT_DEVICE_CONSTANTS_H_
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+inline constexpr char kTfMlrtCustomDevice[] = "tf_mlrt.custom_device";
+inline constexpr char kTpuHostDevice[] = "tpu_host_device";
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_MLRT_DEVICE_CONSTANTS_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/parallelization.cc

@@ -0,0 +1,835 @@
+/* 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/tfrt/transforms/mlrt/parallelization.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SetVector.h"
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/analysis/cost_analysis.h"
+#include "tensorflow/compiler/mlir/tfrt/constants.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tfrt/compiler/stream_analysis.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace mlrt_compiler {
+namespace {
+
+using tensorflow::tfrt_compiler::CostAnalysis;
+using tfrt::compiler::Stream;
+using tfrt::compiler::StreamAnalysis;
+
+std::string GetStreamFunctionName(absl::string_view func_name,
+                                  const Stream& stream) {
+  return absl::StrCat(func_name, "_stream_", stream.id());
+}
+
+bool IsConstant(mlir::Operation* op) {
+  return op && llvm::isa<mlir::TF::ConstOp, mlir::TF::_TfrtGetResourceOp>(op);
+}
+
+// StreamInfo is a bookkeeping for inputs, futures, and promises for a stream.
+struct StreamInfo {
+  const Stream* parent = nullptr;
+
+  // The values that are produced by constant ops. Instead of using
+  // promise/await to pass these values between streams, we can just copying
+  // these ops to the streams that use these constants.
+  llvm::SetVector<mlir::Value> constants;
+  // The values that are the inputs to the stream.
+  llvm::SetVector<mlir::Value> inputs;
+  // The values that will be the futures to the stream.
+  llvm::SetVector<mlir::Value> futures;
+  // The values that will be the control futures (i.e. futures with no data) to
+  // the stream.
+  llvm::SetVector<mlir::Operation*> control_futures;
+  // The values that will be the promises to the stream.
+  llvm::SetVector<mlir::Value> promises;
+  // The values that will be the control promises (i.e., promises with no data)
+  // to the stream.
+  llvm::SetVector<mlir::Operation*> control_promises;
+  // The values that are defined by the operations in the stream. Note that all
+  // values in `futures` will also be in `results`.
+  llvm::DenseSet<mlir::Value> results;
+
+  bool contains_only_constants = true;
+
+  bool IsRoot() const { return parent == nullptr; }
+};
+
+// Preprocess the block to produce StreamInfo for every stream.
+llvm::DenseMap<const Stream*, StreamInfo> PreprocessStreamInfo(
+    mlir::Block& block,
+    const llvm::DenseMap<mlir::Operation*,
+                         llvm::SmallSetVector<mlir::Operation*, 4>>&
+        control_predecessors,
+    const StreamAnalysis& stream_analysis) {
+  llvm::DenseMap<const Stream*, StreamInfo> stream_map;
+
+  // All values that will be promises in the block.
+  llvm::DenseSet<mlir::Value> promises;
+
+  // All operations that will be control promises in the block.
+  llvm::DenseSet<mlir::Operation*> control_promises;
+
+  // Keep track of all available values and controls as we traverse the stream
+  // tree in depth-first order.
+  llvm::DenseSet<mlir::Value> available_values;
+  llvm::DenseSet<mlir::Operation*> available_controls;
+
+  struct Entry {
+    explicit Entry(const Stream* stream) : stream(stream) {}
+
+    const Stream* stream = nullptr;
+
+    // Keep track of the next operation to be processed. If all operations are
+    // processed, we can pop this stream from the DFS stack.
+    int op_idx = 0;
+  };
+
+  std::vector<Entry> stack;
+  stack.reserve(stream_analysis.GetNumStreams());
+
+  // We first push the entry for the root stream.
+  const auto& root_stream = stream_analysis.GetRootStream();
+  auto& root_stream_info = stream_map[&root_stream];
+  available_values.insert(block.getArguments().begin(),
+                          block.getArguments().end());
+  root_stream_info.results.insert(block.getArguments().begin(),
+                                  block.getArguments().end());
+  stack.push_back(Entry(&root_stream));
+
+  // The root stream's first operation a dummy operation that defines all block
+  // arguments.
+  for (auto* child_stream : root_stream.GetChildStreamsForRootOp()) {
+    stream_map[child_stream].parent = &root_stream;
+    stack.push_back(Entry(child_stream));
+  }
+
+  // The first DFS traveral populates inputs and futures for every stream but
+  // not promises. We only know whether a value definition is a promise only
+  // after traversing all streams, so it is not possible to know it in the first
+  // pass.
+  while (!stack.empty()) {
+    auto& [stream, op_idx] = stack.back();
+    auto& stream_info = stream_map[stream];
+
+    auto ops = stream->ops();
+
+    // If we finish processing all operations in the stream, we can pop this
+    // stream, as well as the values defined by its operations.
+    if (op_idx == ops.size()) {
+      for (auto* op : stream->ops()) {
+        // Erase the values and controls produced by the current stream.
+        for (auto result : op->getResults()) {
+          available_values.erase(result);
+        }
+        available_controls.erase(op);
+      }
+      // Futures and control futures will also be available, so we erase them as
+      // well.
+      for (auto future : stream_info.futures) {
+        available_values.erase(future);
+      }
+      for (auto* control_future : stream_info.control_futures) {
+        available_controls.erase(control_future);
+      }
+
+      if (!stream_info.IsRoot()) {
+        // Merge inputs, futures, and promises into the parent stream, as they
+        // will be passed down from the root in the output program.
+        DCHECK_GT(stream_map.count(stream_info.parent), 0);
+        auto& parent_info = stream_map[stream_info.parent];
+
+        for (const auto& input : stream_info.inputs) {
+          DCHECK(available_values.contains(input));
+          if (!parent_info.results.contains(input)) {
+            // An input in the current stream will be an input in the parent
+            // stream only if it is not a result in the parent stream.
+            parent_info.inputs.insert(input);
+          }
+        }
+
+        for (auto future : stream_info.futures) {
+          DCHECK(!available_values.contains(future));
+          parent_info.futures.insert(future);
+        }
+        for (auto* control_future : stream_info.control_futures) {
+          DCHECK(!available_controls.contains(control_future));
+          parent_info.control_futures.insert(control_future);
+        }
+      }
+
+      // Update the global promise set.
+      promises.insert(stream_info.futures.begin(), stream_info.futures.end());
+      control_promises.insert(stream_info.control_futures.begin(),
+                              stream_info.control_futures.end());
+
+      stack.pop_back();
+      continue;
+    }
+
+    // We process the operations one by one. If the operation has child streams,
+    // we process the child streams first before continuing to the next
+    // operation.
+    bool has_child_streams = false;
+    for (; op_idx < ops.size() && !has_child_streams; ++op_idx) {
+      auto* op = ops[op_idx];
+
+      stream_info.contains_only_constants &= IsConstant(op);
+
+      // Check every operand to see whether it is a future or input.
+      for (mlir::Value operand : op->getOperands()) {
+        // If the value is defined in the current stream, nothing needs to be
+        // done.
+        if (!stream_info.results.contains(operand)) {
+          if (available_values.insert(operand).second) {
+            // If the operand is not available in the current stream or any
+            // parent stream, it will be a future and then become a result.
+            if (IsConstant(operand.getDefiningOp())) {
+              stream_info.constants.insert(operand);
+            } else {
+              stream_info.futures.insert(operand);
+            }
+            stream_info.results.insert(operand);
+          } else {
+            // If the operand is not available in the current stream but
+            // available in the parent stream, it is an input.
+            if (IsConstant(operand.getDefiningOp())) {
+              stream_info.constants.insert(operand);
+            } else {
+              stream_info.inputs.insert(operand);
+            }
+          }
+        }
+      }
+
+      // Insert mlrt.await_control if this op has control deps on other ops.
+      if (auto ctrl_iter = control_predecessors.find(op);
+          ctrl_iter != control_predecessors.end()) {
+        const auto& ctrl_deps = ctrl_iter->second;
+
+        for (mlir::Operation* control_dep : ctrl_deps) {
+          if (available_controls.insert(control_dep).second) {
+            // If the control is not already available, it will be a control
+            // future and then become available.
+            stream_info.control_futures.insert(control_dep);
+          }
+        }
+      }
+
+      // Update results of this operations.
+      for (mlir::Value result : op->getResults()) {
+        available_values.insert(result);
+        stream_info.results.insert(result);
+      }
+
+      // Update this op as an available control.
+      available_controls.insert(op);
+
+      // Pause processing the current stream to process the child streams first.
+      const auto& child_streams = stream->GetChildStreams(op);
+      has_child_streams = !child_streams.empty();
+      for (auto* child_stream : child_streams) {
+        stream_map[child_stream].parent = stream;
+        stack.push_back(Entry(child_stream));
+      }
+    }
+  }
+
+  // The second pass populates promises for each stream. We also need to merge
+  // promises in a child to its parent stream. We can do this by traversing the
+  // operation in reverse program order.
+  for (auto& op : llvm::reverse(block)) {
+    const auto& stream = stream_analysis.GetStream(&op);
+    auto& stream_info = stream_map[&stream];
+
+    for (mlir::Value result : op.getResults()) {
+      if (promises.contains(result)) {
+        stream_info.promises.insert(result);
+      }
+    }
+
+    if (control_promises.contains(&op)) {
+      stream_info.control_promises.insert(&op);
+    }
+
+    for (const auto* child_stream : stream.GetChildStreams(&op)) {
+      const auto& child_info = stream_map[child_stream];
+
+      stream_info.promises.insert(child_info.promises.begin(),
+                                  child_info.promises.end());
+      stream_info.control_promises.insert(child_info.control_promises.begin(),
+                                          child_info.control_promises.end());
+    }
+  }
+
+  // Special handling for the dummy operation in the root.
+  auto& root_info = stream_map[&root_stream];
+  for (const auto* child_stream : root_stream.GetChildStreamsForRootOp()) {
+    const auto& child_info = stream_map[child_stream];
+
+    root_info.promises.insert(child_info.promises.begin(),
+                              child_info.promises.end());
+    root_info.control_promises.insert(child_info.control_promises.begin(),
+                                      child_info.control_promises.end());
+  }
+
+  return stream_map;
+}
+
+// A custom struct that groups mappings for values, futures and promises for a
+// stream during creating the corresponding stream function.
+struct Mapping {
+  // This is the mappings for the SSA values used in the original and new
+  // operations.
+  mlir::IRMapping value_mapping;
+
+  // Maps the original tensor value that will be a future to the corresponding
+  // !mlrt.future value.
+  mlir::IRMapping future_mapping;
+
+  // Maps the original tensor value that will be a promise to the corresponding
+  // !mlrt.promise value.
+  mlir::IRMapping promise_mapping;
+
+  // In addition to value mappings, we also need mappings for input control
+  // dependencies to the corresponding !mlrt.future and !mlrt.promise values.
+  llvm::DenseMap<mlir::Operation*, mlir::Value> future_control_mapping;
+  llvm::DenseMap<mlir::Operation*, mlir::Value> promise_control_mapping;
+};
+
+mlrt::compiler::AsyncOp CreateAsyncOp(
+    mlir::OpBuilder& builder, absl::string_view function_name,
+    const llvm::DenseMap<const Stream*, StreamInfo>& stream_map,
+    const Stream& stream, const Mapping& mapping, mlir::Location loc) {
+  auto iter = stream_map.find(&stream);
+  DCHECK(iter != stream_map.end());
+  const auto& stream_info = iter->second;
+
+  if (stream_info.contains_only_constants) return nullptr;
+
+  const auto& [value_mapping, future_mapping, promise_mapping,
+               future_control_mapping, promise_control_mapping] = mapping;
+
+  llvm::SmallVector<mlir::Value> async_operands;
+
+  for (auto input : stream_info.inputs) {
+    async_operands.push_back(value_mapping.lookup(input));
+    DCHECK(async_operands.back());
+  }
+
+  for (auto future : stream_info.futures) {
+    async_operands.push_back(future_mapping.lookup(future));
+    DCHECK(async_operands.back());
+  }
+
+  for (auto* control_future : stream_info.control_futures) {
+    DCHECK_GT(future_control_mapping.count(control_future), 0);
+    async_operands.push_back(future_control_mapping.lookup(control_future));
+    DCHECK(async_operands.back());
+  }
+
+  for (auto promise : stream_info.promises) {
+    async_operands.push_back(promise_mapping.lookup(promise));
+    DCHECK(async_operands.back());
+  }
+
+  for (auto* control_promise : stream_info.control_promises) {
+    DCHECK_GT(promise_control_mapping.count(control_promise), 0);
+    async_operands.push_back(promise_control_mapping.lookup(control_promise));
+    DCHECK(async_operands.back());
+  }
+
+  return mlrt::compiler::AsyncOp::create(
+      builder, loc, builder.getType<mlrt::compiler::AsyncHandleType>(),
+      async_operands,
+      mlir::SymbolRefAttr::get(builder.getContext(),
+                               GetStreamFunctionName(function_name, stream)));
+}
+
+mlir::func::FuncOp CreateStreamFunction(
+    mlir::OpBuilder& builder, Mapping& mapping, absl::string_view name,
+    const Stream& stream, const StreamInfo& stream_info, mlir::Location loc) {
+  if (stream_info.contains_only_constants) return nullptr;
+
+  auto& [value_mapping, future_mapping, promise_mapping, future_control_mapping,
+         promise_control_mapping] = mapping;
+
+  llvm::SmallVector<mlir::Type> arg_types;
+  for (mlir::Value input : stream_info.inputs) {
+    arg_types.push_back(input.getType());
+  }
+
+  arg_types.append(
+      stream_info.futures.size() + stream_info.control_futures.size(),
+      builder.getType<mlrt::compiler::FutureType>());
+  arg_types.append(
+      stream_info.promises.size() + stream_info.control_promises.size(),
+      builder.getType<mlrt::compiler::PromiseType>());
+
+  // The stream function has no result.
+  auto func_type = builder.getFunctionType(arg_types, /*results=*/{});
+
+  auto func = mlir::func::FuncOp::create(
+      builder, loc, GetStreamFunctionName(name, stream), func_type);
+  func.setVisibility(mlir::func::FuncOp::Visibility::Private);
+
+  // Populate the body of the stream function by copying over the operations
+  // in the stream.
+  auto* new_block = func.addEntryBlock();
+
+  // Replace inputs with the function arguments.
+  for (int i = 0; i < stream_info.inputs.size(); ++i) {
+    value_mapping.map(stream_info.inputs[i], new_block->getArgument(i));
+  }
+
+  // Maps the original tensor value that will be a future or a promise to
+  // the corresponding !mlrt.future or !mlrt.promise value.
+  size_t start = stream_info.inputs.size();
+  for (int i = 0; i < stream_info.futures.size(); ++i) {
+    future_mapping.map(stream_info.futures[i],
+                       new_block->getArgument(i + start));
+  }
+
+  start += stream_info.futures.size();
+  for (int i = 0; i < stream_info.control_futures.size(); ++i) {
+    future_control_mapping[stream_info.control_futures[i]] =
+        new_block->getArgument(i + start);
+  }
+
+  start += stream_info.control_futures.size();
+  for (int i = 0; i < stream_info.promises.size(); ++i) {
+    promise_mapping.map(stream_info.promises[i],
+                        new_block->getArgument(i + start));
+  }
+
+  start += stream_info.promises.size();
+  for (int i = 0; i < stream_info.control_promises.size(); ++i) {
+    promise_control_mapping[stream_info.control_promises[i]] =
+        new_block->getArgument(i + start);
+  }
+
+  return func;
+}
+
+void CreateAllocateFuturesOp(mlir::OpBuilder& builder, Mapping& mapping,
+                             const StreamInfo& stream_info,
+                             mlir::Location loc) {
+  auto& [value_mapping, future_mapping, promise_mapping, future_control_mapping,
+         promise_control_mapping] = mapping;
+
+  DCHECK_EQ(stream_info.futures.size(), stream_info.promises.size());
+
+  llvm::SmallVector<mlir::Type> promise_types(
+      stream_info.promises.size(),
+      builder.getType<mlrt::compiler::PromiseType>());
+  llvm::SmallVector<mlir::Type> future_types(
+      stream_info.futures.size(),
+      builder.getType<mlrt::compiler::FutureType>());
+
+  if (!stream_info.futures.empty()) {
+    auto allocate_futures = tf_mlrt::AllocateFuturesOp::create(
+        builder, loc, promise_types, future_types, stream_info.futures.size());
+    for (int i = 0; i < stream_info.futures.size(); ++i) {
+      future_mapping.map(stream_info.futures[i],
+                         allocate_futures.getFutures()[i]);
+    }
+
+    for (int i = 0; i < stream_info.futures.size(); ++i) {
+      // Use the original values in `futures` to make sure futures[i] shares the
+      // state with promises[i].
+      DCHECK(stream_info.promises.contains(stream_info.futures[i]));
+      promise_mapping.map(stream_info.futures[i],
+                          allocate_futures.getPromises()[i]);
+    }
+  }
+
+  DCHECK_EQ(stream_info.control_futures.size(),
+            stream_info.control_promises.size());
+  if (!stream_info.control_futures.empty()) {
+    promise_types.resize(stream_info.control_promises.size(),
+                         builder.getType<mlrt::compiler::PromiseType>());
+    future_types.resize(stream_info.control_futures.size(),
+                        builder.getType<mlrt::compiler::FutureType>());
+
+    auto allocate_control_futures =
+        mlrt::compiler::AllocateControlFuturesOp::create(
+            builder, loc, promise_types, future_types,
+            stream_info.control_futures.size());
+    for (int i = 0; i < stream_info.control_futures.size(); ++i) {
+      future_control_mapping[stream_info.control_futures[i]] =
+          allocate_control_futures.getFutures()[i];
+    }
+    for (int i = 0; i < stream_info.control_futures.size(); ++i) {
+      // Use the original operations in `control_futures` to make sure
+      // control_futures[i] shares the state with control_promises[i].
+      DCHECK(stream_info.control_promises.contains(
+          stream_info.control_futures[i]));
+      promise_control_mapping[stream_info.control_futures[i]] =
+          allocate_control_futures.getPromises()[i];
+    }
+  }
+}
+
+class TensorflowCostModel : public StreamAnalysis::CostModelInterface {
+ public:
+  explicit TensorflowCostModel(CostAnalysis* cost_analysis)
+      : cost_analysis_(*cost_analysis) {}
+
+  std::optional<int64_t> GetOperationCost(mlir::Operation* op) const override {
+    return cost_analysis_.GetCost(op);
+  }
+
+ private:
+  const CostAnalysis& cost_analysis_;
+};
+
+bool SkipControlDep(mlir::Operation* op) {
+  // TODO(chky): Consider define side effects more properly for these ops.
+  return llvm::isa<mlir::TF::TPUCompileMlirAndExecuteOp, mlir::TF::AssertOp>(
+      op);
+}
+
+void ParallelizeBlock(
+    absl::string_view name, mlir::Block& block,
+    const mlir::TF::SideEffectAnalysis::Info& side_effect_analysis,
+    const tfrt_stub::CostRecorder* cost_recorder) {
+  // First, we use SideEffectAnalysis to find out control predecessors for each
+  // operation. We use this map later to insert control futures.
+  llvm::DenseMap<mlir::Operation*, llvm::SmallSetVector<mlir::Operation*, 4>>
+      control_predecessors;
+  for (auto& op : block) {
+    auto& deps = control_predecessors[&op];
+    for (auto* dep : side_effect_analysis.DirectControlPredecessors(&op)) {
+      // If we skip the control deps of `op`, then we need to use the control
+      // deps of these control deps instead.
+      if (SkipControlDep(dep)) {
+        for (auto* d : control_predecessors[dep]) {
+          DCHECK(!SkipControlDep(d));
+          deps.insert(d);
+        }
+      } else {
+        deps.insert(dep);
+      }
+    }
+  }
+
+  // Remove skipped control deps.
+  for (auto& op : block) {
+    if (SkipControlDep(&op)) {
+      control_predecessors.erase(&op);
+    }
+  }
+
+  // Perform stream analysis.
+  CostAnalysis cost_analysis(
+      llvm::cast<mlir::func::FuncOp>(block.getParentOp()), cost_recorder);
+  TensorflowCostModel cost_model(&cost_analysis);
+  StreamAnalysis stream_analysis(block, &cost_model);
+
+  // Preprocess all streams to gather StreamInfos for all streams, without
+  // modifying the program.
+  llvm::DenseMap<const Stream*, StreamInfo> stream_map =
+      PreprocessStreamInfo(block, control_predecessors, stream_analysis);
+
+  // Then we perform a DFS traversal to create stream functions and insert async
+  // operations.
+  std::vector<const Stream*> stack;
+  stack.reserve(stream_analysis.GetNumStreams());
+
+  const auto& root_stream = stream_analysis.GetRootStream();
+  stack.push_back(&root_stream);
+
+  llvm::SmallVector<mlir::Operation*> to_remove;
+
+  mlir::OpBuilder builder(block.getParentOp());
+
+  while (!stack.empty()) {
+    const auto* stream = stack.back();
+    stack.pop_back();
+    DCHECK(stream);
+
+    DCHECK_GT(stream_map.count(stream), 0);
+    const auto& stream_info = stream_map[stream];
+
+    Mapping mapping;
+    auto& [value_mapping, future_mapping, promise_mapping,
+           future_control_mapping, promise_control_mapping] = mapping;
+
+    // `async_handles` keeps the !mlrt.async_handle created in the stream. A
+    // mlrt.await_handle op will be inserted at the end of the stream function
+    // for each async handle.
+    llvm::SmallVector<mlir::Value> async_handles;
+
+    mlir::func::FuncOp stream_func;
+    if (!stream_info.IsRoot()) {
+      // If it is not a root stream, we need to create a new function for this
+      // stream. And futures and promises are also passed as parameters. For the
+      // root stream, futures and promises are allocated in the body.
+
+      // Insert the stream function before the original function.
+      builder.setInsertionPoint(block.getParentOp());
+
+      stream_func =
+          CreateStreamFunction(builder, mapping, name, *stream, stream_info,
+                               block.getParentOp()->getLoc());
+
+      if (stream_func) {
+        // Set the insertion point to the start of the new block in the
+        // function.
+        builder.setInsertionPointToStart(&stream_func.front());
+      }
+    } else {
+      stream_func = llvm::cast<mlir::func::FuncOp>(block.getParentOp());
+
+      DCHECK_EQ(stream, &root_stream);
+      // If it is the root stream, we insert new operations in the original
+      // function. And we need to allocate all the futures used here.
+      builder.setInsertionPointToStart(&block);
+
+      // The block arguments of the root stream are in the `results`. There will
+      // be no additional inputs in `inputs`.
+      DCHECK(stream_info.inputs.empty());
+
+      // Put the original arguments in the mapping as they are not changed.
+      for (auto arg : block.getArguments()) {
+        value_mapping.map(arg, arg);
+      }
+
+      // Insert a tf_mlrt.allocate_futures op to allocate all futures used.
+      CreateAllocateFuturesOp(builder, mapping, stream_info,
+                              block.getParentOp()->getLoc());
+
+      // Lastly for the root stream, we need to handle the dummy op that defines
+      // the arguments.
+      for (const auto* child_stream : stream->GetChildStreamsForRootOp()) {
+        stack.push_back(child_stream);
+        if (auto async =
+                CreateAsyncOp(builder, name, stream_map, *child_stream, mapping,
+                              block.getParentOp()->getLoc())) {
+          async_handles.push_back(async);
+        }
+      }
+    }
+
+    for (auto* op : stream->ops()) {
+      to_remove.push_back(op);
+    }
+
+    // Skip empty streams.
+    if (!stream_func) continue;
+
+    mlir::Operation* return_op = nullptr;
+
+    // Cloning the operations in the stream. If the operand is a future, a
+    // tf_mlrt.Await op will be inserted. If the result is a promise, a
+    // tf_mlrt.Promise will be inserted. Similar to control futures and control
+    // promises.
+    for (auto* op : stream->ops()) {
+      // Clone the current op into the function of this stream, using the
+      // new operands, which can be futures.
+      for (mlir::Value operand : op->getOperands()) {
+        if (stream_info.constants.contains(operand) &&
+            !value_mapping.contains(operand)) {
+          builder.clone(*operand.getDefiningOp(), value_mapping);
+        } else if (stream_info.futures.contains(operand) &&
+                   !value_mapping.contains(operand)) {
+          // Insert Await op if it is a future.
+          auto future_value = tf_mlrt::TFAwaitOp::create(
+              builder, op->getLoc(), operand.getType(),
+              future_mapping.lookup(operand));
+
+          // Now this future is available in the current stream, so it can be a
+          // normal value.
+          value_mapping.map(operand, future_value);
+        }
+      }
+
+      if (auto ctrl_iter = control_predecessors.find(op);
+          ctrl_iter != control_predecessors.end()) {
+        const auto& ctrl_deps = ctrl_iter->second;
+
+        for (mlir::Operation* control_dep : ctrl_deps) {
+          // This control may be available in the ancestors or in a previous
+          // AwaitControl, we only insert a new AwaitControl if it is not.
+          if (stream_info.control_futures.contains(control_dep)) {
+            if (auto iter = future_control_mapping.find(control_dep);
+                iter != future_control_mapping.end()) {
+              mlrt::compiler::AwaitControlOp::create(
+                  builder, control_dep->getLoc(), iter->second);
+
+              // Now we no longer need this control dep in this stream.
+              future_control_mapping.erase(iter);
+            }
+          }
+        }
+      }
+
+      // Clone the op using the value mapping that includes values from futures.
+      auto* new_op = builder.clone(*op, value_mapping);
+
+      // TODO(chky): Ensure the original return op is in the root stream. This
+      // is currently an implicit guarantee in stream analysis.
+      if (llvm::isa<mlir::func::ReturnOp>(op)) {
+        DCHECK(stream_info.IsRoot()) << name << " " << stream->id();
+        return_op = new_op;
+      }
+
+      for (mlir::Value result : op->getResults()) {
+        if (stream_info.promises.contains(result)) {
+          // Insert Promise op if the result is a promise.
+          tf_mlrt::TFPromiseOp::create(builder, op->getLoc(),
+                                       promise_mapping.lookup(result),
+                                       value_mapping.lookup(result));
+        }
+      }
+
+      if (stream_info.control_promises.contains(op)) {
+        // Insert Promise op if this op produce a control dependency to ops in
+        // other streams.
+        mlrt::compiler::PromiseControlOp::create(builder, op->getLoc(),
+                                                 promise_control_mapping[op]);
+      }
+
+      // If this op has child streams, insert mlrt.async ops.
+      for (auto* child_stream : stream->GetChildStreams(op)) {
+        stack.push_back(child_stream);
+        if (auto async = CreateAsyncOp(builder, name, stream_map, *child_stream,
+                                       mapping, op->getLoc())) {
+          async_handles.push_back(async);
+        }
+      }
+    }
+
+    // Create the return op for non-root streams.
+    //
+    // TODO(chky): Ensure the original return op is in the root stream. This is
+    // currently an implicit guarantee in stream analysis.
+    if (!return_op) {
+      DCHECK(!stream_info.IsRoot()) << name << " " << stream->id();
+      return_op =
+          mlir::func::ReturnOp::create(builder, block.getParentOp()->getLoc());
+    }
+
+    // We need to wait for async executions at the end of the stream function,
+    // in order to manage resource lifetime and handle errors properly. These
+    // mlrt.await_handle ops are inserted before the return op.
+    builder.setInsertionPoint(return_op);
+    for (auto handle : async_handles) {
+      mlrt::compiler::AwaitHandleOp::create(
+          builder, block.getParentOp()->getLoc(), handle);
+    }
+  }
+
+  // Remove the operations in the original block.
+  for (auto* op : llvm::reverse(to_remove)) {
+    op->dropAllDefinedValueUses();
+    op->erase();
+  }
+}
+
+class ParallelizationPass
+    : public mlir::PassWrapper<ParallelizationPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  ParallelizationPass() = default;
+  ParallelizationPass(uint64_t cost_threshold,
+                      bool merge_inter_dependent_streams,
+                      const tfrt_stub::CostRecorder* cost_recorder) {
+    cost_threshold_ = cost_threshold;
+    merge_inter_dependent_streams_ = merge_inter_dependent_streams;
+    cost_recorder_ = cost_recorder;
+  }
+  ParallelizationPass(const ParallelizationPass&) {}
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ParallelizationPass)
+
+ private:
+  void getDependentDialects(mlir::DialectRegistry& registry) const override {
+    registry.insert<tensorflow::tf_mlrt::TensorflowMlrtDialect>();
+    registry.insert<mlrt::compiler::MlrtDialect>();
+  }
+
+  llvm::StringRef getArgument() const final {
+    return "tf-mlrt-parallelization";
+  }
+
+  llvm::StringRef getDescription() const final {
+    return "Parallelize tf graphs by inserting mlrt async operations.";
+  }
+
+  void runOnOperation() override {
+    auto module = getOperation();
+
+    mlir::Builder builder(module);
+    module->setAttr("tfrt.cost_threshold",
+                    builder.getI64IntegerAttr(cost_threshold_));
+    module->setAttr("tfrt.merge_inter_dependent_streams",
+                    builder.getBoolAttr(merge_inter_dependent_streams_));
+
+    mlir::TF::SideEffectAnalysis side_effect_analysis(module);
+
+    for (auto func_op :
+         llvm::make_early_inc_range(module.getOps<mlir::func::FuncOp>())) {
+      ParallelizeBlock(func_op.getSymName(), func_op.front(),
+                       side_effect_analysis.GetAnalysisForFunc(func_op),
+                       cost_recorder_);
+    }
+  }
+
+  Option<uint64_t> cost_threshold_{
+      *this, "tfrt-cost-threshold",
+      llvm::cl::desc("If a sequence of operations has a cost lower than the "
+                     "cost-threshold, the sequence will be executed as a block "
+                     "in the same thread."),
+      llvm::cl::init(1)};
+  Option<bool> merge_inter_dependent_streams_{
+      *this, "tfrt-merge-inter-dependent-streams",
+      llvm::cl::desc("If true, streams with inter data depenedencies will be "
+                     "preferred to be merged for inline execution."),
+      llvm::cl::init(false)};
+  const tfrt_stub::CostRecorder* cost_recorder_ = nullptr;
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateParallelizationPass(
+    uint64_t cost_threshold, bool merge_inter_dependent_streams,
+    const tfrt_stub::CostRecorder* cost_recorder) {
+  return std::make_unique<ParallelizationPass>(
+      cost_threshold, merge_inter_dependent_streams, cost_recorder);
+}
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateParallelizationPass() {
+  return std::make_unique<ParallelizationPass>();
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/parallelization.h

@@ -0,0 +1,37 @@
+/* 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_TFRT_TRANSFORMS_MLRT_PARALLELIZATION_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_PARALLELIZATION_H_
+
+#include <memory>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateParallelizationPass(
+    uint64_t cost_threshold, bool merge_inter_dependent_streams,
+    const tfrt_stub::CostRecorder* cost_recorder = nullptr);
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateParallelizationPass();
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_PARALLELIZATION_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/passes.cc

@@ -0,0 +1,83 @@
+/* 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/tfrt/transforms/mlrt/passes.h"
+
+#include "absl/log/check.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/assign_op_key.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/async_while.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/fuse_mlrt_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/ifrt_set_tpu_host_allocator.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/parallelization.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/rewrite_ifrt_load_variable.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/tf_to_mlrt.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/while_to_map_fn.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+void RegisterMlrtPasses() {
+  mlir::registerPass([]() { return CreateAssignOpKeyPass(); });
+  mlir::registerPass([]() { return CreateAsyncWhilePass(); });
+  mlir::registerPass([]() { return CreateParallelizationPass(); });
+  mlir::registerPass([]() { return CreateWhileToMapFnPass(); });
+  mlir::registerPass([]() { return CreateRewriteIfrtLoadVariablePass(); });
+  mlir::registerPass([]() { return CreateIfrtSetTpuHostAllocatorPass(); });
+  mlir::registerPass(
+      []() { return CreateTfToMlrtPreParallelizationConversionPass({}); });
+  mlir::registerPass([]() { return CreateTfToMlrtConversionPass({}); });
+  mlir::registerPass([]() { return CreateFuseMlrtOpPass(); });
+}
+
+void CreateTfToMlrtPipeline(mlir::OpPassManager &pm,
+                            const TfrtPipelineOptions &options,
+                            const tfrt_stub::FallbackState *fallback_state,
+                            const tfrt_stub::CostRecorder *cost_recorder) {
+  pm.addPass(
+      mlrt_compiler::CreateTfToMlrtPreParallelizationConversionPass(options));
+
+  if (options.use_tpu_host_allocator_for_inputs) {
+    pm.addNestedPass<mlir::func::FuncOp>(
+        mlrt_compiler::CreateIfrtSetTpuHostAllocatorPass());
+  }
+  pm.addPass(mlrt_compiler::CreateRewriteIfrtLoadVariablePass());
+
+  if (options.enable_while_parallel_iterations) {
+    pm.addPass(mlrt_compiler::CreateAsyncWhilePass());
+  }
+
+  pm.addPass(mlrt_compiler::CreateParallelizationPass(
+      options.cost_threshold, options.merge_inter_dependent_streams,
+      cost_recorder));
+
+  DCHECK(fallback_state);
+  pm.addPass(
+      mlrt_compiler::CreateTfToMlrtConversionPass(options, fallback_state));
+
+  // Perform optimizations in the lowered MLIR.
+  pm.addNestedPass<mlir::func::FuncOp>(mlrt_compiler::CreateFuseMlrtOpPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+  pm.addPass(mlir::createInlinerPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCSEPass());
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/passes.h

@@ -0,0 +1,38 @@
+/* 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_TFRT_TRANSFORMS_MLRT_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_PASSES_H_
+
+#include "mlir/Pass/PassOptions.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+void RegisterMlrtPasses();
+
+// Creates a pipeline of passes that lowers MLIR TF dialect to MLRT dialects.
+// The op costs from `cost_recorder` (if non-null) are used for Stream Analysis.
+void CreateTfToMlrtPipeline(
+    mlir::OpPassManager& pm, const TfrtPipelineOptions& options,
+    const tfrt_stub::FallbackState* fallback_state,
+    const tfrt_stub::CostRecorder* cost_recorder = nullptr);
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_PASSES_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/rewrite_ifrt_load_variable.cc

@@ -0,0 +1,117 @@
+/* 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/tfrt/transforms/mlrt/rewrite_ifrt_load_variable.h"
+
+#include <memory>
+#include <vector>
+
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
+#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/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+namespace {
+
+class RewriteIfrtLoadVariablePass
+    : public mlir::PassWrapper<RewriteIfrtLoadVariablePass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  RewriteIfrtLoadVariablePass() = default;
+  RewriteIfrtLoadVariablePass &operator=(const RewriteIfrtLoadVariablePass &) =
+      delete;
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RewriteIfrtLoadVariablePass)
+
+ private:
+  void getDependentDialects(mlir::DialectRegistry &registry) const override {
+    registry.insert<tensorflow::tf_mlrt::TensorflowMlrtDialect>();
+    registry.insert<mlrt::compiler::MlrtDialect>();
+  }
+
+  llvm::StringRef getArgument() const final {
+    return "tf-mlrt-rewrite-ifrt-load-variable";
+  }
+
+  llvm::StringRef getDescription() const final {
+    return "Convert tf.IfrtLoadVariable to tf_mlrt.TFIfrtLoadVariable";
+  }
+
+  void runOnOperation() override {
+    mlir::ModuleOp module = getOperation();
+    mlir::OpBuilder builder(module);
+
+    module->walk([&](mlir::TF::IfrtLoadVariableOp load_variable_op) {
+      builder.setInsertionPoint(load_variable_op);
+
+      std::vector<mlir::Type> result_types;
+      result_types.push_back(load_variable_op.getArrayKey().getType());
+      result_types.push_back(builder.getType<mlrt::compiler::FutureType>());
+      auto mlrt_load_variable_op = tf_mlrt::TFIfrtLoadVariableOp::create(
+          builder, load_variable_op->getLoc(), result_types,
+          load_variable_op->getOperands(), load_variable_op->getAttrs());
+      tf_mlrt::TFAwaitOp await_op;
+      for (auto user : llvm::make_early_inc_range(
+               load_variable_op.getTensorFuture().getUsers())) {
+        // Materialize the future for the first use. Reuse it for the rest of
+        // the uses.
+        if (!await_op) {
+          builder.setInsertionPoint(user);
+          await_op = tf_mlrt::TFAwaitOp::create(
+              builder, user->getLoc(),
+              load_variable_op.getTensorFuture().getType(),
+              mlrt_load_variable_op.getTensorFuture());
+        } else {
+          if (user->isBeforeInBlock(await_op)) {
+            await_op->moveBefore(user);
+          }
+        }
+        user->replaceUsesOfWith(load_variable_op.getTensorFuture(),
+                                await_op.getResult());
+      }
+
+      for (auto user : llvm::make_early_inc_range(
+               load_variable_op.getArrayKey().getUsers())) {
+        user->replaceUsesOfWith(load_variable_op.getArrayKey(),
+                                mlrt_load_variable_op.getArrayKey());
+      }
+
+      load_variable_op->erase();
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateRewriteIfrtLoadVariablePass() {
+  return std::make_unique<RewriteIfrtLoadVariablePass>();
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/rewrite_ifrt_load_variable.h

@@ -0,0 +1,36 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_REWRITE_IFRT_LOAD_VARIABLE_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_REWRITE_IFRT_LOAD_VARIABLE_H_
+
+#include <memory>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+// Creates a pass that converts tf.IfrtLoadVariableOp to
+// tf_mlrt.TFIfrtLoadVariableOp and inserts tf_mlrt.Await on the returned future
+// from tf_mlrt.TFIfrtLoadVariableOp if it is used by CPU ops.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateRewriteIfrtLoadVariablePass();
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_REWRITE_IFRT_LOAD_VARIABLE_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/tf_to_mlrt.h

@@ -0,0 +1,48 @@
+/* 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_TFRT_TRANSFORMS_MLRT_TF_TO_MLRT_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_TF_TO_MLRT_H_
+#include <memory>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+// The conversion pass that is run before 'tf-mlrt-parallelization' passes. The
+// parallelization pass changes the graph content, so any rewrite/conversion
+// that depends on the graph instead of individual ops should be done before
+// parallelization.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateTfToMlrtPreParallelizationConversionPass(
+    const TfrtPipelineOptions& options);
+
+// The conversion pass that is run after 'tf-mlrt-parallelization' passes. The
+// parallelization pass changes the graph content, so this pass should only
+// contain conversion that depends on individual ops.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateTfToMlrtConversionPass(const TfrtPipelineOptions& options);
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateTfToMlrtConversionPass(const TfrtPipelineOptions& options,
+                             const tfrt_stub::FallbackState* fallback_state);
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_TF_TO_MLRT_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/tpu_conversion_patterns.cc

@@ -0,0 +1,169 @@
+/* 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/tfrt/transforms/mlrt/tpu_conversion_patterns.h"
+
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/execute_op_registry.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/mlrt_device_constants.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/utils.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+namespace {
+
+class TPUCompileMlirAndExecuteOpPreParallelizationConversion
+    : public mlir::OpConversionPattern<mlir::TF::TPUCompileMlirAndExecuteOp> {
+ public:
+  TPUCompileMlirAndExecuteOpPreParallelizationConversion(
+      mlir::MLIRContext* context, bool use_tpu_host_allocator_for_inputs)
+      : OpConversionPattern(context),
+        use_tpu_host_allocator_for_inputs_(use_tpu_host_allocator_for_inputs) {}
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::TF::TPUCompileMlirAndExecuteOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    llvm::SmallVector<int> constant_operand_indices;
+    llvm::SmallVector<int> non_constant_operand_indices;
+
+    for (int i = 0; i < adaptor.getArgs().size(); ++i) {
+      auto operand = adaptor.getOperands()[i];
+      auto original_operand = op.getOperand(i);
+      if (IsResultVariable(original_operand, operand)) {
+        // NOTE: It's important to populate constant_operand_indices in
+        // ascending order.
+        constant_operand_indices.push_back(i);
+      } else {
+        non_constant_operand_indices.push_back(i);
+      }
+    }
+
+    llvm::SmallVector<mlir::Value> operands = adaptor.getArgs();
+
+    size_t tensor_operands_size = operands.size();
+    operands.append(adaptor.getStaticShapes().begin(),
+                    adaptor.getStaticShapes().end());
+
+    auto producer_name = op->getAttrOfType<mlir::StringAttr>("producer_name");
+
+    llvm::SmallVector<int32_t> operands_with_static_shapes;
+    if (adaptor.getOperandsWithStaticShape().has_value()) {
+      for (auto attr : adaptor.getOperandsWithStaticShapeAttr()
+                           .getAsRange<mlir::IntegerAttr>()) {
+        operands_with_static_shapes.push_back(
+            static_cast<int32_t>(attr.getInt()));
+      }
+    }
+
+    if (use_tpu_host_allocator_for_inputs_) {
+      llvm::DenseMap<mlir::Operation*, mlir::Operation*> replaced_ops;
+
+      for (int i : non_constant_operand_indices) {
+        DCHECK_LT(i, op.getNumOperands());
+        auto old_value = operands[i];
+        mlir::Operation* def = old_value.getDefiningOp();
+
+        if (def && llvm::isa<mlir::TF::TensorFlowDialect>(def->getDialect())) {
+          auto*& op_with_device = replaced_ops[def];
+          if (!op_with_device) {
+            mlir::ConversionPatternRewriter::InsertionGuard guard(rewriter);
+            rewriter.setInsertionPoint(def);
+
+            op_with_device = rewriter.clone(*def);
+            op_with_device->setAttr(kTfMlrtCustomDevice,
+                                    rewriter.getStringAttr(kTpuHostDevice));
+            rewriter.replaceOp(def, op_with_device->getResults());
+          }
+        }
+      }
+    }
+
+    auto compile_and_execute_op = tf_mlrt::TFTPUCompileAndExecuteOp::create(
+        rewriter, op.getLoc(), op.getResultTypes(), operands,
+        rewriter.getDenseI32ArrayAttr(constant_operand_indices),
+        op.getMetadataAttr(), op.getMlirModuleAttr(),
+        rewriter.getUI32IntegerAttr(tensor_operands_size),
+        rewriter.getDenseI32ArrayAttr(operands_with_static_shapes),
+        producer_name);
+
+    rewriter.replaceOp(op, compile_and_execute_op->getResults());
+
+    return mlir::success();
+  }
+
+ private:
+  bool use_tpu_host_allocator_for_inputs_ = false;
+};
+
+class TPUCompileMlirAndExecuteOpConversion
+    : public mlir::OpConversionPattern<tf_mlrt::TFTPUCompileAndExecuteOp> {
+ public:
+  TPUCompileMlirAndExecuteOpConversion(mlir::TypeConverter* type_converter,
+                                       mlir::MLIRContext* context,
+                                       ExecuteOpRegistry* execute_op_registry)
+      : OpConversionPattern(*type_converter, context) {}
+
+  mlir::LogicalResult matchAndRewrite(
+      tf_mlrt::TFTPUCompileAndExecuteOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    llvm::SmallVector<mlir::Value> operands =
+        adaptor.getOperandsAndStaticShapes();
+    llvm::SmallVector<mlir::Type> result_types;
+    result_types.push_back(rewriter.getType<tf_mlrt::TFTensorType>());
+    result_types.append(op.getResults().size(),
+                        rewriter.getType<mlrt::compiler::FutureType>());
+
+    auto compile_and_execute_op = tf_mlrt_tpu::CompileAndExecuteOp::create(
+        rewriter, op.getLoc(), result_types, operands,
+        op.getConstantOperandIndices(), op.getMetadataAttr(),
+        op.getMlirModuleAttr(), op.getNumOperands(),
+        op.getOperandsWithStaticShape(), op.getProducerName());
+
+    rewriter.replaceOp(op, compile_and_execute_op->getResults());
+
+    return mlir::success();
+  }
+};
+
+}  // namespace
+
+void PopulateTpuPreParallelizationConversionPatterns(
+    mlir::ConversionTarget& target, mlir::RewritePatternSet& patterns,
+    const TfrtPipelineOptions& options) {
+  target.addIllegalOp<mlir::TF::TPUCompileMlirAndExecuteOp>();
+  patterns.add<TPUCompileMlirAndExecuteOpPreParallelizationConversion>(
+      patterns.getContext(), options.use_tpu_host_allocator_for_inputs);
+}
+
+void PopulateTpuConversionPatterns(mlir::ConversionTarget& target,
+                                   mlir::RewritePatternSet& patterns,
+                                   mlir::TypeConverter& type_converter,
+                                   ExecuteOpRegistry& execute_op_registry,
+                                   const TfrtPipelineOptions& options) {
+  target.addIllegalOp<tf_mlrt::TFTPUCompileAndExecuteOp>();
+  target.addLegalDialect<tf_mlrt_tpu::TensorflowMlrtTpuDialect>();
+
+  patterns.add<TPUCompileMlirAndExecuteOpConversion>(
+      &type_converter, patterns.getContext(), &execute_op_registry);
+}
+
+void RegisterTpuDialect(mlir::DialectRegistry& registry) {
+  registry.insert<tf_mlrt_tpu::TensorflowMlrtTpuDialect>();
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/tpu_conversion_patterns.h

@@ -0,0 +1,42 @@
+/* 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_TFRT_TRANSFORMS_MLRT_TPU_CONVERSION_PATTERNS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_TPU_CONVERSION_PATTERNS_H_
+
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/mlrt/execute_op_registry.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+void RegisterTpuDialect(mlir::DialectRegistry& registry);
+
+void PopulateTpuPreParallelizationConversionPatterns(
+    mlir::ConversionTarget& target, mlir::RewritePatternSet& patterns,
+    const TfrtPipelineOptions& options);
+
+void PopulateTpuConversionPatterns(mlir::ConversionTarget& target,
+                                   mlir::RewritePatternSet& patterns,
+                                   mlir::TypeConverter& type_converter,
+                                   ExecuteOpRegistry& execute_op_registry,
+                                   const TfrtPipelineOptions& options);
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_TPU_CONVERSION_PATTERNS_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/util.cc

@@ -0,0 +1,45 @@
+/* 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/tfrt/transforms/mlrt/util.h"
+
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h.inc"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_a_m.h.inc"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops_n_z.h.inc"
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+bool UseFallback(mlir::Operation* op) {
+  if (!llvm::isa<mlir::TF::TensorFlowDialect>(op->getDialect())) return false;
+
+  // TODO(b/173017701): have a centralized place to hold the information
+  // whether a TF op should be lowered to FallbackExecute op.
+  // TODO(b/319045348): Define trait to reflect that IfrtLoadVariableOp has no
+  // TF kernels so that we don't need to check every op here.
+  return !llvm::isa<
+      mlir::TF::_TfrtSetResourceOp, mlir::TF::_TfrtGetResourceOp,
+      mlir::TF::BatchFunctionOp, mlir::TF::CaseOp,
+      mlir::TF::IfrtRestoreVariableOp, mlir::TF::StatefulPartitionedCallOp,
+      mlir::TF::PartitionedCallOp, mlir::TF::LegacyCallOp, mlir::TF::IfOp,
+      mlir::TF::WhileOp, mlir::TF::TPUCompileMlirAndExecuteOp,
+      mlir::TF::AsyncIfrtCallOp>(op);
+}
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/mlrt/util.h

@@ -0,0 +1,30 @@
+/* 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_TFRT_TRANSFORMS_MLRT_UTIL_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_UTIL_H_
+
+#include "mlir/IR/Operation.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+// Use fallback by default for anything that does not have a native kernel
+// with some exceptions.
+bool UseFallback(mlir::Operation* op);
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_UTIL_H_

tensorflow/compiler/mlir/tfrt/transforms/mlrt/while_to_map_fn.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_TFRT_TRANSFORMS_MLRT_WHILE_TO_MAP_FN_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_WHILE_TO_MAP_FN_H_
+
+#include <memory>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace mlrt_compiler {
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateWhileToMapFnPass();
+
+}  // namespace mlrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_MLRT_WHILE_TO_MAP_FN_H_

tensorflow/compiler/mlir/tfrt/transforms/optimize.cc

@@ -0,0 +1,181 @@
+/* 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 <cstdint>
+#include <memory>
+#include <utility>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Block.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/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.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/Rewrite/FrozenRewritePatternSet.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/core/util/device_name_utils.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+// Fold tf.DeviceIndex to tf.Const if it has device assigned.
+class FoldDeviceIndex : public mlir::OpRewritePattern<mlir::TF::DeviceIndexOp> {
+ public:
+  using mlir::OpRewritePattern<mlir::TF::DeviceIndexOp>::OpRewritePattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::TF::DeviceIndexOp op,
+      mlir::PatternRewriter &rewriter) const override {
+    auto device = op->getAttrOfType<mlir::StringAttr>("device");
+    if (!device) return mlir::failure();
+
+    DeviceNameUtils::ParsedName parsed_name;
+    if (!DeviceNameUtils::ParseFullName(device.getValue().str(),
+                                        &parsed_name) ||
+        !parsed_name.has_type)
+      return mlir::failure();
+
+    int32_t i = 0;
+    mlir::ArrayAttr device_names = op.getDeviceNames();
+    for (; i < device_names.size(); ++i) {
+      auto device_name =
+          mlir::cast<mlir::StringAttr>(device_names[i]).getValue();
+      if (device_name == parsed_name.type) break;
+    }
+
+    rewriter.replaceOpWithNewOp<mlir::TF::ConstOp>(
+        op,
+        mlir::DenseIntElementsAttr::get(
+            mlir::RankedTensorType::get(/*shape=*/{}, rewriter.getI32Type()),
+            i));
+
+    return mlir::success();
+  }
+};
+
+// A custom hash and compare function for finding out common ops.
+struct SimpleOperationInfo : public llvm::DenseMapInfo<mlir::Operation *> {
+  static unsigned getHashValue(const mlir::Operation *opC) {
+    return mlir::OperationEquivalence::computeHash(
+        const_cast<mlir::Operation *>(opC),
+        /*hashOperands=*/mlir::OperationEquivalence::directHashValue,
+        /*hashResults=*/mlir::OperationEquivalence::ignoreHashValue,
+        mlir::OperationEquivalence::IgnoreLocations);
+  }
+  static bool isEqual(const mlir::Operation *lhsC,
+                      const mlir::Operation *rhsC) {
+    auto *lhs = const_cast<mlir::Operation *>(lhsC);
+    auto *rhs = const_cast<mlir::Operation *>(rhsC);
+    if (lhs == rhs) return true;
+    return mlir::OperationEquivalence::isEquivalentTo(
+        const_cast<mlir::Operation *>(lhsC),
+        const_cast<mlir::Operation *>(rhsC),
+        mlir::OperationEquivalence::IgnoreLocations);
+  }
+};
+
+void EliminateCommonMultinomialOps(mlir::Block &block) {
+  llvm::SmallDenseMap<mlir::Operation *,
+                      llvm::SmallVector<mlir::TF::MultinomialOp>, 2,
+                      SimpleOperationInfo>
+      multinomial_to_eliminate;
+
+  auto eliminate = [&]() {
+    auto &list = multinomial_to_eliminate.begin()->second;
+    auto first = list.front();
+    for (auto op : llvm::drop_begin(list)) {
+      op.getOutput().replaceAllUsesWith(first.getOutput());
+      op->erase();
+    }
+    multinomial_to_eliminate.clear();
+  };
+
+  for (auto &op : block) {
+    auto multinomial_op = llvm::dyn_cast<mlir::TF::MultinomialOp>(&op);
+    // Conservatively, we only eliminate back-to-back tf.Multinomial ops.
+    if (multinomial_op) {
+      if (multinomial_to_eliminate.find(multinomial_op) ==
+              multinomial_to_eliminate.end() &&
+          !multinomial_to_eliminate.empty()) {
+        // If the current op is a tf.Multinomial but it is different from the
+        // preiously found tf.Multinomial, then we eliminate the prviously found
+        // tf.Multinomial.
+        eliminate();
+      }
+      multinomial_to_eliminate[multinomial_op].push_back(multinomial_op);
+    } else if (!multinomial_to_eliminate.empty()) {
+      // If the current op is not a tf.Multinomial, then we eliminate previously
+      // found tf.Multinomial
+      eliminate();
+    }
+  }
+}
+
+// Optimization pass for TFRT-specific rewrite patterns.
+class OptimizeTfForTfrt
+    : public mlir::PassWrapper<OptimizeTfForTfrt,
+                               mlir::OperationPass<mlir::func::FuncOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OptimizeTfForTfrt)
+
+  llvm::StringRef getArgument() const final { return "optimize-tf-for-tfrt"; }
+
+  llvm::StringRef getDescription() const final {
+    return "optmize TF MLIR for TFRT workflow.";
+  }
+
+  mlir::LogicalResult initialize(mlir::MLIRContext *context) override {
+    mlir::RewritePatternSet pattern_list(context);
+    pattern_list.add<FoldDeviceIndex>(context);
+    patterns_ = std::move(pattern_list);
+    return mlir::success();
+  }
+
+  void runOnOperation() override {
+    auto func = getOperation();
+
+    EliminateCommonMultinomialOps(func.getBody().front());
+
+    if (mlir::failed(mlir::applyPatternsGreedily(func, patterns_)))
+      signalPassFailure();
+  }
+
+ private:
+  mlir::FrozenRewritePatternSet patterns_;
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateOptimizeTfForTfrtPass() {
+  return std::make_unique<OptimizeTfForTfrt>();
+}
+
+static mlir::PassRegistration<OptimizeTfForTfrt> register_pass;
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/optimize_tf_control_flow_side_effect.cc

@@ -0,0 +1,186 @@
+/* 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 <memory>
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/analysis/tensor_array_side_effect_analysis.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/utils.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+bool IsOpNonSideEffectingOrReadOnly(mlir::Operation* op,
+                                    bool include_read_only = true) {
+  // Though tf.Assert and tf.Timestamp are side-effecting, they do not
+  // interfere with any other side-effecting ops. For now, if control flow
+  // ops' callee functions contain them, we treat them as non-side-effecting.
+  if (llvm::isa<mlir::TF::AssertOp, mlir::TF::TimestampOp>(op)) return true;
+
+  // TensorArray ops have explicit operand/result for handling side effects.
+  if (IsTensorArrayOp(op)) return true;
+
+  if (include_read_only) {
+    if (mlir::hasSingleEffect<mlir::MemoryEffects::Allocate>(op) ||
+        mlir::hasSingleEffect<mlir::MemoryEffects::Read>(op) ||
+        mlir::hasSingleEffect<mlir::MemoryEffects::Free>(op))
+      return true;
+  }
+
+  return mlir::isMemoryEffectFree(op);
+}
+
+bool FunctionHasSideEffect(
+    mlir::func::FuncOp func_op, bool include_read_only,
+    llvm::DenseMap<mlir::func::FuncOp, bool>& function_side_effect) {
+  auto iter = function_side_effect.find(func_op);
+  if (iter != function_side_effect.end()) return iter->second;
+
+  auto& block = func_op.front();
+
+  auto op_has_side_effect = [&](mlir::Operation* op) {
+    if (auto while_op = llvm::dyn_cast<mlir::TF::WhileOp>(op)) {
+      if (while_op.getIsStateless()) return false;
+
+      return FunctionHasSideEffect(while_op.cond_function(), include_read_only,
+                                   function_side_effect) ||
+             FunctionHasSideEffect(while_op.body_function(), include_read_only,
+                                   function_side_effect);
+    }
+
+    if (auto if_op = llvm::dyn_cast<mlir::TF::IfOp>(op)) {
+      if (if_op.getIsStateless()) return false;
+
+      return FunctionHasSideEffect(if_op.else_function(), include_read_only,
+                                   function_side_effect) ||
+             FunctionHasSideEffect(if_op.then_function(), include_read_only,
+                                   function_side_effect);
+    }
+
+    return !IsOpNonSideEffectingOrReadOnly(op, include_read_only);
+  };
+
+  // Speculatively setting the function to have no side effect to avoid infinite
+  // recursion. The correct side effect will be updated later once more
+  // operations in the block are checked.
+  function_side_effect[func_op] = false;
+
+  for (mlir::Operation& op : block) {
+    if (op_has_side_effect(&op)) {
+      function_side_effect[func_op] = true;
+      return true;
+    }
+  }
+
+  function_side_effect[func_op] = false;
+  return false;
+}
+
+// This pass sets `is_stateless` attribute of tf.If and tf.While ops to true if
+// their callee functions contains only non-side-effecting ops.
+class OptimizeTfControlFlowSideEffectPass
+    : public mlir::PassWrapper<OptimizeTfControlFlowSideEffectPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      OptimizeTfControlFlowSideEffectPass)
+
+ private:
+  llvm::StringRef getArgument() const final {
+    return "tfrt-optimize-tf-control-flow-side-effect";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Set tf control flow ops to stateless if their callee functions "
+           "contains only non-side-effecting ops";
+  }
+  void runOnOperation() override {
+    auto module = getOperation();
+
+    llvm::SmallVector<mlir::func::FuncOp> functions;
+
+    bool include_read_only = true;
+    for (auto func : module.getOps<mlir::func::FuncOp>()) {
+      if (func.isPublic() && !IsSessionInitializer(func)) {
+        functions.push_back(func);
+
+        for (auto& op : func.front()) {
+          // Skip control flow ops in the first pass.
+          if (llvm::isa<mlir::TF::WhileOp, mlir::TF::IfOp>(&op)) continue;
+
+          if (!IsOpNonSideEffectingOrReadOnly(&op)) {
+            include_read_only = false;
+          }
+        }
+      }
+    }
+
+    llvm::DenseMap<mlir::func::FuncOp, bool> function_side_effect;
+
+    mlir::Builder builder(module.getContext());
+
+    for (auto func : functions) {
+      func.walk([&](mlir::Operation* op) {
+        if (auto while_op = llvm::dyn_cast<mlir::TF::WhileOp>(op)) {
+          if (while_op.getIsStateless()) return;
+
+          if (!FunctionHasSideEffect(while_op.cond_function(),
+                                     include_read_only, function_side_effect) &&
+              !FunctionHasSideEffect(while_op.body_function(),
+                                     include_read_only, function_side_effect)) {
+            while_op->setAttr("is_stateless", builder.getBoolAttr(true));
+          }
+        }
+
+        if (auto if_op = llvm::dyn_cast<mlir::TF::IfOp>(op)) {
+          if (if_op.getIsStateless()) return;
+
+          if (!FunctionHasSideEffect(if_op.else_function(), include_read_only,
+                                     function_side_effect) &&
+              !FunctionHasSideEffect(if_op.then_function(), include_read_only,
+                                     function_side_effect)) {
+            if_op->setAttr("is_stateless", builder.getBoolAttr(true));
+          }
+        }
+      });
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateOptimizeTfControlFlowSideEffectPass() {
+  return std::make_unique<OptimizeTfControlFlowSideEffectPass>();
+}
+
+static mlir::PassRegistration<OptimizeTfControlFlowSideEffectPass>
+    register_pass(CreateOptimizeTfControlFlowSideEffectPass);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/passes.cc

@@ -0,0 +1,282 @@
+
+/* 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/tfrt/transforms/passes.h"
+
+#include <cassert>
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/PassOptions.h"
+#include "mlir/Transforms/Passes.h"
+#include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
+#include "absl/status/status.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_saved_model_asset_sinking_pass.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/bridge_logger.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/set_shape_invariant_in_while_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/util/device_name_utils.h"
+#include "tsl/platform/errors.h"
+
+namespace tensorflow {
+namespace {
+
+// Assigns devices so that later passes can utilize device information.
+// Device assignment might have not been done by the upstream pipeline, or get
+// removed by previous passes. However, we assume most of the device assignment
+// has been done by the upstream pipeline, so we simply assign the default
+// device to unassigned ops. Specifically, we do assignment for ConstOp first to
+// place it on the same device as its user operation, instead of placing it on
+// the default device blindly.
+// TODO(b/221297389): Figure out a more robust way to handle dropped device
+// assignment.
+void AddTfDeviceAssignmentPasses(mlir::OpPassManager &pm,
+                                 const TfrtPipelineOptions &options) {
+  pm.addPass(mlir::TF::CreateConstantOpDeviceAssignmentPass());
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::TF::CreateTFDeviceAssignmentByFuncAttrPass());
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::TF::CreateSimpleTFDeviceAssignmentPass(options.default_device));
+}
+
+}  // namespace
+
+void CreateTFExecutorToTFPreInvariantOptimizationPipelineHelper(
+    mlir::OpPassManager &pm, const TfrtPipelineOptions &options) {
+  // Due to b/191304670, functionalized while ops might not have the
+  // shape_invariant attribute set correctly, which leads to failure in shape
+  // inference. As a workaround, we conservatively (e.g., we place less
+  // restrictions on tf.while which will avoid failures but lead to potentially
+  // less exact shape inference) set the shape_invariant attribute in all
+  // tf.While ops before performing shape inference.
+  //
+  // Note that this pass might not work well with TF XLA bridge, but this is
+  // fine as TF XLA bridge is run before this pipeline. For CPU ops, less exact
+  // shape inference may lead to fewer optimizations but it should be fine as it
+  // is limited to while ops currently.
+  //
+  // TODO(b/191304670): Remove this pass once the shape_invariant attribute is
+  // set correctly in the upstream.
+  pm.addNestedPass<mlir::func::FuncOp>(
+      tfrt_compiler::CreateSetShapeInvariantInWhileOps());
+
+  // We pass the MLIR module through the TF standard pipeline, which for
+  // instances does shape inference, canonicalization, inlining, etc.
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::tf_executor::CreateTFExecutorGraphPruningPass());
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::tf_executor::CreateTFExecutorIslandCoarseningPass());
+
+  AddTfDeviceAssignmentPasses(pm, options);
+
+  if (options.allow_xla_cpu) {
+    pm.addPass(tfrt_compiler::CreateTfrtXlaRewritePass());
+  }
+
+  // Here we perform TFRT specific optimization before standard TF optimization,
+  // as TFRT-specific optimization may create more opportunities.
+  pm.addNestedPass<mlir::func::FuncOp>(
+      tfrt_compiler::CreateOptimizeTfForTfrtPass());
+
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::CreateExecutorDialectToFunctionalConversionPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+  // Guarantee all functions have one use, which enables more exact shape
+  // inference.
+  pm.addPass(mlir::TF::CreateGuaranteeAllFuncsOneUsePass());
+  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
+  pm.addPass(mlir::createInlinerPass());
+  pm.addPass(mlir::createSymbolDCEPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::TF::CreateTFOptimizePass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCSEPass());
+
+  AddTfDeviceAssignmentPasses(pm, options);
+
+  // After the standard pass, we now have MLIR in TF dialect, and now we convert
+  // reference variable to resource variables, which is best effort.
+  pm.addPass(CreateConvertReferenceVariableToResourceVariablePass());
+
+  // Move the tf.Assert op to the end of the function, so that it does not
+  // impose unnecessary control dependencies on other ops.
+  pm.addPass(tfrt_compiler::CreateReorderTfAssertPass());
+
+  // Optimize the side-effects of control flow ops by examining the ops in its
+  // callees.
+  pm.addPass(tfrt_compiler::CreateOptimizeTfControlFlowSideEffectPass());
+
+  // Remove tf.If ops' operands that are produced by tf.Const ops.
+  pm.addPass(tfrt_compiler::CreateRemoveTfIfConstArgsPass());
+
+  // Merge non-side-effecting tf.If ops if their operands are the same.
+  pm.addPass(tfrt_compiler::CreateMergeTfIfOpsPass());
+
+  pm.addPass(tfrt_compiler::CreateReconfigBatchOpPass({
+      .min_num_batch_threads = options.min_num_batch_threads,
+      .min_max_enqueued_batches = options.min_max_enqueued_batches,
+      .batch_padding_policy = options.batch_padding_policy,
+      .num_batch_threads = options.num_batch_threads,
+      .max_batch_size = options.max_batch_size,
+      .batch_timeout_micros = options.batch_timeout_micros,
+      .allowed_batch_sizes = options.allowed_batch_sizes,
+      .max_enqueued_batches = options.max_enqueued_batches,
+      .low_priority_max_batch_size = options.low_priority_max_batch_size,
+      .low_priority_batch_timeout_micros =
+          options.low_priority_batch_timeout_micros,
+      .low_priority_allowed_batch_sizes =
+          options.low_priority_allowed_batch_sizes,
+      .low_priority_max_enqueued_batches =
+          options.low_priority_max_enqueued_batches,
+      .num_warmup_batch_threads = options.num_warmup_batch_threads,
+      .enable_large_batch_splitting = options.enable_large_batch_splitting,
+      .mixed_priority_batching_policy = options.mixed_priority_batching_policy,
+      .batch_queue_global_prioritization_num_threads =
+          options.batch_queue_global_prioritization_num_threads,
+      .enable_priority_aware_batch_scheduler =
+          options.enable_priority_aware_batch_scheduler,
+      .enable_priority_aware_batch_scheduler_resplit =
+          options.enable_priority_aware_batch_scheduler_resplit,
+      .enable_batching_task_lazy_cancellation =
+          options.enable_batching_task_lazy_cancellation,
+  }));
+
+  // Deduplicate functions invoked by tf.BatchFunction with the same
+  // shared_name
+  pm.addPass(
+      tfrt_compiler::CreateDeduplicateFunctionsInovkedByBatchFunctionPass());
+
+  // RemoveUnusedWhileResultsPass operates on the region-based control flow, so
+  // the functional control flow is first converted to region-based control
+  // flow, which is converted back after the optimization passes are performed.
+  pm.addPass(mlir::TF::CreateTFFunctionalControlFlowToRegions());
+  pm.addPass(mlir::createInlinerPass());
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::TF::CreateRemoveUnusedWhileResultsPass());
+
+  // Apply standard optimization after optimizing control flow ops.
+  pm.addPass(mlir::createInlinerPass());
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCSEPass());
+
+  // TODO(b/187876545): An extra shape inference pass is added because it does
+  // not work well with tf.Identity op that remove ref type. So we work around
+  // by performing shape inference again after reference variable to resource
+  // variable conversion. We should remove this after b/187876545 is fixed.
+  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
+  pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());
+
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::TFDevice::CreateLaunchToDeviceAttributePass());
+
+  // After all standard passes run layout optimization to assign optimal data
+  // format for all layout sensitive operations.
+  mlir::TF::LayoutOptimizationPipelineOptions layout_optimization_options;
+  layout_optimization_options.force_data_format =
+      options.force_data_format.getValue();
+  // TODO(b/191304261): Folding transpose in ops is buggy in the layout
+  // optimization pass. Disable it to avoid errors in b/191304261. This should
+  // not affect CPU performance as it does not change the number of ops, nor
+  // does it change the types of the ops.
+  layout_optimization_options.skip_fold_transpose_in_ops = true;
+  mlir::TF::CreateLayoutOptimizationPipeline(pm.nest<mlir::func::FuncOp>(),
+                                             layout_optimization_options);
+
+  // Run canonicalization pipeline to remove unused constants and bypassed
+  // transpose operations left in the IR after layout optimization.
+  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+
+  // Decompose resource ops as resource variables will be converted to tensors
+  // directly.
+  if (options.decompose_resource_ops)
+    pm.addNestedPass<mlir::func::FuncOp>(
+        mlir::TFDevice::CreateDecomposeResourceOpsPass());
+
+  AddTfDeviceAssignmentPasses(pm, options);
+
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::TF::CreateTensorDeviceCopyConversionPass());
+
+  // Rewriter operation sequences to device specific fusions.
+  DeviceNameUtils::ParsedName parsed_name;
+
+  // Ignore error.
+  bool success =
+      DeviceNameUtils::ParseFullName(options.default_device, &parsed_name);
+  assert(success && "default device is invalid");
+  (void)success;
+
+  if (parsed_name.has_type && parsed_name.type == DEVICE_GPU)
+    pm.addNestedPass<mlir::func::FuncOp>(mlir::TF::CreateGpuOpFusionPass());
+
+  if (parsed_name.has_type && parsed_name.type == DEVICE_CPU)
+    pm.addNestedPass<mlir::func::FuncOp>(
+        mlir::TF::CreateFusedKernelMatcherPass());
+
+  if (options.tpu_fuse_ops) {
+    pm.addNestedPass<mlir::func::FuncOp>(
+        tfrt_compiler::CreateFuseTpuCompileAndExecutePass());
+    // Remove ops for the input to _TPUCompileMlirOp, which are no longer needed
+    // after CreateFuseTpuCompileAndExecutePass
+    pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+  }
+
+  AddTfDeviceAssignmentPasses(pm, options);
+}
+
+void CreateTFInvariantOptimizationPipelineHelper(
+    mlir::OpPassManager& pm, const TfrtPipelineOptions& options) {
+  if (options.sink_in_invariant_ops) {
+    pm.addPass(CreateSinkInInvariantOpsPass());
+  }
+
+  if (!options.saved_model_dir.empty()) {
+    pm.addPass(
+        mlir::tf_saved_model::CreateAssetSinkingPass(options.saved_model_dir));
+  }
+
+  pm.addPass(CreateLowerTFSavedModelPass(
+      options.hoist_invariant_ops, options.fuse_get_resource_ops_in_hoisting));
+}
+
+absl::Status ValidateTfrtPipelineOptions(const TfrtPipelineOptions &options) {
+  if (options.target_tpurt && options.target_gpu) {
+    return absl::InternalError(
+        "Invalid pipeline options. Targeting both TPU and GPU is not "
+        "supported.");
+  }
+  return absl::OkStatus();
+}
+
+absl::Status CreateTFExecutorToTFPreInvariantOptimizationPipeline(
+    mlir::PassManager &pm, const TfrtPipelineOptions &options) {
+  TF_RETURN_IF_ERROR(ValidateTfrtPipelineOptions(options));
+  if (VLOG_IS_ON(1)) {
+    // Print the whole module after each pass, which requires disabling
+    // multi-threading as well.
+    pm.getContext()->disableMultithreading();
+    pm.enableIRPrinting(std::make_unique<tensorflow::BridgeLoggerConfig>(
+        /*print_module_scope=*/true));
+  }
+  CreateTFExecutorToTFPreInvariantOptimizationPipelineHelper(pm, options);
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

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

@@ -0,0 +1,182 @@
+/* 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_TFRT_TRANSFORMS_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_PASSES_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "absl/status/status.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/CommandLine.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassOptions.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/analysis/side_effect_analysis.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/tpu_passes.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace mlir {
+class PassManager;
+}
+
+namespace tensorflow {
+
+namespace tfrt_compiler {
+
+// Create a pass to insert kernels that copy fallback tensors when they are
+// passed to multiple threads, to avoid atomic contention on their refcounts.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateInsertFallbackTensorCopyPass();
+
+// Create a pass to reorder tf.Assert ops or tf.If ops that contains only
+// tf.Assert ops to the end of the function, to avoid unnecessary control
+// dependencies to other ops.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateReorderTfAssertPass();
+
+// Create a pass to optimize the side-effect of control flow ops. eg. if both
+// branches of a tf.If op contains only non-side-effecting ops, its
+// `is_stateless` attribute will be set to true.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateOptimizeTfControlFlowSideEffectPass();
+
+// Create a pass to remove tf.If ops' operands that are produced by tf.Const
+// ops.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateRemoveTfIfConstArgsPass();
+
+// Create a pass to merge non-side-effecting tf.If ops that have the same
+// operands.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateMergeTfIfOpsPass();
+
+// Create a pass to deduplicate the function invoked by tf.BatchFunction with
+// the same shared_name.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateDeduplicateFunctionsInovkedByBatchFunctionPass();
+
+// Create a pass to lower bound the number of threads in tf.BatchFunction.
+struct ReconfigBatchOpPassOptions {
+  int64_t min_num_batch_threads = 1;
+  int64_t min_max_enqueued_batches = 1;
+  std::string batch_padding_policy = "";
+  int64_t num_batch_threads = 0;
+  int64_t max_batch_size = 0;
+  int64_t batch_timeout_micros = 0;
+  llvm::ArrayRef<int64_t> allowed_batch_sizes = {};
+  int64_t max_enqueued_batches = 0;
+  int64_t low_priority_max_batch_size = 0;
+  int64_t low_priority_batch_timeout_micros = 0;
+  llvm::ArrayRef<int64_t> low_priority_allowed_batch_sizes = {};
+  int64_t low_priority_max_enqueued_batches = 0;
+  int64_t num_warmup_batch_threads = 0;
+  bool enable_large_batch_splitting = false;
+  std::string mixed_priority_batching_policy = "";
+  int64_t batch_queue_global_prioritization_num_threads = 0;
+  bool enable_priority_aware_batch_scheduler = false;
+  bool enable_priority_aware_batch_scheduler_resplit = false;
+  bool enable_batching_task_lazy_cancellation = false;
+};
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateReconfigBatchOpPass(
+    ReconfigBatchOpPassOptions options);
+
+// Create a pass to fuse the TPU Ops for TFRT.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateFuseTpuCompileAndExecutePass();
+
+// Create a pass to optimize TF dialect for TFRT workflow.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateOptimizeTfForTfrtPass();
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateTfrtXlaRewritePass();
+
+// Create a pass to deduplicate results of tf.If ops.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateDeduplicateIfResultPass();
+
+}  // namespace tfrt_compiler
+
+class CoreRTConverter;
+
+// Create a pass that sink in the var handle op to the callee function when
+// proper.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateSinkInInvariantOpsPass();
+
+// Create a pass that rewrites tf_saved_model dialect's ops according to TFRT's
+// requirements.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateLowerTFSavedModelPass(bool hoist_invariant_ops,
+                            bool fuse_get_resource_ops);
+
+// Create a pass that converts ref variables to resource variables in a limited
+// number of cases.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateConvertReferenceVariableToResourceVariablePass();
+
+// Run *ToCoreRTConversionPassRun as free functions. Useful for
+// reusing the pass logic in a custom pass with additional conversions.
+mlir::LogicalResult TFSavedModelToCoreRTConversionPassRun(
+    mlir::MLIRContext* context, mlir::func::FuncOp func,
+    mlir::ConversionTarget* target, mlir::RewritePatternSet* patterns,
+    CoreRTConverter* corert_converter);
+
+// Create an operation pass that removes the device attribute from every
+// corert.executeop.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateRemoveDeviceAttributePass();
+
+// Create an operation pass that inserts corert.transfer op to make sure any
+// argument of any op is on the same device of the op itself.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateCrossDeviceTransferPass();
+
+// Create a pass that converts MLIR TF dialect to MLIR TFRT dialect.
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateTfToTfrtConversionPass(const TfrtPipelineOptions& options);
+
+// Creates a pipeline of passes that lowers MLIR TF dialect to TFRT dialects.
+void CreateTfToTfrtPipeline(mlir::OpPassManager& pm,
+                            const TfrtPipelineOptions& options);
+
+// Creates a pipeline of passes that lowers MLIR TF dialect from tf.function to
+// TFRT dialect. SavedModel related conversions are not included.
+absl::Status CreateTfExecutorToTfrtPipeline(mlir::PassManager& pm,
+                                            const TfrtPipelineOptions& options);
+
+// Creates a pipeline of passes that lowers MLIR TF Executor dialect to TF
+// dialect for CoreRT purposes.
+absl::Status CreateTFExecutorToTFPipeline(mlir::PassManager& pm,
+                                          const TfrtPipelineOptions& options);
+
+// TODO(deqiangc): refactor below helpers once mlrt is OSSed.
+void CreateTFExecutorToTFPreInvariantOptimizationPipelineHelper(
+    mlir::OpPassManager& pm, const TfrtPipelineOptions& options);
+void CreateTFInvariantOptimizationPipelineHelper(
+    mlir::OpPassManager& pm, const TfrtPipelineOptions& options);
+
+absl::Status CreateTFExecutorToTFPreInvariantOptimizationPipeline(
+    mlir::PassManager& pm, const TfrtPipelineOptions& options);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_PASSES_H_

tensorflow/compiler/mlir/tfrt/transforms/reconfig_batch_op.cc

@@ -0,0 +1,273 @@
+/* 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 <algorithm>
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/CommandLine.h"
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/core/platform/cpu_info.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+class ReconfigBatchOpPass
+    : public mlir::PassWrapper<ReconfigBatchOpPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  explicit ReconfigBatchOpPass(ReconfigBatchOpPassOptions options)
+      : mlir::PassWrapper<ReconfigBatchOpPass,
+                          mlir::OperationPass<mlir::ModuleOp>>() {
+    min_num_batch_threads_ = options.min_num_batch_threads;
+    min_max_enqueued_batches_ = options.min_max_enqueued_batches;
+    batch_padding_policy_ = options.batch_padding_policy;
+    num_batch_threads_ = options.num_batch_threads;
+    max_batch_size_ = options.max_batch_size;
+    batch_timeout_micros_ = options.batch_timeout_micros;
+    allowed_batch_sizes_ = options.allowed_batch_sizes;
+    max_enqueued_batches_ = options.max_enqueued_batches;
+    low_priority_max_batch_size_ = options.low_priority_max_batch_size;
+    low_priority_batch_timeout_micros_ =
+        options.low_priority_batch_timeout_micros;
+    low_priority_allowed_batch_sizes_ =
+        options.low_priority_allowed_batch_sizes;
+    low_priority_max_enqueued_batches_ =
+        options.low_priority_max_enqueued_batches;
+    num_warmup_batch_threads_ = options.num_warmup_batch_threads;
+    enable_large_batch_splitting_ = options.enable_large_batch_splitting;
+    mixed_priority_batching_policy_ = options.mixed_priority_batching_policy;
+    batch_queue_global_prioritization_num_threads_ =
+        options.batch_queue_global_prioritization_num_threads;
+    enable_priority_aware_batch_scheduler_ =
+        options.enable_priority_aware_batch_scheduler;
+    enable_priority_aware_batch_scheduler_resplit_ =
+        options.enable_priority_aware_batch_scheduler_resplit;
+    enable_batching_task_lazy_cancellation_ =
+        options.enable_batching_task_lazy_cancellation;
+  }
+  ReconfigBatchOpPass()
+      : mlir::PassWrapper<ReconfigBatchOpPass,
+                          mlir::OperationPass<mlir::ModuleOp>>() {}
+  ReconfigBatchOpPass(const ReconfigBatchOpPass& other)
+      : mlir::PassWrapper<ReconfigBatchOpPass,
+                          mlir::OperationPass<mlir::ModuleOp>>(other) {}
+
+  ReconfigBatchOpPass& operator=(const ReconfigBatchOpPass& other) = delete;
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReconfigBatchOpPass)
+
+ private:
+  llvm::StringRef getArgument() const final { return "tfrt-reconfig-batch-op"; }
+
+  llvm::StringRef getDescription() const final {
+    return "Reconfig batch op such as num_batch_threads and "
+           "max_enqueued_batches.";
+  }
+
+  void runOnOperation() override {
+    if (min_num_batch_threads_ == 0 && min_max_enqueued_batches_ == 0 &&
+        batch_padding_policy_.empty() && num_batch_threads_ == 0 &&
+        max_batch_size_ == 0 && batch_timeout_micros_ == 0 &&
+        allowed_batch_sizes_.empty() && max_enqueued_batches_ == 0 &&
+        low_priority_max_batch_size_ == 0 &&
+        low_priority_batch_timeout_micros_ == 0 &&
+        low_priority_allowed_batch_sizes_.empty() &&
+        low_priority_max_enqueued_batches_ == 0 &&
+        num_warmup_batch_threads_ == 0 &&
+        !enable_priority_aware_batch_scheduler_ &&
+        !enable_priority_aware_batch_scheduler_resplit_ &&
+        !enable_batching_task_lazy_cancellation_) {
+      return;
+    }
+    mlir::ModuleOp module = getOperation();
+    module.walk([&](mlir::TF::BatchFunctionOp batch_op) {
+      int64_t num_batch_threads = batch_op.getNumBatchThreads();
+      num_batch_threads =
+          std::max(num_batch_threads, min_num_batch_threads_.getValue());
+      batch_op.setNumBatchThreads(num_batch_threads);
+
+      int64_t max_enqueued_batches = batch_op.getMaxEnqueuedBatches();
+      max_enqueued_batches =
+          std::max(max_enqueued_batches, min_max_enqueued_batches_.getValue());
+      batch_op.setMaxEnqueuedBatches(max_enqueued_batches);
+
+      if (!batch_padding_policy_.empty()) {
+        batch_op.setBatchPaddingPolicy(batch_padding_policy_);
+      }
+      if (num_batch_threads_ > 0) {
+        batch_op.setNumBatchThreads(num_batch_threads_);
+      }
+      if (max_batch_size_ > 0) {
+        batch_op.setMaxBatchSize(max_batch_size_);
+      }
+      if (batch_timeout_micros_ > 0) {
+        batch_op.setBatchTimeoutMicros(batch_timeout_micros_);
+      }
+      if (!allowed_batch_sizes_.empty()) {
+        batch_op.setAllowedBatchSizesAttr(
+            mlir::Builder(module.getContext())
+                .getI64ArrayAttr(allowed_batch_sizes_));
+      }
+      if (max_enqueued_batches_ > 0) {
+        batch_op.setMaxEnqueuedBatches(max_enqueued_batches_);
+      }
+      if (low_priority_max_batch_size_ > 0) {
+        batch_op.setLowPriorityMaxBatchSize(low_priority_max_batch_size_);
+      }
+      if (low_priority_batch_timeout_micros_ > 0) {
+        batch_op.setLowPriorityBatchTimeoutMicros(
+            low_priority_batch_timeout_micros_);
+      }
+      if (!low_priority_allowed_batch_sizes_.empty()) {
+        batch_op.setLowPriorityAllowedBatchSizesAttr(
+            mlir::Builder(module.getContext())
+                .getI64ArrayAttr(low_priority_allowed_batch_sizes_));
+      }
+      if (low_priority_max_enqueued_batches_ > 0) {
+        batch_op.setLowPriorityMaxEnqueuedBatches(
+            low_priority_max_enqueued_batches_);
+      }
+      if (num_warmup_batch_threads_ > 0) {
+        batch_op.setNumWarmupBatchThreads(num_warmup_batch_threads_);
+      }
+      if (enable_large_batch_splitting_) {
+        batch_op.setEnableLargeBatchSplittingAttr(
+            mlir::Builder(module.getContext()).getBoolAttr(true));
+      }
+      if (!mixed_priority_batching_policy_.empty()) {
+        batch_op.setMixedPriorityPolicy(mixed_priority_batching_policy_);
+      }
+      if (batch_queue_global_prioritization_num_threads_ > 0) {
+        batch_op.setNumBatchThreads(
+            batch_queue_global_prioritization_num_threads_);
+        batch_op.setNumWarmupBatchThreads(port::MaxParallelism());
+        batch_op.setEnablePriorityAwareBatchScheduler(true);
+      }
+      if (enable_priority_aware_batch_scheduler_) {
+        batch_op.setEnablePriorityAwareBatchScheduler(true);
+      }
+      if (enable_priority_aware_batch_scheduler_resplit_) {
+        batch_op.setEnablePriorityAwareBatchSchedulerResplit(true);
+      }
+      if (enable_batching_task_lazy_cancellation_) {
+        batch_op.setEnableBatchingTaskLazyCancellation(true);
+      }
+    });
+  }
+
+ protected:
+  mlir::Pass::Option<int64_t> min_num_batch_threads_{
+      *this, "tfrt-min-num-batch-threads", llvm::cl::init(1),
+      llvm::cl::desc("Minimum number of batch threads")};
+  mlir::Pass::Option<int64_t> min_max_enqueued_batches_{
+      *this, "tfrt-min-max-enqueued-batches", llvm::cl::init(1),
+      llvm::cl::desc(
+          "Minimum of the maximum number of outstanding enqueued batches")};
+  mlir::Pass::Option<std::string> batch_padding_policy_{
+      *this, "tfrt-batch-padding-policy", llvm::cl::init(""),
+      llvm::cl::desc("The policy used when padding (or splitting) batches.")};
+  mlir::Pass::Option<int64_t> num_batch_threads_{
+      *this, "tfrt-num-batch-threads", llvm::cl::init(0),
+      llvm::cl::desc(
+          "The number of threads for processing batches in parallel")};
+  mlir::Pass::Option<int64_t> max_batch_size_{
+      *this, "tfrt-max-batch-size", llvm::cl::init(0),
+      llvm::cl::desc("The maximum allowed batch size")};
+  mlir::Pass::Option<int64_t> batch_timeout_micros_{
+      *this, "tfrt-batch-timeout-micros", llvm::cl::init(0),
+      llvm::cl::desc("The maximum number of microseconds before outputting an "
+                     "incomplete batch")};
+  mlir::Pass::ListOption<int64_t> allowed_batch_sizes_{
+      *this, "tfrt-allowed-batch-sizes",
+      llvm::cl::desc("Allowed sizes for padding (or splitting) batches")};
+  mlir::Pass::Option<int64_t> max_enqueued_batches_{
+      *this, "tfrt-max-enqueued-batches", llvm::cl::init(0),
+      llvm::cl::desc("The maximum number of batches enqueued for processing "
+                     "before requests are failed fast")};
+  mlir::Pass::Option<int64_t> low_priority_max_batch_size_{
+      *this, "tfrt-low-priority-max-batch-size", llvm::cl::init(0),
+      llvm::cl::desc(
+          "The maximum allowed batch size for low priority requests")};
+  mlir::Pass::Option<int64_t> low_priority_batch_timeout_micros_{
+      *this, "tfrt-low-priority-batch-timeout-micros", llvm::cl::init(0),
+      llvm::cl::desc("The maximum number of microseconds before outputting an "
+                     "incomplete batch for low priority requests")};
+  mlir::Pass::ListOption<int64_t> low_priority_allowed_batch_sizes_{
+      *this, "tfrt-low-priority-allowed-batch-sizes",
+      llvm::cl::desc("Allowed sizes for padding (or splitting) batches for low "
+                     "priority requests")};
+  mlir::Pass::Option<int64_t> low_priority_max_enqueued_batches_{
+      *this, "tfrt-low-priority-max-enqueued-batches", llvm::cl::init(0),
+      llvm::cl::desc("The maximum number of batches enqueued for processing "
+                     "before low priority requests are failed fast")};
+  // TODO(b/516818455): Remove warmup threads after in-lining warmup requests.
+  mlir::Pass::Option<int64_t> num_warmup_batch_threads_{
+      *this, "tfrt-num-warmup-batch-threads", llvm::cl::init(0),
+      llvm::cl::desc(
+          "The number of threads for processing warmup requests. "
+          "Useful to process warmup requests without starving the "
+          "regular batch threads when global scheduler is enabled.")};
+  mlir::Pass::Option<bool> enable_large_batch_splitting_{
+      *this, "tfrt-enable-large-batch-splitting", llvm::cl::init(false),
+      llvm::cl::desc("If true, enables large batch splitting to reduce "
+                     "padding inefficiency")};
+  mlir::Pass::Option<std::string> mixed_priority_batching_policy_{
+      *this, "tfrt-mixed-priority-batching-policy", llvm::cl::init(""),
+      llvm::cl::desc("Policy for mixed priority batching")};
+  mlir::Pass::Option<int64_t> batch_queue_global_prioritization_num_threads_{
+      *this, "tfrt-batch-queue-global-prioritization-num-threads",
+      llvm::cl::init(0),
+      llvm::cl::desc("If non-zero, all models on this server are switched to "
+                     "use a prioritized batching function using this number of "
+                     "global threads.")};
+  mlir::Pass::Option<bool> enable_priority_aware_batch_scheduler_{
+      *this, "tfrt-enable-priority-aware-batch-scheduler",
+      llvm::cl::init(false),
+      llvm::cl::desc("If true, the queue implementation will have a separate "
+                     "subqueue for each criticality.")};
+  mlir::Pass::Option<bool> enable_priority_aware_batch_scheduler_resplit_{
+      *this, "tfrt-enable-priority-aware-batch-scheduler-resplit",
+      llvm::cl::init(false),
+      llvm::cl::desc("If true, the queue implementation will allow task "
+                     "resplit for priority aware batch scheduler.")};
+  mlir::Pass::Option<bool> enable_batching_task_lazy_cancellation_{
+      *this, "tfrt-enable-batching-task-lazy-cancellation",
+      llvm::cl::init(false),
+      llvm::cl::desc("If true, enable lazy cancellation filtering in the "
+                     "priority-aware batch scheduler.")};
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateReconfigBatchOpPass(
+    ReconfigBatchOpPassOptions options) {
+  return std::make_unique<ReconfigBatchOpPass>(options);
+}
+
+static mlir::PassRegistration<ReconfigBatchOpPass> register_pass;
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/remove_device_attribute.cc

@@ -0,0 +1,87 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This pass removes the device attribute from every corert.executeop.
+
+#include <memory>
+#include <utility>
+
+#include "llvm/ADT/StringRef.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Visitors.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/TypeID.h"  // from @llvm-project
+#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+namespace {
+
+constexpr const char* kDevice = "device";
+
+struct RemoveDeviceAttributePass
+    : public PassWrapper<RemoveDeviceAttributePass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RemoveDeviceAttributePass)
+
+  llvm::StringRef getArgument() const final {
+    return "tfrt-remove-device-attribute";
+  }
+  llvm::StringRef getDescription() const final {
+    return "This pass removes the device attribute from every corert.executeop";
+  }
+
+  void runOnOperation() override;
+};
+
+void RemoveDeviceAttributePass::runOnOperation() {
+  ModuleOp module = getOperation();
+
+  module.walk([&](tfrt::corert::ExecuteOp execute_op) {
+    SmallVector<std::pair<StringRef, Attribute>, 4> op_func_attrs;
+    SmallVector<std::pair<StringRef, Attribute>, 4> op_attrs;
+    SmallVector<std::pair<StringRef, Attribute>, 4> new_op_attrs;
+    execute_op.getOpFuncAttrs(&op_func_attrs);
+    execute_op.getOpAttrs(&op_attrs);
+    for (std::pair<StringRef, Attribute> attr : op_attrs) {
+      if (attr.first != kDevice) {
+        new_op_attrs.push_back(attr);
+      }
+    }
+    if (op_attrs.size() == new_op_attrs.size()) return WalkResult::advance();
+
+    OpBuilder builder(execute_op);
+    auto new_execute_op = tfrt::corert::ExecuteOp::create(
+        builder, execute_op.getLoc(), execute_op.getResultTypes(),
+        execute_op.getOpHandler(), execute_op.getArguments(), new_op_attrs,
+        op_func_attrs, execute_op.getOpName());
+    execute_op.replaceAllUsesWith(new_execute_op.getResults());
+    execute_op.erase();
+    return WalkResult::advance();
+  });
+}
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> CreateRemoveDeviceAttributePass() {
+  return std::make_unique<RemoveDeviceAttributePass>();
+}
+
+static PassRegistration<RemoveDeviceAttributePass> pass;
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/remove_tf_if_const_args.cc

@@ -0,0 +1,196 @@
+/* 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 <cassert>
+#include <memory>
+#include <string>
+
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.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/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+// This pass removes tf.If ops' operands that are produced by tf.Const ops.
+// These constants can be moved into branches' function body for further
+// optimziation.
+class RemoveTfIfConstArgs
+    : public mlir::PassWrapper<RemoveTfIfConstArgs,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RemoveTfIfConstArgs)
+
+ private:
+  llvm::StringRef getArgument() const final {
+    return "tfrt-remove-tf-if-const-args";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Remove const args from tf.If ops";
+  }
+
+  void runOnOperation() override {
+    auto module = getOperation();
+    for (auto func_op :
+         llvm::make_early_inc_range(module.getOps<mlir::func::FuncOp>())) {
+      ProcessFunction(func_op);
+    }
+  }
+
+  void ProcessFunction(mlir::func::FuncOp op) {
+    // Set the insertion point to the current function, as we will insert new
+    // functions here.
+    mlir::OpBuilder builder(op);
+    for (mlir::Operation &op : op.front()) {
+      auto if_op = llvm::dyn_cast<mlir::TF::IfOp>(&op);
+      if (!if_op) continue;
+
+      // Record the operands that are produced by tf.Const ops.
+      llvm::SmallVector<mlir::TF::ConstOp, 2> const_args;
+      // Record these operands's corresponding operand indices.
+      llvm::SmallVector<unsigned, 2> const_arg_indices;
+      // Record the remaining operands that won't be removed.
+      llvm::SmallVector<mlir::Value, 2> remaining_args;
+      for (auto iter : llvm::enumerate(if_op.getInput())) {
+        mlir::Value operand = iter.value();
+        if (auto const_op = operand.getDefiningOp<mlir::TF::ConstOp>()) {
+          const_args.push_back(const_op);
+          const_arg_indices.push_back(iter.index());
+        } else {
+          remaining_args.push_back(operand);
+        }
+      }
+
+      if (const_args.empty()) continue;
+
+      RemoveConstArgsFromTfIfOp(builder, if_op, const_args, const_arg_indices,
+                                remaining_args);
+    }
+  }
+
+  void RemoveConstArgsFromTfIfOp(mlir::OpBuilder &builder, mlir::TF::IfOp if_op,
+                                 llvm::ArrayRef<mlir::TF::ConstOp> const_args,
+                                 llvm::ArrayRef<unsigned> const_arg_indices,
+                                 llvm::ArrayRef<mlir::Value> remaining_args) {
+    auto branch_suffix = absl::StrCat("_removed_const_args_", id_++);
+
+    // Create wrapper functions with the new arguments (as const args are
+    // removed) for both then function and else function.
+    auto new_then_function_name =
+        CreateBranchFunction(builder, if_op.then_function(), branch_suffix,
+                             const_args, const_arg_indices);
+    auto new_else_function_name =
+        CreateBranchFunction(builder, if_op.else_function(), branch_suffix,
+                             const_args, const_arg_indices);
+
+    // Change the if_op's argumetns to the new arguments, branches to new
+    // branches. Note that the outputs are not changed.
+    if_op.getInputMutable().assign(remaining_args);
+    if_op.setThenBranchAttr(
+        mlir::SymbolRefAttr::get(builder.getContext(), new_then_function_name));
+    if_op.setElseBranchAttr(
+        mlir::SymbolRefAttr::get(builder.getContext(), new_else_function_name));
+  }
+
+  llvm::StringRef CreateBranchFunction(
+      mlir::OpBuilder &builder, mlir::func::FuncOp branch,
+      absl::string_view branch_suffix,
+      llvm::ArrayRef<mlir::TF::ConstOp> const_args,
+      llvm::ArrayRef<unsigned> const_arg_indices) {
+    // Get the new function type as const args are removed.
+    llvm::BitVector const_arg_indices_bv(branch.getNumArguments());
+    for (auto i : const_arg_indices) const_arg_indices_bv.set(i);
+    auto new_branch_type = branch.getFunctionType().getWithoutArgsAndResults(
+        const_arg_indices_bv, {});
+    std::string new_branch_name =
+        absl::StrCat(branch.getSymName().str(), branch_suffix);
+    // Create the wrapper function with the new arguments that calls the
+    // original branch.
+    auto new_branch = mlir::func::FuncOp::create(
+        builder, branch.getLoc(), new_branch_name, new_branch_type);
+    new_branch.setVisibility(mlir::func::FuncOp::Visibility::Private);
+
+    // In its function body, we will add the corresponding const ops and call
+    // the original branch.
+
+    mlir::OpBuilder::InsertionGuard guard(builder);
+    auto *block = new_branch.addEntryBlock();
+    builder.setInsertionPointToStart(block);
+
+    // Prepare the function arguments of the original branch.
+    llvm::SmallVector<mlir::Value, 4> call_args(branch.getNumArguments());
+
+    // For those removed const args, we copy the tf.Const op, and use that as
+    // the corresponding argument when calling the original branch.
+    for (const auto &iter : llvm::zip(const_args, const_arg_indices)) {
+      auto const_op =
+          llvm::cast<mlir::TF::ConstOp>(builder.clone(*std::get<0>(iter)));
+      unsigned index = std::get<1>(iter);
+      call_args[index] = const_op;
+    }
+
+    // For the rest, they are now coming from the wrapper function's arguments
+    // in the original order.
+    for (int i = 0, j = 0; i < call_args.size(); ++i) {
+      if (!call_args[i]) {
+        assert(j < block->getNumArguments());
+        call_args[i] = block->getArgument(j++);
+      }
+    }
+
+    // Now create the call op to the original branch.
+    auto call_op = mlir::TF::StatefulPartitionedCallOp::create(
+        builder, new_branch.getLoc(), new_branch_type.getResults(), call_args,
+        /*args_attrs=*/nullptr, /*res_attrs=*/nullptr, branch.getSymName(), "",
+        "", "");
+    // Note that the outputs are not changed.
+    mlir::func::ReturnOp::create(builder, new_branch.getLoc(),
+                                 call_op.getOutput());
+
+    return new_branch.getSymName();
+  }
+
+  int id_ = 0;
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateRemoveTfIfConstArgsPass() {
+  return std::make_unique<RemoveTfIfConstArgs>();
+}
+
+static mlir::PassRegistration<RemoveTfIfConstArgs> register_pass(
+    CreateRemoveTfIfConstArgsPass);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/reorder_assert.cc

@@ -0,0 +1,114 @@
+/* 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 <memory>
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+// Reorder tf.Assert ops or tf.If ops that contains only tf.Assert ops to the
+// end of the function, in order to avoid unnecessary control dependencies
+// between tf.Assert and other ops.
+class ReorderTfAssertPass
+    : public mlir::PassWrapper<ReorderTfAssertPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ReorderTfAssertPass)
+
+  llvm::StringRef getArgument() const final { return "tfrt-reorder-tf-assert"; }
+  llvm::StringRef getDescription() const final {
+    return "Move tf.Assert to the end of the function to avoid unnecessary "
+           "control dependencies";
+  }
+
+  void runOnOperation() override {
+    auto module = getOperation();
+    for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
+      ProcessFunction(func_op);
+    }
+  }
+
+  void ProcessFunction(mlir::func::FuncOp func_op) {
+    auto& block = func_op.front();
+
+    llvm::SmallVector<mlir::Operation*, 2> assert_ops;
+    for (mlir::Operation& op : block) {
+      if (auto assert_op = llvm::dyn_cast<mlir::TF::AssertOp>(&op)) {
+        assert_ops.push_back(assert_op);
+      }
+
+      if (auto if_op = llvm::dyn_cast<mlir::TF::IfOp>(&op)) {
+        if (IsAssertOnlyIfOp(if_op)) {
+          assert_ops.push_back(if_op);
+        }
+      }
+    }
+
+    auto& return_op = block.back();
+
+    for (auto assert_op : assert_ops) {
+      assert_op->moveBefore(&return_op);
+    }
+  }
+
+  bool IsAssertOnlyIfOp(mlir::TF::IfOp op) {
+    // If the results of the if op are used by some other ops, we cannot reorder
+    // it.
+    if (!op->use_empty()) return false;
+
+    // Only reorder if both branches are non-side-effecting or containing only
+    // Assert ops.
+    if (IsFunctionNonSideEffectingOrAssert(op.then_function()) &&
+        IsFunctionNonSideEffectingOrAssert(op.else_function()))
+      return true;
+
+    return false;
+  }
+
+  bool IsFunctionNonSideEffectingOrAssert(mlir::func::FuncOp func_op) {
+    auto& block = func_op.front();
+    for (mlir::Operation& op : block) {
+      if (!llvm::isa<mlir::TF::AssertOp>(&op) && !mlir::isMemoryEffectFree(&op))
+        return false;
+    }
+    return true;
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateReorderTfAssertPass() {
+  return std::make_unique<ReorderTfAssertPass>();
+}
+
+static mlir::PassRegistration<ReorderTfAssertPass> register_pass(
+    CreateReorderTfAssertPass);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/set_shape_invariant_in_while_ops.cc

@@ -0,0 +1,57 @@
+/* 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/tfrt/transforms/set_shape_invariant_in_while_ops.h"
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+class SetShapeInvariantInWhileOps
+    : public mlir::PassWrapper<SetShapeInvariantInWhileOps,
+                               mlir::OperationPass<mlir::func::FuncOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SetShapeInvariantInWhileOps)
+
+  void runOnOperation() override {
+    mlir::func::FuncOp func_op = getOperation();
+
+    auto shape_invariant = mlir::UnitAttr::get(&getContext());
+
+    func_op.walk([&](mlir::TF::WhileOp op) {
+      // Skip tf.While op on TPU.
+      if (!op->hasAttr("_tpu_replicate")) {
+        op.setShapeInvariantAttr(shape_invariant);
+      }
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateSetShapeInvariantInWhileOps() {
+  return std::make_unique<SetShapeInvariantInWhileOps>();
+}
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/set_shape_invariant_in_while_ops.h

@@ -0,0 +1,35 @@
+/* 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_TFRT_TRANSFORMS_SET_SHAPE_INVARIANT_IN_WHILE_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_SET_SHAPE_INVARIANT_IN_WHILE_OPS_H_
+
+#include <memory>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+// Create a pass to set shape_invariant attribute for all tf.While ops except
+// those are on TPU.
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateSetShapeInvariantInWhileOps();
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_SET_SHAPE_INVARIANT_IN_WHILE_OPS_H_

tensorflow/compiler/mlir/tfrt/transforms/sink_in_invariant_ops.cc

@@ -0,0 +1,239 @@
+/* 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 <memory>
+#include <utility>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace tensorflow {
+namespace {
+
+// TODO(b/262610234): Generalize the sinking conditions.
+// Check if the op qualifies to sink to the callee.
+bool IsSinkCandidate(mlir::Operation *op) {
+  return op && llvm::isa<mlir::TF::VarHandleOp, mlir::TF::ConstOp,
+                         mlir::TF::HashTableV2Op>(op);
+}
+
+// Check if the op is allowed to be sinked. We are being conservative here to
+// whilelist very limited set of ops here.
+struct AllowSinkHelper {
+  explicit AllowSinkHelper(mlir::Operation* sinked_op, mlir::Operation* user,
+                           int arg_index) {
+    if (llvm::isa<mlir::TF::BatchFunctionOp,
+                  mlir::TF::StatefulPartitionedCallOp>(user)) {
+      allow_sink_to = true;
+      callee_arg_index = arg_index;
+      return;
+    }
+
+    // We tend to limit this support on WhileOp to only VarHandleOp to satisfy
+    // IFRT lowering requirements.
+    // Sinking other invariants like ConstOp is error-prone because it requires
+    // non-trivial effort to avoid sinking Consts when they are used by cond
+    // function and we don't need such support.
+    if (llvm::isa<mlir::TF::VarHandleOp>(sinked_op) &&
+        llvm::isa<mlir::TF::WhileOp>(user)) {
+      allow_sink_to = true;
+      callee_arg_index = arg_index;
+      return;
+    }
+
+    if (llvm::isa<mlir::TF::IfOp>(user) && arg_index > 0) {
+      allow_sink_to = true;
+      callee_arg_index = arg_index - 1;
+      return;
+    }
+  }
+
+  bool allow_sink_to = false;
+  int callee_arg_index = 0;
+};
+
+llvm::SmallVector<mlir::Value> FindValueInCallees(
+    const mlir::SymbolTable &symbol_table,
+    const mlir::SymbolUserMap &symbol_users, mlir::Operation *caller,
+    int arg_index) {
+  llvm::SmallVector<mlir::Value> values;
+  llvm::SmallDenseSet<llvm::StringRef> callees;
+  for (const auto &named_attr : caller->getAttrs()) {
+    if (auto symbol_attr =
+            mlir::dyn_cast<mlir::FlatSymbolRefAttr>(named_attr.getValue())) {
+      auto symbol = symbol_attr.getValue();
+
+      auto callee = symbol_table.lookup<mlir::func::FuncOp>(symbol);
+      if (!callee) continue;
+
+      // One callee invoked by multiple caller is skipped for simplicity.
+      // Consider adding support if more usage are observed from production.
+      if (llvm::ArrayRef<mlir::Operation *> users =
+              symbol_users.getUsers(callee);
+          users.size() > 1)
+        continue;
+
+      // Invoked by same caller multiple times, only process the first one.
+      if (!callees.insert(symbol).second) continue;
+
+      values.push_back(callee.getArgument(arg_index));
+    }
+  }
+  return values;
+}
+
+void FindSinkTarget(
+    const mlir::SymbolTable &symbol_table,
+    const mlir::SymbolUserMap &symbol_users, mlir::OpResult original,
+    mlir::Value value,
+    llvm::DenseMap<mlir::OpOperand *, llvm::SmallDenseSet<mlir::OpResult>>
+        &targets) {
+  for (mlir::OpOperand &use : value.getUses()) {
+    auto *user = use.getOwner();
+
+    AllowSinkHelper helper(original.getDefiningOp(), user,
+                           use.getOperandNumber());
+
+    if (helper.allow_sink_to) {
+      auto values = FindValueInCallees(symbol_table, symbol_users, user,
+                                       helper.callee_arg_index);
+      for (auto value : values) {
+        FindSinkTarget(symbol_table, symbol_users, original, value, targets);
+      }
+    } else if (value != original) {
+      // If the sinked op is directly used by ReturnOp, we don't sink it.
+      // One example is for tf.WhileOp, the input and output of the cond
+      // function and the body function must be the same. If the cond function
+      // has an input of type tf.VarHandleOp and it just return the VarHandleOp,
+      // we don't need to sink it.
+      if (llvm::isa<mlir::func::ReturnOp>(user)) {
+        continue;
+      }
+      targets[&use].insert(original);
+    }
+  }
+}
+
+// Sink in invariant ops like tf.Const, tf.VarHandleOp and tf.HashTableV2 ops
+// into sinkable calls like tf.BatchFunction and tf.If. If there are nested
+// calls, the invariant ops will only be copied at the target.
+void SinkInInvariantOps(mlir::ModuleOp module) {
+  mlir::SymbolTable symbol_table(module);
+  mlir::SymbolTableCollection symbol_table_collection;
+  mlir::SymbolUserMap symbol_users(symbol_table_collection, module);
+
+  // TODO(b/263191534): Replace with CallOpInterface to handle callees.
+  // Identify the invariant Op, Caller, Callee FuncOp to update.
+
+  llvm::DenseMap<mlir::OpOperand *, llvm::SmallDenseSet<mlir::OpResult>>
+      targets;
+  module.walk([&](mlir::Operation *op) {
+    if (IsSinkCandidate(op)) {
+      for (auto value : op->getOpResults()) {
+        FindSinkTarget(symbol_table, symbol_users, value, value, targets);
+      }
+    }
+  });
+
+  // Clone the sinkable op associated with the func op to the func op
+  mlir::OpBuilder builder(module);
+  for (const auto &p : targets) {
+    if (p.second.size() != 1) continue;
+
+    auto *use = p.first;
+
+    builder.setInsertionPointToStart(use->getOwner()->getBlock());
+
+    mlir::OpResult original = *p.second.begin();
+    auto *new_op = builder.clone(*original.getDefiningOp());
+    // Use `use->set()` instead of `replaceAllUsesWith()` on the block argument.
+    // `replaceAllUsesWith()` would replace ALL uses of the block argument in
+    // the region, including those that were explicitly skipped (e.g.,
+    // `ReturnOp` uses). Replacing a `ReturnOp` use with a cloned ranked
+    // resource/constant would change the type of the returned value without
+    // updating the function signature, leading to a verifier crash.
+    //
+    // For example, in the following IR:
+    //
+    // func.func private @body(
+    //     %arg0: tensor<*x!tf_type.resource>
+    // ) -> tensor<*x!tf_type.resource> {
+    //   %0 = "tf.ReadVariableOp"(%arg0)
+    //     : (tensor<*x!tf_type.resource>) -> tensor<i32>
+    //   func.return %arg0 : tensor<*x!tf_type.resource>
+    // }
+    //
+    // If %arg0 is a constant with known rank, we only want to sink the use in
+    // `ReadVariableOp`. Using `use->set()` ensures only that specific use is
+    // replaced with the cloned ranked constant, while the `ReturnOp` continues
+    // to use the unranked `%arg0`, preserving signature compliance.
+    //
+    // Since the `targets` map is constructed by recursively traversing all uses
+    // of the invariant value (via `FindSinkTarget`), it is guaranteed to
+    // capture every single safe leaf use. By iterating over `targets` and
+    // updating them individually, we ensure all eligible uses are successfully
+    // sunk without missing any, while safely leaving the ineligible ones (like
+    // `ReturnOp`) untouched.
+    use->set(new_op->getResult(original.getResultNumber()));
+  }
+}
+
+class SinkInInvariantOpsPass
+    : public mlir::PassWrapper<SinkInInvariantOpsPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SinkInInvariantOpsPass)
+
+  llvm::StringRef getArgument() const final {
+    return "tfrt-sink-in-invariant-ops";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Sink in the invariant ops to facilitate invariant ops hoisting.";
+  }
+
+  void runOnOperation() override {
+    auto module = getOperation();
+    SinkInInvariantOps(module);
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateSinkInInvariantOpsPass() {
+  return std::make_unique<SinkInInvariantOpsPass>();
+}
+
+static mlir::PassRegistration<SinkInInvariantOpsPass>
+    sink_in_invariant_ops_pass;
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h

@@ -0,0 +1,269 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_TFRT_PIPELINE_OPTIONS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_TFRT_PIPELINE_OPTIONS_H_
+
+#include <cstdint>
+#include <string>
+
+#include "llvm/Support/CommandLine.h"
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+
+namespace tensorflow {
+
+struct TfrtPipelineOptions
+    : public mlir::PassPipelineOptions<TfrtPipelineOptions> {
+  Option<std::string> saved_model_dir{*this, "saved-model-dir",
+                                      llvm::cl::desc(""), llvm::cl::init("")};
+  Option<std::string> default_device{
+      *this, "default-device", llvm::cl::desc("default device assignment"),
+      llvm::cl::init("/job:localhost/replica:0/task:0/device:CPU:0")};
+  Option<bool> enable_optimizer{
+      *this, "enable-optimizer",
+      llvm::cl::desc("run optimization passes on corert dialect"),
+      llvm::cl::init(false)};
+  Option<bool> decompose_resource_ops{
+      *this, "decompose-resource-ops",
+      llvm::cl::desc("decompose composite resource ops into ReadVariableOp and "
+                     "non-resource ops. This is currently used in TFRT "
+                     "savedmodel pipeline."),
+      llvm::cl::init(false)};
+  Option<std::string> force_data_format{
+      *this, "force-data-format",
+      llvm::cl::desc("force data format for all layout sensitive operations")};
+  // TODO(tfrt-devs): consider making compiler to figure out whether to fold
+  // transpose or not instead of exposing the specific option.
+  Option<bool> skip_fold_transpose_in_ops{
+      *this, "skip-fold-transpose-in-ops",
+      llvm::cl::desc("Skip folding transpose operands in Ops which can support "
+                     "different layouts.")};
+  Option<bool> target_tpurt{*this, "target-tpurt",
+                            llvm::cl::desc("target TPURT dialect if true"),
+                            llvm::cl::init(false)};
+  Option<bool> tpu_use_core_selector{
+      *this, "tpu-use-core-selector",
+      llvm::cl::desc("If true, use ServingCoreSelector to pick TPU core. "
+                     "Otherwise, use the assigned core. Currently we use "
+                     "core selector for Servo serving use cases."),
+      llvm::cl::init(true)};
+  Option<bool> tpu_use_bundled_transfer{
+      *this, "tpu-use-bundled-transfer",
+      llvm::cl::desc("If true, use BundledTransferToTpuOp to transfer "
+                     "variables and input tensors to TPU."),
+      llvm::cl::init(true)};
+  Option<bool> tpu_lower_to_fallback{
+      *this, "tpu-lower-to-fallback",
+      llvm::cl::desc("If true, lower an TF op that's placed on TPU device "
+                     "to be executed by tfrt_fallback.execute."),
+      llvm::cl::init(true)};
+  Option<bool> tpu_fuse_ops{
+      *this, "tpu-fuse-ops",
+      llvm::cl::desc("If true, use the TPU fused compile_and_execute kernel"),
+      llvm::cl::init(false)};
+  // TODO(b/194081364): remove this option once we unify servo TPU serving
+  // result transfer behavior.
+  Option<bool> tpu_transfer_result_to_host{
+      *this, "tpu-transfer-result-to-host",
+      llvm::cl::desc("If true, transfer the result of tpurt.execute from TPU "
+                     "to host."),
+      llvm::cl::init(true)};
+  Option<bool> use_tpu_host_allocator_for_inputs{
+      *this, "use-tpu-host-allocator-for-inputs",
+      llvm::cl::desc("If true, fallback executeops that produce inputs to tpu "
+                     "program will use tpu host allocator."),
+      llvm::cl::init(false)};
+  Option<TfrtCompileOptions::TpuAllowUnpaddedBatch> tpu_allow_unpadded_batch{
+      *this, "tpu-allow-unpadded-batch",
+      llvm::cl::desc("To allow unpadded batch for TPU execution."),
+      llvm::cl::values(
+          clEnumValN(TfrtCompileOptions::TpuAllowUnpaddedBatch::kDisabled,
+                     "disabled", "Disable this feature."),
+          clEnumValN(TfrtCompileOptions::TpuAllowUnpaddedBatch::kAuto, "auto",
+                     "Enable this feature when in-graph batching is detected."),
+          clEnumValN(TfrtCompileOptions::TpuAllowUnpaddedBatch::kEnforced,
+                     "enforced", "Force to enable this feature.")),
+      llvm::cl::init(TfrtCompileOptions::TpuAllowUnpaddedBatch::kDisabled)};
+
+  Option<bool> target_gpu{
+      *this, "target-gpu",
+      llvm::cl::desc("If true, target GPU compiler passes."),
+      llvm::cl::init(false)};
+
+  // TODO(b/294895431): Remove the flag and default to the fused op.
+  Option<bool> use_gpu_compile_and_execute_op{
+      *this, "use-gpu-compile-and-execute-op",
+      llvm::cl::desc("If true, gpurt.compile_and_execute is used for GPU"),
+      llvm::cl::init(false)};
+
+  Option<bool> enable_while_parallel_iterations{
+      *this, "enable-while-parallel-iterations",
+      llvm::cl::desc("If true, tf.While op will be parallelized. This is "
+                     "currently experimental."),
+      llvm::cl::init(false)};
+
+  Option<bool> hoist_invariant_ops{
+      *this, "hoist-invariant-ops",
+      llvm::cl::desc("If true, invariant ops in savedmodels will be hoisted "
+                     "out to run during loading."),
+      llvm::cl::init(false)};
+
+  Option<bool> fuse_get_resource_ops_in_hoisting{
+      *this, "fuse-get-resource-ops-in-hoisting",
+      llvm::cl::desc("If true, get_resource_op will be fused during hoisting"),
+      llvm::cl::init(true)};
+
+  Option<bool> sink_in_invariant_ops{
+      *this, "sink-in-invariant-ops",
+      llvm::cl::desc("If true, sink the selected invariant ops in to the "
+                     "nested functions to facilitate invariant ops hoisting."),
+      llvm::cl::init(false)};
+
+  Option<uint64_t> cost_threshold{
+      *this, "tfrt-cost-threshold",
+      llvm::cl::desc(
+          "The cost threshold to decide whether a sequence of operations is "
+          "cheap, and then whether it can be executed inline."),
+      llvm::cl::init(1)};
+
+  Option<int64_t> min_num_batch_threads{
+      *this, "tfrt-min-num-batch-threads",
+      llvm::cl::desc("The minimum number of batch threads"), llvm::cl::init(1)};
+
+  Option<int64_t> min_max_enqueued_batches{
+      *this, "tfrt-min-max-enqueued-batches",
+      llvm::cl::desc(
+          "The minimum of the maximum number of outstanding enqueued batches"),
+      llvm::cl::init(1)};
+
+  Option<int64_t> batch_queue_global_prioritization_num_threads{
+      *this, "tfrt-batch-queue-global-prioritization-num-threads",
+      llvm::cl::desc(
+          "If non-zero, all models on this server are switched to use a "
+          "prioritized batching function using this number of global threads."),
+      llvm::cl::init(0)};
+
+  Option<bool> enable_priority_aware_batch_scheduler{
+      *this, "tfrt-enable-priority-aware-batch-scheduler",
+      llvm::cl::desc("If true, the queue implementation will have a separate "
+                     "subqueue for each criticality."),
+      llvm::cl::init(false)};
+
+  Option<bool> enable_priority_aware_batch_scheduler_resplit{
+      *this, "tfrt-enable-priority-aware-batch-scheduler-resplit",
+      llvm::cl::desc("If true, the queue implementation will allow task "
+                     "resplit for priority aware batch scheduler."),
+      llvm::cl::init(false)};
+
+  Option<bool> enable_batching_task_lazy_cancellation{
+      *this, "tfrt-enable-batching-task-lazy-cancellation",
+      llvm::cl::desc("If true, enable lazy cancellation filtering in the "
+                     "priority-aware batch scheduler."),
+      llvm::cl::init(false)};
+
+  Option<std::string> batch_padding_policy{
+      *this, "tfrt-batch-padding-policy",
+      llvm::cl::desc("The policy used when padding (or splitting) batches."),
+      llvm::cl::init("")};
+
+  Option<int64_t> num_batch_threads{
+      *this, "tfrt-num-batch-threads",
+      llvm::cl::desc(
+          "The number of threads for processing batches in parallel"),
+      llvm::cl::init(0)};
+
+  Option<int64_t> max_batch_size{
+      *this, "tfrt-max-batch-size",
+      llvm::cl::desc("The maximum allowed batch size"), llvm::cl::init(0)};
+
+  Option<int64_t> batch_timeout_micros{
+      *this, "tfrt-batch-timeout-micros",
+      llvm::cl::desc("The maximum number of microseconds before outputting an "
+                     "incomplete batch"),
+      llvm::cl::init(0)};
+
+  ListOption<int64_t> allowed_batch_sizes{
+      *this, "tfrt-allowed-batch-sizes",
+      llvm::cl::desc("Allowed sizes for padding (or splitting) batches")};
+
+  Option<int64_t> max_enqueued_batches{
+      *this, "tfrt-max-enqueued-batches",
+      llvm::cl::desc("The maximum number of batches enqueued for processing "
+                     "before requests are failed fast"),
+      llvm::cl::init(0)};
+
+  /**experimental options to override legacy priority batching parameters**/
+
+  // The following options are for legacy priority batching, and are only for
+  // experimentation.
+  // TODO(b/488097129): Remove these options after the experiment is done.
+  Option<int64_t> low_priority_max_batch_size{
+      *this, "tfrt-low-priority-max-batch-size",
+      llvm::cl::desc("The maximum allowed batch size for low priority "
+                     "requests"),
+      llvm::cl::init(0)};
+
+  Option<int64_t> low_priority_batch_timeout_micros{
+      *this, "tfrt-low-priority-batch-timeout-micros",
+      llvm::cl::desc("The maximum number of microseconds before outputting an "
+                     "incomplete batch for low priority requests"),
+      llvm::cl::init(0)};
+
+  ListOption<int64_t> low_priority_allowed_batch_sizes{
+      *this, "tfrt-low-priority-allowed-batch-sizes",
+      llvm::cl::desc("Allowed sizes for padding (or splitting) batches for low "
+                     "priority requests")};
+
+  Option<int64_t> low_priority_max_enqueued_batches{
+      *this, "tfrt-low-priority-max-enqueued-batches",
+      llvm::cl::desc("The maximum number of batches enqueued for processing "
+                     "before low priority requests are failed fast"),
+      llvm::cl::init(0)};
+
+  /*experimental options end*/
+
+  Option<int64_t> num_warmup_batch_threads{
+      *this, "tfrt-num-warmup-batch-threads",
+      llvm::cl::desc("The number of threads for processing warmup requests. "
+                     "Useful to process warmup requests without starving the "
+                     "regular batch threads when global scheduler is enabled."),
+      llvm::cl::init(0)};
+
+  Option<bool> enable_large_batch_splitting{
+      *this, "tfrt-enable-large-batch-splitting",
+      llvm::cl::desc("If true, enables large batch splitting to reduce padding "
+                     "inefficiency"),
+      llvm::cl::init(false)};
+
+  Option<std::string> mixed_priority_batching_policy{
+      *this, "tfrt-mixed-priority-batching-policy",
+      llvm::cl::desc("Policy for mixed priority batching"), llvm::cl::init("")};
+
+  Option<bool> merge_inter_dependent_streams{
+      *this, "tfrt-merge-inter-dependent-streams",
+      llvm::cl::desc("If true, streams with inter data depenedencies will be "
+                     "preferred to be merged for inline execution."),
+      llvm::cl::init(false)};
+  Option<bool> allow_xla_cpu{
+      *this,
+      "allow-xla-cpu",
+      llvm::cl::desc("If true, allow XLA:CPU for CPU computations."),
+      llvm::cl::init(true),
+  };
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_TFRT_PIPELINE_OPTIONS_H_

tensorflow/compiler/mlir/tfrt/transforms/tpu_passes.h

@@ -0,0 +1,87 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_TPU_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_TPU_PASSES_H_
+
+// This file contains stub implementations for Google internal TPU APIs.
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Pass/PassOptions.h"
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+
+namespace tensorflow {
+
+class CoreRTConverter;
+
+namespace tfrt_compiler {
+
+class FallbackConverter;
+
+}
+
+struct TfrtTpuCompileOptions
+    : mlir::PassPipelineOptions<TfrtTpuCompileOptions> {
+  Option<bool> move_resource_gather_to_host{
+      *this, "move-resource-gather-to-host",
+      llvm::cl::desc("Move resource gather ops to host"),
+      llvm::cl::init(false)};
+  Option<int64_t> gather_table_width_threshold_bytes{
+      *this, "gather-table-width-threshold-bytes",
+      llvm::cl::desc(
+          "The threshold to control whether a TPU resource gather op should be "
+          "moved to host. A negative values means all are moved."),
+      llvm::cl::init(-1)};
+};
+
+struct TfrtTpuExecuteOpConversionOptions {
+  bool use_core_selector = false;
+  bool use_bundled_transfer = false;
+  bool transfer_result_to_host = false;
+  bool use_tpu_host_allocator_for_inputs = false;
+  TfrtCompileOptions::TpuAllowUnpaddedBatch allow_unpadded_batch =
+      TfrtCompileOptions::TpuAllowUnpaddedBatch::kDisabled;
+};
+
+// Registers a set of dialects used in TFRT TPU lowering.
+inline void RegisterTPUDialects(mlir::DialectRegistry *registry) {}
+
+// Adds a target dialect and a set of rewrite patterns for TFRT TPU lowering.
+inline void AddTPUTargetDialectAndPatterns(
+    mlir::ConversionTarget *target, mlir::RewritePatternSet *patterns,
+    mlir::MLIRContext *context, CoreRTConverter *corert_converter,
+    tfrt_compiler::FallbackConverter *fallback_converter,
+    const TfrtTpuExecuteOpConversionOptions &tpu_exec_conv_opts,
+    bool tpu_lower_to_fallback) {}
+
+// Rewrites specific TF TPU ops to equivalent TF ops in a module.
+inline mlir::LogicalResult RunTPUBackwardCompatConversion(
+    mlir::ModuleOp module, const TfrtTpuCompileOptions &options) {
+  return mlir::failure();
+}
+
+// The rewrite rules to support the fallback execution of TPUPartitionedCallOp.
+inline mlir::LogicalResult RunTPUPartitionedCallFallbackCompatConversion(
+    mlir::ModuleOp module) {
+  return mlir::failure();
+}
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_TPU_PASSES_H_

tensorflow/compiler/mlir/tfrt/transforms/update_op_cost_in_tfrt_mlir.cc

@@ -0,0 +1,53 @@
+/* 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/tfrt/transforms/update_op_cost_in_tfrt_mlir.h"
+
+#include <cstdint>
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/analysis/cost_analysis.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+constexpr char kCostAttrName[] = "_tfrt_cost";
+constexpr char kOpKeyAttrName[] = "op_key";
+
+void UpdateOpCostInTfrtMlir(mlir::ModuleOp op,
+                            const tfrt_stub::CostRecorder& cost_recorder) {
+  mlir::Builder builder(op);
+  op.walk([&](mlir::Operation* op) {
+    // TODO(b/259602527): Add unit test for the precedence.
+    // Registered cost function has higher priority than online cost analysis.
+    if (HasCostFunctionRegistered(op->getName().getStringRef())) return;
+    // Only update ops with existing cost attr.
+    const auto cost_attr = op->getAttrOfType<mlir::IntegerAttr>(kCostAttrName);
+    if (!cost_attr) return;
+    // Only fallback ops have `op_key`s.
+    const auto op_key_attr =
+        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
+    if (!op_key_attr) return;
+    // Set the cost attr with a new value.
+    const int64_t op_key = op_key_attr.getInt();
+    op->setAttr(kCostAttrName, builder.getI64IntegerAttr(
+                                   cost_recorder.GetCost(op_key)));
+  });
+}
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/transforms/update_op_cost_in_tfrt_mlir.h

@@ -0,0 +1,32 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_UPDATE_OP_COST_IN_TFRT_MLIR_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_UPDATE_OP_COST_IN_TFRT_MLIR_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+// Updates the existing costs for all the fallback ops with the records in
+// `cost_recorder`.
+void UpdateOpCostInTfrtMlir(mlir::ModuleOp op,
+                            const tfrt_stub::CostRecorder& cost_recorder);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_UPDATE_OP_COST_IN_TFRT_MLIR_H_

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

@@ -0,0 +1,112 @@
+/* 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/tfrt/transforms/utils.h"
+
+#include <optional>
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.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/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/host_runtime/tfrt_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
+
+namespace tensorflow {
+
+bool IsResourceArgument(mlir::Value value) {
+  auto arg = mlir::dyn_cast<mlir::BlockArgument>(value);
+  if (!arg) return false;
+
+  auto func = llvm::cast<mlir::func::FuncOp>(arg.getOwner()->getParentOp());
+
+  return func.getArgAttr(arg.getArgNumber(), "tf.resource_name") != nullptr;
+}
+
+bool IsResultVariable(const mlir::Value &original_operand,
+                      const mlir::Value &operand) {
+  if (mlir::isa<mlir::OpResult>(original_operand)) {
+    auto defining_op = original_operand.getDefiningOp();
+
+    // TODO(b/174753886): When device assignment is properly done, we
+    // should check that TF::ReadVariableOp is for TPU device here.
+    if (llvm::isa<mlir::TF::ReadVariableOp>(defining_op) &&
+        defining_op->getNumOperands() == 1) {
+      return true;
+    } else if (llvm::isa<mlir::TF::_TfrtGetResourceOp>(defining_op)) {
+      return true;
+    }
+    return false;
+  }
+  return IsResourceArgument(operand);
+}
+
+std::optional<std::string> CanonicalizeTensorflowFunctionName(
+    const mlir::SymbolTable &symbol_table, absl::string_view mlir_func_name,
+    bool use_mlir_func_name) {
+  if (use_mlir_func_name) {
+    return std::string(mlir_func_name);
+  }
+
+  // Currently in TF graph to MLIR importing, a "0" is appended to the original
+  // function name. The renaming is for TF/XLA v1 bridge use cases. Refer to
+  // b/142268695, b/141617294 for more context.
+  //
+  // TFRT currently uses the original function library. Hence, we retrieve the
+  // original function name from the function attributes. Longer term, we
+  // probably want to export the MLIR functions.
+  auto callee =
+      symbol_table.lookup<mlir::func::FuncOp>(std::string(mlir_func_name));
+  if (!callee) return std::nullopt;
+
+  mlir::StringAttr original_func_name =
+      callee->getAttrOfType<mlir::StringAttr>("tf._original_func_name");
+  if (!original_func_name) {
+    // If there is no function attribute "tf._original_func_name" in the callee,
+    // we use the workaround to recover the original function name by removing
+    // the last char of the MLIR function name.
+    // TODO(b/259138201): Remove this workwaround after we make sure
+    // "tf._original_func_name" is present in callees in all code paths.
+    mlir_func_name.remove_suffix(1);
+    return std::string(mlir_func_name);
+  }
+
+  return original_func_name.str();
+}
+
+bool IsSessionInitializer(mlir::func::FuncOp op) {
+  auto session_initializer_op = mlir::tf_saved_model::GetSessionInitializerOp(
+      op->getParentOfType<mlir::ModuleOp>());
+  if (!session_initializer_op) return false;
+
+  for (auto sym_ref : session_initializer_op.getInitializers()) {
+    if (op.getSymName() ==
+        mlir::cast<mlir::FlatSymbolRefAttr>(sym_ref).getValue())
+      return true;
+  }
+
+  return false;
+}
+
+}  // namespace tensorflow

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

@@ -0,0 +1,46 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_UTILS_H_
+
+#include <optional>
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+
+namespace tensorflow {
+
+// Checks if the given `value` is a resource argument.
+bool IsResourceArgument(mlir::Value value);
+
+// Checks if an operand is the value of a variable.
+bool IsResultVariable(const mlir::Value &original_operand,
+                      const mlir::Value &operand);
+
+// Canonicalize the symbol attr to the original TF function name.
+std::optional<std::string> CanonicalizeTensorflowFunctionName(
+    const mlir::SymbolTable &symbol_table, absl::string_view mlir_func_name,
+    bool use_mlir_func_name = false);
+
+// Returns true if the function is a session initializer in tf_saved_model
+// dialect.
+bool IsSessionInitializer(mlir::func::FuncOp op);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_TRANSFORMS_UTILS_H_

tensorflow/compiler/mlir/tfrt/transforms/xla_rewrite_pass.cc

@@ -0,0 +1,121 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_interfaces.h"
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/core/ir/types/dialect.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+template <typename OpType>
+struct RewriteFunctionCallToXlaLaunchOnCpu
+    : public mlir::OpRewritePattern<OpType> {
+ public:
+  using mlir::OpRewritePattern<OpType>::OpRewritePattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      OpType op, mlir::PatternRewriter& rewriter) const override {
+    if (auto xla_must_compile =
+            op->template getAttrOfType<mlir::BoolAttr>("_XlaMustCompile");
+        !xla_must_compile || !xla_must_compile.getValue()) {
+      return mlir::failure();
+    }
+
+    llvm::StringRef device = GetDeviceOrEmpty(op);
+
+    if (device.empty() || !device.contains("CPU")) return mlir::failure();
+
+    llvm::SmallVector<int64_t> constants;
+    llvm::SmallVector<int64_t> resources;
+
+    for (int i = 0; i < op.getNumOperands(); ++i) {
+      auto value = op.getOperand(i);
+      if (llvm::isa<mlir::tf_type::ResourceType>(
+              llvm::cast<mlir::TensorType>(value.getType()).getElementType())) {
+        resources.push_back(i);
+      } else if (auto* def = value.getDefiningOp();
+                 def && llvm::isa<mlir::TF::ConstOp>(def)) {
+        constants.push_back(i);
+      }
+    }
+
+    rewriter.replaceOpWithNewOp<mlir::TF::XlaLaunchV2Op>(
+        op, op.getResultTypes(), op.getOperands(),
+        rewriter.getI64ArrayAttr(constants),
+        rewriter.getI64ArrayAttr(resources), op.getFAttr());
+
+    return mlir::success();
+  }
+};
+
+struct TfrtXlaRewritePass
+    : public mlir::PassWrapper<TfrtXlaRewritePass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+  llvm::StringRef getArgument() const override { return "tfrt-xla-rewrite"; }
+
+  llvm::StringRef getDescription() const override {
+    return "rewrites for XLA host ops.";
+  }
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TfrtXlaRewritePass)
+
+  void runOnOperation() override {
+    mlir::RewritePatternSet patterns(&getContext());
+
+    patterns
+        .add<RewriteFunctionCallToXlaLaunchOnCpu<mlir::TF::PartitionedCallOp>>(
+            &getContext());
+    patterns.add<RewriteFunctionCallToXlaLaunchOnCpu<
+        mlir::TF::StatefulPartitionedCallOp>>(&getContext());
+
+    if (mlir::failed(
+            mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+      signalPassFailure();
+      return;
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+CreateTfrtXlaRewritePass() {
+  return std::make_unique<TfrtXlaRewritePass>();
+}
+
+static mlir::PassRegistration<TfrtXlaRewritePass> register_pass(
+    CreateTfrtXlaRewritePass);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow