chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfrt/BUILD

@@ -0,0 +1,766 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "if_google", "tf_cc_binary", "tf_cc_test")
+
+# Note: keep the following lines separate due to the way copybara works
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+)
+
+# TF to TFRT kernels conversion.
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    packages = [
+        "//tensorflow/compiler/...",
+        "//tensorflow/core/tfrt/experimental/data/...",
+        "//tensorflow/core/tfrt/graph_executor/...",
+        "//tensorflow/core/tfrt/ifrt/...",
+        "//tensorflow/core/tfrt/kernels/...",
+        "//tensorflow/core/tfrt/mlrt/...",
+        "//tensorflow/core/tfrt/saved_model/...",
+        "//tensorflow/core/tfrt/tfrt_session/...",
+    ] + if_google([
+        "//learning/brain/experimental/mlir/tflite/tfmrt/...",
+        "//learning/brain/experimental/mlir/tfrt_compiler/...",
+        "//learning/brain/experimental/tfrt/...",
+        "//learning/brain/tfrt/...",
+        "//learning/infra/mira/...",
+        "//learning/serving/contrib/tfrt/mlir/...",
+        "//learning/serving/testing/functional/...",
+        # Allow visibility from the mlir language server.
+        "//learning/brain/mlir/mlir_lsp_server/...",
+        "//smartass/brain/ops/...",
+        "//third_party/mlir_edge/tpgen/...",
+        "//tensorflow_serving/servables/tensorflow/...",
+        "//tensorflow_serving/servables/tensorflow/google/...",
+        "//third_party/tf_runtime_google/...",
+    ]),
+)
+
+exports_files(["run_lit.sh"])
+
+td_library(
+    name = "runtime_fallback_ops_td_files",
+    srcs = [
+        "runtime_fallback/runtime_fallback_ops.td",
+    ],
+    deps = [
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+        "@tf_runtime//:OpBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "runtime_fallback_ops_inc_gen",
+    tbl_outs = {
+        "runtime_fallback_ops.h.inc": ["-gen-op-decls"],
+        "runtime_fallback_ops.cc.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "runtime_fallback/runtime_fallback_ops.td",
+    deps = [":runtime_fallback_ops_td_files"],
+)
+
+cc_library(
+    name = "runtime_fallback_opdefs",
+    srcs = [
+        "runtime_fallback/runtime_fallback_combine.cc",
+        "runtime_fallback/runtime_fallback_ops.cc",
+    ],
+    hdrs = [
+        "runtime_fallback/runtime_fallback_ops.h",
+    ],
+    deps = [
+        ":runtime_fallback_ops_inc_gen",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:tensor_opdefs",
+    ],
+)
+
+cc_library(
+    name = "runtime_fallback_executor",
+    testonly = True,
+    srcs = [
+        "runtime_fallback/runtime_fallback_executor.cc",
+    ],
+    hdrs = [
+        "runtime_fallback/runtime_fallback_executor.h",
+    ],
+    deps = [
+        ":host_context_util",
+        ":tf_to_tfrt",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/platform:env",
+        "//tensorflow/core/platform:threadpool_interface",
+        "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_execute_compat_eager",
+        "//tensorflow/core/runtime_fallback/runtime:kernel_utils",
+        "//tensorflow/core/runtime_fallback/runtime:runtime_fallback_alwayslink",
+        "//tensorflow/core/tfrt/utils:fallback_tensor",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@tf_runtime//:basic_kernels_alwayslink",
+        "@tf_runtime//:bef",
+        "@tf_runtime//:befexecutor",
+        "@tf_runtime//:core_runtime_alwayslink",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:mlirtobef",
+        "@tf_runtime//:support",
+    ],
+)
+
+cc_library(
+    name = "corert_converter",
+    srcs = [
+        "transforms/corert_converter.cc",
+    ],
+    hdrs = [
+        "transforms/corert_converter.h",
+    ],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    deps = [
+        ":attr_lowering_utils",
+        ":transform_utils",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/core:framework",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:DerivedAttributeOpInterface",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:core_runtime_opdefs",
+    ],
+)
+
+cc_library(
+    name = "fallback_converter",
+    srcs = [
+        "transforms/fallback_converter.cc",
+    ],
+    hdrs = [
+        "transforms/fallback_converter.h",
+    ],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_opdefs",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:core_runtime_opdefs",
+    ],
+)
+
+cc_library(
+    name = "transforms/gpu_passes",
+    srcs = ["transforms/gpu_passes.cc"],
+    hdrs = ["transforms/gpu_passes.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_gpu_opdefs",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+cc_library(
+    name = "tfrt_pipeline_options",
+    hdrs = ["transforms/tfrt_pipeline_options.h"],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    deps = [
+        ":tfrt_compile_options",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Pass",
+    ],
+)
+
+cc_library(
+    name = "tf_to_tfrt",
+    srcs = [
+        "transforms/cross_device_transfer.cc",
+        "transforms/deduplicate_batch_function.cc",
+        "transforms/deduplicate_if_result_pass.cc",
+        "transforms/fuse_tpu_compile_and_execute_ops.cc",
+        "transforms/insert_tensor_copy.cc",
+        "transforms/lower_saved_model.cc",
+        "transforms/merge_tf_if_ops.cc",
+        "transforms/optimize.cc",
+        "transforms/optimize_tf_control_flow_side_effect.cc",
+        "transforms/passes.cc",
+        "transforms/reconfig_batch_op.cc",
+        "transforms/remove_device_attribute.cc",
+        "transforms/remove_tf_if_const_args.cc",
+        "transforms/reorder_assert.cc",
+        "transforms/sink_in_invariant_ops.cc",
+        "transforms/tf_to_tfrt.cc",
+        "transforms/xla_rewrite_pass.cc",
+    ],
+    hdrs = [
+        "transforms/passes.h",
+    ],
+    deps = [
+        ":attr_lowering_utils",
+        ":corert_converter",
+        ":cost_analysis",
+        ":fallback_converter",
+        ":tensor_array_side_effect_analysis",
+        ":tfrt_compile_options",
+        ":tfrt_pipeline_options",
+        ":tpu_passes",
+        ":transform_utils",
+        ":transforms/gpu_passes",
+        ":transforms/set_shape_invariant_in_while_ops",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:bridge_logger",
+        "//tensorflow/compiler/mlir/tensorflow:convert_tensor",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_analysis",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_op_interfaces",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/compiler/mlir/tensorflow/ir/host_runtime:tensorflow_tfrt_ops",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_asset_sinking_pass",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_gpu_opdefs",
+        "//tensorflow/compiler/tf2xla:tf2xla_defs",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_lite",
+        "//tensorflow/core/ir/types:Dialect",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Rewrite",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:core_runtime_opdefs",
+        "@tf_runtime//:stream_analysis",
+        "@tf_runtime//:test_kernels_opdefs",
+        "@tsl//tsl/platform:errors",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "host_context_util",
+    srcs = ["utils/host_context.cc"],
+    hdrs = ["utils/host_context.h"],
+    deps = [
+        "//tensorflow/core/platform:logging",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/log",
+        "@tf_runtime//:hostcontext",
+    ],
+)
+
+cc_library(
+    name = "function",
+    srcs = [
+        "function/function.cc",
+    ],
+    hdrs = [
+        "function/function.h",
+    ],
+    deps = [
+        ":tf_to_tfrt",
+        ":tfrt_compile_options",
+        ":tfrt_pipeline_options",
+        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:translate_lib",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:vlog_is_on",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:bef",
+        "@tf_runtime//:core_runtime",
+        "@tf_runtime//:mlirtobef",
+    ],
+)
+
+cc_library(
+    name = "saved_model",
+    srcs = [
+        "saved_model/saved_model.cc",
+    ],
+    hdrs = [
+        "saved_model/saved_model.h",
+    ],
+    deps = [
+        ":tf_to_tfrt",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:convert_type",
+        "//tensorflow/compiler/mlir/tensorflow:translate_lib",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_dialect_passes",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@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_runtime//:bef",
+        "@tf_runtime//:core_runtime",
+        "@tf_runtime//:mlirtobef",
+        "@tsl//tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "import_model",
+    srcs = [
+        "translate/import_model.cc",
+    ],
+    hdrs = [
+        "translate/import_model.h",
+    ],
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt/tests/saved_model:__pkg__",
+        "//tensorflow/compiler/mlir/tfrt/transforms/mlrt:__pkg__",
+        "//tensorflow/core/tfrt/graph_executor:__pkg__",
+        "//tensorflow/core/tfrt/saved_model:__pkg__",
+        "//tensorflow/core/tfrt/saved_model/utils:__pkg__",
+    ],
+    deps = [
+        ":backend_compiler",
+        ":function",
+        ":tf_to_tfrt",
+        ":tfrt_compile_options",
+        ":tfrt_pipeline_options",
+        ":tpu_passes",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
+        "//tensorflow/compiler/mlir/tensorflow:serialize_mlir_module_utils",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tf_saved_model_asset_sinking_pass",
+        "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:cluster_tf",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_dialect_to_executor",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
+        "//tensorflow/compiler/tf2xla:xla_op_registry",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/common_runtime:function_def_utils",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/tfrt/fallback:fallback_state",
+        "//tensorflow/core/tfrt/runtime",
+        "//tensorflow/core/tpu:tpu_defs",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/functional:function_ref",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:vlog_is_on",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:bef",
+        "@tf_runtime//:mlirtobef",
+        "@tf_runtime//:support",
+        "@tsl//tsl/platform:env",
+        "@tsl//tsl/platform:errors",
+    ],
+)
+
+tf_cc_test(
+    name = "import_model_test",
+    srcs = [
+        "translate/import_model_test.cc",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":import_model",
+        ":tfrt_compile_options",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "tfrt_compile_options",
+    srcs = ["translate/tfrt_compile_options.cc"],
+    hdrs = ["translate/tfrt_compile_options.h"],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_library(
+    name = "cost_analysis",
+    srcs = ["analysis/cost_analysis.cc"],
+    hdrs = ["analysis/cost_analysis.h"],
+    deps = [
+        ":constants",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/tfrt/fallback:cost_recorder",
+        "//tensorflow/core/tfrt/fallback:op_cost_map_proto_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:compiler_tfrt_op_interfaces",
+    ],
+)
+
+cc_library(
+    name = "test_cost_analysis_pass",
+    srcs = ["analysis/test_cost_analysis_pass.cc"],
+    deps = [
+        ":cost_analysis",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:Pass",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "tensor_array_side_effect_analysis",
+    srcs = ["analysis/tensor_array_side_effect_analysis.cc"],
+    hdrs = ["analysis/tensor_array_side_effect_analysis.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SideEffectInterfaces",
+    ],
+)
+
+cc_library(
+    name = "test_tensor_array_side_effect_analysis",
+    srcs = ["analysis/test_tensor_array_side_effect_analysis.cc"],
+    deps = [
+        ":tensor_array_side_effect_analysis",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+    ],
+    alwayslink = 1,
+)
+
+tf_proto_library(
+    name = "analysis/analysis_proto",
+    srcs = ["analysis/analysis.proto"],
+)
+
+cc_library(
+    name = "passes",
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt:tf_to_tfrt",
+    ],
+)
+
+cc_library(
+    name = "transforms/set_shape_invariant_in_while_ops",
+    srcs = ["transforms/set_shape_invariant_in_while_ops.cc"],
+    hdrs = ["transforms/set_shape_invariant_in_while_ops.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+cc_library(
+    name = "tf_tfrt_opt_lib",
+    testonly = True,
+    srcs = ["tf-tfrt-opt.cc"],
+    deps = [
+        ":passes",
+        ":test_cost_analysis_pass",
+        ":test_opkernels",
+        ":test_tensor_array_side_effect_analysis",
+        ":tf_to_tfrt",
+        ":tpu_passes",
+        ":transforms/gpu_passes",
+        "//tensorflow/compiler/mlir:init_mlir",
+        "//tensorflow/compiler/mlir:passes",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "//tensorflow/compiler/mlir/tensorflow/transforms/host_runtime:lower_cluster_to_runtime_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_sync_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/ir/mlrt:tf_mlrt_ops",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:pack_inputs_pass",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf_ifrt_passes",
+        "//tensorflow/compiler/mlir/tfrt/transforms/mlrt:passes",
+        "//tensorflow/core:tensorflow",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:MlirOptLib",
+        "@llvm-project//mlir:RegisterAllDialects",  # buildcleaner: keep
+        "@llvm-project//mlir:RegisterAllPasses",  # buildcleaner: keep
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@stablehlo//:register",
+        "@tf_runtime//:init_tfrt_dialects",
+        "@tf_runtime//:print_stream_pass",
+        "@xla//xla/mlir_hlo",
+    ],
+)
+
+tf_cc_binary(
+    name = "tf-tfrt-opt",
+    testonly = True,
+    deps = [":tf_tfrt_opt_lib"],
+)
+
+tf_cc_binary(
+    name = "tfrt_translate",
+    srcs = ["tools/tfrt_translate/static_registration.cc"],
+    visibility = [":friends"],
+    deps = [
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:TranslateLib",
+        "@tf_runtime//:beftomlir_translate",
+        "@tf_runtime//:init_tfrt_dialects",
+        "@tf_runtime//:mlirtobef_translate",
+    ] + if_google(
+        [
+            "//third_party/tf_runtime_llvm:tfrt_translate_main",
+        ],
+        ["@tf_runtime//third_party/llvm_derived:tfrt_translate_main"],
+    ),
+)
+
+tf_cc_binary(
+    name = "bef_executor",
+    testonly = True,
+    visibility = [":friends"],
+    deps = [
+        "@tf_runtime//:dtype",
+        "@tf_runtime//:simple_tracing_sink",
+        "@tf_runtime//tools:bef_executor_expensive_kernels",
+        "@tf_runtime//tools:bef_executor_lib",
+        "@tf_runtime//tools:bef_executor_lightweight_kernels",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_registration",
+    srcs = [
+        "tfrt_fallback_registration.cc",
+    ],
+    hdrs = [
+        "tfrt_fallback_registration.h",
+    ],
+    visibility = [":friends"] + if_google([
+        "//learning/brain/experimental/tfrt/visualization:__pkg__",
+        # Allow visibility from the mlir language server.
+        "//learning/brain/mlir/mlir_lsp_server:__pkg__",
+    ]),
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_sync_opdefs",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+tf_cc_binary(
+    name = "tfrt_fallback_translate",
+    srcs = [
+        "tfrt_fallback_translate_registration.cc",
+    ],
+    visibility = [":friends"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tfrt:tfrt_fallback_registration",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_opdefs",
+        "@llvm-project//mlir:TranslateLib",
+        "@tf_runtime//:init_tfrt_dialects",
+        "@tf_runtime//:mlirtobef_translate",
+    ] + if_google(
+        [
+            "//third_party/tf_runtime_llvm:tfrt_translate_main",
+        ],
+        ["@tf_runtime//third_party/llvm_derived:tfrt_translate_main"],
+    ),
+)
+
+cc_library(
+    name = "attr_lowering_utils",
+    srcs = [
+        "transforms/attr_lowering_utils.cc",
+    ],
+    hdrs = [
+        "transforms/attr_lowering_utils.h",
+    ],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    visibility = [":friends"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_attributes",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:core_runtime_opdefs",
+    ],
+)
+
+cc_library(
+    name = "transform_utils",
+    srcs = [
+        "transforms/utils.cc",
+    ],
+    hdrs = [
+        "transforms/utils.h",
+    ],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    visibility = [":friends"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops",
+        "//tensorflow/compiler/mlir/tensorflow/ir/host_runtime:tensorflow_tfrt_ops",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:core_runtime_opdefs",
+    ],
+)
+
+cc_library(
+    name = "transforms/update_op_cost_in_tfrt_mlir",
+    srcs = ["transforms/update_op_cost_in_tfrt_mlir.cc"],
+    hdrs = ["transforms/update_op_cost_in_tfrt_mlir.h"],
+    deps = [
+        ":cost_analysis",
+        "//tensorflow/core/tfrt/fallback:cost_recorder",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "tpu_passes",
+    hdrs = ["transforms/tpu_passes.h"],
+    visibility = [
+        ":friends",
+    ] + if_google([
+        "//learning/brain/tfrt/ifrt/pjrt/__subpackages__",
+        "//learning/serving/servables/tfrt:__subpackages__",
+        "//smartass/brain/util:__pkg__",
+    ]),
+    deps = [
+        ":fallback_converter",
+        ":tfrt_compile_options",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+cc_library(
+    name = "constants",
+    hdrs = ["constants.h"],
+)
+
+cc_library(
+    name = "test_opkernels",
+    testonly = True,
+    srcs = ["test_opkernels.cc"],
+    deps = ["//tensorflow/core:framework"],
+    alwayslink = True,
+)
+
+cc_library(
+    name = "backend_compiler",
+    srcs = ["backend_compiler.cc"],
+    hdrs = ["backend_compiler.h"],
+    deps = [
+        "//tensorflow/core/tfrt/runtime",
+        "@com_google_absl//absl/status",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+cc_library(
+    name = "export",
+    srcs = ["utils/export.cc"],
+    hdrs = ["utils/export.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
+        "//tensorflow/compiler/mlir/tensorflow/transforms:tensorflow_passes",
+        "//tensorflow/compiler/mlir/tf2xla/api/v1:tf_dialect_to_executor",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tsl//tsl/platform:errors",
+        "@tsl//tsl/profiler/lib:traceme",
+    ],
+)

tensorflow/compiler/mlir/tfrt/analysis/analysis.proto

@@ -0,0 +1,41 @@
+// 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 mlir.tfrt;
+
+message CompatibilityAnalysisReportProto {
+  bool unknown_dialect = 1;
+  bool ref_variable = 2;
+  bool incompatible_variable = 3;
+  bool incompatible_attribute = 4;
+  bool control_flow_v1 = 5;
+  string method_name = 6;
+
+  // TODO(chky): add more checks, eg. tensor datatypes.
+}
+
+message CompatibilityAnalysisProto {
+  CompatibilityAnalysisReportProto summary = 1;
+
+  message OpInfo {
+    int32 count = 1;
+
+    CompatibilityAnalysisReportProto report = 2;
+  }
+
+  map<string, OpInfo> ops = 2;
+}

tensorflow/compiler/mlir/tfrt/analysis/cost_analysis.cc

@@ -0,0 +1,235 @@
+/* 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/analysis/cost_analysis.h"
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <string>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/string_view.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/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#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_types.h"
+#include "tensorflow/compiler/mlir/tfrt/constants.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tfrt/compiler/opdefs/tfrt_op_interfaces.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+constexpr int64_t kDefaultCheapCost = 1;
+
+int64_t GetRankedTensorSize(mlir::TensorType type) {
+  auto shape = type.getShape();
+
+  int64_t size = 1;
+  for (int64_t dim : shape) {
+    // For unknown dimensions, use 1 as the size because it is usually the batch
+    // dimension.
+    //
+    // TODO(chky): Find out a better default number for this case.
+    size *= std::max(kDefaultCheapCost, dim);
+  }
+
+  return size;
+}
+
+int64_t InferTensorSize(const CostContext& context, mlir::TensorType type) {
+  if (type.hasRank()) return GetRankedTensorSize(type);
+  return context.default_unranked_tensor_size;
+}
+
+// The cost function for tf.LookupTableFindV2.
+int64_t InferLookupTableFindV2Cost(const CostContext& context,
+                                   mlir::TF::LookupTableFindV2Op op) {
+  // tf.LookupTableFindV2 ops are usually more costly than tf.AddV2 with the
+  // same input size, as it involves more operations like hashing, map lookup,
+  // etc.
+  constexpr int64_t kLookupTableFindCostScale = 8;
+  constexpr int64_t kLookupTableFindStringKeyCostScale = 16;
+
+  auto value_type = mlir::cast<mlir::TensorType>(op.getValues().getType());
+  auto key_type = mlir::cast<mlir::TensorType>(op.getKeys().getType());
+
+  int64_t output_size = InferTensorSize(context, value_type);
+
+  int64_t cost = kLookupTableFindCostScale * output_size;
+
+  if (mlir::isa<mlir::TF::StringType>(key_type.getElementType()))
+    cost *= kLookupTableFindStringKeyCostScale;
+
+  return cost;
+}
+
+// The cost function for tf.GatherV2.
+int64_t InferGatherV2Cost(const CostContext& context, mlir::TF::GatherV2Op op) {
+  return InferTensorSize(
+      context, mlir::cast<mlir::TensorType>(op.getOutput().getType()));
+}
+
+// The cost function for tf.SparseSegmentSumOp.
+template <typename OpType>
+int64_t InferSparseSegmentOpCost(const CostContext& context, OpType op) {
+  return InferTensorSize(
+      context, mlir::cast<mlir::TensorType>(op.getOutput().getType()));
+}
+
+// CostFunctionRegistry is a map from op names to their cost functions.
+using CostFunctionRegistry = absl::flat_hash_map<std::string, CostFunction>;
+
+void RegisterCostFunction(CostFunctionRegistry& registry,
+                          absl::string_view op_name,
+                          CostFunction cost_function) {
+  auto r = registry.try_emplace(op_name, std::move(cost_function));
+  assert(r.second);
+  (void)r;
+}
+
+template <typename OpType, typename F>
+void RegisterCostFunction(CostFunctionRegistry& registry, F f) {
+  RegisterCostFunction(
+      registry, OpType::getOperationName().str(),
+      [f = std::move(f)](const CostContext& context, mlir::Operation* op) {
+        return f(context, llvm::cast<OpType>(op));
+      });
+}
+
+CostFunctionRegistry& GetCostFunctionRegistry() {
+  static auto* const registry = []() {
+    auto* registry = new CostFunctionRegistry;
+    // TODO(chky): Find a more scalable way to register cost functions. One
+    // option is to incorporate it is TF MLIR ODS.
+    RegisterCostFunction<mlir::TF::GatherV2Op>(*registry, InferGatherV2Cost);
+    RegisterCostFunction<mlir::TF::SparseSegmentSumOp>(
+        *registry, InferSparseSegmentOpCost<mlir::TF::SparseSegmentSumOp>);
+    RegisterCostFunction<mlir::TF::SparseSegmentMeanOp>(
+        *registry, InferSparseSegmentOpCost<mlir::TF::SparseSegmentMeanOp>);
+    RegisterCostFunction<mlir::TF::SparseSegmentSqrtNOp>(
+        *registry, InferSparseSegmentOpCost<mlir::TF::SparseSegmentSqrtNOp>);
+    RegisterCostFunction<mlir::TF::LookupTableFindV2Op>(
+        *registry, InferLookupTableFindV2Cost);
+    return registry;
+  }();
+  return *registry;
+}
+
+}  // namespace
+
+void RegisterCostFunction(absl::string_view op_name,
+                          CostFunction cost_function) {
+  RegisterCostFunction(GetCostFunctionRegistry(), op_name,
+                       std::move(cost_function));
+}
+
+bool HasCostFunctionRegistered(absl::string_view op_name) {
+  return GetCostFunctionRegistry().contains(op_name);
+}
+
+int64_t CostAnalysis::GetCost(mlir::Operation* op) const {
+  assert(cost_map_.count(op) > 0);
+  return cost_map_.lookup(op);
+}
+
+void CostAnalysis::AnalyzeArguments(mlir::func::FuncOp func_op) {
+  // Use the max size among function inputs as the default size of dynamic
+  // shaped tensors in the function.
+  for (auto arg : func_op.getArguments()) {
+    if (!mlir::isa<mlir::TensorType>(arg.getType())) continue;
+    auto type = mlir::cast<mlir::TensorType>(arg.getType());
+    if (type.hasRank()) {
+      max_arg_size_ = std::max(max_arg_size_, GetRankedTensorSize(type));
+    }
+  }
+}
+
+void CostAnalysis::AnalyzeBlock(mlir::Block* block) {
+  for (auto& op : *block) {
+    EvaluateCost(&op);
+  }
+}
+
+void CostAnalysis::EvaluateCost(mlir::Operation* op) {
+  if (auto cost_function =
+          mlir::dyn_cast<tfrt::compiler::CostFunctionInterface>(op)) {
+    cost_map_[op] = cost_function.cost();
+    return;
+  }
+
+  if (!llvm::isa<mlir::TF::TensorFlowDialect>(op->getDialect())) {
+    cost_map_[op] = max_arg_size_;
+    return;
+  }
+
+  // Try to use its cost function if it is registered.
+  const auto& registry = GetCostFunctionRegistry();
+  absl::string_view op_name = op->getName().getStringRef();
+  auto iter = registry.find(op_name);
+  if (iter != registry.end()) {
+    CostContext context;
+    context.default_unranked_tensor_size = max_arg_size_;
+    cost_map_[op] = iter->second(context, op);
+    return;
+  }
+
+  // Try to use the recorded cost if any.
+  if (cost_recorder_ != nullptr) {
+    const auto op_key_attr =
+        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
+    if (op_key_attr) {
+      cost_map_[op] = cost_recorder_->GetCost(op_key_attr.getInt());
+      return;
+    }
+  }
+
+  // These ops are cheap regardless of their input sizes.
+  //
+  // TODO(chky): Find a more scalable way to figure out cheap ops.
+  if (llvm::isa<mlir::TF::ShapeOp, mlir::TF::StridedSliceOp,
+                mlir::TF::ReshapeOp, mlir::TF::ExpandDimsOp>(op)) {
+    cost_map_[op] = kDefaultCheapCost;
+    return;
+  }
+
+  // For other ops, use the sum of input sizes as its cost.
+  int64_t cost = kDefaultCheapCost;
+  for (auto operand : op->getOperands()) {
+    auto type = mlir::cast<mlir::TensorType>(operand.getType());
+    if (type.hasRank()) {
+      cost += GetRankedTensorSize(type);
+    } else {
+      // For unranked tensors, use the max size among the input tensors. This is
+      // because the only dynamic information of the function should be the
+      // input, so the size of dynamic tensors should be usually capped by
+      // inputs' sizes.
+      cost += max_arg_size_;
+    }
+  }
+
+  cost_map_[op] = cost;
+}
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/analysis/cost_analysis.h

@@ -0,0 +1,98 @@
+/* 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_ANALYSIS_COST_ANALYSIS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_ANALYSIS_COST_ANALYSIS_H_
+
+#include <cstdint>
+#include <functional>
+
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tensorflow/core/tfrt/fallback/op_cost_map.pb.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+// Analyze costs for tensorflow operations.
+//
+// The current heuristic used is quite simple, which is to calculate the total
+// size of input tensors. The exception is that ops whose cost is irrelevant to
+// input sizes, such as tf.Shape and tf.Reshape, are whitelisted to have cheap
+// cost. This cost analysis is expected to be used conservatively (eg. use a low
+// threshold to decide whether a cost is cheap or expensive), as it might not be
+// accurate in some cases.
+//
+class CostAnalysis {
+ public:
+  explicit CostAnalysis(
+      mlir::func::FuncOp func_op,
+      const tfrt_stub::CostRecorder* cost_recorder = nullptr) {
+    cost_recorder_ = cost_recorder;
+    AnalyzeArguments(func_op);
+    AnalyzeBlock(&func_op.front());
+  }
+
+  int64_t GetCost(mlir::Operation* op) const;
+
+ private:
+  void AnalyzeArguments(mlir::func::FuncOp func_op);
+  void AnalyzeBlock(mlir::Block* block);
+  void EvaluateCost(mlir::Operation* op);
+
+  int64_t max_arg_size_ = 1;
+  llvm::DenseMap<mlir::Operation*, int64_t> cost_map_;
+  const tfrt_stub::CostRecorder* cost_recorder_;
+};
+
+struct CostContext {
+  int64_t default_unranked_tensor_size;
+};
+
+using CostFunction =
+    std::function<int64_t(const CostContext&, mlir::Operation*)>;
+
+void RegisterCostFunction(absl::string_view op_name,
+                          CostFunction cost_function);
+
+template <typename OpType, typename F>
+void RegisterCostFunction(F f) {
+  RegisterCostFunction(
+      OpType::getOperationName().str(),
+      [f = std::move(f)](const CostContext& context, mlir::Operation* op) {
+        return f(context, llvm::cast<OpType>(op));
+      });
+}
+
+template <typename OpType>
+struct CostFunctionRegistration {
+  explicit CostFunctionRegistration(
+      std::function<int64_t(const CostContext&, OpType)> cost_function) {
+    RegisterCostFunction<OpType>(std::move(cost_function));
+  }
+};
+
+bool HasCostFunctionRegistered(absl::string_view op_name);
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_ANALYSIS_COST_ANALYSIS_H_

tensorflow/compiler/mlir/tfrt/analysis/tensor_array_side_effect_analysis.cc

@@ -0,0 +1,53 @@
+/* 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/analysis/tensor_array_side_effect_analysis.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 "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+bool IsTensorArrayOp(mlir::Operation* op) {
+  return llvm::isa<mlir::TF::TensorArrayV3Op, mlir::TF::TensorArrayScatterV3Op,
+                   mlir::TF::TensorArrayGatherV3Op,
+                   mlir::TF::TensorArrayReadV3Op,
+                   mlir::TF::TensorArrayWriteV3Op>(op);
+}
+
+static bool FunctionContainsOnlyNoSideEffectOpOrTensorArrayOp(
+    mlir::func::FuncOp func_op) {
+  for (mlir::Operation& op : func_op.front()) {
+    if (!mlir::isMemoryEffectFree(&op) && !IsTensorArrayOp(&op)) return false;
+  }
+
+  return true;
+}
+
+TensorArraySideEffectAnalysis::TensorArraySideEffectAnalysis(
+    mlir::ModuleOp module) {
+  for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
+    if (FunctionContainsOnlyNoSideEffectOpOrTensorArrayOp(func_op)) {
+      set_.insert(func_op);
+    }
+  }
+}
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/analysis/tensor_array_side_effect_analysis.h

@@ -0,0 +1,52 @@
+/* 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_ANALYSIS_TENSOR_ARRAY_SIDE_EFFECT_ANALYSIS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_ANALYSIS_TENSOR_ARRAY_SIDE_EFFECT_ANALYSIS_H_
+
+#include "llvm/ADT/DenseSet.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
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+// Return true if it is a TensorArrayOp, eg. TensorArrayV3Op.
+bool IsTensorArrayOp(mlir::Operation* op);
+
+// This class provides utilities for analyzing side effects for TensorArray ops
+// in the graph. mlir::TF::SideEffectAnalysis currently produces suboptimal
+// side-effect analysis for TensorArray ops. On the other hand, control
+// dependencies are already sorted out for TensorArray ops in the original TF
+// graph. Each TensorArray op will take or produce a `flow` value and they are
+// already properly chained in the origninal TF graph.
+class TensorArraySideEffectAnalysis {
+ public:
+  explicit TensorArraySideEffectAnalysis(mlir::ModuleOp module);
+
+  // Return if the function contains only non-side-effecting ops or TensorArray
+  // ops.
+  bool HasAtMostTensorArrayEffect(mlir::func::FuncOp func_op) const {
+    return set_.count(func_op) > 0;
+  }
+
+ private:
+  llvm::DenseSet<mlir::func::FuncOp> set_;
+};
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_ANALYSIS_TENSOR_ARRAY_SIDE_EFFECT_ANALYSIS_H_

tensorflow/compiler/mlir/tfrt/analysis/test_cost_analysis_pass.cc

@@ -0,0 +1,46 @@
+/* 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 "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Pass/PassRegistry.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/analysis/cost_analysis.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+class TestCostAnalysis
+    : public mlir::PassWrapper<TestCostAnalysis,
+                               mlir::OperationPass<mlir::func::FuncOp>> {
+  llvm::StringRef getArgument() const final {
+    return "tfrt-test-cost-analysis";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Add remarks based on cost analysis for testing purpose.";
+  }
+  void runOnOperation() override {
+    const auto& cost_analysis = getAnalysis<CostAnalysis>();
+
+    auto func_op = getOperation();
+    for (auto& op : func_op.front()) {
+      op.emitRemark() << "Cost: " << cost_analysis.GetCost(&op);
+    }
+  }
+};
+
+static mlir::PassRegistration<TestCostAnalysis> pass;
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/analysis/test_tensor_array_side_effect_analysis.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 "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
+#include "mlir/Support/TypeID.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/analysis/tensor_array_side_effect_analysis.h"
+
+namespace tensorflow {
+namespace tfrt_compiler {
+namespace {
+
+class TestTensorArraySideEffectAnalysis
+    : public mlir::PassWrapper<TestTensorArraySideEffectAnalysis,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+ public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+      TestTensorArraySideEffectAnalysis)
+
+ private:
+  llvm::StringRef getArgument() const final {
+    return "tfrt-test-tensor-array-effect";
+  }
+  llvm::StringRef getDescription() const final {
+    return "Test TensorArraySideEffectAnalysis";
+  }
+  void runOnOperation() override {
+    auto module = getOperation();
+    TensorArraySideEffectAnalysis tensor_array_side_effect_analysis(module);
+
+    for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
+      func_op.emitRemark() << "HasAtMostTensorArrayEffect: "
+                           << tensor_array_side_effect_analysis
+                                  .HasAtMostTensorArrayEffect(func_op);
+    }
+  }
+};
+
+mlir::PassRegistration<TestTensorArraySideEffectAnalysis> pass;
+
+}  // namespace
+}  // namespace tfrt_compiler
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/backend_compiler.cc

@@ -0,0 +1,22 @@
+/* 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/backend_compiler.h"
+
+namespace tensorflow {
+
+BackendCompiler::~BackendCompiler() = default;
+
+}

tensorflow/compiler/mlir/tfrt/backend_compiler.h

@@ -0,0 +1,41 @@
+/* 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_BACKEND_COMPILER_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_BACKEND_COMPILER_H_
+
+#include "absl/status/status.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+
+namespace tensorflow {
+
+class BackendCompiler {
+ public:
+  virtual ~BackendCompiler();
+
+  virtual void GetDependentDialects(mlir::DialectRegistry& registry) const {}
+
+  // Compile the `module` in TF dialect. The result module should be also in TF
+  // dialect.
+  virtual absl::Status CompileTensorflow(
+      tfrt_stub::ModelRuntimeContext& model_context,
+      mlir::ModuleOp module) const = 0;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_BACKEND_COMPILER_H_

tensorflow/compiler/mlir/tfrt/constants.h

@@ -0,0 +1,28 @@
+/* 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_CONSTANTS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_CONSTANTS_H_
+
+namespace tensorflow {
+namespace tfrt_compiler {
+
+// Use __ prefix to indicate this is internal attribute.
+inline constexpr char kOpKeyAttrName[] = "__op_key";
+
+}  // namespace tfrt_compiler
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_CONSTANTS_H_

tensorflow/compiler/mlir/tfrt/function/function.cc

@@ -0,0 +1,100 @@
+/* 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/function/function.h"
+
+#include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "mlir/IR/OperationSupport.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/translate/tf_mlir_translate.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/compiler/mlir/tfrt/transforms/tfrt_pipeline_options.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/bef_converter/mlir_to_bef.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+absl::Status CompileTFMLIRToBEF(const TfrtFunctionCompileOptions& options,
+                                mlir::ModuleOp module,
+                                tfrt::BefBuffer* bef_buffer) {
+  mlir::OpPrintingFlags print_flags;
+  print_flags.elideLargeElementsAttrs();
+
+  if (VLOG_IS_ON(1)) {
+    VLOG(1) << "Input TF Executor dialect:";
+    DumpMlirOpToFile("tf_to_tfrt_tf_executor_dialect", module);
+  }
+
+  mlir::StatusScopedDiagnosticHandler diag_handler(module.getContext());
+
+  // Lower MLIR TF Dialect to MLIR TFRT CoreRT dialect.
+  mlir::PassManager pm(module.getContext());
+  tensorflow::applyTensorflowAndCLOptions(pm);
+
+  tensorflow::TfrtPipelineOptions pass_options;
+  if (!options.default_device.empty()) {
+    pass_options.default_device = options.default_device;
+  }
+  if (!options.force_data_format.empty()) {
+    pass_options.force_data_format = options.force_data_format;
+  }
+  // TODO(tfrt-devs): Current MaxPoolingOp only supports NHWC on device type
+  // CPU. Enable this layout optimization after we introduce TFRT native ops
+  // for training.
+  if (absl::StrContains(pass_options.default_device, "CPU")) {
+    pass_options.skip_fold_transpose_in_ops = true;
+  }
+  pass_options.enable_optimizer = options.enable_optimizer;
+  // Use TFRT TPU OpKernel for training.
+  pass_options.target_tpurt = false;
+  pass_options.tpu_use_core_selector = options.tpu_use_core_selector;
+  pass_options.tpu_use_bundled_transfer = options.tpu_use_bundled_transfer;
+  pass_options.tpu_lower_to_fallback = options.tpu_lower_to_fallback;
+  pass_options.tpu_fuse_ops = options.tpu_fuse_ops;
+  pass_options.tpu_transfer_result_to_host =
+      options.tpu_transfer_result_to_host;
+  absl::Status status =
+      tensorflow::CreateTfExecutorToTfrtPipeline(pm, pass_options);
+  if (!status.ok()) {
+    return diag_handler.Combine(status);
+  }
+
+  if (mlir::failed(pm.run(module)))
+    return diag_handler.Combine(
+        absl::InternalError("failed to lower TF Dialect to CoreRT dialect."));
+
+  if (VLOG_IS_ON(1)) {
+    VLOG(1) << "TFRT dialect: ";
+    DumpMlirOpToFile("tf_to_tfrt_tfrt_dialect", module);
+  }
+
+  *bef_buffer =
+      tfrt::ConvertMLIRToBEF(module, /* disable_optional_sections = */ true);
+  if (bef_buffer->empty())
+    return diag_handler.Combine(
+        absl::InternalError("failed to convert MLIR to BEF."));
+
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/function/function.h

@@ -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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_FUNCTION_FUNCTION_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_FUNCTION_FUNCTION_H_
+
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+#include "tensorflow/core/platform/status.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/core_runtime/tensor_handle.h"  // from @tf_runtime
+
+namespace tfrt {
+class CoreRuntime;
+}
+
+namespace mlir {
+class ModuleOp;
+}
+
+namespace tensorflow {
+
+struct TfrtFunctionCompileOptions : public TfrtCompileOptions {
+  // Currently only SavedModel API inference uses the tpu_fuse_ops option
+  TfrtFunctionCompileOptions() {
+    tpu_fuse_ops = false;
+    // Currently grappler is not correctly applied in the eager execution of TF
+    // functions, as it may sometimes remove arguments and results.
+    enable_grappler = false;
+  }
+
+  // If true, use ServingCoreSelector to pick TPU core. Otherwise, obtain core
+  // location from assigned device name.
+  // Currently we don't use core_selector for training use cases.
+  bool tpu_use_core_selector = false;
+
+  // If true, use BundledTransferToTpuOp to transfer variables and input tensors
+  // to TPU.
+  bool tpu_use_bundled_transfer = false;
+
+  // If true, lower an TF op that's placed on TPU device to be executed with
+  // tfrt_fallback.execute.
+  // Currently for training use cases we need to lower the op to corert.execute
+  // to execute with TPU OpHandler, and with TFRT's native implementation.
+  // TODO(b/188940204): remove this config after we clear up the TPU variable
+  // implementation.
+  bool tpu_lower_to_fallback = false;
+  // If true, transfer the result of TPUExecuteOp from TPU to host.
+  // Currently for training and Python bulk inference use cases, we don't need
+  // to proactively transfer the result to host since the consumer op (or
+  // function) of the result may still be on TPU.
+  // TODO(b/194081364): remove this option once we unify servo TPU serving
+  // result transfer behavior.
+  bool tpu_transfer_result_to_host = false;
+};
+
+// Compile MLIR generated by tf.function in TF dialect into BEF.
+absl::Status CompileTFMLIRToBEF(const TfrtFunctionCompileOptions& options,
+                                mlir::ModuleOp module,
+                                tfrt::BefBuffer* bef_buffer);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_FUNCTION_FUNCTION_H_

tensorflow/compiler/mlir/tfrt/ir/BUILD

@@ -0,0 +1,235 @@
+# BUILD targets for runtime fallback opdefs.
+
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+# Note: keep the following lines separate due to the way copybara works
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "tfrt_fallback_opdefs",
+    srcs = [
+        "tfrt_fallback.cc",
+    ],
+    hdrs = ["tfrt_fallback.h"],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    visibility = [
+        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
+        # copybara:uncomment "//learning/brain/tfrt/tpu/compiler:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/runtime_fallback:internal",
+        "//tensorflow/core/tfrt/saved_model:friends",
+        # copybara:uncomment "//third_party/tf_runtime_google:__subpackages__",
+    ],
+    deps = [
+        ":tfrt_fallback_opdefs_inc_gen",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_async_opdefs",
+    srcs = [
+        "tfrt_fallback_async.cc",
+    ],
+    hdrs = ["tfrt_fallback_async.h"],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    visibility = [
+        # copybara:uncomment "//learning/brain/experimental/tfrt:__subpackages__",
+        # copybara:uncomment "//learning/brain/tfrt/tpu/compiler/mlir:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/tfrt/saved_model:friends",
+    ],
+    deps = [
+        ":tfrt_fallback_async_opdefs_inc_gen",
+        ":tfrt_fallback_common",
+        ":tfrt_fallback_opdefs",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InliningUtils",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:compiler_tfrt_op_interfaces",
+        "@tf_runtime//:compiler_tfrt_traits",
+        "@tf_runtime//:core_runtime_opdefs",
+        "@tf_runtime//:core_runtime_sync_opdefs",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_sync_opdefs",
+    srcs = [
+        "tfrt_fallback_sync.cc",
+    ],
+    hdrs = ["tfrt_fallback_sync.h"],
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/runtime_fallback:internal",
+    ],
+    deps = [
+        ":tfrt_fallback_common",
+        ":tfrt_fallback_opdefs",
+        ":tfrt_fallback_sync_opdefs_inc_gen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@tf_runtime//:basic_kernels_opdefs",
+        "@tf_runtime//:core_runtime_opdefs",
+        "@tf_runtime//:tensor_opdefs",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_common",
+    srcs = ["tfrt_fallback_common.cc"],
+    hdrs = ["tfrt_fallback_common.h"],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:basic_kernels_opdefs",
+    ],
+)
+
+cc_library(
+    name = "tfrt_fallback_util",
+    srcs = ["tfrt_fallback_util.cc"],
+    hdrs = ["tfrt_fallback_util.h"],
+    visibility = ["//tensorflow/compiler/mlir/tfrt:__subpackages__"],
+    deps = [
+        ":tfrt_fallback_async_opdefs",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+    ],
+)
+
+td_library(
+    name = "tfrt_fallback_td_files",
+    srcs = [
+        "tfrt_fallback.td",
+        "tfrt_fallback_async.td",
+        "tfrt_fallback_sync.td",
+    ],
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    includes = ["."],
+    visibility = [
+        "//learning/brain/experimental/tfrt:__subpackages__",
+        "//learning/brain/tfrt/tpu/compiler:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+        "@tf_runtime//:CoreRTTdFiles",
+        "@tf_runtime//:OpBaseTdFiles",
+        "@tf_runtime//:compiler_td_files",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tfrt_fallback_opdefs_inc_gen",
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    tbl_outs = {
+        "tfrt_fallback.h.inc": ["-gen-op-decls"],
+        "tfrt_fallback.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tfrt_fallback.td",
+    deps = [":tfrt_fallback_td_files"],
+)
+
+gentbl_cc_library(
+    name = "tfrt_fallback_async_opdefs_inc_gen",
+    compatible_with = get_compatible_with_portable(),  # copybara: comment
+    tbl_outs = {
+        "tfrt_fallback_async.h.inc": ["-gen-op-decls"],
+        "tfrt_fallback_async.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tfrt_fallback_async.td",
+    deps = [":tfrt_fallback_td_files"],
+)
+
+gentbl_cc_library(
+    name = "tfrt_fallback_sync_opdefs_inc_gen",
+    tbl_outs = {
+        "tfrt_fallback_sync.h.inc": ["-gen-op-decls"],
+        "tfrt_fallback_sync.cpp.inc": ["-gen-op-defs"],
+        "tfrt_fallback_sync_dialect.h.inc": [
+            "-gen-dialect-decls",
+            "-dialect=tfrt_fallback_sync",
+        ],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tfrt_fallback_sync.td",
+    test = True,
+    visibility = [
+        "//learning/brain/experimental/tfrt:__subpackages__",
+    ],
+    deps = [":tfrt_fallback_td_files"],
+)
+
+td_library(
+    name = "gpu_ops_td_file",
+    srcs = [
+        "gpu_ops.td",
+    ],
+    includes = ["."],
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [
+        ":tfrt_fallback_td_files",
+        "@tf_runtime//:OpBaseTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tfrt_gpu_opdefs_inc_gen",
+    tbl_outs = {
+        "gpu_ops.h.inc": ["-gen-op-decls"],
+        "gpu_ops.cpp.inc": ["-gen-op-defs"],
+        "gpurt_dialect.h.inc": [
+            "-gen-dialect-decls",
+            "-dialect=gpurt",
+        ],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "gpu_ops.td",
+    test = True,
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [":gpu_ops_td_file"],
+)
+
+cc_library(
+    name = "tfrt_gpu_opdefs",
+    srcs = [
+        "gpu_ops.cc",
+    ],
+    hdrs = ["gpu_ops.h"],
+    visibility = [
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [
+        ":tfrt_fallback_opdefs",
+        ":tfrt_gpu_opdefs_inc_gen",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:IR",
+    ],
+)

tensorflow/compiler/mlir/tfrt/ir/gpu_ops.cc

@@ -0,0 +1,41 @@
+/* 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/ir/gpu_ops.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
+
+namespace tfrt {
+namespace gpu {
+
+GpuRuntimeDialect::GpuRuntimeDialect(MLIRContext *context)
+    : Dialect(/*name=*/"gpurt", context, TypeID::get<GpuRuntimeDialect>()) {
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/gpu_ops.cpp.inc"
+      >();
+}
+
+}  // namespace gpu
+}  // namespace tfrt
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/gpu_ops.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/gpu_ops.h

@@ -0,0 +1,41 @@
+/* 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_IR_GPU_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_GPU_OPS_H_
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+
+using namespace mlir;  // NOLINT
+
+namespace tfrt {
+namespace gpu {
+
+// Dialect for TFRT GPU operations.
+class GpuRuntimeDialect : public Dialect {
+ public:
+  explicit GpuRuntimeDialect(MLIRContext *context);
+  static StringRef getDialectNamespace() { return "gpurt"; }
+};
+
+}  // namespace gpu
+}  // namespace tfrt
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/gpu_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_GPU_OPS_H_

tensorflow/compiler/mlir/tfrt/ir/gpu_ops.td

@@ -0,0 +1,114 @@
+/* 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.
+==============================================================================*/
+
+#ifdef TFRT_GPU_OPS
+#else
+#define TFRT_GPU_OPS
+
+include "tfrt/tfrt_op_base.td"
+include "tfrt/compiler/opdefs/tfrt_op_interfaces.td"
+include "tfrt/compiler/opdefs/tfrt_traits.td"
+include "tfrt_fallback.td"
+
+def TFRT_GPU_Dialect : Dialect {
+  let name = "gpurt";
+
+  let description = [{
+    The TFRT GPU Dialect.
+  }];
+
+  let cppNamespace = "::tfrt::gpu";
+}
+
+class Gpu_Op<string mnemonic, list<Trait> traits = []> :
+    Op<TFRT_GPU_Dialect, mnemonic, traits> {
+}
+
+// TODO(b/260267885): We may add a device argument when we want to support
+// GPU MIG.
+def TransferToDeviceOp: Gpu_Op<"transfer_to_device"> {
+  let summary = "Transfer a CPU tensor to device.";
+
+  let description = [{
+    Transfer a CPU tensor to device.
+
+    Example:
+      %device_tensor = gpurt.transfer_to_device %cpu_tensor
+  }];
+
+  let arguments = (ins TFTensorType);
+  let results = (outs TFTensorType);
+  let assemblyFormat = "operands attr-dict";
+}
+
+// TODO(b/260267885): We may add a device argument when we want to support
+// GPU MIG.
+def TransferFromDeviceOp: Gpu_Op<"transfer_from_device"> {
+  let summary = "Transfer a tensor from device.";
+
+  let description = [{
+    Transfer a tensor from device.
+
+    Example:
+      %cpu_tensor = gpurt.transfer_from_device %device_tensor
+  }];
+
+  let arguments = (ins TFTensorType);
+  let results = (outs TFTensorType);
+  let assemblyFormat = "operands attr-dict";
+}
+
+// TODO(b/260267885): We may add a device argument when we want to support
+// GPU MIG.
+def MaybeTransferVariableOp: Gpu_Op<"maybe_transfer_variable"> {
+  let summary = "Transfer a CPU variable tensor to device.";
+  let description = [{
+    Transfer a CPU variable tensor to device if the variable has not been
+    transferred before.
+
+    Example:
+      %device_var = gpurt.maybe_transfer_variable %cpu_var
+  }];
+
+  let arguments = (ins TFTensorType);
+  let results = (outs TFTensorType);
+  let assemblyFormat = "operands attr-dict";
+}
+
+def CompileAndExecuteOp: Gpu_Op<"compile_and_execute"> {
+  let summary = "GPU compile and execute operation.";
+  let description = [{
+    The op compiles and executes a GPU cluster function.
+
+    func_name is the name of the function to be executed on GPU.
+    resource_indices are the indices of inputs that are resources.
+    used_output_indices are the indices of outputs that have users.
+
+    Example:
+      %results = gpurt.compile_and_execute {func_name = "xla_func_0", resource_indices = [1] ...}
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$operands,
+    StrAttr:$func_name,
+    I64ArrayAttr:$resource_indices,
+    I64ArrayAttr:$used_output_indices
+  );
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+}
+
+#endif // TFRT_GPU_OPS

tensorflow/compiler/mlir/tfrt/ir/mlrt/BUILD

@@ -0,0 +1,157 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+td_library(
+    name = "mlrt_td_files",
+    srcs = [
+        "mlrt_dialect.td",
+        "mlrt_ops.td",
+    ],
+    includes = ["."],
+    visibility = ["//visibility:private"],
+    deps = [
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "mlrt_ops_inc_gen",
+    tbl_outs = {
+        "mlrt_ops.h.inc": ["-gen-op-decls"],
+        "mlrt_ops.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "mlrt_ops.td",
+    deps = [":mlrt_td_files"],
+)
+
+cc_library(
+    name = "mlrt_ops",
+    srcs = [
+        "mlrt_dialect.cc",
+        "mlrt_ops.cc",
+    ],
+    hdrs = [
+        "mlrt_dialect.h",
+        "mlrt_ops.h",
+    ],
+    visibility = ["//tensorflow/compiler/mlir/tfrt:__subpackages__"],
+    deps = [
+        ":mlrt_ops_inc_gen",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InliningUtils",
+        "@llvm-project//mlir:Support",
+    ],
+)
+
+td_library(
+    name = "tf_mlrt_td_files",
+    srcs = [
+        "tf_mlrt_dialect.td",
+        "tf_mlrt_ops.td",
+        "tf_ops.td",
+    ],
+    includes = ["."],
+    deps = [
+        ":mlrt_td_files",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+        "@tf_runtime//:compiler_td_files",
+    ],
+)
+
+td_library(
+    name = "tf_mlrt_tpu_td_files",
+    srcs = [
+        "tf_mlrt_tpu_ops.td",
+    ],
+    includes = ["."],
+    visibility = ["//visibility:private"],
+    deps = [
+        ":mlrt_td_files",
+        ":tf_mlrt_td_files",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_ops_td_files",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tf_mlrt_ops_inc_gen",
+    tbl_outs = {
+        "tf_mlrt_ops.h.inc": ["-gen-op-decls"],
+        "tf_mlrt_ops.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tf_mlrt_ops.td",
+    deps = [":tf_mlrt_td_files"],
+)
+
+gentbl_cc_library(
+    name = "tf_mlrt_tpu_ops_inc_gen",
+    tbl_outs = {
+        "tf_mlrt_tpu_ops.h.inc": ["-gen-op-decls"],
+        "tf_mlrt_tpu_ops.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tf_mlrt_tpu_ops.td",
+    deps = [":tf_mlrt_tpu_td_files"],
+)
+
+gentbl_cc_library(
+    name = "tf_ops_inc_gen",
+    tbl_outs = {
+        "tf_ops.h.inc": ["-gen-op-decls"],
+        "tf_ops.cpp.inc": ["-gen-op-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tf_ops.td",
+    deps = [":tf_mlrt_td_files"],
+)
+
+cc_library(
+    name = "tf_mlrt_ops",
+    srcs = ["tf_mlrt_ops.cc"],
+    hdrs = ["tf_mlrt_ops.h"],
+    visibility = [
+        # copybara:uncomment "//learning/brain/experimental/tfrt/mlrt/application/tensorflow/tests:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+    ],
+    deps = [
+        ":mlrt_ops",
+        ":tf_mlrt_ops_inc_gen",
+        ":tf_ops_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:tensorflow_side_effects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InliningUtils",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:Transforms",
+        "@tf_runtime//:compiler_tfrt_op_interfaces",
+        "@tf_runtime//:compiler_tfrt_traits",
+    ],
+)
+
+cc_library(
+    name = "tf_mlrt_tpu_ops",
+    srcs = ["tf_mlrt_tpu_ops.cc"],
+    hdrs = ["tf_mlrt_tpu_ops.h"],
+    visibility = ["//tensorflow/compiler/mlir/tfrt/transforms/mlrt:__subpackages__"],
+    deps = [
+        ":mlrt_ops",
+        ":tf_mlrt_ops",
+        ":tf_mlrt_tpu_ops_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SideEffectInterfaces",
+    ],
+)

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.cc

@@ -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.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.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/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Region.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h"
+
+namespace mlrt {
+namespace compiler {
+
+namespace {
+
+struct MlrtInlinerInterface : public mlir::DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+  bool isLegalToInline(mlir::Operation *op, mlir::Region *dest,
+                       bool would_be_cloned,
+                       mlir::IRMapping &mapping) const final {
+    // All mlrt dialect ops can be inlined.
+    return true;
+  }
+};
+
+}  // namespace
+
+MlrtDialect::MlrtDialect(mlir::MLIRContext *context)
+    : mlir::Dialect(/*name=*/"mlrt", context,
+                    mlir::TypeID::get<MlrtDialect>()) {
+  addTypes<FutureType>();
+  addTypes<PromiseType>();
+  addTypes<AsyncHandleType>();
+  addInterfaces<MlrtInlinerInterface>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.cpp.inc"
+      >();
+}
+
+// Parse a type registered to this dialect.
+mlir::Type MlrtDialect::parseType(mlir::DialectAsmParser &parser) const {
+  llvm::StringRef keyword;
+  if (parser.parseKeyword(&keyword)) return mlir::Type();
+
+  if (keyword == "future") return FutureType::get(getContext());
+  if (keyword == "promise") return PromiseType::get(getContext());
+  if (keyword == "async_handle") return AsyncHandleType::get(getContext());
+
+  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
+  return mlir::Type();
+}
+
+// Print a type registered to this dialect.
+void MlrtDialect::printType(mlir::Type type,
+                            mlir::DialectAsmPrinter &os) const {
+  if (mlir::isa<FutureType>(type)) {
+    os << "future";
+    return;
+  }
+
+  if (mlir::isa<PromiseType>(type)) {
+    os << "promise";
+    return;
+  }
+
+  if (mlir::isa<AsyncHandleType>(type)) {
+    os << "async_handle";
+    return;
+  }
+
+  llvm_unreachable("unexpected mlrt type kind");
+}
+
+}  // namespace compiler
+}  // namespace mlrt

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h

@@ -0,0 +1,62 @@
+/* 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_IR_MLRT_MLRT_DIALECT_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_DIALECT_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+
+namespace mlrt {
+namespace compiler {
+
+class MlrtDialect : public mlir::Dialect {
+ public:
+  explicit MlrtDialect(mlir::MLIRContext *context);
+  static llvm::StringRef getDialectNamespace() { return "mlrt"; }
+
+  mlir::Type parseType(mlir::DialectAsmParser &parser) const override;
+  void printType(mlir::Type type, mlir::DialectAsmPrinter &os) const override;
+};
+
+// The MLIR type represents a C++ mlrt::Future.
+class FutureType
+    : public mlir::Type::TypeBase<FutureType, mlir::Type, mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "mlrt.compiler.future";
+};
+
+// The MLIR type represents a C++ mlrt::Promise.
+class PromiseType
+    : public mlir::Type::TypeBase<PromiseType, mlir::Type, mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "mlrt.compiler.promise";
+};
+
+// The MLIR type represents a C++ mlrt::AsyncHandle.
+class AsyncHandleType : public mlir::Type::TypeBase<AsyncHandleType, mlir::Type,
+                                                    mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "mlrt.compiler.async_handle";
+};
+
+}  // namespace compiler
+}  // namespace mlrt
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_DIALECT_H_

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td

@@ -0,0 +1,55 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifdef MLRT_DIALECT
+#else
+#define MLRT_DIALECT
+
+include "mlir/IR/OpBase.td"
+
+def Mlrt_Dialect : Dialect {
+  let name = "mlrt";
+
+  let description = [{
+    The MLRT Dialect.
+  }];
+
+  let cppNamespace = "::mlrt::compiler";
+}
+
+def MlrtFutureType : DialectType<Mlrt_Dialect,
+    CPred<"::llvm::isa<::mlrt::compiler::FutureType>($_self)">, "!mlrt.future type">,
+    BuildableType<"$_builder.getType<::mlrt::compiler::FutureType>()"> {
+  let description = [{
+    `!mlrt.future type` represents a C++ mlrt::Future.
+  }];
+}
+
+def MlrtPromiseType : DialectType<Mlrt_Dialect,
+    CPred<"::llvm::isa<::mlrt::compiler::PromiseType>($_self)">, "!mlrt.promise type">,
+    BuildableType<"$_builder.getType<::mlrt::compiler::PromiseType>()"> {
+  let description = [{
+    `!mlrt.promise type` represents a C++ mlrt::Promise.
+  }];
+}
+
+def MlrtAsyncHandleType : DialectType<Mlrt_Dialect,
+    CPred<"::llvm::isa<::mlrt::compiler::AsyncHandleType>($_self)">, "!mlrt.async_handle type">,
+    BuildableType<"$_builder.getType<::mlrt::compiler::AsyncHandleType>()"> {
+  let description = [{
+    `!mlrt.async_handle type` represents a C++ mlrt::AsyncHandle.
+  }];
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.cc

@@ -0,0 +1,28 @@
+/* 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/ir/mlrt/mlrt_ops.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h

@@ -0,0 +1,27 @@
+/* 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_IR_MLRT_MLRT_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_OPS_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_MLRT_OPS_H_

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_ops.td

@@ -0,0 +1,240 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifdef MLRT_OPS
+#else
+#define MLRT_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+
+class Mlrt_Op<string mnemonic, list<Trait> traits = []> :
+    Op<Mlrt_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+def CondOp: Mlrt_Op<"cond", []> {
+  let summary = "mlrt.cond op";
+
+  let description = [{
+    Execute $a_true_fn with $args if $cond is true; otherwise, %b_false_fn is
+    executed.
+  }];
+
+  let arguments = (ins
+    I1:$cond,
+    Variadic<AnyType>:$args,
+    SymbolRefAttr:$a_true_fn,
+    SymbolRefAttr:$b_false_fn
+  );
+
+  let results = (outs
+    Variadic<AnyType>:$results
+  );
+
+  let assemblyFormat = [{
+    $cond $a_true_fn $b_false_fn `(` $args `)` attr-dict `:` `(` type($args) `)` `->` `(` type($results) `)`
+  }];
+}
+
+def AsyncOp: Mlrt_Op<"async", []> {
+  let summary = "Launches a function asynchronously.";
+
+  let description = [{
+    Launch a function asynchronously.
+
+    $args: a list of arguments to be passed.
+    $callee: The function to be launched. Its return op must not have operands.
+
+    $handle: This op returns a handle object that manages the context of the async execution.
+  }];
+
+  let arguments = (ins
+    Variadic<AnyType>:$args,
+    SymbolRefAttr:$callee
+  );
+
+  let results = (outs
+    MlrtAsyncHandleType:$handle
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $handle)";
+}
+
+def AwaitHandleOp: Mlrt_Op<"await_handle", []> {
+  let summary = "Awaits an async execution ";
+
+  let description = [{
+    Awaits an async execution.
+
+    $handle: The handle returned by mlrt.async op.
+  }];
+
+  let arguments = (ins
+    MlrtAsyncHandleType:$handle
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def AwaitAllHandleOp: Mlrt_Op<"await_all_handle", []> {
+  let summary = "Awaits multiple async executions";
+
+  let description = [{
+    Awaits multiple async execution.
+
+    $handles: A list of handles returned by mlrt.async ops.
+  }];
+
+  let arguments = (ins
+    Variadic<MlrtAsyncHandleType>:$handles
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($handles)";
+}
+
+def AwaitControlOp: Mlrt_Op<"await_control", []> {
+  let summary = "Await a signal from a future";
+
+  let description = [{
+    Await a signal, instead of a value, from a future.
+
+    $future: A value of !mlrt.future type.
+  }];
+
+  let arguments = (ins
+    MlrtFutureType:$future
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def AwaitAllControlOp: Mlrt_Op<"await_all_control", []> {
+  let summary = "Awaits multiple signals";
+
+  let description = [{
+    Awaits multiple signals
+
+    $futures: A list of !mlrt.futures
+  }];
+
+  let arguments = (ins
+    Variadic<MlrtFutureType>:$futures
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($futures)";
+}
+
+def PromiseControlOp: Mlrt_Op<"promise_control", []> {
+  let summary = "Set a control promise";
+
+  let description = [{
+    Set a control promise.
+
+    $promise: A value of !mlrt.promise type.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def CaseOp : Mlrt_Op<"case"> {
+  let summary = "An n-way switch statement which calls a single branch function.";
+  let description = [{
+    An n-way switch statement, implementing the following:
+        ```
+        switch (branch_index) {
+          case 0:
+            outputs = branches[0](inputs);
+            break;
+          case 1:
+            outputs = branches[1](inputs);
+            break;
+          ...
+          case [[nbranches-1]]:
+          default:
+            outputs = branches[nbranches-1](inputs);
+            break;
+        }
+        ```
+    Example: %res = mlrt.case %branch_idx [@branch0, @branch1] (%arg0, %arg1) : (i32, i32) -> (i32)
+  }];
+
+  let arguments = (ins I32:$branch_index,
+                       ConfinedAttr<SymbolRefArrayAttr, [ArrayMinCount<1>]>:$branches,
+                       Variadic<AnyType>:$branch_operands);
+
+  let results = (outs Variadic<AnyType>:$branch_outputs);
+  let assemblyFormat = [{
+    $branch_index $branches `(` $branch_operands `)` attr-dict `:` `(` type($branch_operands) `)` `->` `(` type($branch_outputs) `)`
+  }];
+}
+
+def AllocateControlFuturesOp: Mlrt_Op<"allocate_control_futures", [AttrSizedResultSegments]> {
+  let summary = "Allocate futures and corresponding promises";
+
+  let description = [{
+    Allocate futures and corresponding promises.
+
+    $num: The number of futures to be allocated.
+
+    $promises: There are $num promises, and promises[i] shares the state with futures[i].
+    $futures: There are $num futures, and futures[i] shares the state with promises[i].
+  }];
+
+  let arguments = (ins
+    I32Attr:$num
+  );
+
+  let results = (outs
+    Variadic<MlrtPromiseType>:$promises,
+    Variadic<MlrtFutureType>:$futures
+  );
+}
+
+def WhileOp : Mlrt_Op<"while", []> {
+  let summary = "while operation";
+  let description = [{
+    cond: The boolean to control whether the first iteration should be
+      executed.
+    operands: The arguments to the first iteration.
+    results: The results of the last iteration. The number and types of results
+      excluding the last one are the same as the number and types of operands. The
+      last element of results is an I1 value that is false.
+    body_fn: The body function that takes the arguments and returns the results
+      that includes an I1 value to indicate whether next iteration should be executed.
+
+    The pseudo code:
+
+    while(cond) {
+      results = body_fn(operands)
+      cond = results#1
+    }
+    return results
+
+  }];
+
+  let arguments = (ins I1:$cond,
+                       Variadic<AnyType>:$arguments,
+                       FlatSymbolRefAttr:$body_fn);
+
+  let results = (outs Variadic<AnyType>);
+
+  let assemblyFormat = [{
+    $cond $body_fn `(` $arguments `)` attr-dict `:` `(` type($arguments) `)` `->` `(` type(results) `)`
+  }];
+}
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td

@@ -0,0 +1,56 @@
+/* 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.
+==============================================================================*/
+#ifdef TF_MLRT_DIALECT
+#else
+#define TF_MLRT_DIALECT
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+
+// TODO(chky,rohitju): Unify this dialect with tfrt_fallback_sync dialect after
+// vrooml is using the new interpreter.
+def TensorflowMlrt_Dialect : Dialect {
+  let name = "tf_mlrt";
+
+  let description = [{
+    The TF MLRT Dialect.
+  }];
+
+  let cppNamespace = "::tensorflow::tf_mlrt";
+}
+
+class TensorflowMlrt_Op<string mnemonic, list<Trait> traits = []> :
+    Op<TensorflowMlrt_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+// This corresponds to tensorflow::Tensor.
+def TFTensorType : DialectType<TensorflowMlrt_Dialect,
+    CPred<"::llvm::isa<::tensorflow::tf_mlrt::TFTensorType>($_self)">, "!tf_mlrt.tensor type">,
+    BuildableType<"$_builder.getType<::tensorflow::tf_mlrt::TFTensorType>()"> {
+  let description = [{
+    `!tf_mlrt.tensor type` represents a tensorflow::Tensor.
+  }];
+}
+
+// This corresponds to tensorflow::Device* .
+def TFDeviceType : DialectType<TensorflowMlrt_Dialect,
+    CPred<"::llvm::isa<::tensorflow::tf_mlrt::TFDeviceType>($_self)">, "!tf_mlrt.device type">,
+    BuildableType<"$_builder.getType<::tensorflow::tf_mlrt::TFDeviceType>()"> {
+  let description = [{
+    `!tf_mlrt.device type` represents a tensorflow::device.
+  }];
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.cc

@@ -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.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.h"
+
+namespace tensorflow {
+namespace tf_mlrt {
+
+namespace {
+
+struct TensorflowMlrtInlinerInterface : public mlir::DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+  bool isLegalToInline(mlir::Operation *op, mlir::Region *dest,
+                       bool would_be_cloned,
+                       mlir::IRMapping &mapping) const final {
+    // All tf_mlrt dialect ops can be inlined.
+    return true;
+  }
+  // Note that CallOp and ReturnOp are handled by func; so need to implement
+  // handleTerminator.
+};
+
+}  // namespace
+
+TensorflowMlrtDialect::TensorflowMlrtDialect(mlir::MLIRContext *context)
+    : mlir::Dialect(/*name=*/"tf_mlrt", context,
+                    mlir::TypeID::get<TensorflowMlrtDialect>()) {
+  addTypes<TFTensorType, TFDeviceType>();
+  addInterfaces<TensorflowMlrtInlinerInterface>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.cpp.inc"
+      >();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_ops.cpp.inc"
+      >();
+}
+
+// Parse a type registered to this dialect.
+mlir::Type TensorflowMlrtDialect::parseType(
+    mlir::DialectAsmParser &parser) const {
+  llvm::StringRef keyword;
+  if (parser.parseKeyword(&keyword)) return mlir::Type();
+
+  if (keyword == "tensor") return TFTensorType::get(getContext());
+
+  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
+  return mlir::Type();
+}
+
+// Print a type registered to this dialect.
+void TensorflowMlrtDialect::printType(mlir::Type type,
+                                      mlir::DialectAsmPrinter &os) const {
+  if (mlir::isa<TFTensorType>(type)) {
+    os << "tensor";
+    return;
+  }
+
+  llvm_unreachable("unexpected tf_mlrt type kind");
+}
+
+}  // namespace tf_mlrt
+}  // namespace tensorflow
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.cpp.inc"
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_ops.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h

@@ -0,0 +1,63 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_OPS_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_side_effects.h"
+#include "tfrt/compiler/opdefs/tfrt_op_interfaces.h"  // from @tf_runtime
+#include "tfrt/compiler/opdefs/tfrt_traits.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tf_mlrt {
+
+class TensorflowMlrtDialect : public mlir::Dialect {
+ public:
+  explicit TensorflowMlrtDialect(mlir::MLIRContext *context);
+  static llvm::StringRef getDialectNamespace() { return "tf_mlrt"; }
+
+  mlir::Type parseType(mlir::DialectAsmParser &parser) const override;
+  void printType(mlir::Type type, mlir::DialectAsmPrinter &os) const override;
+};
+
+// The MLIR type represents a tensorflow::Tensor.
+class TFTensorType
+    : public mlir::Type::TypeBase<TFTensorType, mlir::Type, mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "tensorflow.tf_mlrt.tf_tensor";
+};
+
+// The MLIR type represents a tensorflow::Device*
+class TFDeviceType
+    : public mlir::Type::TypeBase<TFDeviceType, mlir::Type, mlir::TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr mlir::StringLiteral name = "tensorflow.tf_mlirt.tf_device";
+};
+
+}  // namespace tf_mlrt
+}  // namespace tensorflow
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.h.inc"
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_OPS_H_

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_ops.td

@@ -0,0 +1,577 @@
+/* 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.
+==============================================================================*/
+#ifdef TF_MLRT_OPS
+#else
+#define TF_MLRT_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+
+def CreateOp: TensorflowMlrt_Op<"createop", []> {
+  let summary = "The Fallback CreateOp";
+
+  let description = [{
+    The CreateOp creates the tensorflow::OpKernel in the fallback context.
+  }];
+
+  let arguments = (ins
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
+def ConstOp: TensorflowMlrt_Op<"constop", []> {
+  let summary = "The tf_mlrt ConstOp";
+
+  let description = [{
+    The ConstOp creates a constant tensorflow::Tensor from serialized proto.
+  }];
+
+  let arguments = (ins
+    StrAttr:$tensor_proto
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
+def ExecuteOp : TensorflowMlrt_Op<"executeop", []> {
+  let summary = "The Fallback ExecuteOp";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def ExecuteOpWithDevice: TensorflowMlrt_Op<"executeop.device", []> {
+  let summary = "The Fallback ExecuteOp with custom device";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device using a custom device.
+  }];
+
+  let arguments = (ins
+    TFDeviceType:$device,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "`(` $device`)` `(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def AsyncExecuteOp : TensorflowMlrt_Op<"async_executeop", []> {
+  let summary = "The Fallback ExecuteOp for tensorflow::AsyncOpKernel";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device asynchronously.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def AsyncExecuteOpWithDevice : TensorflowMlrt_Op<"async_executeop.device", []> {
+  let summary = "The Fallback ExecuteOp for tensorflow::AsyncOpKernel";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device asynchronously.
+  }];
+
+  let arguments = (ins
+    TFDeviceType:$device,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = "`(` $device`)` `(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def SetResourceOp : TensorflowMlrt_Op<"set_resource", []> {
+  let summary = "Set a tensor in resource array";
+
+  let description = [{
+    Set a tensor in resource array.
+
+    arg: the tensor to be set in the resource array.
+    index: the index in the resource array
+  }];
+
+  let arguments = (ins
+    TFTensorType:$arg,
+    I64Attr:$index
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def GetResourceOp : TensorflowMlrt_Op<"get_resource", []> {
+  let summary = "get a tensor in resource array";
+
+  let description = [{
+    Get a tensor in resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "attr-dict `:` type($results)";
+}
+
+def AwaitOp: TensorflowMlrt_Op<"await", [Pure]> {
+  let summary = "Await a tensor from a !mlrt.future";
+
+  let description = [{
+    Await a tensor from a !mlrt.future.
+
+    $future: A value of type !mlrt.future. The underlying value must be a tensorflow tensor.
+
+    $result: a tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtFutureType:$future
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def AwaitAllOp: TensorflowMlrt_Op<"await_all", [Pure]> {
+  let summary = "Await tensors from a list of !mlrt.future";
+
+  let description = [{
+    Await tensors from a list of !mlrt.future.
+
+    $futures: A list of !mlrt.future. The underlying value must be tensorflow tensors.
+
+    $results: A list of tensorflow tensors.
+  }];
+
+  let arguments = (ins
+    Variadic<MlrtFutureType>:$futures
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "$futures attr-dict `:` type($results)";
+}
+
+def PromiseOp: TensorflowMlrt_Op<"promise", []> {
+  let summary = "Set a tensor in a promise";
+
+  let description = [{
+    Set a tensor in a promise.
+
+    $promise: A value of type !mlrt.promise. The underlying value must be a tensorflow tensor.
+    $tensor: A tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise,
+    TFTensorType:$tensor
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def PromiseFutureOp: TensorflowMlrt_Op<"promise_future", []> {
+  let summary = "Set a tensor future in a promise";
+
+  let description = [{
+    Set a tensor future in a promise.
+
+    $promise: A value of type !mlrt.promise. The underlying value must be a tensorflow tensor.
+    $future: A value of type !mlrt.future. Must represents a tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise,
+    MlrtFutureType:$tensor
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def PromiseReturnOp: TensorflowMlrt_Op<"promise_return", [Terminator]> {
+  let summary = "Set a tensor in a promise and return from a function with no output";
+
+  let description = [{
+    Set a tensor in a promise and return from a function with no output.
+
+    $promise: A value of type !mlrt.promise. The underlying value must be a tensorflow tensor.
+    $tensor: A tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise,
+    TFTensorType:$tensor
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def AllocateFuturesOp: TensorflowMlrt_Op<"allocate_futures", [AttrSizedResultSegments]> {
+  let summary = "Allocate futures and promsies for tensorflow tensors";
+
+  let description = [{
+    Allocate futures and promsies for tensorflow tensors.
+
+    $num_futures: The number of futures to be allocated.
+
+    $promises: There are $num_futures promises. promises[i] shares the state with futures[i].
+    $futures: There are $num_futures futures. futures[i] shares the state with promises[i].
+  }];
+
+  let arguments = (ins
+    I32Attr:$num_futures
+  );
+
+  let results = (outs
+    Variadic<MlrtPromiseType>:$promises,
+    Variadic<MlrtFutureType>:$futures
+  );
+}
+
+def TensorToIntOp : TensorflowMlrt_Op<"tensor_to_int32", [Pure]> {
+  let summary = "Cast a Tensor to int32.";
+  let description = [{
+    Cast a Tensor to int32.
+
+    Example:
+      %one = tf_mlrt.tensor_to_int32 %src_tenosr
+  }];
+
+  let arguments = (ins TFTensorType:$src);
+  let results = (outs I32:$result);
+  let assemblyFormat = "operands attr-dict";
+}
+
+def BoolToTensorOp : TensorflowMlrt_Op<"bool_to_tensor", [Pure]> {
+  let summary = "Cast a boolean to a tensor.";
+  let description = [{
+    Cast a boolean to a tensor.
+
+    Example:
+      %one = tf_mlrt.bool_to_tensor %src_bool
+  }];
+
+  let arguments = (ins I1: $src);
+  let results = (outs TFTensorType:$result);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def PredicateOp : TensorflowMlrt_Op<"predicate", [Pure]> {
+  let summary = "Converts a fallback tensor to a bool";
+
+  let description = [{
+    Note: this kernel is used for CPU tensors.
+
+    Converts a fallback tensor to a bool with the following rules:
+
+    - For 0D tensors, truthiness is determined by comparing against a "zero"
+      value. For numerical types it is the obvious zero. For strings it is the
+      empty string.
+
+    - For >0D tensors, truthiness is determined by looking at the number of
+      elements. If has zero elements, then the result is false. Otherwise the
+      result is true.
+
+    input: a fallback tensor representing the condition.
+    device: the name of the tensorflow device that is associated with the
+      input fallback tensor.
+
+    output: the converted bool.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$input
+  );
+
+  let results = (outs
+    I1:$output
+  );
+
+  let assemblyFormat = "$input attr-dict";
+}
+
+def BatchFunctionOp : TensorflowMlrt_Op<"batch_function", [Pure]> {
+  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";
+
+  let description = [{
+    This kernel executes a variant tf.BatchFunction kernel that supports having
+    the `f` attribute as a bytecode function.
+
+    Example:
+      %res = tf_mlrt.batch_function(%input, %captured_input)  {
+          device = "/device:CPU:0",
+          f = @batch_function,
+          node_def = "..."
+        } : (!tf_mlrt.tensor,!tf_mlrt.tensor) -> (!tf_mlrt.tensor)
+
+    Note that the trailing number indicates the number of results.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    SymbolRefAttr:$f,
+    StrAttr:$node_def
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+def BatchFunctionWithDeviceOp : TensorflowMlrt_Op<"batch_function.device", [Pure]> {
+  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction with custom device.";
+
+  let description = [{
+    This kernel executes a variant tf.BatchFunction kernel that supports having
+    the `f` attribute as a bytecode function with a custom device.
+
+    Example:
+      %res = tf_mlrt.batch_function.device(%custom_device, %input, %captured_input)  {
+          device = "/device:CPU:0",
+          f = @batch_function,
+          node_def = "..."
+        } : (!tf_mlrt.tensor,!tf_mlrt.tensor) -> (!tf_mlrt.tensor)
+
+    Note that the trailing number indicates the number of results.
+    Also note that the custom device argument is not the same as the device attribute. 
+    The latter is just the name in which the enclosure graph is executed.
+  }];
+
+  let arguments = (ins
+    TFDeviceType:$custom_device,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    SymbolRefAttr:$f,
+    StrAttr:$node_def
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = " `(` $custom_device`)` `(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+
+def CancelOp: TensorflowMlrt_Op<"cancel", []> {
+  let summary = "Handle cancellation request.";
+
+  let description = [{
+    This kernel will early terminate the program upon cancellation request (e.g. time out).
+  }];
+}
+
+def MapFnOp : TensorflowMlrt_Op<"map_fn", [AttrSizedOperandSegments, Pure]> {
+  let summary = "The Parallel Map for tf_mlrt dialect";
+  let description = [{
+    The Pmap executes body function in parallel for all ranges up to $max_iterations.
+
+    The pseudo code:
+      for(int i = 0; i < $max_iterations; i++) {
+        body_fn(MlrtFture($tensor_list_or_flow_in[i]),
+                MlrtPromise($tensor_list_or_flow_in[i+1]),
+                i, i, $invariant_args);
+      }
+
+      return $tensor_list_or_flow_in[$max_iterations]
+  }];
+
+  let arguments = (ins
+    TFTensorType:$max_iterations,
+    Variadic<TFTensorType>:$tensor_list_or_flow_in,
+    Variadic<TFTensorType>:$invariant_args,
+    FlatSymbolRefAttr:$body_fn,
+    I32Attr:$num_tensor_list_or_flow_in
+  );
+
+  let results = (outs
+     Variadic<TFTensorType>:$result
+  );
+
+  let assemblyFormat = "`(`$max_iterations`,` $tensor_list_or_flow_in`,` $invariant_args `)` attr-dict `:` functional-type(operands, results)";
+}
+
+def AsyncWhileOp : TensorflowMlrt_Op<"async_while", [Pure]> {
+  let summary = "Asynchronously execution of while op for tf_mlrt";
+  let description = [{
+    cond: The boolean to control whether the first iteration should be executed.
+    arguments: the last $immutable_size elements are invariants between iterations.
+    results: a list of futures.
+    body_fn: its input are [predicate_promise, arg0_future, arg0_promise, arg1_future, arg1_promise, ..., invariant_args] and has no returned results.
+  }];
+
+  let arguments = (ins TFTensorType:$cond,
+                       Variadic<TFTensorType>:$arguments,
+                       FlatSymbolRefAttr:$body_fn,
+                       I32Attr:$invariant_size);
+
+  let results = (outs Variadic<MlrtFutureType>);
+
+  let assemblyFormat = [{
+    $body_fn `(` $cond `,` $arguments `)` attr-dict `:` `(` type($cond) `,` type($arguments) `)` `->` `(` type(results) `)`
+  }];
+}
+
+def IfrtLoadVariableOp: TensorflowMlrt_Op<"ifrt_load_variable", [Pure]> {
+  let summary = "Loads a variable tensor as an IFRT array for mlrt";
+
+  let description = [{
+    This op loads a restored variable tensor as a tensor future. It is a
+    replacement of `tf.ReadVariableOp`.
+
+    This op returns a scalar string tensor containing the restored variable name, which can be
+    used as a key within the runtime, as well as a future for the tensor.
+
+    The `tf.IfrtCall` kernel uses the output $array_key.
+    Other ops executed by TFRT may make use of $tensor_future.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$variable,
+    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
+  );
+
+  let results = (outs
+    TFTensorType:$array_key,
+    MlrtFutureType: $tensor_future
+  );
+}
+
+def IfrtRestoreVariableOp: TensorflowMlrt_Op<"ifrt_restore_variable", []> {
+  let summary = "Restore variable tensors";
+  let description = [{
+    This is the MLRT version of tf.IfrtRestoreVariableOp.
+
+    This op restores tensors from a checkpoint asynchronously, similar to tf.RestoreV2.
+    It fuses variable assignment: tensors that are consumed by tf.AssignVariableOp
+    are handled internally via the $var_handles operands and do not appear in
+    the results.
+
+    Tensors that are consumed by other operations (e.g. tf.LookupTableImportV2)
+    will be returned as results in $tensors. Consequently, the number of
+    results is dynamic, containing only those tensors not assigned to variables
+    but consumed directly by other operations. A restored tensor with no
+    consumer will cause an error in the kernel.
+    This Op is specific to MLRT runtime and is not a stable interface for
+    serialization.
+
+  }];
+
+  let arguments = (ins
+    TFTensorType:$prefix,
+    TFTensorType:$tensor_names,
+    TFTensorType:$shape_and_slices,
+    Variadic<TFTensorType>:$var_handles,
+    TypeArrayAttr: $restored_dtypes,
+    DenseBoolArrayAttr: $truncate_in_cast,
+     // The names of the returned tensors. The order of names corresponds to the
+    // order of the returned tensors.
+    StrArrayAttr: $returned_tensor_names
+  );
+
+  let results = (outs
+    Res<Variadic<MlrtFutureType>, [{The restored tensor futures as requested by the $returned_Tensor_names attribute.}]>:$tensors
+  );
+}
+
+def MlrtIfrtResourceDeserializeOp: TensorflowMlrt_Op<"ifrt_resource_deserialize", []> {
+  let summary = "Deserialize resource vars.";
+  let description = [{
+    This is the MLRT version of the IfrtResourceDeserialize op.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$resource_var,
+    TFTensorType:$input_dir,
+    StrAttr:$tensor_name
+  );
+}
+
+def AsyncIfrtCallOp : TensorflowMlrt_Op<"async_ifrt_call", []> {
+  let summary = "IFRT execution op for mlrt";
+  let description = [{
+    This op executes an IFRT program asynchronously.
+  }];
+
+  let arguments = (ins
+    I64Attr:$program_id,
+    I32ArrayAttr:$variable_arg_indices,
+    Variadic<TFTensorType>:$args
+  );
+
+  let results = (outs
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = "`(` $args `)` attr-dict `:` functional-type($args, $results)";
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.cc

@@ -0,0 +1,46 @@
+/* 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/ir/mlrt/tf_mlrt_tpu_ops.h"
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#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 tf_mlrt_tpu {
+
+TensorflowMlrtTpuDialect::TensorflowMlrtTpuDialect(mlir::MLIRContext *context)
+    : mlir::Dialect(/*name=*/"tf_mlrt_tpu", context,
+                    mlir::TypeID::get<TensorflowMlrtTpuDialect>()) {
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.cpp.inc"
+      >();
+}
+
+}  // namespace tf_mlrt_tpu
+}  // namespace tensorflow
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.h

@@ -0,0 +1,39 @@
+/* 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_IR_MLRT_TF_MLRT_TPU_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_TPU_OPS_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace tf_mlrt_tpu {
+
+class TensorflowMlrtTpuDialect : public mlir::Dialect {
+ public:
+  explicit TensorflowMlrtTpuDialect(mlir::MLIRContext *context);
+  static llvm::StringRef getDialectNamespace() { return "tf_mlrt_tpu"; }
+};
+
+}  // namespace tf_mlrt_tpu
+}  // namespace tensorflow
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_MLRT_TF_MLRT_TPU_OPS_H_

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td

@@ -0,0 +1,82 @@
+/* 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.
+==============================================================================*/
+#ifdef TF_MLRT_TPU_OPS
+#else
+#define TF_MLRT_TPU_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+
+def TensorflowMlrtTpu_Dialect : Dialect {
+  let name = "tf_mlrt_tpu";
+
+  let description = [{
+    The TF MLRT TPU Dialect.
+  }];
+
+  let cppNamespace = "::tensorflow::tf_mlrt_tpu";
+}
+
+class TensorflowMlrtTpu_Op<string mnemonic, list<Trait> traits = []> :
+    Op<TensorflowMlrtTpu_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+def GetTpuHostDeviceOp : TensorflowMlrtTpu_Op<"get_tpu_host_device", [Pure]> {
+  let summary = "get the tpu host allocator that implements tensorflow::Device";
+
+  let results = (outs
+    TFDeviceType:$device
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
+def CompileAndExecuteOp : TensorflowMlrtTpu_Op<"compile_and_execute"> {
+  let summary = "tpu compile and execute operation";
+  let description = [{
+    tf_mlrt_tpu.compile_and_execute compiles a mlir tpu program and executes the compiled tpu program.
+
+    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.
+    $metadata is a serialized TPUCompileMetadataProto describing the shapes and types of the inputs to the computation, as well as a mapping onto the TPU pod topology.
+    $constant_operand_indices are the indices of the inputs that are constant to the TPU program (e.g. weights in inference), the rest of the inputs are input tensors.
+    constant_operand_indices is sorted in ascending order.
+    $operands_with_static_shape are indices of operands that are tagged with a maximum static shape.
+    $producer_name is a string describing the name of the framework that added support for running this portion of the model on TPUs.
+
+    Example:
+      %rendezvous_key_base, %result = tf_mlrt_tpu.compile_and_execute (%operands) constant_operand_indices = [1, 3] metadata = "metadata..." mlir_module = "mlir_module..."
+  }];
+  let arguments = (ins
+    Variadic<TFTensorType>:$operands_and_static_shapes,
+    DenseI32ArrayAttr:$constant_operand_indices,
+    StrAttr:$metadata,
+    StrAttr:$mlir_module,
+    UI32Attr:$num_operands,
+    DenseI32ArrayAttr:$operands_with_static_shape,
+    StrAttr:$producer_name
+  );
+
+  let results = (outs
+    TFTensorType:$rendezvous_key_base,
+    Variadic<MlrtFutureType>:$results
+  );
+
+  let assemblyFormat = [{
+    `(` $operands_and_static_shapes `)` attr-dict `:` functional-type($operands_and_static_shapes, results)
+  }];
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_ops.td

@@ -0,0 +1,182 @@
+/* Copyright 2015 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.
+==============================================================================*/
+#ifdef MLRT_TF_OPS
+#else
+#define MLRT_TF_OPS
+
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td"
+include "tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
+include "tfrt/compiler/opdefs/tfrt_op_interfaces.td"  // from @tf_runtime
+include "tfrt/compiler/opdefs/tfrt_traits.td"  // from @tf_runtime
+
+// tf_mlrt.tf_await returns a tensorflow Tensor. It is a fake op that is only
+// used during parallelization and has no runtime implementation.
+def TFAwaitOp: TensorflowMlrt_Op<"tf_await", [Pure, TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Await a tensor from a !mlrt.future";
+
+  let description = [{
+    Await a tensor from a !mlrt.future.
+
+    $future: A value of type !mlrt.future. The underlying value must be a tensorflow tensor.
+
+    $result: a tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtFutureType:$future
+  );
+
+  let results = (outs
+    TF_Tensor:$result
+  );
+}
+
+// tf_mlrt.tf_promise takes a tensorflow Tensor. It is a fake op that is only
+// used during parallelization and has no runtime implementation.
+def TFPromiseOp: TensorflowMlrt_Op<"tf_promise", [TF_MustExecute, TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Set a tensor in a promise";
+
+  let description = [{
+    Set a tensor in a promise.
+
+    $promise: A value of type !mlrt.promise. The underlying value will always be a tensorflow tensor.
+    $tensor: A tensorflow tensor.
+  }];
+
+  let arguments = (ins
+    MlrtPromiseType:$promise,
+    TF_Tensor:$tensor
+  );
+}
+
+def TFMapFnOp : TensorflowMlrt_Op<"tf_map_fn", [AttrSizedOperandSegments, Pure]> {
+  let summary = "The Parallel Map for tf_mlrt dialect";
+  let description = [{
+    The Pmap executes body function in parallel for all ranges up to $max_iterations.
+
+    The pseudo code:
+      for(int i = 0; i < $max_iterations; i++) {
+        body_fn(MlrtFture($tensor_list_or_flow_in[i]),
+                MlrtPromise($tensor_list_or_flow_in[i+1]),
+                i, i, $invariant_args);
+      }
+
+      return $tensor_list_or_flow_in[$max_iterations]
+  }];
+
+  let arguments = (ins
+    TF_Tensor:$max_iterations,
+    Variadic<TF_Tensor>:$tensor_list_or_flow_in,
+    Variadic<TF_Tensor>:$invariant_args,
+    FlatSymbolRefAttr:$body_fn,
+    I32Attr:$num_tensor_list_or_flow_in
+  );
+
+  let results = (outs
+     Variadic<TF_Tensor>:$result
+  );
+
+  let assemblyFormat = "`(`$max_iterations`,` $tensor_list_or_flow_in`,` $invariant_args `)` attr-dict `:` functional-type(operands, results)";
+}
+
+def TFAsyncWhileOp : TensorflowMlrt_Op<"tf_async_while", [Pure]> {
+  let summary = "Asynchronously execution of while op for tf_mlrt";
+  let description = [{
+    cond: The boolean to control whether the first iteration should be executed.
+    arguments: The last $invariant_size elements are invariants between iterations.
+    results: a list of futures.
+    body_fn: input is [predicate_promise, arg0_future, arg0_promise, arg1_future, arg1_promise, ..., invariant_args] and has no returned results.
+  }];
+
+  let arguments = (ins TF_Tensor:$cond,
+                       Variadic<TF_Tensor>:$arguments,
+                       FlatSymbolRefAttr:$body_fn,
+                       I32Attr:$invariant_size);
+
+  let results = (outs Variadic<MlrtFutureType>);
+
+  let assemblyFormat = [{
+    $body_fn `(` $cond `,` $arguments `)` attr-dict `:` `(` type($cond) `,` type($arguments) `)` `->` `(` type(results) `)`
+  }];
+}
+
+def TFTPUCompileAndExecuteOp : TensorflowMlrt_Op<"tf_tpu_compile_and_execute", [TF_MustExecute]> {
+  let summary = "tpu compile and execute operation";
+  let description = [{
+    tf_mlrt_tpu.compile_and_execute compiles a mlir tpu program and executes the compiled tpu program.
+
+    $mlir_module is a serialized MLIR module with a `main` function that contains target computation.
+    $metadata is a serialized TPUCompileMetadataProto describing the shapes and types of the inputs to the computation, as well as a mapping onto the TPU pod topology.
+    $constant_operand_indices are the indices of the inputs that are constant to the TPU program (e.g. weights in inference), the rest of the inputs are input tensors.
+    constant_operand_indices is sorted in ascending order.
+    $operands_with_static_shape are indices of operands that are tagged with a maximum static shape.
+    $producer_name is a string describing the name of the framework that added support for running this portion of the model on TPUs.
+
+    Example:
+      %rendezvous_key_base, %result = tf_mlrt_tpu.compile_and_execute (%operands) constant_operand_indices = [1, 3] metadata = "metadata..." mlir_module = "mlir_module..."
+  }];
+  let arguments = (ins
+    Variadic<TF_Tensor>:$operands_and_static_shapes,
+    DenseI32ArrayAttr:$constant_operand_indices,
+    StrAttr:$metadata,
+    StrAttr:$mlir_module,
+    UI32Attr:$num_operands,
+    DenseI32ArrayAttr:$operands_with_static_shape,
+    StrAttr:$producer_name
+  );
+
+  let results = (outs
+    TF_Tensor:$rendezvous_key_base,
+    Variadic<TF_Tensor>:$results
+  );
+
+  let assemblyFormat = [{
+    `(` $operands_and_static_shapes `)` attr-dict `:` functional-type($operands_and_static_shapes, results)
+  }];
+}
+
+def TFIfrtLoadVariableOp: TensorflowMlrt_Op<"tf_ifrt_load_variable", [Pure]> {
+  let summary = "Loads a variable tensor as an IFRT array for mlrt";
+
+  let description = [{
+    This is the MLRT version of tf.IfrtLoadVariableOp.
+
+    This op loads a variable tensor as an IFRT array and binds it with the specified name.
+
+    This op is an replacement of `tf.ReadVariableOp` in the case that a constant
+    variable tensor is an input to the tpu program invoked by `tf.IfrtCall`.
+
+    After a `tf.ReadVariableOp` is lowered into `tf.IfrtLoadVariableOp`, the `tf.IfrtCall` kernel
+    will bind the loaded IFRT array by name with the tpu program's input.
+
+    This op returns a scalar string tensor as a key for user to look for the loaded array
+    and a future containing the restored tensor.
+  }];
+
+  let arguments = (ins
+    TF_Tensor:$variable,
+    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
+  );
+
+  let results = (outs
+    TF_Tensor:$array_key,
+    MlrtFutureType: $tensor_future
+  );
+}
+
+
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.cc

@@ -0,0 +1,72 @@
+/* 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/ir/tfrt_fallback.h"
+
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace tfrt {
+namespace fallback {
+
+FallbackDialect::FallbackDialect(MLIRContext *context)
+    : Dialect(/*name=*/"tfrt_fallback", context,
+              TypeID::get<FallbackDialect>()) {
+  addTypes<TFTensorType, TFAllocatorType>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.cpp.inc"
+      >();
+}
+
+/// Parse a type registered to this dialect.
+Type FallbackDialect::parseType(DialectAsmParser &parser) const {
+  StringRef keyword;
+  if (parser.parseKeyword(&keyword)) return Type();
+
+  if (keyword == "tf_tensor") return TFTensorType::get(getContext());
+  if (keyword == "tf_allocator") return TFAllocatorType::get(getContext());
+
+  parser.emitError(parser.getNameLoc(), "unknown type: ") << keyword;
+  return Type();
+}
+
+/// Print a type registered to this dialect.
+void FallbackDialect::printType(Type type, DialectAsmPrinter &os) const {
+  if (mlir::isa<TFTensorType>(type)) {
+    os << "tf_tensor";
+    return;
+  }
+
+  if (mlir::isa<TFAllocatorType>(type)) {
+    os << "tf_allocator";
+    return;
+  }
+
+  llvm_unreachable("unexpected fallback type kind");
+}
+
+}  // namespace fallback
+}  // namespace tfrt
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.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_IR_TFRT_FALLBACK_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+
+using namespace mlir;  // NOLINT
+
+namespace tfrt {
+namespace fallback {
+
+// Dialect for fallback operations.
+class FallbackDialect : public Dialect {
+ public:
+  explicit FallbackDialect(MLIRContext *context);
+  static StringRef getDialectNamespace() { return "tfrt_fallback"; }
+
+  Type parseType(DialectAsmParser &parser) const override;
+  void printType(Type type, DialectAsmPrinter &os) const override;
+};
+
+// The MLIR type represents a tensorflow::Tensor.
+class TFTensorType : public Type::TypeBase<TFTensorType, Type, TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr StringLiteral name = "tfrt.tf_tensor";
+};
+
+// The MLIR type represents a tensorflow::Allocator.
+class TFAllocatorType
+    : public Type::TypeBase<TFAllocatorType, Type, TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr StringLiteral name = "tfrt.tf_allocator";
+};
+
+}  // namespace fallback
+}  // namespace tfrt
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_H_

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.td

@@ -0,0 +1,50 @@
+/* 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.
+==============================================================================*/
+
+#ifdef TFRT_FALLBACK_OPS
+#else
+#define TFRT_FALLBACK_OPS
+
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+def Fallback_Dialect : Dialect {
+  let name = "tfrt_fallback";
+
+  let description = [{
+    The TFRT Fallback Dialect.
+  }];
+
+  let cppNamespace = "::tfrt::fallback";
+}
+
+// This corresponds to tensorflow::Tensor.
+def TFTensorType : DialectType<Fallback_Dialect,
+    CPred<"::llvm::isa<::tfrt::fallback::TFTensorType>($_self)">, "!tfrt_fallback.tf_tensor type">,
+    BuildableType<"$_builder.getType<::tfrt::fallback::TFTensorType>()"> {
+  let description = [{
+    `!tfrt_fallback.tf_tensor type` represents a tensorflow::Tensor.
+  }];
+}
+
+// This corresponds to tensorflow::Allocator.
+def TFAllocatorType : DialectType<Fallback_Dialect,
+    CPred<"::llvm::isa<::tfrt::fallback::TFAllocatorType>($_self)">, "!tfrt_fallback.tf_allocator type">,
+    BuildableType<"$_builder.getType<::tfrt::fallback::TFAllocatorType>()"> {
+  let description = [{
+    `!tfrt_fallback.tf_tensor type` represents a tensorflow::Tensor.
+  }];
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.cc

@@ -0,0 +1,407 @@
+/* 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/ir/tfrt_fallback_async.h"
+
+#include <utility>
+
+#include "llvm/ADT/STLExtras.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
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/Matchers.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Transforms/InliningUtils.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.h"
+#include "tfrt/basic_kernels/opdefs/tfrt_base.h"  // from @tf_runtime
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/attributes.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/core_runtime.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/sync/core_runtime.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/types.h"  // from @tf_runtime
+
+namespace tfrt {
+namespace fallback_async {
+
+namespace {
+
+struct FallbackInlinerInterface : public mlir::DialectInlinerInterface {
+  using DialectInlinerInterface::DialectInlinerInterface;
+
+  bool isLegalToInline(Operation *op, Region *dest, bool would_be_cloned,
+                       IRMapping &) const final {
+    return true;
+  }
+};
+
+}  // namespace
+
+FallbackAsyncDialect::FallbackAsyncDialect(MLIRContext *context)
+    : Dialect(/*name=*/"tfrt_fallback_async", context,
+              TypeID::get<FallbackAsyncDialect>()) {
+  context->getOrLoadDialect<tfrt::fallback::FallbackDialect>();
+  context->getOrLoadDialect<compiler::TFRTDialect>();
+  context->getOrLoadDialect<corert::CoreRTDialect>();
+
+  allowUnknownTypes();
+
+  allowUnknownOperations();
+
+  addInterfaces<FallbackInlinerInterface>();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.cpp.inc"
+      >();
+}
+
+static Type GetChainType(Builder *builder) {
+  return builder->getType<compiler::ChainType>();
+}
+
+LogicalResult CreateOp::verify() {
+  return fallback_common::VerifyFallbackExecuteOp(*this);
+}
+LogicalResult ExecuteOp::verify() {
+  return fallback_common::VerifyFallbackExecuteOp(*this);
+}
+LogicalResult ExecuteOpSeq::verify() {
+  return fallback_common::VerifyFallbackExecuteOp(*this);
+}
+LogicalResult ExecuteOpWithAllocator::verify() {
+  return fallback_common::VerifyExecuteOpCommon(*this);
+}
+LogicalResult ExecuteOpSeqWithAllocator::verify() {
+  return fallback_common::VerifyExecuteOpCommon(*this);
+}
+LogicalResult BatchFunctionOp::verify() {
+  return fallback_common::VerifyExecuteOpCommon(*this);
+}
+
+ParseResult CreateOp::parse(OpAsmParser &parser, OperationState &result) {
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = true;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = false;
+
+  auto &builder = parser.getBuilder();
+  if (mlir::failed(fallback_common::ParseExecuteOpCommon(
+          parser, builder, result, builder.getType<fallback::TFTensorType>(),
+          parse_options)))
+    return mlir::failure();
+
+  mlir::IntegerAttr num_args;
+  if (parser.parseKeyword("num_args") || parser.parseLParen() ||
+      parser.parseAttribute(num_args, "num_args", result.attributes) ||
+      parser.parseRParen())
+    return mlir::failure();
+
+  return mlir::success();
+}
+ParseResult ExecuteOp::parse(OpAsmParser &parser, OperationState &result) {
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = false;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = true;
+
+  auto &builder = parser.getBuilder();
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+ParseResult ExecuteOpSeq::parse(OpAsmParser &parser, OperationState &result) {
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = true;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = true;
+
+  auto &builder = parser.getBuilder();
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+ParseResult ExecuteOpWithAllocator::parse(OpAsmParser &parser,
+                                          OperationState &result) {
+  auto &builder = parser.getBuilder();
+  llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 1> allocator;
+  if (parser.parseOperandList(allocator,
+                              /*requiredOperandCount=*/1,
+                              mlir::OpAsmParser::Delimiter::Paren))
+    return mlir::failure();
+
+  if (parser.resolveOperand(allocator.front(),
+                            builder.getType<fallback::TFAllocatorType>(),
+                            result.operands))
+    return mlir::failure();
+
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = false;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = true;
+
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+ParseResult ExecuteOpSeqWithAllocator::parse(OpAsmParser &parser,
+                                             OperationState &result) {
+  auto &builder = parser.getBuilder();
+  llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 2>
+      chain_and_allocator;
+  if (parser.parseOperandList(chain_and_allocator,
+                              /*requiredOperandCount=*/2,
+                              mlir::OpAsmParser::Delimiter::Paren))
+    return mlir::failure();
+
+  auto &chain = chain_and_allocator[0];
+  auto &allocator = chain_and_allocator[1];
+
+  if (parser.resolveOperand(chain, builder.getType<compiler::ChainType>(),
+                            result.operands))
+    return mlir::failure();
+
+  if (parser.resolveOperand(allocator,
+                            builder.getType<fallback::TFAllocatorType>(),
+                            result.operands))
+    return mlir::failure();
+
+  // The first result is a chain.
+  result.types.push_back(builder.getType<compiler::ChainType>());
+
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = false;
+  parse_options.has_key = true;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = true;
+  parse_options.has_cost = true;
+
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+
+ParseResult BatchFunctionOp::parse(OpAsmParser &parser,
+                                   OperationState &result) {
+  fallback_common::ParseExecuteOpOptions parse_options;
+  parse_options.has_chain = false;
+  parse_options.has_key = false;
+  parse_options.has_device = true;
+  parse_options.has_func_attr = false;
+  parse_options.has_cost = false;
+  parse_options.has_op_name = false;
+  parse_options.has_symbol_ref = true;
+
+  auto &builder = parser.getBuilder();
+  return fallback_common::ParseExecuteOpCommon(
+      parser, builder, result, builder.getType<fallback::TFTensorType>(),
+      parse_options);
+}
+
+void CreateOp::print(OpAsmPrinter &p) {
+  CreateOp op = *this;
+  p << "(" << op.getInCh() << ") key("
+    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") device("
+    << op->getAttr("device") << ") " << op->getAttr("op_name") << "()";
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+
+  p << " num_args(" << op->getAttrOfType<mlir::IntegerAttr>("num_args").getInt()
+    << ')';
+}
+
+void ExecuteOp::print(OpAsmPrinter &p) {
+  ExecuteOp op = *this;
+  p << " key(" << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt()
+    << ") cost(" << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()
+    << ") device(" << op->getAttr("device") << ") " << op->getAttr("op_name")
+    << '(' << op.getArgs() << ')';
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+  if (!op.getResults().empty()) p << " : " << op.getResults().size();
+}
+
+void ExecuteOpSeq::print(OpAsmPrinter &p) {
+  ExecuteOpSeq op = *this;
+  p << "(" << op.getInOpChain() << ") key("
+    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") cost("
+    << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()
+    << ") device(" << op->getAttr("device") << ") " << op->getAttr("op_name")
+    << '(' << op.getArgs() << ')';
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+  if (!op.getResults().empty()) p << " : " << op.getResults().size();
+}
+
+void ExecuteOpWithAllocator::print(OpAsmPrinter &p) {
+  ExecuteOpWithAllocator op = *this;
+  p << "(" << op.getAllocator() << ") key("
+    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") cost("
+    << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()
+    << ") device(" << op->getAttr("device") << ") " << op->getAttr("op_name")
+    << '(' << op.getArgs() << ')';
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+  if (!op.getResults().empty()) p << " : " << op.getResults().size();
+}
+
+void ExecuteOpSeqWithAllocator::print(OpAsmPrinter &p) {
+  ExecuteOpSeqWithAllocator op = *this;
+  p << "(" << op.getInOpChain() << ", " << op.getAllocator() << ") key("
+    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") cost("
+    << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()
+    << ") device(" << op->getAttr("device") << ") " << op->getAttr("op_name")
+    << '(' << op.getArgs() << ')';
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+  fallback_common::PrintExecuteOpFuncAttribute(p, op);
+  if (!op.getResults().empty()) p << " : " << op.getResults().size();
+}
+
+void BatchFunctionOp::print(OpAsmPrinter &p) {
+  BatchFunctionOp op = *this;
+  p << " device(" << op->getAttr("device") << ") " << op->getAttr("f") << " ("
+    << op.getArgs() << ")";
+
+  fallback_common::PrintExecuteOpCommon(p, op);
+}
+
+void ExecuteOp::getOpAttrs(
+    SmallVectorImpl<std::pair<StringRef, Attribute>> *op_attrs) {
+  fallback_common::GetExecuteOpAttrsCommon(
+      this->getContext(), this->getOpAttrs().getValue(), op_attrs);
+}
+
+//===----------------------------------------------------------------------===//
+// ConstDenseTensorOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult ConstDenseTensorOp::fold(FoldAdaptor) { return getValue(); }
+
+//===----------------------------------------------------------------------===//
+// CoreRTTensorHandleToFallbackTensorOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+// Simplifies pattern containing a corert const tensor op followed by a
+// `tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor` op to a single
+// tfrt_fallback_async const tensor.
+struct ConstCoreRTTensorHandleToFallbackTensorCanonicalization
+    : public OpRewritePattern<CoreRTTensorHandleToFallbackTensorOp> {
+  using OpRewritePattern<
+      CoreRTTensorHandleToFallbackTensorOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CoreRTTensorHandleToFallbackTensorOp op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<Value, 1> new_values;
+    bool should_rewrite = false;
+    for (auto operand : op.getArgs()) {
+      if (auto corert_const_dense_tensor_op =
+              operand.getDefiningOp<corert::ConstDenseTensorOp>()) {
+        new_values.push_back(fallback_async::ConstDenseTensorOp::create(
+            rewriter, op.getLoc(), rewriter.getType<fallback::TFTensorType>(),
+            corert_const_dense_tensor_op.getValue()));
+        should_rewrite = true;
+        continue;
+      }
+      if (auto corert_const_string_tensor_op =
+              operand.getDefiningOp<corert::ConstStringTensorOp>()) {
+        new_values.push_back(fallback_async::ConstStringTensorOp::create(
+            rewriter, op.getLoc(), rewriter.getType<fallback::TFTensorType>(),
+            corert_const_string_tensor_op.getShape(),
+            corert_const_string_tensor_op.getValue()));
+        should_rewrite = true;
+        continue;
+      }
+      // To guarantee that the new values are in the same order as the old
+      // ones, we create individual ops for the non-canonicalizable operands.
+      // For simplicity, we don't consolidate these ops when all the
+      // non-canonicalizable operands are adjacent.
+      new_values.push_back(
+          fallback_async::CoreRTTensorHandleToFallbackTensorOp::create(
+              rewriter, op.getLoc(), rewriter.getType<fallback::TFTensorType>(),
+              operand, op->getAttrOfType<mlir::StringAttr>("device"))
+              .getResult(0));
+    }
+
+    if (!should_rewrite) return failure();
+    rewriter.replaceOp(op, new_values);
+    return success();
+  }
+};
+
+// Removes the following double tensor conversion:
+//  %1 = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %0
+//  %2 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %1
+struct RemoveDoubleTensorConversion
+    : mlir::OpRewritePattern<CoreRTTensorHandleToFallbackTensorOp> {
+  using OpRewritePattern<
+      CoreRTTensorHandleToFallbackTensorOp>::OpRewritePattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      CoreRTTensorHandleToFallbackTensorOp op,
+      mlir::PatternRewriter &rewriter) const override {
+    // Currently only handles the case where there is only one value in the
+    // conversion op. This should be enough for most of the cases.
+    if (op.getNumOperands() > 1) return mlir::failure();
+
+    auto def =
+        op.getOperand(0).getDefiningOp<FallbackTensorToCoreRTTensorHandleOp>();
+    if (!def) return mlir::failure();
+
+    if (def.getNumResults() > 1) return mlir::failure();
+
+    rewriter.replaceOp(op, def.getOperand(0));
+
+    return mlir::success();
+  }
+};
+
+}  // namespace
+
+void CoreRTTensorHandleToFallbackTensorOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *context) {
+  results.add<ConstCoreRTTensorHandleToFallbackTensorCanonicalization,
+              RemoveDoubleTensorConversion>(context);
+}
+
+}  // namespace fallback_async
+}  // namespace tfrt
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h

@@ -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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_ASYNC_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_ASYNC_H_
+
+#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tfrt/compiler/opdefs/tfrt_op_interfaces.h"  // from @tf_runtime
+#include "tfrt/compiler/opdefs/tfrt_traits.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/traits.h"  // from @tf_runtime
+
+using namespace mlir;  // NOLINT
+
+namespace tfrt {
+namespace fallback_async {
+
+// Dialect for fallback async operations.
+class FallbackAsyncDialect : public Dialect {
+ public:
+  explicit FallbackAsyncDialect(MLIRContext *context);
+  static StringRef getDialectNamespace() { return "tfrt_fallback_async"; }
+};
+
+}  // namespace fallback_async
+}  // namespace tfrt
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_ASYNC_H_

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td

@@ -0,0 +1,531 @@
+/* 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.
+==============================================================================*/
+
+#ifdef TFRT_FALLBACK_ASYNC_OPS
+#else
+#define TFRT_FALLBACK_ASYNC_OPS
+
+include "tfrt/tfrt_op_base.td"
+include "tfrt/compiler/opdefs/tfrt_op_interfaces.td"
+include "tfrt/compiler/opdefs/tfrt_traits.td"
+include "tfrt/core_runtime/opdefs/corert_base.td"
+include "tfrt/core_runtime/opdefs/corert_traits.td"
+include "tfrt_fallback.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+def FallbackAsync_Dialect : Dialect {
+  let name = "tfrt_fallback_async";
+
+  let description = [{
+    The TFRT Fallback Async Dialect.
+  }];
+
+  let cppNamespace = "::tfrt::fallback_async";
+}
+
+class FallbackAsync_Op<string mnemonic, list<Trait> traits = []> :
+    Op<FallbackAsync_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+def SetResourceOp : FallbackAsync_Op<"set_resource", [CoreRT_TypedAttributeTrait]> {
+  let summary = "Set a tensor in resource array";
+
+  let description = [{
+    Set a tensor in resource array.
+
+    arg: the tensor to be set in the resource array.
+    index: the index in the resource array
+
+    in_ch and out_ch are for control dependencies.
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_ch,
+    TFTensorType:$arg,
+    StrAttr:$device,
+    I64Attr:$index
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_ch
+  );
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def GetResourceOp : FallbackAsync_Op<"get_resource",
+    [CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "get a tensor in resource array";
+
+  let description = [{
+    Get a tensor in resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+
+    in_ch and out_ch are for control dependencies.
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_ch,
+    StrAttr:$device,
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_ch,
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "operands attr-dict `:` `(` type($in_ch) `)` `->` `(` type($out_ch) `,` type($results) `)`";
+}
+
+def ConstTensorProtoOp : FallbackAsync_Op<"const_tensor_proto",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Create a tensorflow tensor from serialized tensor proto";
+
+  let description = [{
+    Create a tensorflow tensor from serilaized tensor proto.
+
+    Example:
+      %tensor = tfrt_fallback_async.const_tensor_proto "<serialized proto content>"
+  }];
+
+  let arguments = (ins
+    StrAttr:$serialized_tensor_proto
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "$serialized_tensor_proto attr-dict";
+}
+
+def ConstDenseTensorOp : FallbackAsync_Op<"const_dense_tensor",
+    [ConstantLike, Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Create a constant dense tensorflow tensor.";
+
+  let description = [{
+    ConstDenseTensorOp creates a constant dense tensorflow tensor.
+
+    Example:
+      %tensor = tfrt_fallback_async.const_dense_tensor dense<[0, 1, 2]> : tensor<3xi32>
+  }];
+
+  let arguments = (ins
+    ElementsAttr:$value
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "$value attr-dict";
+
+  let hasFolder = 1;
+}
+
+// TODO(b/181257202): This should be `ConstantLike`.
+def ConstStringTensorOp : FallbackAsync_Op<"const_string_tensor",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Create a constant string tensorflow tensor.";
+  let description = [{
+    ConstStringTensorOp creates a constant string tensorflow tensor.
+
+    Example:
+      %tensor = tfrt_fallback_async.const_string_tensor {shape = [1, 2], value = ["const", "string"]}
+
+    Note: if the entries in `value` are all the same, we can just keep one copy.
+      In other words, the following two lines are effectively the same.
+      %tensor = tfrt_fallback_async.const_string_tensor {shape = [1, 2], value = ["string", "string"]}
+      %tensor = tfrt_fallback_async.const_string_tensor {shape = [1, 2], value = ["string"]}
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$shape,
+    StrArrayAttr:$value
+  );
+
+  let results = (outs
+    TFTensorType:$result
+  );
+
+  let assemblyFormat = "attr-dict";
+}
+
+def FallbackTensorToCoreRTTensorHandleOp : FallbackAsync_Op<
+                              "fallback_tensor_to_corert_tensorhandle",
+                              [Pure, CoreRT_TypedAttributeTrait,
+                               TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Convert fallback tensor to corert tensorhandle";
+
+  let description = [{
+    Convert the fallback tensor to corert tensorhandle.
+
+    args: a list of fallback tensors.
+    device: the name of the tensorflow device that is associated with the
+      input fallback tensors.
+
+    results: the converted corert tensorhandles.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device
+  );
+
+  let results = (outs
+    Variadic<CoreRT_TensorHandleType>:$results
+  );
+
+  let assemblyFormat = "operands attr-dict `:` `(` type($args) `)` `->` `(` type($results) `)`";
+}
+
+def CoreRTTensorHandleToFallbackTensorOp : FallbackAsync_Op<
+                              "corert_tensorhandle_to_fallback_tensor",
+                              [Pure, CoreRT_TypedAttributeTrait,
+                               TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Convert corert tensorhandle to fallback tensor";
+
+  let description = [{
+    Convert the corert tensorhandle to the fallback tensor.
+
+    args: a list of corert tensorhandles.
+    device: the name of the tensorflow device that is associated with the
+      output fallback tensor.
+
+    results: the converted fallback tensors.
+  }];
+
+  let arguments = (ins
+    Variadic<CoreRT_TensorHandleType>:$args,
+    StrAttr:$device
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "operands attr-dict `:` `(` type($args) `)` `->` `(` type($results) `)`";
+
+  let hasCanonicalizer = 1;
+}
+
+def PredicateOp : FallbackAsync_Op<"predicate",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Converts a fallback tensor to a bool";
+
+  let description = [{
+    Note: this kernel is used for CPU tensors.
+
+    Converts a fallback tensor to a bool with the following rules:
+
+    - For 0D tensors, truthiness is determined by comparing against a "zero"
+      value. For numerical types it is the obvious zero. For strings it is the
+      empty string.
+
+    - For >0D tensors, truthiness is determined by looking at the number of
+      elements. If has zero elements, then the result is false. Otherwise the
+      result is true.
+
+    input: a fallback tensor representing the condition.
+    device: the name of the tensorflow device that is associated with the
+      input fallback tensor.
+
+    output: the converted bool.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$input
+  );
+
+  let results = (outs
+    I1:$output
+  );
+
+  let assemblyFormat = "$input attr-dict";
+}
+
+def CopyIfSmallOp : FallbackAsync_Op<"copy_if_small",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
+  let summary = "Copy the CPU fallback tensor if it is small";
+
+  let description = [{
+    This kernel performs deep copy on the input tensor if it is small, to avoid
+    atomic contention on its refcount.
+
+    Note that this kernel always create a new AsyncValue for each result to
+    avoid atomic contention on AsyncValue's refcount.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$operand
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "$operand attr-dict `:` `(` type($operand) `)` `->` `(` type($results) `)`";
+}
+
+def CreateOp: FallbackAsync_Op<"createop", [CoreRT_TypedAttributeTrait]> {
+  let summary = "The Fallback CreateOp";
+
+  let description = [{
+    The CreateOp creates the tensorflow::OpKernel in the fallback context.
+
+    op_key: The unique key to identify this op.
+    num_args: The number of inputs.
+    device: The tensorflow device. eg. "/CPU:0"
+    op_attrs: The tensorflow attributes excluding the func attrs.
+    op_func_attrs: The func attrs.
+    op_name: The tensorflow op name. eg. "tf.AddV2"
+
+    Example:
+      %out_ch = tfrt_fallback_async.createop(%in_ch) key(0) device("/CPU:0")
+        "some.op"() {attr1 = value, attr2 = value2} {f_attr1 = "fn1", f_attr2 = "fn2"} num_args(1)
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_ch,
+    I64Attr:$num_args,
+    StrAttr:$device,
+    ArrayAttr:$op_attrs,
+    // TODO(b/173025975): consider using DictionaryAttr after we support
+    // BEF conversion for this type.
+    ArrayAttr:$op_func_attrs,
+    I64Attr:$op_key,
+    StrAttr:$op_name
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_ch
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def ExecuteOp : FallbackAsync_Op<"executeop",
+    [Pure, CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_AttrCostTrait]> {
+  let summary = "The Fallback ExecuteOp";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device.
+
+    Example:
+      %res = tfrt_fallback_async.executeop key(0) cost(100) device("/CPU:0") "some.op"(%arg) : 1
+
+    Note that the trailing number indicates the number of results.
+
+    The func attrs and op attrs are placed in two separate regions.
+
+    Example:
+      %res = tfrt_fallback_async.executeop key(0) device("/CPU:0") "some.op"(%arg)
+          {attr1 = value, attr2 = value2} {f_attr1 = "fn1", f_attr2 = "fn2"} : 1
+
+    Note that in this case, the first bracket contains op attributes other than
+    function attributes, and the second braket contains function attributes,
+    which represent function names.
+
+    Refer to b/172597167 for approaches considered for supporting function
+    attributes.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    ArrayAttr:$op_attrs,
+    // TODO(b/173025975): consider using DictionaryAttr after we support
+    // BEF conversion for this type.
+    ArrayAttr:$op_func_attrs,
+    I64Attr:$op_key,
+    StrAttr:$op_name,
+    I64Attr:$_tfrt_cost
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let extraClassDeclaration = [{
+    void getOpAttrs(SmallVectorImpl<std::pair<::llvm::StringRef, ::mlir::Attribute>>* op_attrs);
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def ExecuteOpSeq : FallbackAsync_Op<"executeop.seq",
+    [CoreRT_TypedAttributeTrait,
+     TFRT_CostFunctionInterface, TFRT_AttrCostTrait]> {
+  let summary = "The sequenced version of Fallback ExecuteOp";
+  let description = [{
+    The ExecuteOpSeq executes an operation on the specified device, similar to
+    ExecuteOp. In addition, it takes an in_op_chain and produces an out_op_chain
+    that can be used to sequence op executions. ExecuteOpSeq is nonstrict on
+    in_op_chain.
+
+    Example:
+      %op_ch_out, %res = tfrt_fallback_async.executeop.seq(%op_ch_in)
+        key(0) cost(100) device("/CPU:0") "some.op"(%arg) : 1
+
+    Note that the trailing number indicates the number of results.
+
+    The func attrs and op attrs are placed in two separate regions.
+
+    Example:
+      %op_ch_out, %res = tfrt_fallback_async.executeop.seq(%op_ch_in)
+        key(0) cost(100) device("/CPU:0") "some.op"(%arg)
+        {fn = function_1} {attr1 = value, attr2 = value2} : 1
+
+    Note that in this case, the first braket contains op attributes other than
+    function attributes, and the second braket contains function attributes,
+    which represent function names.
+
+    Refer to b/172597167 for approaches considered for supporting function
+    attributes.
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_op_chain,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    ArrayAttr:$op_attrs,
+    // TODO(b/173025975): consider using DictionaryAttr after we support
+    // BEF conversion for this type.
+    ArrayAttr:$op_func_attrs,
+    I64Attr:$op_key,
+    StrAttr:$op_name,
+    I64Attr:$_tfrt_cost
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_op_chain,
+    Variadic<TFTensorType>:$results
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def BatchFunctionOp : FallbackAsync_Op<"batch_function",
+    [Pure, CoreRT_TypedAttributeTrait]> {
+  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";
+  let description = [{
+    This kernel executes a variant tf.BatchFunction kernel that supports having
+    the `f` attribute as a BEF function.
+
+    Example:
+      %res = tfrt_fallback_async.batch_function device("/CPU:0") @f(%input, %captured_input)  {
+          num_batch_threads = 1,
+          max_batch_size = 4,
+          allowed_batch_sizes = [2 , 4],
+          batch_timeout_micros = 50,
+          container = "container",
+          shared_name = "shared_name",
+          batching_queue = "batching_queue",
+          enable_large_batch_splitting = false,
+          Tin = [f32],
+          Tcaptured = [f32],
+          Tout = [f32]} : 1
+
+    Note that the trailing number indicates the number of results.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    SymbolRefAttr:$f,
+    ArrayAttr:$op_attrs
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def ExecuteOpWithAllocator : FallbackAsync_Op<"executeop.allocator",
+    [Pure, CoreRT_TypedAttributeTrait, TFRT_CostFunctionInterface, TFRT_AttrCostTrait]> {
+  let summary = "The Fallback ExecuteOp with custom allocator";
+  let description = [{
+    Similar to ExecuteOp but takes a custom allocator for allocating output tensors.
+  }];
+
+  let arguments = (ins
+    TFAllocatorType:$allocator,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    ArrayAttr:$op_attrs,
+    ArrayAttr:$op_func_attrs,
+    I64Attr:$op_key,
+    StrAttr:$op_name,
+    I64Attr:$_tfrt_cost
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def ExecuteOpSeqWithAllocator : FallbackAsync_Op<"executeop.seq.allocator",
+    [CoreRT_TypedAttributeTrait, TFRT_CostFunctionInterface, TFRT_AttrCostTrait]> {
+  let summary = "The sequenced version of Fallback ExecuteOp with custom allocator";
+  let description = [{
+    Similar to ExecuteOpSeq but takes a custom allocator for allocating output tensors.
+
+    Example:
+      %op_ch_out, %res = tfrt_fallback_async.executeop.seq.allocator(%op_ch_in, %allocator)
+        key(0) cost(100) device("/CPU:0") "some.op"(%arg) : 1
+  }];
+
+  let arguments = (ins
+    TFRT_ChainType:$in_op_chain,
+    TFAllocatorType:$allocator,
+    Variadic<TFTensorType>:$args,
+    StrAttr:$device,
+    ArrayAttr:$op_attrs,
+    // TODO(b/173025975): consider using DictionaryAttr after we support
+    // BEF conversion for this type.
+    ArrayAttr:$op_func_attrs,
+    I64Attr:$op_key,
+    StrAttr:$op_name,
+    I64Attr:$_tfrt_cost
+  );
+
+  let results = (outs
+    TFRT_ChainType:$out_op_chain,
+    Variadic<TFTensorType>:$results
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.cc

@@ -0,0 +1,145 @@
+/* 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/ir/tfrt_fallback_common.h"
+
+#include <cassert>
+#include <cstdint>
+#include <utility>
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace tfrt {
+namespace fallback_common {
+
+void GetExecuteOpAttrsCommon(
+    mlir::MLIRContext *context, llvm::ArrayRef<mlir::Attribute> op_attr_array,
+    llvm::SmallVectorImpl<std::pair<llvm::StringRef, mlir::Attribute>>
+        *op_attrs) {
+  assert(op_attrs);
+  op_attrs->clear();
+
+  mlir::Builder builder(context);
+  for (auto iter : op_attr_array) {
+    auto key_value = mlir::cast<mlir::ArrayAttr>(iter).getValue();
+    llvm::StringRef key = mlir::cast<mlir::StringAttr>(key_value[0]).getValue();
+    mlir::Attribute value = key_value[1];
+    op_attrs->push_back({key, value});
+  }
+}
+
+mlir::ParseResult ParseExecuteOpCommon(mlir::OpAsmParser &parser,
+                                       mlir::Builder &builder,
+                                       mlir::OperationState &result,
+                                       mlir::Type tensor_type,
+                                       const ParseExecuteOpOptions &options) {
+  auto chain_type = builder.getType<compiler::ChainType>();
+
+  mlir::IntegerAttr op_key;
+  mlir::IntegerAttr cost;
+  mlir::StringAttr device;
+  mlir::StringAttr op_name;
+  mlir::SymbolRefAttr f;
+  llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 4> in_chains;
+  llvm::SmallVector<mlir::OpAsmParser::UnresolvedOperand, 4> operands;
+  mlir::NamedAttrList op_attrs;
+  mlir::NamedAttrList op_func_attrs;
+  auto loc = parser.getNameLoc();
+
+  if (options.has_chain &&
+      parser.parseOperandList(in_chains,
+                              /*requiredOperandCount=*/1,
+                              mlir::OpAsmParser::Delimiter::Paren))
+    return mlir::failure();
+
+  if (options.has_key &&
+      (parser.parseKeyword("key") || parser.parseLParen() ||
+       parser.parseAttribute(op_key, "op_key", result.attributes) ||
+       parser.parseRParen()))
+    return mlir::failure();
+
+  if (options.has_cost &&
+      (parser.parseKeyword("cost") || parser.parseLParen() ||
+       parser.parseAttribute(cost, "_tfrt_cost", result.attributes) ||
+       parser.parseRParen()))
+    return mlir::failure();
+
+  if (options.has_device &&
+      (parser.parseKeyword("device") || parser.parseLParen() ||
+       parser.parseAttribute(device, "device", result.attributes) ||
+       parser.parseRParen()))
+    return mlir::failure();
+
+  if (options.has_op_name &&
+      parser.parseAttribute(op_name, "op_name", result.attributes))
+    return mlir::failure();
+
+  if (options.has_symbol_ref &&
+      parser.parseAttribute(f, "f", result.attributes))
+    return mlir::failure();
+
+  if (parser.parseOperandList(operands, mlir::OpAsmParser::Delimiter::Paren) ||
+      parser.parseOptionalAttrDict(op_attrs) ||
+      parser.parseOptionalAttrDict(op_func_attrs))
+    return mlir::failure();
+
+  int64_t num_results = 0;
+  if (succeeded(parser.parseOptionalColon())) {
+    mlir::IntegerAttr attr;
+    mlir::NamedAttrList attrs;
+    if (failed(parser.parseAttribute(attr, "num_results", attrs)))
+      return mlir::failure();
+    num_results = attr.getValue().getSExtValue();
+  }
+  if (num_results < 0) return mlir::failure();
+
+  llvm::SmallVector<mlir::Type, 4> operand_types;
+  if (options.has_chain) operand_types.push_back(chain_type);
+  if (parser.resolveOperands(in_chains, operand_types, loc, result.operands) ||
+      parser.resolveOperands(operands, tensor_type, result.operands))
+    return mlir::failure();
+
+  if (options.has_chain) result.types.push_back(chain_type);
+  result.types.append(num_results, tensor_type);
+
+  llvm::SmallVector<mlir::Attribute, 4> op_attr_array;
+  for (const auto &key_value : op_attrs) {
+    auto key = key_value.getName();
+    auto value = key_value.getValue();
+    op_attr_array.push_back(builder.getArrayAttr({key, value}));
+  }
+
+  result.attributes.push_back(
+      builder.getNamedAttr("op_attrs", builder.getArrayAttr(op_attr_array)));
+
+  // TODO(tfrt-devs): support func attributes in tfrt_fallback_sync.
+  if (options.has_func_attr) {
+    llvm::SmallVector<mlir::Attribute, 4> op_func_attr_array;
+    for (const auto &key_value : op_func_attrs) {
+      auto key = key_value.getName();
+      auto value = key_value.getValue();
+      op_func_attr_array.push_back(builder.getArrayAttr({key, value}));
+    }
+
+    result.attributes.push_back(builder.getNamedAttr(
+        "op_func_attrs", builder.getArrayAttr(op_func_attr_array)));
+  }
+
+  return mlir::success();
+}
+
+}  // namespace fallback_common
+}  // namespace tfrt

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.h

@@ -0,0 +1,127 @@
+/* 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_IR_TFRT_FALLBACK_COMMON_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_COMMON_H_
+
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+
+namespace tfrt {
+namespace fallback_common {
+
+template <typename OpTy>
+mlir::LogicalResult VerifyExecuteOpCommon(OpTy op) {
+  auto op_attr_array = op.getOpAttrs().getValue();
+  for (auto op_attr : op_attr_array) {
+    auto key_value = mlir::dyn_cast<mlir::ArrayAttr>(op_attr);
+    if (!key_value || key_value.getValue().size() != 2 ||
+        !mlir::isa<mlir::StringAttr>(key_value.getValue()[0]))
+      return op.emitOpError() << "each op_attr should be a key-value pair, "
+                                 "where the key is a string";
+  }
+  return mlir::success();
+}
+
+template <typename OpTy>
+mlir::LogicalResult VerifyFallbackExecuteOp(OpTy op) {
+  auto result = VerifyExecuteOpCommon(op);
+  if (failed(result)) return result;
+
+  // Verify function attributes.
+  auto op_func_attr_array = op.getOpFuncAttrs().getValue();
+  for (auto op_attr : op_func_attr_array) {
+    auto key_value = mlir::dyn_cast<mlir::ArrayAttr>(op_attr);
+    if (!key_value || key_value.getValue().size() != 2 ||
+        !mlir::isa<mlir::StringAttr>(key_value.getValue()[0]) ||
+        !mlir::isa<mlir::StringAttr>(key_value.getValue()[1]))
+      return op.emitOpError() << "each op_func_attr should be a key-value "
+                                 "pair, where both the key and the value are "
+                                 "strings";
+  }
+  return mlir::success();
+}
+
+template <typename OpTy>
+void PrintExecuteOpFuncAttribute(mlir::OpAsmPrinter &p, OpTy op) {
+  auto op_func_attrs = op.getOpFuncAttrs();
+  if (!op_func_attrs.empty()) {
+    auto print_key_value = [&](mlir::Attribute attr) {
+      auto key_value = mlir::cast<mlir::ArrayAttr>(attr).getValue();
+      auto key = key_value[0];
+      auto value = key_value[1];
+
+      p << mlir::cast<mlir::StringAttr>(key).getValue();
+      p << " = ";
+      p << value;
+    };
+
+    auto op_func_attr_array = op_func_attrs.getValue();
+    p << " {";
+    llvm::interleaveComma(op_func_attr_array, p, print_key_value);
+    p << '}';
+  }
+}
+
+template <typename OpTy>
+void PrintExecuteOpCommon(mlir::OpAsmPrinter &p, OpTy op) {
+  auto op_attrs = op.getOpAttrs();
+  if (!op_attrs.empty()) {
+    auto print_key_value = [&](mlir::Attribute attr) {
+      auto key_value = mlir::cast<mlir::ArrayAttr>(attr).getValue();
+      auto key = key_value[0];
+      auto value = key_value[1];
+
+      p << mlir::cast<mlir::StringAttr>(key).getValue();
+      p << " = ";
+      p << value;
+    };
+
+    auto op_attr_array = op_attrs.getValue();
+    p << " {";
+    llvm::interleaveComma(op_attr_array, p, print_key_value);
+    p << '}';
+  }
+}
+
+void GetExecuteOpAttrsCommon(
+    mlir::MLIRContext *context, llvm::ArrayRef<mlir::Attribute> op_attr_array,
+    llvm::SmallVectorImpl<std::pair<llvm::StringRef, mlir::Attribute>>
+        *op_attrs);
+
+struct ParseExecuteOpOptions {
+  bool has_chain = false;
+  bool has_key = false;
+  bool has_device = false;
+  bool has_func_attr = false;
+  bool has_cost = false;
+  bool has_op_name = true;
+  bool has_symbol_ref = false;
+};
+
+mlir::ParseResult ParseExecuteOpCommon(mlir::OpAsmParser &parser,
+                                       mlir::Builder &builder,
+                                       mlir::OperationState &result,
+                                       mlir::Type tensor_type,
+                                       const ParseExecuteOpOptions &options);
+}  // namespace fallback_common
+}  // namespace tfrt
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_COMMON_H_

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cc

@@ -0,0 +1,65 @@
+/* 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/ir/tfrt_fallback_sync.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_common.h"
+#include "tfrt/basic_kernels/opdefs/tfrt_base.h"  // from @tf_runtime
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+#include "tfrt/core_runtime/opdefs/attributes.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
+#include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime
+
+namespace tfrt {
+namespace fallback_sync {
+
+FallbackSyncDialect::FallbackSyncDialect(MLIRContext *context)
+    : Dialect(/*name=*/"tfrt_fallback_sync", context,
+              TypeID::get<FallbackSyncDialect>()) {
+  context->getOrLoadDialect<tfrt::fallback::FallbackDialect>();
+  context->getOrLoadDialect<compiler::TFRTDialect>();
+  context->getOrLoadDialect<corert::CoreRTDialect>();
+
+  allowUnknownTypes();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cpp.inc"
+      >();
+}
+
+static Type GetTensorType(Builder *builder) {
+  return tfrt::tfrt_tensor::TensorType::get(builder->getContext());
+}
+
+}  // namespace fallback_sync
+}  // namespace tfrt
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.cpp.inc"

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.h

@@ -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.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_SYNC_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_SYNC_H_
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tfrt/core_runtime/opdefs/traits.h"  // from @tf_runtime
+#include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime
+
+using namespace mlir;  // NOLINT
+
+namespace tfrt {
+namespace fallback_sync {
+
+// Dialect for fallback operations.
+class FallbackSyncDialect : public Dialect {
+ public:
+  explicit FallbackSyncDialect(MLIRContext *context);
+  static StringRef getDialectNamespace() { return "tfrt_fallback_sync"; }
+};
+
+}  // namespace fallback_sync
+}  // namespace tfrt
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_SYNC_H_

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td

@@ -0,0 +1,286 @@
+/* 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.
+==============================================================================*/
+
+#ifdef TFRT_FALLBACK_SYNC_OPS
+#else
+#define TFRT_FALLBACK_SYNC_OPS
+
+include "tfrt/tfrt_op_base.td"
+include "tfrt/core_runtime/opdefs/corert_base.td"
+include "tfrt/core_runtime/opdefs/corert_traits.td"
+include "tfrt_fallback.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
+
+def FallbackSync_Dialect : Dialect {
+  let name = "tfrt_fallback_sync";
+
+  let description = [{
+    The TFRT Fallback Sync Dialect.
+  }];
+
+  let cppNamespace = "::tfrt::fallback_sync";
+}
+
+class FallbackSync_Op<string mnemonic, list<Trait> traits = []> :
+    Op<FallbackSync_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
+}
+
+// TODO(rohitju): Add a verifier if there is anything that should be verified
+// compile time.
+def SetResourceOp : FallbackSync_Op<"set_resource", [CoreRT_TypedAttributeTrait]> {
+  let summary = "Set a tensor in resource array";
+
+  let description = [{
+    Set a tensor in resource array.
+
+    arg: the tensor to be set in the resource array.
+    index: the index in the resource array
+  }];
+
+  let arguments = (ins
+    TFTensorType:$arg,
+    I64Attr:$index
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def SetResourceDhtOp : FallbackSync_Op<"set_resource_dht", [CoreRT_TypedAttributeTrait]> {
+  let summary = "Set a DHT in resource array";
+
+  let description = [{
+    Set a DHT in resource array.
+
+    arg: the tensor to be set in the resource array.
+    index: the index in the resource array
+  }];
+
+  let arguments = (ins
+    TensorType:$arg,
+    I64Attr:$index
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def GetResourceOp : FallbackSync_Op<"get_resource",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "get a tensor in resource array";
+
+  let description = [{
+    Get a tensor in resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "attr-dict `:` type($results)";
+}
+
+def GetResourceDhtOp : FallbackSync_Op<"get_resource_dht",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "get a DHT in resource array";
+
+  let description = [{
+    Get a tensor in resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    Variadic<TensorType>:$results
+  );
+
+  let assemblyFormat = "attr-dict `:` type($results)";
+}
+
+def CreateOp: FallbackSync_Op<"createop", []> {
+  let summary = "The Fallback CreateOp";
+
+  let description = [{
+    The CreateOp creates the tensorflow::OpKernel in the fallback context.
+
+    op_key: The unique key to identify this op.
+    num_args: The number of inputs.
+    device: The tensorflow device. eg. "/CPU:0"
+    op_attrs: The tensorflow attributes excluding the func attrs.
+    op_name: The tensorflow op name. eg. "tf.AddV2"
+
+    Example:
+      tfrt_fallback_sync.createop() key(0) device("/CPU:0")
+        "some.op"() {attr1 = value, attr2 = value2} {f_attr1 = "fn1", f_attr2 = "fn2"} num_args(1)
+  }];
+
+  let arguments = (ins
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "attr-dict";
+}
+
+def SyncExecuteOp : FallbackSync_Op<"executeop", [Pure]> {
+  let summary = "The Fallback Sync ExecuteOp";
+  let description = [{
+    The ExecuteOp executes an operation on the specified device.
+
+    Example:
+      %res = tfrt_fallback_sync.executeop "some.op"(%arg) : 1
+
+    Note that the trailing number indicates the number of results.
+  }];
+
+  let arguments = (ins
+    Variadic<TFTensorType>,
+    StrAttr:$node_def,
+    I32Attr:$op_key
+  );
+
+  let results = (outs Variadic<TFTensorType>);
+
+  let assemblyFormat = "`(`operands`)` attr-dict `:` functional-type(operands, results)";
+}
+
+
+def ConvertDhtToFallbackTensorOp : FallbackSync_Op<"convert_dht_to_fallback_tensor",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "convert a DHT to a fallback tensor";
+
+  let description = [{
+    Convert a TFRT DHT to a fallback tensor.
+
+    dht: the input DHT.
+    fallback_tensor: The output fallback tensor
+  }];
+
+  let arguments = (ins
+    TensorType:$dht
+  );
+
+  let results = (outs
+    TFTensorType:$fallback_tensor
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($fallback_tensor)";
+}
+
+def ConvertShtToFallbackTensorOp : FallbackSync_Op<"convert_sht_to_fallback_tensor",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "convert a SHT to a fallback tensor";
+
+  let description = [{
+    Convert a TFRT SHT to a fallback tensor.
+
+    sht: the input SHT.
+    fallback_tensor: The output fallback tensor
+  }];
+
+  let arguments = (ins
+    TensorType:$sht
+  );
+
+  let results = (outs
+    TFTensorType:$fallback_tensor
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($fallback_tensor)";
+}
+
+def ConvertFallbackTensorToDhtOp : FallbackSync_Op<"convert_fallback_tensor_to_dht",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "convert fallback tensor to a DHT";
+
+  let description = [{
+    Convert a fallback tensor to a TFRT DHT.
+
+    fallback_tensor: The input fallback tensor
+    dht: the output DHT.
+  }];
+
+  let arguments = (ins
+    TFTensorType:$fallback_tensor
+  );
+
+  let results = (outs
+    TensorType:$dht
+  );
+
+  let assemblyFormat = "operands attr-dict `:` type($dht)";
+}
+
+// TODO(rohitju): This is Ads specific, need to find an appropriate home for it.
+def SetSparseMatrixResourceOp : FallbackSync_Op<"set_sparse_matrix_resource", [CoreRT_TypedAttributeTrait]> {
+  let summary = "Set a Sparse matrix in resource array";
+
+  let description = [{
+    Set a sparse matrix in resource array.
+
+    arg: the matrix to be set in the resource array.
+    index: the index in the resource array
+  }];
+
+  let arguments = (ins
+    TFTensorType:$arg,
+    I64Attr:$index
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = "operands attr-dict";
+}
+
+def GetSparseMatrixResourceOp : FallbackSync_Op<"get_sparse_matrix_resource",
+    [CoreRT_TypedAttributeTrait]> {
+  let summary = "get a sparse matrix from resource array";
+
+  let description = [{
+    Get a sparse matrix from resource array.
+
+    indices: the indices in the resource array.
+    results: the tensor values for the corresponding indices.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr:$indices
+  );
+
+  let results = (outs
+    Variadic<TFTensorType>:$results
+  );
+
+  let assemblyFormat = "attr-dict `:` type($results)";
+}
+
+#endif

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_util.cc

@@ -0,0 +1,56 @@
+/* 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/ir/tfrt_fallback_util.h"
+
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h"
+
+namespace tfrt {
+namespace fallback_async {
+
+bool IsArgConsumedByFallback(mlir::func::FuncOp func, int arg_index) {
+  auto arg = func.getArgument(arg_index);
+
+  // Return true if any user is a fallback op. It is more interesting to know
+  // whether it is consumed by any fallback op than whether it is only consumed
+  // by fallback ops. For example, the conversion from a DenseHostTensor to a
+  // fallback tensor is more costly than the conversion from fallback tensor to
+  // a DenseHostTensor. So as long as one of the users of a captured variable is
+  // a fallback op, we should keep this variable as a fallback tensor.
+  for (mlir::Operation *user : arg.getUsers()) {
+    if (llvm::isa<FallbackAsyncDialect>(user->getDialect())) return true;
+  }
+  return false;
+}
+
+void ForEachArgConsumedByFallback(
+    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action) {
+  for (int arg_index = 0; arg_index < func.getNumArguments(); ++arg_index) {
+    if (IsArgConsumedByFallback(func, arg_index)) action(arg_index);
+  }
+}
+
+void ForEachArgConsumedByFallback(
+    mlir::ModuleOp module,
+    llvm::function_ref<void(llvm::StringRef func_name, int arg_index)> action) {
+  for (auto func : module.getOps<mlir::func::FuncOp>()) {
+    ForEachArgConsumedByFallback(
+        func, [func_name = func.getName(), action](int arg_index) {
+          action(func_name, arg_index);
+        });
+  }
+}
+
+}  // namespace fallback_async
+}  // namespace tfrt

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_util.h

@@ -0,0 +1,36 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_UTIL_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_UTIL_H_
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+
+namespace tfrt {
+namespace fallback_async {
+
+bool IsArgConsumedByFallback(mlir::func::FuncOp func, int arg_index);
+
+void ForEachArgConsumedByFallback(
+    mlir::func::FuncOp func, llvm::function_ref<void(int arg_index)> action);
+
+void ForEachArgConsumedByFallback(
+    mlir::ModuleOp module,
+    llvm::function_ref<void(llvm::StringRef func_name, int arg_index)> action);
+
+}  // namespace fallback_async
+}  // namespace tfrt
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_IR_TFRT_FALLBACK_UTIL_H_

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_combine.cc

@@ -0,0 +1,80 @@
+/* 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 file implements a set of simple combiners for optimizing operations in
+// the Runtime Fallback dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.h"
+
+// This optimizes the following scenario:
+// %tft0, %c2 = "tfd.move_dht_to_tft"(%dht0, %c1)
+//     : (!dht.host_tensor, !tfrt.chain) -> (!tfd.tf_tensor, !tfrt.chain)
+// %dht1, %c3 = "tfd.convert_tft_to_dht"(%tft0, %c2)
+//     : (!tfd.tf_tensor, !tfrt.chain) -> (!dht.host_tensor, !tfrt.chain)
+// some_op %dht1, %c3
+//
+// becomes
+// some_op %dht0, %c1
+
+struct SimplifyDoubleConversion
+    : public mlir::OpRewritePattern<mlir::tfd::ConvertTftToDhtOp> {
+  // We register this pattern to match every tfd.move_dht_to_tft op.
+  // The "benefit" is used by the framework to order the patterns and process
+  // them in order of profitability.
+  explicit SimplifyDoubleConversion(mlir::MLIRContext* context)
+      : mlir::OpRewritePattern<mlir::tfd::ConvertTftToDhtOp>(context,
+                                                             /*benefit=*/1) {}
+
+  // This method attempts to match a pattern and rewrite it. The rewriter
+  // argument is the orchestrator of the sequence of rewrites. The pattern is
+  // expected to interact with it to perform any changes to the IR from here.
+  mlir::LogicalResult matchAndRewrite(
+      mlir::tfd::ConvertTftToDhtOp op,
+      mlir::PatternRewriter& rewriter) const override {
+    // Look through the inputs of the ConvertTftToDhtOp.
+    mlir::Value convert_op_input_0 = op.getOperand(0);
+    mlir::Value convert_op_input_1 = op.getOperand(1);
+    mlir::tfd::MoveDhtToTftOp move_input_op_0 =
+        llvm::dyn_cast_or_null<mlir::tfd::MoveDhtToTftOp>(
+            convert_op_input_0.getDefiningOp());
+    mlir::tfd::MoveDhtToTftOp move_input_op_1 =
+        llvm::dyn_cast_or_null<mlir::tfd::MoveDhtToTftOp>(
+            convert_op_input_1.getDefiningOp());
+
+    // The inputs should be MoveDhtToTftOp.
+    if (!move_input_op_0 || !move_input_op_1) return mlir::failure();
+    // Both inputs are the same MoveDhtToTftOp.
+    if (move_input_op_0 != move_input_op_1) return mlir::failure();
+
+    // Use the rewriter to replace the ConvertTftToDhtOp's users with the
+    // operands of MoveDhtToTftOp.
+    rewriter.replaceOp(
+        op, {move_input_op_0.getOperand(0), move_input_op_0.getOperand(1)});
+    return mlir::success();
+  }
+};
+
+// Register rewrite pattern as "canonicalization" patterns on the MoveDhtToTftOp
+// so that they can be picked up by the Canonicalization framework.
+void mlir::tfd::ConvertTftToDhtOp::getCanonicalizationPatterns(
+    RewritePatternSet& results, MLIRContext* context) {
+  results.add<SimplifyDoubleConversion>(context);
+}

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_executor.cc

@@ -0,0 +1,230 @@
+/* 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/runtime_fallback/runtime_fallback_executor.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/SourceMgr.h"
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tfrt/utils/host_context.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/threadpool.h"
+#include "tensorflow/core/platform/threadpool_interface.h"
+#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_execute_compat_eager.h"
+#include "tensorflow/core/runtime_fallback/runtime/kernel_utils.h"
+#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/bef_converter/mlir_to_bef.h"  // from @tf_runtime
+#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
+#include "tfrt/host_context/async_value.h"  // from @tf_runtime
+#include "tfrt/host_context/chain.h"  // from @tf_runtime
+#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
+#include "tfrt/host_context/function.h"  // from @tf_runtime
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
+#include "tfrt/support/ref_count.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+using ::tfrt::AsyncValue;
+using ::tfrt::BEFFile;
+using ::tfrt::ExecutionContext;
+using ::tfrt::Function;
+using ::tfrt::MakeAvailableAsyncValueRef;
+using ::tfrt::RCReference;
+using ::tfrt::RequestContextBuilder;
+
+using ::tensorflow::Env;
+using ::tensorflow::thread::ThreadPool;
+using ::tensorflow::thread::ThreadPoolInterface;
+
+using ::tensorflow::tfrt_stub::FallbackTensor;
+
+// -------------------------------------------------------------------------- //
+// Run function via the TF->TFRT fallback lowering.
+// -------------------------------------------------------------------------- //
+
+namespace {
+// Thread pool for running `intra-op` tasks scheduled by the fallback kernels.
+class IntraOpThreadPool : public ThreadPoolInterface {
+ public:
+  explicit IntraOpThreadPool(int64_t num_threads)
+      : tpool_(Env::Default(), "intra-op",
+               std::max(1, static_cast<int32_t>(num_threads))) {}
+
+  void Schedule(std::function<void()> fn) override {
+    tpool_.Schedule(std::move(fn));
+  }
+
+  int NumThreads() const override { return tpool_.NumThreads(); }
+  int CurrentThreadId() const override { return tpool_.CurrentThreadId(); }
+  void Cancel() override {}
+
+ private:
+  ThreadPool tpool_;
+};
+}  // namespace
+
+RuntimeFallbackExecutor::RuntimeFallbackExecutor(int64_t num_threads)
+    : intra_op_(std::make_unique<IntraOpThreadPool>(num_threads)) {
+  // Create a HostContext for running TFRT functions. Concurrent work queue acts
+  // similar to the Tensorflow `inter-op` thread pool, so we'll match the size.
+  host_context_ = num_threads ? CreateMultiThreadedHostContext(num_threads)
+                              : CreateSingleThreadedHostContext();
+  tfrt::RegisterStaticKernels(host_context_->GetMutableRegistry());
+
+  // Build an ExecutionContext from the HostContext.
+  auto builder = RequestContextBuilder(host_context_.get(), &resource_context_);
+
+  // Get tensorflow::EagerContext for the kernel fallback.
+  auto* eager_context_resource =
+      resource_context_
+          .GetOrCreateResource<tensorflow::tfd::EagerContextResource>(
+              tensorflow::tfd::kEagerContextResourceName);
+  auto expected_eager_context = eager_context_resource->GetTFEagerContext();
+  auto* eager_context = expected_eager_context.get();
+
+  // Initialize fallback kernels state with a custom intra-op thread pool.
+  auto status = tensorflow::tfd::SetUpKernelFallbackCompatRequestContext(
+      &builder, /*runner_table=*/nullptr, eager_context, intra_op_.get());
+  CHECK(status.ok()) << "Failed to setup request context: " << status.message();
+
+  auto req_ctx = std::move(builder).build();
+  if (auto err = req_ctx.takeError())
+    LOG(FATAL) << "Failed to build a request context";
+
+  exec_ctx_ = std::make_unique<tfrt::ExecutionContext>(std::move(*req_ctx));
+}
+
+void RuntimeFallbackExecutor::Prepare(llvm::StringRef mlir_input) {
+  // We only support IR written in the Tensorflow dialect.
+  mlir::DialectRegistry registry;
+  mlir::RegisterAllTensorFlowDialects(registry);
+  mlir::MLIRContext context(registry);
+
+  llvm::SourceMgr source_mgr;
+  source_mgr.AddNewSourceBuffer(
+      llvm::MemoryBuffer::getMemBuffer(mlir_input, "test_ir"), llvm::SMLoc());
+
+  // Parse a kernel source code into the MLIR Module.
+  mlir::OwningOpRef<mlir::ModuleOp> module(
+      mlir::parseSourceFile<mlir::ModuleOp>(source_mgr, &context));
+  CHECK(module) << "failed to parse mlir module";
+
+  // Collect all diagnostics emitted while lowering parsed kernel module.
+  std::string diagnostic_str;
+  llvm::raw_string_ostream os(diagnostic_str);
+  mlir::SourceMgrDiagnosticHandler handler(source_mgr, module->getContext(),
+                                           os);
+
+  // Convert TF to TFRT fallback dialect.
+  TfrtPipelineOptions pipeline_opts;
+  pipeline_opts.default_device = kDefaultHostDeviceName;
+  pipeline_opts.hoist_invariant_ops = true;
+  pipeline_opts.sink_in_invariant_ops = false;
+  pipeline_opts.cost_threshold = 1024;
+  pipeline_opts.merge_inter_dependent_streams = true;
+
+  mlir::PassManager pm(module->getContext());
+  pm.addPass(CreateTfToTfrtConversionPass(pipeline_opts));
+
+  CHECK(mlir::succeeded(pm.run(*module)))
+      << "Failed to lower module to TFRT: " << os.str();
+
+  // Convert module to BEF.
+  bef_buffer_ =
+      tfrt::ConvertMLIRToBEF(*module, /*disable_optional_sections=*/false);
+  CHECK(!bef_buffer_.empty()) << "Failed to convert module to BEF";
+
+  bef_file_ =
+      BEFFile::Open(bef_buffer_, host_context_->GetKernelRegistry(),
+                    host_context_->diag_handler(), host_context_->allocator());
+  CHECK(bef_file_) << "Failed to open BEF";
+
+  // Run TFRT initialization function to pre-instantiate fallback kernels.
+  RunTfrtInitializer();
+}
+
+llvm::SmallVector<Tensor> RuntimeFallbackExecutor::Execute(
+    llvm::StringRef function_name, llvm::ArrayRef<Tensor> arguments) {
+  // Get the kernel entrypoint function.
+  const Function* compute = bef_file_->GetFunction(function_name);
+  CHECK(compute) << "Entrypoint function not found";
+  CHECK_EQ(arguments.size() + 1, compute->num_arguments())
+      << "Wrong number of arguments for function " << function_name.str();
+
+  // Prepare function arguments from ready Chain and input Tensors.
+  llvm::SmallVector<tfrt::AsyncValue*> exec_arguments;
+  exec_arguments.reserve(compute->num_arguments());
+  exec_arguments.push_back(tfrt::GetReadyChain().release());
+  for (const Tensor& input_tensor : arguments) {
+    auto av = MakeAvailableAsyncValueRef<FallbackTensor>(input_tensor);
+    exec_arguments.push_back(av.release());
+  }
+
+  // Space for returned values.
+  llvm::SmallVector<RCReference<AsyncValue>> results(compute->num_results());
+
+  compute->Execute(*exec_ctx_, exec_arguments, results);
+
+  // Wait for the function execution to finish, as well as the side-effects.
+  host_context_->Await(results);
+
+  // Check that all results are available.
+  llvm::SmallVector<Tensor> ret_values;
+  for (unsigned i = 1; i < results.size(); ++i) {
+    if (auto* error = results[i]->GetErrorIfPresent())
+      LOG(FATAL) << "Failed to execute a function: " << error->message();
+    ret_values.push_back(results[i]->get<tfrt_stub::FallbackTensor>().tensor());
+  }
+
+  // Deallocate arguments.
+  for (auto* argument : exec_arguments) argument->DropRef();
+  return ret_values;
+}
+
+// Run TFRT fallback initialization function to instantiate all fallback
+// kernels ahead of executing the compute function.
+void RuntimeFallbackExecutor::RunTfrtInitializer() {
+  const Function* func = bef_file_->GetFunction("_tfrt_fallback_init");
+  CHECK(func) << "TFRT initialization function was not found";
+  CHECK_EQ(func->argument_types().size(), 1);
+
+  llvm::SmallVector<RCReference<AsyncValue>, 1> results;
+  results.resize(func->result_types().size());
+  CHECK_EQ(results.size(), 1);
+
+  func->Execute(*exec_ctx_, tfrt::GetReadyChain().GetAsyncValue(), results);
+
+  host_context_->Await(results);
+
+  CHECK(!results[0]->IsError()) << "Failed to run TFRT initialization function";
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_executor.h

@@ -0,0 +1,63 @@
+/* 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_RUNTIME_FALLBACK_RUNTIME_FALLBACK_EXECUTOR_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_EXECUTOR_H_
+
+#include <cstdint>
+#include <memory>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/threadpool_interface.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
+#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
+#include "tfrt/support/ref_count.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+class RuntimeFallbackExecutor {
+ public:
+  explicit RuntimeFallbackExecutor(int64_t num_threads);
+
+  // Prepare() needs to be called once before calling Execute(). It sets up all
+  // things necessary to execute the given 'mlir_input' with the fallback to
+  // tensorflow.
+  void Prepare(llvm::StringRef mlir_input);
+
+  // Execute() can be called several times after the call to Prepare() (e.g. for
+  // benchmarking).
+  llvm::SmallVector<Tensor> Execute(llvm::StringRef function_name,
+                                    llvm::ArrayRef<Tensor> arguments);
+
+ private:
+  void RunTfrtInitializer();
+
+  std::unique_ptr<thread::ThreadPoolInterface> intra_op_;
+  std::unique_ptr<tfrt::HostContext> host_context_;
+  tfrt::ResourceContext resource_context_;
+  std::unique_ptr<tfrt::ExecutionContext> exec_ctx_;
+  tfrt::BefBuffer bef_buffer_;
+  tfrt::RCReference<tfrt::BEFFile> bef_file_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_EXECUTOR_H_

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.cc

@@ -0,0 +1,47 @@
+/* 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/runtime_fallback/runtime_fallback_ops.h"
+
+#include "tfrt/basic_kernels/opdefs/types.h"  // from @tf_runtime
+
+namespace mlir {
+namespace tfd {
+
+//===----------------------------------------------------------------------===//
+// TfrtDelegate Dialect
+//===----------------------------------------------------------------------===//
+
+RuntimeFallbackDialect::RuntimeFallbackDialect(MLIRContext *context)
+    : Dialect(/*name=*/"tfd", context, TypeID::get<RuntimeFallbackDialect>()) {
+  allowUnknownTypes();
+
+  allowUnknownOperations();
+
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback_ops.cc.inc"
+      >();
+}
+
+}  // namespace tfd
+}  // namespace mlir
+
+//===----------------------------------------------------------------------===//
+// TableGen'd op method definitions
+//===----------------------------------------------------------------------===//
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback_ops.cc.inc"

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.h

@@ -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.
+==============================================================================*/
+
+// This file defines the operations used in the Runtime Fallback dialect.
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_OPS_H_
+
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/TypeUtilities.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "tfrt/tensor/opdefs/tensor.h"  // from @tf_runtime
+
+namespace mlir {
+namespace tfd {
+
+// Dialect for TFRT delegate operations.
+class RuntimeFallbackDialect : public Dialect {
+ public:
+  explicit RuntimeFallbackDialect(MLIRContext* context);
+  static StringRef getDialectNamespace() { return "tfd"; }
+};
+
+}  // namespace tfd
+}  // namespace mlir
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfrt/runtime_fallback_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_RUNTIME_FALLBACK_RUNTIME_FALLBACK_OPS_H_

tensorflow/compiler/mlir/tfrt/runtime_fallback/runtime_fallback_ops.td

@@ -0,0 +1,158 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is the definition file for the Runtime Fallback Dialect.
+
+#ifdef TFRT_DELEGATE_DIALECT
+#else
+#define TFRT_DELEGATE_DIALECT
+
+include "tfrt/tfrt_op_base.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+//===----------------------------------------------------------------------===//
+// Type definitions
+//===----------------------------------------------------------------------===//
+def TfTensorType : OpaqueType<"tfd", "tf_tensor", "!tfd.tf_tensor type">;
+
+//===----------------------------------------------------------------------===//
+// Runtime Fallback Dialect definitions
+//===----------------------------------------------------------------------===//
+
+def RuntimeFallback_Dialect : Dialect {
+  let name = "tfd";
+
+  let description = [{
+    The Runtime Fallback dialect.
+
+    This dialect contains operations to run existing TF kernels on TFRT by
+    invoking TF Eager API.
+  }];
+
+  let cppNamespace = "::mlir::tfd";
+}
+
+//===----------------------------------------------------------------------===//
+// Runtime Fallback Dialect Ops definitions
+//===----------------------------------------------------------------------===//
+
+// Base class for the operation in this dialect.
+class RuntimeFallbackDialect_Op<string mnemonic, list<Trait> traits = []> :
+    Op<RuntimeFallback_Dialect, mnemonic, traits> { }
+
+def InitEagerContextOp : RuntimeFallbackDialect_Op<"init_eager_context"> {
+  let summary = "eager context initialization operation";
+  let description = [{
+    The "tfd.init_eager_context" operation takes an input chain, creates and
+    initializes the TF EagerContext and returns an output chain.
+
+    Example:
+      %c1 = "tfd.init_eager_context"(%c0): (!tfrt.chain) -> !tfrt.chain
+  }];
+
+  let arguments = (ins TFRT_ChainType);
+  let results = (outs TFRT_ChainType);
+}
+
+def DelegateKernelOp : RuntimeFallbackDialect_Op<"delegate_kernel"> {
+  let summary = "delegate kernel operation";
+  let description = [{
+    The "tfd.delegate_kernel" operation takes an input chain, and arbitrary
+    number of input arguments, and runs a specified TF op via TFE C API. It
+    returns an output chain and variable number of outputs from the TF op.
+
+    The input arguments and attributes are passed to the TF op. The outputs are
+    outputs of the TF op.
+
+    Note that `_name` is a required attribute specifying the TF op to run.
+    TFRT attributes are sorted alphabetically, passed in as positional
+    attributes to the TFRT kernel, rather than as named attributes.
+
+    Example:
+      To run "tf.MatMul" op, which has two boolean attributes,
+        1. Set _name = "MatMul"
+        2. For each TF attribute, split it into two attributes, one for name of
+           the TF attribute, and the other for the type and value of the
+           attribute value. Attribute value is a string with the format of
+           "type$val", where type can be "bool", "string", "tfdtype", "tfshape",
+           "tftensor".
+           The value serialization format can be found in attr_util.h.
+
+      %out_c, %out_tensor = "tfd.delegate_kernel"(
+        %in_c, %in1_tensor, %in2_tensor) {
+        _name = "MatMul",
+        attr1_name = "transpose_a", attr1_value = "bool$false",
+        attr2_name = "transpose_b", attr2_value = "bool$false"
+      } : (!tfrt.chain, !tfd.tf_tensor, !tfd.tf_tensor) -> (
+        !tfrt.chain, !tfd.tf_tensor)
+  }];
+
+  let arguments = (ins TFRT_ChainType, Variadic<AnyType>);
+  let results = (outs TFRT_ChainType, Variadic<AnyType>);
+}
+
+def PrintTftOp : RuntimeFallbackDialect_Op<"print_tft"> {
+  let summary = "print TF tensor operation";
+  let description = [{
+    The "tfd.print_tft" operation prints the input TF tensor. It takes an input
+    TF tensor to be printed and an input chain, and returns an output chain.
+
+    Example:
+      %c1 = "tfd.print_tft"(%t, %c) : (!tfd.tf_tensor, !tfrt.chain) -> !tfrt.chain
+
+  }];
+
+  let arguments = (ins TfTensorType, TFRT_ChainType);
+  let results = (outs TFRT_ChainType);
+}
+
+def ConvertTftToDhtOp : RuntimeFallbackDialect_Op<"convert_tft_to_dht", [Pure]> {
+  let summary = "convert TF tensor to TFRT DHT tensor operation";
+  let description = [{
+    The "tfd.convert_tft_to_dht" operation converts a TF tensor to a TFRT
+    DenseHostTensor.
+
+    It takes as input a TF Tensor and an input chain, and returns a converted
+    TFRT DHT tensor and an output chain.
+
+    Example:
+      %dht, %c0 = "tfd.convert_tft_to_dht"(%tft, %c)
+      : (!tfd.tf_tensor, !tfrt.chain) -> (!dht.host_tensor, !tfrt.chain)
+  }];
+
+  let arguments = (ins TfTensorType, TFRT_ChainType);
+  // Enable registering canonicalization patterns with this operation.
+  let hasCanonicalizer = 1;
+  let results = (outs TensorType, TFRT_ChainType);
+}
+
+def MoveDhtToTftOp : RuntimeFallbackDialect_Op<"move_dht_to_tft", [Pure]> {
+  let summary = "convert TFRT DHT tensor to DHT tensor operation";
+  let description = [{
+    The "tfd.move_dht_to_tft" operation moves a TFRT tensor into a TF Tensor.
+
+    It takes as input a TFRT Tensor and an input chain, and returns a TF tensor
+    with the same underlying buffer and an output chain.
+
+    Example:
+      %dht, %c0 = "tfd.convert_tft_to_dht"(%tft, %c)
+        : (!tfd.tf_tensor, !tfrt.chain) -> (!dht.host_tensor, !tfrt.chain)
+  }];
+
+  let arguments = (ins TensorType, TFRT_ChainType);
+  let results = (outs TfTensorType, TFRT_ChainType);
+}
+
+#endif // TFRT_DELEGATE_DIALECT

tensorflow/compiler/mlir/tfrt/saved_model/saved_model.cc

@@ -0,0 +1,154 @@
+/* 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/saved_model/saved_model.h"
+
+#include <cstdint>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/SymbolTable.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/ir/tf_saved_model.h"
+#include "tensorflow/compiler/mlir/tensorflow/utils/convert_type.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tsl/platform/errors.h"
+#include "tfrt/bef_converter/mlir_to_bef.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace {
+
+using ::mlir::tf_saved_model::kTfSavedModelIndexPathAttr;
+
+llvm::StringRef ProcessIndexPath(mlir::ArrayAttr index_path) {
+  if (index_path.size() == 1 && mlir::isa<mlir::StringAttr>(index_path[0])) {
+    // TODO(chky): Support cases where index_path is not a single string.
+    return mlir::cast<mlir::StringAttr>(index_path[0]).getValue();
+  }
+  return "";
+}
+
+absl::StatusOr<std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
+ProcessTensorSpec(mlir::TensorType type) {
+  tensorflow::DataType dtype;
+  TF_RETURN_IF_ERROR(
+      ConvertScalarTypeToDataType(type.getElementType(), &dtype));
+
+  if (!type.hasRank())
+    return std::make_pair(dtype, tensorflow::PartialTensorShape());
+
+  auto shape = type.getShape();
+  llvm::SmallVector<int64_t, 4> dims;
+  dims.assign(shape.begin(), shape.end());
+  return std::make_pair(dtype, tensorflow::PartialTensorShape(dims));
+}
+
+}  // namespace
+
+absl::Status MapFunctionSignaturesFromTFSavedModelMLIR(
+    mlir::ModuleOp module,
+    llvm::function_ref<void(const TFRTSavedModelSignatureInfo&)> map_fn) {
+  // Create bound inputs for each functions.
+  mlir::SymbolTable symbol_table(module);
+  absl::Status status = absl::OkStatus();
+  module.walk([&symbol_table, map_fn, &status](mlir::func::FuncOp func) {
+    // Use the exported name as the function name, and skip non-exported
+    // functions.
+    auto func_names = mlir::tf_saved_model::GetExportedNames(func);
+    if (func_names.empty()) return mlir::WalkResult::advance();
+
+    auto func_type = func.getFunctionType();
+
+    // Here we walk through each arguments and find out the input/output names,
+    // and input devices, variables used by this function.
+    llvm::SmallVector<llvm::StringRef, 4> input_names;
+    llvm::SmallVector<
+        std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>, 4>
+        input_specs;
+    llvm::SmallVector<llvm::StringRef, 4> input_devices;
+    llvm::SmallVector<mlir::Operation*, 4> bound_inputs;
+    for (unsigned i = 0, e = func.getNumArguments(); i != e; ++i) {
+      if (auto input_index_path = func.getArgAttrOfType<mlir::ArrayAttr>(
+              i, kTfSavedModelIndexPathAttr)) {
+        input_names.push_back(ProcessIndexPath(input_index_path));
+        auto statusor_spec = ProcessTensorSpec(
+            mlir::cast<mlir::TensorType>(func_type.getInput(i)));
+        if (!statusor_spec.ok()) {
+          status = std::move(statusor_spec).status();
+          return mlir::WalkResult::interrupt();
+        }
+        input_specs.push_back(std::move(statusor_spec).value());
+        if (auto input_device =
+                func.getArgAttrOfType<mlir::StringAttr>(i, "tf.device")) {
+          input_devices.push_back(input_device.getValue());
+        } else {
+          input_devices.push_back("");
+        }
+      }
+      if (auto* bound_input =
+              mlir::tf_saved_model::LookupBoundInput(func, i, symbol_table)) {
+        bound_inputs.push_back(bound_input);
+      }
+    }
+
+    llvm::SmallVector<llvm::StringRef, 4> output_names;
+    llvm::SmallVector<
+        std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>, 4>
+        output_specs;
+    for (unsigned i = 0, e = func.getNumResults(); i != e; ++i) {
+      if (auto output_index_path = func.getResultAttrOfType<mlir::ArrayAttr>(
+              i, kTfSavedModelIndexPathAttr)) {
+        output_names.push_back(ProcessIndexPath(output_index_path));
+        auto statusor_spec = ProcessTensorSpec(
+            mlir::cast<mlir::TensorType>(func_type.getResult(i)));
+        if (!statusor_spec.ok()) {
+          status = std::move(statusor_spec).status();
+          return mlir::WalkResult::interrupt();
+        }
+        output_specs.push_back(std::move(statusor_spec).value());
+      }
+    }
+
+    for (auto func_name : func_names) {
+      TFRTSavedModelSignatureInfo sig_info;
+      sig_info.func_name = func_name;
+      sig_info.input_names = input_names;
+      sig_info.input_specs = input_specs;
+      sig_info.input_devices = input_devices;
+      sig_info.output_names = output_names;
+      sig_info.output_specs = output_specs;
+      sig_info.bound_inputs = bound_inputs;
+      map_fn(sig_info);
+    }
+
+    return mlir::WalkResult::advance();
+  });
+
+  return status;
+}
+
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/saved_model/saved_model.h

@@ -0,0 +1,81 @@
+/* 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_SAVED_MODEL_SAVED_MODEL_H_
+#define TENSORFLOW_COMPILER_MLIR_TFRT_SAVED_MODEL_SAVED_MODEL_H_
+
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/core_runtime/tensor_handle.h"  // from @tf_runtime
+
+namespace tfrt {
+class CoreRuntime;
+}
+
+namespace mlir {
+class ModuleOp;
+}
+
+namespace tensorflow {
+
+// TFRTSavedModelSignatureInfo contains the metadata for a signature in the
+// savedmodel such as function name, inputs/outputs' names and types. This can
+// be used to retrieve these information in a tf_saved_model module.
+struct TFRTSavedModelSignatureInfo {
+  llvm::StringRef func_name;
+
+  // The following are metadata for inputs.
+  llvm::ArrayRef<llvm::StringRef> input_names;
+  llvm::ArrayRef<
+      std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
+      input_specs;
+  llvm::ArrayRef<llvm::StringRef> input_devices;
+
+  // The following are metadata for outputs.
+  llvm::ArrayRef<llvm::StringRef> output_names;
+  llvm::ArrayRef<
+      std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>>
+      output_specs;
+
+  // The following are metadata for bound_inputs, ie. captures.
+  llvm::ArrayRef<mlir::Operation*> bound_inputs;
+};
+
+// Apply `map_fn` on every exported function in the module with the
+// corresponding signature metadata populated in TFRTSavedModelSignatureInfo for
+// the function.
+absl::Status MapFunctionSignaturesFromTFSavedModelMLIR(
+    mlir::ModuleOp module,
+    llvm::function_ref<void(const TFRTSavedModelSignatureInfo&)> map_fn);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFRT_SAVED_MODEL_SAVED_MODEL_H_

tensorflow/compiler/mlir/tfrt/test_opkernels.cc

@@ -0,0 +1,52 @@
+/* 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/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace tensorflow {
+namespace tf_mlrt {
+
+REGISTER_OP("TestAsyncIdentity")
+    .Input("in: T")
+    .Output("out: T")
+    .Attr(
+        "T: {bfloat16, half, float, double, uint8, int8, int16, uint32, int32, "
+        "int64, complex64, complex128}")
+    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);
+
+class TestAsyncIdentityKernel : public AsyncOpKernel {
+ public:
+  explicit TestAsyncIdentityKernel(OpKernelConstruction* context)
+      : AsyncOpKernel(context) {}
+
+  void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+    const Tensor& in = ctx->input(0);
+    ctx->set_output(0, in);
+    done();
+  }
+
+ private:
+  TestAsyncIdentityKernel(const TestAsyncIdentityKernel&) = delete;
+  void operator=(const TestAsyncIdentityKernel&) = delete;
+};
+
+REGISTER_KERNEL_BUILDER(Name("TestAsyncIdentity").Device(DEVICE_CPU),
+                        TestAsyncIdentityKernel);
+
+}  // namespace tf_mlrt
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/tests/BUILD

@@ -0,0 +1,29 @@
+load("//tensorflow:tensorflow.bzl", "if_oss")
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    features = if_oss(["--path=org_tensorflow/tensorflow/compiler/mlir/tfrt"]),
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir:tf-mlir-translate",
+        "//tensorflow/compiler/mlir:tf-opt",
+        "//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:not",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+)

tensorflow/compiler/mlir/tfrt/tests/analysis/BUILD

@@ -0,0 +1,50 @@
+load("//tensorflow:tensorflow.bzl", "if_oss", "tf_cc_test")
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    exclude = ["testdata/**"],
+    features = if_oss(["--path=org_tensorflow/tensorflow/compiler/mlir/tfrt"]),
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir:tf-mlir-translate",
+        "//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:not",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+)
+
+tf_cc_test(
+    name = "update_op_cost_in_tfrt_mlir_test",
+    srcs = ["update_op_cost_in_tfrt_mlir_test.cc"],
+    data = [
+        "testdata/test.mlir",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt:transforms/update_op_cost_in_tfrt_mlir",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_sync_opdefs",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:resource_loader",
+        "//tensorflow/core/tfrt/fallback:cost_recorder",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@tf_runtime//:init_tfrt_dialects",
+    ],
+)

tensorflow/compiler/mlir/tfrt/tests/analysis/cost_analysis.mlir

@@ -0,0 +1,165 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -tfrt-test-cost-analysis -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: test_cheap_ops_0
+func.func @test_cheap_ops_0(%arg: tensor<?x!tf_type.string>) -> (tensor<?x8xf32>) {
+    // expected-remark@+1 {{Cost: 1}}
+    %0 = "tf.Const"() {value = dense<> : tensor<0xi64>} : () -> tensor<0xi64>
+    // expected-remark@+1 {{Cost: 1}}
+    %1 = "tf.Const"() {value = dense<"has_login_page_feature"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %2 = "tf.Const"() {value = dense<"num_terms_inside_postform"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %3 = "tf.Const"() {value = dense<"num_terms_outside_postform"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %4 = "tf.Const"() {value = dense<"num_terms_outside_postform_without_bp"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %5 = "tf.Const"() {value = dense<"password_not_in_bp_area"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %6 = "tf.Const"() {value = dense<"query_params_contains_url"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %7 = "tf.Const"() {value = dense<"title_with_login_phase"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %8 = "tf.Const"() {value = dense<"url_contains_login_terms"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %9 = "tf.Const"() {value = dense<> : tensor<0x!tf_type.string>} : () -> tensor<0x!tf_type.string>
+    // expected-remark@+1 {{Cost: 1}}
+    %10 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+    // expected-remark@+1 {{Cost: 1}}
+    %11 = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+    // expected-remark@+1 {{Cost: 19}}
+    %dense_values:8 = "tf.ParseExample"(%arg, %9, %1, %2, %3, %4, %5, %6, %7, %8, %0, %0, %0, %0, %0, %0, %0, %0) {dense_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>], device = "/job:localhost/replica:0/task:0/device:CPU:0", operandSegmentSizes = array<i32: 1, 1, 0, 8, 8>, resultSegmentSizes = array<i32: 0, 0, 0, 8>} : (tensor<?x!tf_type.string>, tensor<0x!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>, tensor<0xi64>) -> (tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>)
+    // expected-remark@+1 {{Cost: 2}}
+    %28 = "tf.Cast"(%dense_values#0) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
+    // expected-remark@+1 {{Cost: 2}}
+    %29 = "tf.Cast"(%dense_values#1) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
+    // expected-remark@+1 {{Cost: 2}}
+    %30 = "tf.Cast"(%dense_values#2) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
+    // expected-remark@+1 {{Cost: 2}}
+    %31 = "tf.Cast"(%dense_values#3) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
+    // expected-remark@+1 {{Cost: 2}}
+    %32 = "tf.Cast"(%dense_values#4) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
+    // expected-remark@+1 {{Cost: 2}}
+    %33 = "tf.Cast"(%dense_values#5) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
+    // expected-remark@+1 {{Cost: 2}}
+    %34 = "tf.Cast"(%dense_values#6) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
+    // expected-remark@+1 {{Cost: 2}}
+    %35 = "tf.Cast"(%dense_values#7) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xi64>) -> tensor<?xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %36 = "tf.ExpandDims"(%28, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %37 = "tf.ExpandDims"(%29, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %38 = "tf.ExpandDims"(%30, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %39 = "tf.ExpandDims"(%31, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %40 = "tf.ExpandDims"(%32, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %41 = "tf.ExpandDims"(%33, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %42 = "tf.ExpandDims"(%34, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %43 = "tf.ExpandDims"(%35, %11) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>) -> tensor<?x1xf32>
+    // expected-remark@+1 {{Cost: 10}}
+    %44 = "tf.ConcatV2"(%36, %37, %38, %39, %40, %41, %42, %43, %10) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<?x1xf32>, tensor<i32>) -> tensor<?x8xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    func.return %44 : tensor<?x8xf32>
+}
+
+// CHECK-LABEL: test_cheap_ops_1
+func.func @test_cheap_ops_1(%arg: tensor<?x8x?x?xf32>) -> (tensor<4xi32>, tensor<?x8x?x?xf32>) {
+    // expected-remark@+1 {{Cost: 1}}
+    %0 = "tf.Const"() {value = dense<8> : tensor<i32>} : () -> tensor<i32>
+    // expected-remark@+1 {{Cost: 1}}
+    %1 = "tf.Const"() {value = dense<4> : tensor<1xi32>} : () -> tensor<1xi32>
+    // expected-remark@+1 {{Cost: 1}}
+    %2 = "tf.Const"() {value = dense<64> : tensor<i32>} : () -> tensor<i32>
+    // expected-remark@+1 {{Cost: 1}}
+    %3 = "tf.Const"() {value = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
+    // expected-remark@+1 {{Cost: 1}}
+    %4 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
+    // expected-remark@+1 {{Cost: 1}}
+    %5 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
+    // expected-remark@+1 {{Cost: 1}}
+    %6 = "tf.Const"() {value = dense<3> : tensor<1xi32>} : () -> tensor<1xi32>
+    // expected-remark@+1 {{Cost: 9}}
+    %7 = "tf.Softmax"(%arg) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x8x?x?xf32>) -> tensor<?x8x?x?xf32>
+    // expected-remark@+1 {{Cost: 1}}
+    %8 = "tf.Shape"(%7) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x8x?x?xf32>) -> tensor<4xi32>
+    // expected-remark@+1 {{Cost: 1}}
+    %9 = "tf.StridedSlice"(%8, %5, %4, %4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
+    // expected-remark@+1 {{Cost: 1}}
+    %10 = "tf.StridedSlice"(%8, %3, %6, %4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
+    // expected-remark@+1 {{Cost: 5}}
+    %11 = "tf.Pack"(%9, %0, %10, %2) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<4xi32>
+    // expected-remark@+1 {{Cost: 1}}
+    %12 = "tf.StridedSlice"(%8, %6, %1, %4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
+    // expected-remark@+1 {{Cost: 5}}
+    %13 = "tf.Pack"(%9, %0, %10, %12) {axis = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<4xi32>
+    // expected-remark@+1 {{Cost: 1}}
+    %14 = "tf.Reshape"(%7, %13) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x8x?x?xf32>, tensor<4xi32>) -> tensor<?x8x?x?xf32>
+    // expected-remark@+1 {{Cost: 8}}
+    func.return %11, %14 : tensor<4xi32>, tensor<?x8x?x?xf32>
+}
+
+// CHECK-LABEL: test_expensive_ops
+func.func @test_expensive_ops(%arg: tensor<?x512xf32>) -> tensor<?x512xf32> {
+    // expected-remark@+1 {{Cost: 1}}
+    %0 = "tf.VarHandleOp"() {allowed_devices = [], container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", shared_name = "var"} : () -> tensor<!tf_type.resource<tensor<512x512xf32>>>
+    // expected-remark@+1 {{Cost: 2}}
+    %1 = "tf.ReadVariableOp"(%0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource<tensor<512x512xf32>>>) -> tensor<512x512xf32>
+    // 262657 = 1 + 512 + 512 * 512
+    // expected-remark@+1 {{Cost: 262657}}
+    %2 = "tf.MatMul"(%arg, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<?x512xf32>, tensor<512x512xf32>) -> tensor<?x512xf32>
+    // expected-remark@+1 {{Cost: 512}}
+    func.return %2 : tensor<?x512xf32>
+}
+
+// CHECK-LABEL: test_dynamic_shape
+func.func @test_dynamic_shape(%key: tensor<?x!tf_type.string>, %value: tensor<8xi64>) -> tensor<*xi1> {
+    // expected-remark@+1 {{Cost: 1}}
+    %default = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i64>} : () -> tensor<i64>
+    // expected-remark@+1 {{Cost: 1}}
+    %0 = "tf.HashTableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", key_dtype = !tf_type.string, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
+    // expected-remark@+1 {{Cost: 1024}}
+    %1 = "tf.LookupTableFindV2"(%0, %key, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
+    // 17 = 1 + 8 + 8
+    // expected-remark@+1 {{Cost: 17}}
+    %2 = "tf.NotEqual"(%1, %value) {device = "/job:localhost/replica:0/task:0/device:CPU:0", incompatible_shape_error = true} : (tensor<*xi64>, tensor<8xi64>) -> tensor<*xi1>
+    // expected-remark@+1 {{Cost: 8}}
+    func.return %2 : tensor<*xi1>
+}
+
+// CHECK-LABEL: test_gather
+func.func @test_gather(%arg0 : tensor<1x2x20xf32>, %arg1 : tensor<3x5xi32>) -> (tensor<1x3x5x20xf32>){
+    // expected-remark@+1 {{Cost: 1}}
+    %0 = "tf.Const"() { value = dense<[1]> : tensor<1xi32> } : () -> tensor<1xi32>
+    // expected-remark@+1 {{Cost: 300}}
+    %1 = "tf.GatherV2"(%arg0, %arg1, %0) : (tensor<1x2x20xf32>, tensor<3x5xi32>, tensor<1xi32>) -> tensor<1x3x5x20xf32>
+    // expected-remark@+1 {{Cost: 40}}
+    func.return %1 : tensor<1x3x5x20xf32>
+}
+
+// CHECK-LABEL: test_sparse_segment_sum
+func.func @test_sparse_segment_sum(%indices: tensor<3xi64>, %segment_ids: tensor<3xi64>) -> (tensor<?x28xf32>){
+    // expected-remark@+1 {{Cost: 1}}
+    %data = "tf.Const"() { value = dense<0.1> : tensor<476x28xf32> } : () -> tensor<476x28xf32>
+    // expected-remark@+1 {{Cost: 28}}
+    %1 = "tf.SparseSegmentSum"(%data, %indices, %segment_ids) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<476x28xf32>, tensor<3xi64>, tensor<3xi64>) -> tensor<?x28xf32>
+    // expected-remark@+1 {{Cost: 3}}
+    func.return %1 : tensor<?x28xf32>
+}

tensorflow/compiler/mlir/tfrt/tests/analysis/tensor_array_side_effect_analysis.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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -tfrt-test-tensor-array-effect -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: @test_tensor_array_effect
+// expected-remark@+1 {{HasAtMostTensorArrayEffect: 1}}
+func.func @test_tensor_array_effect(%index: tensor<i32>, %size: tensor<i32>, %flow_0: tensor<f32>, %flow_1: tensor<f32>, %handle_0: tensor<2x!tf_type.resource<tensor<?x100xf32>>>, %handle_1: tensor<2x!tf_type.resource<tensor<?x512xf32>>>) -> (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>) {
+  %cst = "tf.Const"() {value = dense<1.1> : tensor<100x512xf32>} : () -> tensor<100x512xf32>
+  %one = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %x = "tf.TensorArrayReadV3"(%handle_0, %index, %flow_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<i32>, tensor<f32>) -> tensor<?x100xf32>
+  %y = "tf.MatMul"(%x, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x100xf32>, tensor<100x512xf32>) -> (tensor<?x512xf32>)
+  %flow_1_out = "tf.TensorArrayWriteV3"(%handle_1, %index, %y, %flow_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<i32>, tensor<?x512xf32>, tensor<f32>) -> tensor<f32>
+  %next_index = "tf.AddV2"(%index, %one) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %next_index, %size, %flow_0, %flow_1_out, %handle_0, %handle_1 : tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<2x!tf_type.resource<tensor<?x100xf32>>>, tensor<2x!tf_type.resource<tensor<?x512xf32>>>
+}

tensorflow/compiler/mlir/tfrt/tests/analysis/testdata/test.mlir

@@ -0,0 +1,22 @@
+// 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.
+// ==============================================================================
+
+func.func @test(%ch: !tfrt.chain, %arg0: !corert.tensorhandle, %arg1_th: !corert.tensorhandle) {
+  %cpu = corert.get_op_handler %ch "cpu"
+  %0 = corert.executeop(%cpu) "tf.Relu"(%arg0) { T = f32 } : 1
+  %arg1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %arg1_th {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
+  %1 = tfrt_fallback_async.executeop key(0) cost(100) device("/CPU:0") "tf.Relu"(%arg1) { T = f32 } : 1
+  tfrt.return
+}

tensorflow/compiler/mlir/tfrt/tests/analysis/update_op_cost_in_tfrt_mlir_test.cc

@@ -0,0 +1,89 @@
+/* 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 <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/container/flat_hash_map.h"
+#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/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tfrt/init_tfrt_dialects.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace {
+
+constexpr char kCostAttrName[] = "_tfrt_cost";
+constexpr char kOpKeyAttrName[] = "op_key";
+
+absl::flat_hash_map<int64_t, uint64_t> GetOpCostMap(mlir::ModuleOp op) {
+  absl::flat_hash_map<int64_t, uint64_t> op_cost_map;
+  op.walk([&](mlir::Operation* op) {
+    const auto cost_attr = op->getAttrOfType<mlir::IntegerAttr>(kCostAttrName);
+    if (!cost_attr) return;
+    const auto op_key_attr =
+        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
+    if (!op_key_attr) return;
+    op_cost_map[op_key_attr.getInt()] = cost_attr.getInt();
+  });
+  return op_cost_map;
+}
+
+TEST(CostUpdateTest, Basic) {
+  std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
+      "tensorflow/compiler/mlir/tfrt/tests/analysis/testdata/test.mlir");
+
+  mlir::DialectRegistry registry;
+  tfrt::RegisterTFRTDialects(registry);
+  registry.insert<tfrt::fallback_async::FallbackAsyncDialect>();
+  registry.insert<tfrt::fallback_sync::FallbackSyncDialect>();
+  mlir::MLIRContext context(registry);
+  auto module =
+      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
+  ASSERT_TRUE(module);
+
+  // Create a cost recorder with fake cost records.
+  auto expected_op_cost_map = GetOpCostMap(module.get());
+  EXPECT_EQ(expected_op_cost_map.size(), 1);
+  unsigned int seed = 23579;
+  for (auto& [op_key, cost] : expected_op_cost_map) {
+    cost = rand_r(&seed) % 1000;
+  }
+  tensorflow::tfrt_stub::CostRecorder cost_recorder;
+  for (const auto& [op_key, cost] : expected_op_cost_map) {
+    cost_recorder.RecordCost(op_key, cost);
+  }
+
+  // Update the TFRT MLIR with the cost recorder.
+  tfrt_compiler::UpdateOpCostInTfrtMlir(module.get(), cost_recorder);
+
+  // Check the updated costs.
+  const auto got_op_cost_map = GetOpCostMap(module.get());
+  EXPECT_THAT(got_op_cost_map, ::testing::ContainerEq(expected_op_cost_map));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/mlir/tfrt/tests/batch_function_fallback_resource_variable_as_captured_tensor.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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -pass-pipeline='builtin.module(tf-executor-to-tfrt-pipeline{target-tpurt=true})' %s | FileCheck %s
+
+module attributes {tf_saved_model.semantics}  {
+  // CHECK-LABEL: func @main
+  func.func @main_func() -> (tensor<*xf32> {tf_saved_model.index_path = ["a"]}) attributes {tf_saved_model.exported_names = ["main_func"]} {
+    %0 = tf_executor.graph {
+      %outputs_0, %control_0 = tf_executor.island wraps "tf.VarHandleOp"() {container = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<501000x128xf32>>>
+      %outputs_1, %control_1 = tf_executor.island wraps "tf.Cast"(%outputs_0) {Truncate = false} : (tensor<!tf_type.resource<tensor<501000x128xf32>>>) -> tensor<*x!tf_type.resource>
+
+      // CHECK: tfrt_fallback_async.batch_function device([[DEVICE:.*]]) @batched_func ([[BATCHED_FUNC_ARG:%.*]])
+      // CHECK-SAME: Tcaptured = [!corert.resource]
+      // CHECK-SAME: Tin = []
+      // CHECK-SAME: Tout = [f32]
+      %outputs_2, %control_2 = tf_executor.island wraps "tf.BatchFunction"(%outputs_1) {batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", f = @batched_func, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 0, 1>, shared_name = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
+      tf_executor.fetch %outputs_2 : tensor<*xf32>
+    }
+    func.return %0 : tensor<*xf32>
+  }
+  func.func private @batched_func(%arg0: tensor<*x!tf_type.resource>) -> tensor<?xf32> {
+    %0 = tf_executor.graph {
+      %outputs_0, %control_0 = tf_executor.island wraps "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource>) -> tensor<?xf32>
+      tf_executor.fetch %outputs_0 : tensor<?xf32>
+    }
+    func.return %0 : tensor<?xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/batch_function_lowering.mlir

@@ -0,0 +1,41 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -tf-executor-to-tfrt-pipeline %s | FileCheck %s --dump-input=always
+
+func.func private @batched_function(%arg0: tensor<1x3xf32> {tf._user_specified_name = "0"}, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32> attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
+  %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<1x3xf32>
+  %1 = "tf.AddV2"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
+  %2 = "tf.Identity"(%1) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "batch/BatchFunction:0"}} {
+  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+  // CHECK: tfrt_fallback_async.batch_function device("/device:CPU:0") @batched_function
+  // CHECK-SAME: Tin = [f32]
+  // CHECK-SAME: Tout = [f32]
+  // CHECK-SAME: allowed_batch_sizes = [6]
+  // CHECK-SAME: batch_timeout_micros = 100000 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = false
+  // CHECK-SAME: max_batch_size = 6 : i64
+  // CHECK-SAME: max_enqueued_batches = 10 : i64
+  // CHECK-SAME: num_batch_threads = 1 : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}

tensorflow/compiler/mlir/tfrt/tests/convert_ref_variables.mlir

@@ -0,0 +1,120 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -tfrt-convert-ref-variables -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// Test the basic cases where all uses of a ref variable can be converted.
+
+// CHECK-LABEL: @init
+func.func @init() {
+  // CHECK-NOT: tf.VariableV2
+  // CHECK-NOT: tf.Assign
+
+  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"
+  // CHECK-SAME: shared_name = "x"
+  // CHECK: "tf.AssignVariableOp"([[handle]], {{%.*}})
+  %0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
+  %1 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %2 = "tf.Assign"(%0, %1) {T = i32, device = "", use_locking = true, validate_shape = true} : (tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<!tf_type.int32ref>
+  func.return
+}
+
+// CHECK-LABEL: @inference
+func.func @inference() -> tensor<i32> {
+  // CHECK-NOT: tf.VariableV2
+
+  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"
+  // CHECK-SAME: shared_name = "x"
+  // CHECK: "tf.ReadVariableOp"([[handle]])
+  %0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
+  %1 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
+  func.return %1 : tensor<i32>
+}
+
+// -----
+
+// Test the cases when there are both reads and writes, the order of the reads and writes are preserved.
+
+// CHECK-LABEL: @init
+func.func @init() -> tensor<i32> {
+  // CHECK-NOT: tf.VariableV2
+
+  // CHECK: [[zero:%.*]] = "tf.Const"
+  // CHECK-SAME: dense<0>
+  %0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+
+  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"
+  // CHECK-SAME: shared_name = "x"
+  // CHECK-NEXT: "tf.AssignVariableOp"([[handle]], [[zero]])
+  // CHECK-NEXT: "tf.ReadVariableOp"([[handle]])
+  %1 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
+  %2 = "tf.Assign"(%1, %0) {T = i32, device = "", use_locking = true, validate_shape = true} : (tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<!tf_type.int32ref>
+  %3 = "tf.Identity"(%1) : (tensor<!tf_type.int32ref>) -> tensor<i32>
+
+  // CHECK: [[one:%.*]] = "tf.Const"
+  // CHECK-SAME: dense<1>
+  // CHECK-NEXT: "tf.AssignVariableOp"([[handle]], [[one]])
+  // CHECK-NEXT: "tf.ReadVariableOp"([[handle]])
+  %4 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %5 = "tf.Assign"(%1, %4) {T = i32, device = "", use_locking = true, validate_shape = true} : (tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<!tf_type.int32ref>
+  %6 = "tf.Identity"(%1) : (tensor<!tf_type.int32ref>) -> tensor<i32>
+
+  func.return %6 : tensor<i32>
+}
+
+// CHECK-LABEL: @inference
+func.func @inference() -> (tensor<i32>, tensor<i32>, tensor<i32>) {
+  // CHECK-NOT: tf.VariableV2
+
+  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"
+  // CHECK-SAME: shared_name = "x"
+  // CHECK: "tf.ReadVariableOp"([[handle]])
+  // CHECK: "tf.ReadVariableOp"([[handle]])
+  // CHECK: "tf.ReadVariableOp"([[handle]])
+  %0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
+  %1 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
+  %2 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
+  %3 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
+  func.return %1, %2, %3 : tensor<i32>, tensor<i32>, tensor<i32>
+}
+
+// -----
+
+// Test report error when the shared_name of the tf.VariableV2 op is empty.
+
+// CHECK-LABEL: @inference
+func.func @inference() -> tensor<i32> {
+  // expected-error @+1 {{unable to convert reference variables with empty shared_names.}}
+  %0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = ""} : () -> tensor<!tf_type.int32ref>
+  %1 = "tf.Identity"(%0) : (tensor<!tf_type.int32ref>) -> tensor<i32>
+  func.return %1 : tensor<i32>
+}
+
+// -----
+
+// Test conversion when the user is a side-effect-free op.
+
+// CHECK-LABEL: @side_effect_free_user
+func.func @side_effect_free_user() -> tensor<2xi32> {
+  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"
+  // CHECK-SAME: shared_name = "x"
+  // CHECK: [[value0:%.*]] = "tf.ReadVariableOp"([[handle]])
+  // CHECK: [[value1:%.*]] = "tf.ReadVariableOp"([[handle]])
+  // CHECK: "tf.ConcatV2"([[value1]], [[value0]]
+  // CHECK-SAME: (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<2xi32>
+  %axis = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
+  %1 = "tf.ConcatV2"(%0, %0, %axis) : (tensor<!tf_type.int32ref>, tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<2xi32>
+  func.return %1 : tensor<2xi32>
+}

tensorflow/compiler/mlir/tfrt/tests/cross_device_transfer.mlir

@@ -0,0 +1,45 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -tfrt-cross-device-transfer %s | FileCheck %s
+
+// CHECK-LABEL: test_transfer_op_result
+func.func @test_transfer_op_result(%arg0: !tfrt.chain) -> () {
+  // CHECK-NEXT: %[[RESULT_0:.*]] = corert.get_op_handler %[[ARG_0:.*]] "cpu"
+  %0 = corert.get_op_handler %arg0 "cpu"
+  // CHECK-NEXT: %[[RESULT_1:.*]] = corert.get_op_handler %[[ARG_0]] "gpu"
+  %1 = corert.get_op_handler %arg0 "gpu"
+  // CHECK-NEXT: %[[RESULT_2:.*]] = corert.create_dense_tensor.i32 {shape = [0], value = []}
+  %2 = corert.create_dense_tensor.i32 {shape = [0], value = []}
+  // CHECK-NEXT: %[[RESULT_3:.*]] = corert.executeop(%[[RESULT_0]]) "tf.AddV2"(%[[RESULT_2]], %[[RESULT_2]])
+  %3 = corert.executeop(%0) "tf.AddV2"(%2, %2) {T = f32, device = "/device:CPU:0"} : 1
+  // CHECK-NEXT: %[[RESULT_4:.*]] = tfrt.get_device %[[ARG_0]] {device_name = "/device:GPU:0"}
+  // CHECK-NEXT: %[[RESULT_5:.*]] = corert.get_dst_tensor_type %[[RESULT_3]], %[[RESULT_4]]
+  // CHECK-NEXT: %[[RESULT_6:.*]] = corert.transfer %[[RESULT_3]], %[[RESULT_4]], %[[RESULT_5]]
+  // CHECK-NEXT: %[[RESULT_7:.*]] = corert.executeop(%[[RESULT_1]]) "tf.AddV2"(%[[RESULT_6]], %[[RESULT_6]])
+  %4 = corert.executeop(%1) "tf.AddV2"(%3, %3) {T = f32, device = "/device:GPU:0"} : 1
+  tfrt.return
+}
+
+// CHECK: func @test_transfer_func_arg(%[[ARG_0:.*]]: !tfrt.chain, %[[ARG_1:.*]]: !corert.tensorhandle
+func.func @test_transfer_func_arg(%arg0: !tfrt.chain, %arg1: !corert.tensorhandle {tfrt.device = "/device:CPU:0"}) -> () {
+  // CHECK-NEXT: %[[RESULT_0:.*]] = corert.get_op_handler %[[ARG_0]] "gpu"
+  %0 = corert.get_op_handler %arg0 "gpu"
+  // CHECK-NEXT: %[[RESULT_1:.*]] = tfrt.get_device %[[ARG_0]] {device_name = "/device:GPU:0"}
+  // CHECK-NEXT: %[[RESULT_2:.*]] = corert.get_dst_tensor_type %[[ARG_1]], %[[RESULT_1]]
+  // CHECK-NEXT: %[[RESULT_3:.*]] = corert.transfer %[[ARG_1]], %[[RESULT_1]], %[[RESULT_2]]
+  // CHECK-NEXT: %[[RESULT_4:.*]] = corert.executeop(%[[RESULT_0]]) "tf.AddV2"(%[[RESULT_3]], %[[RESULT_3]])
+  %1 = corert.executeop(%0) "tf.AddV2"(%arg1, %arg1) {T = f32, device = "/device:GPU:0"} : 1
+  tfrt.return
+}

tensorflow/compiler/mlir/tfrt/tests/deduplicate_if_results.mlir

@@ -0,0 +1,61 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tfrt-deduplicate-if-result %s | FileCheck %s -dump-input=fail
+
+func.func private @then(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
+  func.return %x, %x : tensor<i32>, tensor<i32>
+}
+
+func.func private @else(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
+  func.return %y, %y : tensor<i32>, tensor<i32>
+}
+
+// CHECK-LABEL: then/tfrt_dedup_results
+// CHECK: return {{%.*}} : tensor<i32>
+
+// CHECK-LABEL: else/tfrt_dedup_results
+// CHECK: return {{%.*}} : tensor<i32>
+
+// CHECK-LABEL: @basic
+func.func @basic(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
+  // CHECK-NEXT: [[r:%.*]] = "tf.If"
+  // CHECK-NEXT: return [[r]], [[r]] : tensor<i32>, tensor<i32>
+  %0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @else, then_branch = @then, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
+  return %0, %1 : tensor<i32>, tensor<i32>
+}
+
+// -----
+
+func.func private @unmatched_then(%x: tensor<*xi32>, %y: tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>) {
+  func.return %x, %x : tensor<*xi32>, tensor<*xi32>
+}
+
+func.func private @unmatched_else(%x: tensor<i32>, %y: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
+  func.return %y, %y : tensor<i32>, tensor<i32>
+}
+
+// CHECK-LABEL: unmatched_then/tfrt_dedup_results
+// CHECK: return {{%.*}} : tensor<*xi32>
+
+// CHECK-LABEL: unmatched_else/tfrt_dedup_results
+// CHECK: return {{%.*}} : tensor<i32>
+
+// CHECK-LABEL: @unmatched_then_else_type
+func.func @unmatched_then_else_type(%cond: tensor<i1>, %x: tensor<*xi32>, %y: tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>) {
+  // CHECK-NEXT: [[r:%.*]] = "tf.If"
+  // CHECK-NEXT: return [[r]], [[r]] : tensor<*xi32>, tensor<*xi32>
+  %0, %1 = "tf.If"(%cond, %x, %y) {else_branch = @unmatched_else, then_branch = @unmatched_then, is_stateless = true} : (tensor<i1>, tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>)
+  return %0, %1 : tensor<*xi32>, tensor<*xi32>
+}

tensorflow/compiler/mlir/tfrt/tests/fuse_tpu_compile_and_execute_ops.mlir

@@ -0,0 +1,190 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -verify-diagnostics -split-input-file -tfrt-fuse-tpu-compile-and-execute-ops -canonicalize %s | FileCheck %s --dump-input=fail --dump-input-filter=all
+
+module attributes {tf_saved_model.semantics} {
+
+// Test fusing _TPUCompileMlirOp and TPUExecuteOp into TPUCompileMlirAndExecuteOp.
+
+// CHECK-LABEL: func private @test_fuse_tpu_ops
+func.func private @test_fuse_tpu_ops(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<*xi32> {
+  // CHECK-NOT: tf._TPUCompileMlirOp
+  // CHECK-NOT: tf.TPUCompileSucceededAssert
+  // CHECK-NOT: tf.TPUExecuteOp
+
+  // CHECK-NEXT: %0 = "tf.ReadVariableOp"(%arg1)
+  // CHECK:      [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, operands_with_static_shape = [], producer_name = "default"}> : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
+  // CHECK-NEXT: return [[exec_result]] : tensor<*xi32>
+
+  %0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
+  %1 = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
+  %2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
+  %compilation_status, %program = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
+  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
+  %3 = "tf.TPUExecute"(%arg0, %0, %program) {device = "/TPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> tensor<*xi32>
+  func.return %3 : tensor<*xi32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test models using Outside Compilation
+
+// CHECK-LABEL: func private @test_outside_compilation
+func.func private @test_outside_compilation(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<*xi32> {
+  // CHECK-NOT: tf._TPUCompileMlirOp
+  // CHECK-NOT: tf.TPUCompileSucceededAssert
+  // CHECK-NOT: tf.TPUExecuteOp
+
+  // CHECK-NEXT: %0 = "tf.ReadVariableOp"(%arg1)
+  // CHECK:      [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, operands_with_static_shape = [], producer_name = "default"}> : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
+  // CHECK-NEXT: "tf._XlaSendFromHost"(%arg0, %0, [[key]]) <{device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"}> {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> ()
+  // CHECK-NEXT: return [[exec_result]] : tensor<*xi32>
+  %0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
+  %1 = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
+  %2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
+  %compilation_status, %program = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
+  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
+  "tf._XlaSendFromHost"(%arg0, %0, %program) {_xla_has_host_transfer = true, device = "/job:localhost/replica:0/task:0/device:CPU:0", device_ordinal = 0 : i64, key = "host_compute_channel_0_retvals"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> ()
+  %3 = "tf.TPUExecute"(%arg0, %0, %program) {device = "/TPU:0"} : (tensor<*xi32>, tensor<*xi32>, tensor<3x!tf_type.string>) -> tensor<*xi32>
+  func.return %3 : tensor<*xi32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test models with dynamic bounds ops.
+
+// CHECK-LABEL: func private @test_fuse_dynamic_dimension_ops
+func.func private @test_fuse_dynamic_dimension_ops(%arg0: tensor<?x?xi32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<2xi64>, %arg3: tensor<?xi64>, %arg4: tensor<?xi64>) -> tensor<*xi32> {
+  // CHECK-NOT: tf._TPUCompileMlirOp
+  // CHECK-NOT: tf.TPUCompileSucceededAssert
+  // CHECK-NOT: tf.TPUExecuteOp
+  // CHECK-NOT: tf.SetStaticDimensionBounds
+
+  // CHECK: [[read_result:%.*]] = "tf.ReadVariableOp"(%arg1)
+  // CHECK: [[shape_result_1:%.*]] = "tf.Shape"(%arg0) {device = "/CPU:0"} : (tensor<?x?xi32>) -> tensor<?xi64>
+  // CHECK: [[shape_result_2:%.*]] = "tf.Shape"([[read_result]]) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
+  // CHECK: [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, [[shape_result_2]], %0, %0, %arg2, %arg4, %arg3) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 4, 3>, operands_with_static_shape = [0 : i32, 1 : i32, 3 : i32], producer_name = "default"}> : (tensor<?x?xi32>, tensor<?xi64>, tensor<*xi32>, tensor<*xi32>, tensor<2xi64>, tensor<?xi64>, tensor<?xi64>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
+  // CHECK: [[key_1:%.*]], [[exec_result_1:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %2, %0, %1) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 4, 0>, operands_with_static_shape = [], producer_name = "default"}> : (tensor<?x?xi32>, tensor<?xi64>, tensor<*xi32>, tensor<?xi64>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
+  // CHECK-NEXT: return [[exec_result]] : tensor<*xi32>
+  %0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
+  %dyn_arg0 = "tf.SetStaticDimensionBounds" (%arg0, %arg2) :(tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+  %dyn_0 = "tf.SetStaticDimensionBounds" (%0, %arg3) :(tensor<*xi32>, tensor<?xi64>) -> tensor<?xi64>
+  %1 = "tf.Shape"(%dyn_arg0) {device = "/CPU:0"} : (tensor<?x?xi32>) -> tensor<?xi64>
+  %2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
+  %dyn_2 = "tf.SetStaticDimensionBounds" (%2, %arg4) :(tensor<?xi64>, tensor<?xi64>) -> tensor<?xi64>
+  %compilation_status, %program = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
+  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
+  %3 = "tf.TPUExecute"(%dyn_arg0, %dyn_2, %0, %dyn_0, %program) {device = "/TPU:0"} : (tensor<?x?xi32>, tensor<?xi64>, tensor<*xi32>, tensor<?xi64>, tensor<3x!tf_type.string>) -> tensor<*xi32>
+  %compilation_status_2, %program_2 = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
+  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
+  %4 = "tf.TPUExecute"(%arg0, %2, %0, %1, %program_2) {device = "/TPU:0"} : (tensor<?x?xi32>, tensor<?xi64>, tensor<*xi32>, tensor<?xi64>, tensor<3x!tf_type.string>) -> tensor<*xi32>
+  func.return %3 : tensor<*xi32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// CHECK-LABEL: func private @reorder_execute_arg_defining_ops
+// CHECK: tf.VarHandleOp
+// CHECK-NEXT: tf.ReadVariableOp
+// CHECK-NEXT: tf.TPUCompileMlirAndExecute
+func.func private @reorder_execute_arg_defining_ops(%arg0: tensor<1x3xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {
+  %compilation_status, %program = "tf._TPUCompileMlir"() {device = "/CPU:0", metadata = "metadata", mlir_module = "propgram"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
+  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
+  %0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+  %1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+  %2 = "tf.TPUExecute"(%arg0, %1, %program) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x3xf32>, tensor<3x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  return %2 : tensor<1x1xf32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+// CHECK-LABEL: func private @spmd_fuse_mulitple_execute_ops
+// CHECK-NEXT: %0 = "tf.VarHandleOp"()
+// CHECK-NEXT: %1 = "tf.ReadVariableOp"(%0)
+// CHECK-NEXT: %rendezvous_key_base, %results = "tf.TPUCompileMlirAndExecute"(%arg0, %1)
+func.func private @spmd_fuse_mulitple_execute_ops(%arg0: tensor<1x4xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {
+  %cst = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %compilation_status, %program:2 = "tf._TPUCompileMlir"() {device = "/CPU:0", metadata = "metadata", mlir_module = "propgram"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>)
+  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
+  %0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<2x1xf32>>>
+  %1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<2x1xf32>>>) -> tensor<2x1xf32>
+  %2:2 = "tf.Split"(%cst, %arg0) {device = "/CPU:0"} : (tensor<i32>,  tensor<1x4xf32>) -> (tensor<1x2xf32>, tensor<1x2xf32>)
+  %3 = "tf.TPUExecute"(%2#0, %1, %program#0) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x2xf32>, tensor<2x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  %4 = "tf.TPUExecute"(%2#1, %1, %program#1) {_producer_name = "UNKNOWN", device = "/TPU:1"} : (tensor<1x2xf32>, tensor<2x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  return %3 : tensor<1x1xf32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+// CHECK-LABEL: func private @spmd_fuse_mulitple_execute_ops_2
+// CHECK-NEXT: %0 = "tf.VarHandleOp"()
+// CHECK-NEXT: %1 = "tf.ReadVariableOp"(%0)
+// CHECK-NEXT: %rendezvous_key_base, %results = "tf.TPUCompileMlirAndExecute"(%arg0, %1)
+func.func private @spmd_fuse_mulitple_execute_ops_2(%arg0: tensor<1x1xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {
+  %cst = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %compilation_status, %program:2 = "tf._TPUCompileMlir"() {device = "/CPU:0", metadata = "metadata", mlir_module = "propgram"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>)
+  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
+  %0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<2x1xf32>>>
+  %1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<2x1xf32>>>) -> tensor<2x1xf32>
+  %2:2 = "tf.Split"(%cst, %1) {device = "/CPU:0"} : (tensor<i32>,  tensor<2x1xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>)
+  %3 = "tf.TPUExecute"(%arg0, %2#0, %program#0) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  %4 = "tf.TPUExecute"(%arg0, %2#1, %program#1) {_producer_name = "UNKNOWN", device = "/TPU:1"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  return %3 : tensor<1x1xf32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+// CHECK-LABEL: func private @spmd_fuse_split_nd_ops
+// CHECK-NEXT: %0 = "tf.VarHandleOp"()
+// CHECK-NEXT: %1 = "tf.ReadVariableOp"(%0)
+// CHECK-NEXT: %rendezvous_key_base, %results = "tf.TPUCompileMlirAndExecute"(%arg0, %1)
+func.func private @spmd_fuse_split_nd_ops(%arg0: tensor<1x4xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {
+  %cst = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %compilation_status, %program:4 = "tf._TPUCompileMlir"() {device = "/CPU:0", metadata = "metadata", mlir_module = "propgram"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>)
+  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
+  %0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<1x1xf32>>>
+  %1 = "tf.ReadVariableOp"(%0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<1x1xf32>>>) -> tensor<1x1xf32>
+  %2:2 = "tf.Split"(%cst, %arg0) {device = "/CPU:0"} : (tensor<i32>,  tensor<1x4xf32>) -> (tensor<1x2xf32>, tensor<1x2xf32>)
+  %3:2 = "tf.Split"(%cst, %2#0) {device = "/CPU:0"} : (tensor<i32>,  tensor<1x2xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>)
+  %4:2 = "tf.Split"(%cst, %2#1) {device = "/CPU:0"} : (tensor<i32>,  tensor<1x2xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>)
+  %5 = "tf.TPUExecute"(%3#0, %1, %program#0) {_producer_name = "UNKNOWN", device = "/TPU:0"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  %6 = "tf.TPUExecute"(%3#1, %1, %program#1) {_producer_name = "UNKNOWN", device = "/TPU:1"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  %7 = "tf.TPUExecute"(%4#0, %1, %program#2) {_producer_name = "UNKNOWN", device = "/TPU:2"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  %8 = "tf.TPUExecute"(%4#1, %1, %program#3) {_producer_name = "UNKNOWN", device = "/TPU:3"} : (tensor<1x1xf32>, tensor<1x1xf32>, tensor<3x!tf_type.string>) -> tensor<1x1xf32>
+  return %5 : tensor<1x1xf32>
+}
+
+}

tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops.mlir

@@ -0,0 +1,329 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tfrt-lower-tf-savedmodel=hoist-invariant-ops=true %s | FileCheck %s --dump-input=fail --dump-input-filter=all
+
+module attributes {tf_saved_model.semantics} {
+
+// Test hoisting varhandle op.
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor<!tf_type.resource<tensor<i32>>>
+// CHECK: [[x:%.*]] = "tf.ReadVariableOp"([[handle]]) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+// CHECK: "tf._TfrtSetResource"([[x]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()
+
+// CHECK-LABEL: func @test_hoist_varhandleop
+func.func @hoist_varhandleop(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
+  attributes {tf_saved_model.exported_names = ["test_hoist_varhandleop"]} {
+  // CHECK-NOT: tf.VarHandleOp
+  // CHECK-NOT: tf.ReadVariableOp
+  // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [0], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor<i32>
+  // CHECK: [[r:%.*]] = "tf.AddV2"({{.*}}, [[v]]) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: return [[r]]
+  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
+  %x = "tf.ReadVariableOp"(%handle) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  %r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %r : tensor<i32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test hoisting hash table op.
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK: [[handle:%.*]] = "tf.HashTableV2"()
+// CHECK-SAME: shared_name = "x"
+// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
+// CHECK: [[x:%.*]] = "tf.LookupTableSizeV2"([[handle]])
+// CHECK: "tf._TfrtSetResource"([[x]]) <{index = [[size_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i64>) -> ()
+
+// CHECK: func @test_hoist_hash_table
+func.func @hoist_hash_table(%arg: tensor<?x!tf_type.string> {tf_saved_model.index_path = ["input"]}, %default: tensor<i64> {tf_saved_model.index_path = ["default"]}) -> (tensor<i64> {tf_saved_model.index_path = ["r"]}, tensor<*xi64> {tf_saved_model.index_path = ["r1"]})
+  attributes {tf_saved_model.exported_names = ["test_hoist_hash_table"]} {
+  // CHECK-NOT: tf.HashTableV2
+  // CHECK-NOT: tf.LookupTableSizeV2
+  // CHECK: [[v:%.*]]:2 = "tf._TfrtGetResource"() <{container = ["", ""], indices = [0, 1], shared_name = [{{.*}}, {{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
+  // CHECK: [[r:%.*]] = "tf.LookupTableFindV2"([[v]]#[[handle_idx]]
+  // CHECK: return [[v]]#[[size_idx]], [[r]]
+  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "x", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
+  %1 = "tf.LookupTableSizeV2"(%0) {device = ""} : (tensor<!tf_type.resource>) -> tensor<i64>
+  %2 = "tf.LookupTableFindV2"(%0, %arg, %default) {device = "/CPU:0"} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
+  func.return %1, %2 : tensor<i64>, tensor<*xi64>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test hoisting const op.
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/CPU:0"} : () -> tensor<i32>
+// CHECK: [[x:%.*]] = "tf.AddV2"([[const]], [[const]]) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+// CHECK: "tf._TfrtSetResource"([[x]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()
+// CHECK: [[const_1:%.*]] = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = "/CPU:0"} : () -> tensor<i32>
+// CHECK: "tf._TfrtSetResource"([[const_1]]) <{index = 1 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()
+
+// CHECK-LABEL: func @test_hoist_const
+func.func @hoist_const(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
+  attributes {tf_saved_model.exported_names = ["test_hoist_const"]} {
+  // CHECK-NOT: tf.Const
+  // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [0], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor<i32>
+  // CHECK-NEXT: "tf.AddV2"({{.*}}, [[v]]) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK-NEXT: return
+  %const = "tf.Const"() {device = "/CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %x = "tf.AddV2"(%const, %const) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %r : tensor<i32>
+}
+
+// CHECK-LABEL: func @test_hoist_const_return
+func.func @hoist_const_return(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
+  attributes {tf_saved_model.exported_names = ["test_hoist_const_return"]} {
+  // CHECK-NOT: tf.Const
+  // CHECK: [[v:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [1], shared_name = [""]}> {device = "/CPU:0"} : () -> tensor<i32>
+  // CHECK-NEXT: return [[v]]
+  %const = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  func.return %const : tensor<i32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test not hoisting `tf.BatchFunction`.
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> {device = "/CPU:0"} : () -> tensor<1xi32>
+// CHECK: "tf._TfrtSetResource"([[const]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor<1xi32>) -> ()
+
+// CHECK-LABEL: func.func private @func_with_batch_function
+func.func private @func_with_batch_function() -> tensor<*xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "StatefulPartitionedCall:0"}} {
+  // CHECK:  "tf._TfrtGetResource"()
+  %cst = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> {device = "/CPU:0"} : () -> tensor<1xi32>
+  // CHECK:  "tf.BatchFunction"
+  %0 = "tf.BatchFunction"(%cst) <{allowed_batch_sizes = [1], batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", enable_large_batch_splitting = true, f = @_batched, low_priority_allowed_batch_sizes = [], low_priority_batch_timeout_micros = 0 : i64, low_priority_max_batch_size = 0 : i64, low_priority_max_enqueued_batches = 0 : i64, max_batch_size = 1 : i64, max_enqueued_batches = 1 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch_function___inference_signature_wrapper_fn_with_defaults_36"}> {device = "/CPU:0"} : (tensor<1xi32>) -> tensor<*xi32>
+  return %0 : tensor<*xi32>
+}
+func.func private @_batched(%arg0: tensor<1xi32>) -> tensor<1xi32> {
+  return %arg0 : tensor<1xi32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test hoisting write side-effect ops.
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> tensor<i32>
+// CHECK: "tf._TfrtSetResource"([[const]]) <{index = [[const_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> ()
+// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor<!tf_type.resource<tensor<i32>>>
+// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_idx:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>) -> ()
+
+// CHECK: func @test_hoist_var_read_write
+func.func @hoist_var_read_write() -> (tensor<i32> {tf_saved_model.index_path = ["x"]}, tensor<i32> {tf_saved_model.index_path = ["r"]})
+  attributes {tf_saved_model.exported_names = ["test_hoist_var_read_write"]} {
+  // CHECK-NOT: tf.Const
+  // CHECK-NOT: tf.VarHandleOp
+  // CHECK: [[v:%.*]]:2 = "tf._TfrtGetResource"() <{container = ["", ""], indices = [0, 1], shared_name = [{{.*}}, {{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> ({{.*}})
+  // CHECK: [[x:%.*]] = "tf.ReadVariableOp"([[v]]#[[handle_idx]]) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  // CHECK-NEXT: "tf.AssignVariable"([[v]]#[[handle_idx]], [[v]]#[[const_idx]]) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
+  // CHECK-NEXT: [[r:%.*]] = "tf.ReadVariableOp"([[v]]#[[handle_idx]]) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  // CHECK-NEXT: return [[x]], [[r]]
+  %const = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
+  %x = "tf.ReadVariableOp"(%handle) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  "tf.AssignVariable"(%handle, %const) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
+  %r = "tf.ReadVariableOp"(%handle) {device = "/CPU:0", dtype = i32} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  func.return %x, %r : tensor<i32>, tensor<i32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test not hoisting read variable op that used by control flow ops if var handle op and read variable op are separated, but still hoists const ops and var handle ops.
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK: [[handle:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor<!tf_type.resource<tensor<i32>>>
+// CHECK: "tf._TfrtSetResource"([[handle]])
+// CHECK-SAME: index = [[handle_index:.*]]
+// CHECK: [[handle1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "x"}> : () -> tensor<!tf_type.resource<tensor<i32>>>
+// CHECK: "tf._TfrtSetResource"([[handle1]])
+// CHECK-SAME: index = [[handle1_index:.*]]
+// CHECK: [[const:%.*]] = "tf.Const"() <{value = dense<true> : tensor<i1>}> {device = "/CPU:0"} : () -> tensor<i1>
+// CHECK: "tf._TfrtSetResource"([[const]])
+// CHECK-SAME: index = [[const_index:.*]]
+func.func private @some_func(
+    %arg: tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32> {
+  // CHECK: tf.ReadVariableOp
+  %0 = "tf.ReadVariableOp"(%arg) {device = "cpu"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  func.return %0 : tensor<i32>
+}
+
+// CHECK-LABEL: func @test_not_hoist_stateful_call
+func.func @not_hoist_stateful_call(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
+  attributes {tf_saved_model.exported_names = ["test_not_hoist_stateful_call"]} {
+  // CHECK-NOT: tf.VarHandleOp
+  // CHECK:  "tf._TfrtGetResource"()
+  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
+  // CHECK: tf.StatefulPartitionedCall
+  %x = "tf.StatefulPartitionedCall"(%handle) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<i32>)
+  %r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %r : tensor<i32>
+}
+
+// CHECK-LABEL: func @test_not_hoist_if
+func.func @not_hoist_if(%arg: tensor<i32> {tf_saved_model.index_path = ["input"]}) -> (tensor<i32> {tf_saved_model.index_path = ["r"]})
+  attributes {tf_saved_model.exported_names = ["test_not_hoist_if"]} {
+  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
+  // CHECK-NOT: tf.Const
+  // CHECK:  "tf._TfrtGetResource"()
+  %cond = "tf.Const"() {device = "/CPU:0", value = dense<true> : tensor<i1>} : () -> tensor<i1>
+  // CHECK: tf.If
+  %x = "tf.If"(%cond, %handle) {then_branch = @some_func, else_branch = @some_func, is_stateless = false} : (tensor<i1>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  %r = "tf.AddV2"(%arg, %x) {device = "/CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %r : tensor<i32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test hoist var handle op and read variable op in the batch function.
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32>
+  attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
+  // CHECK-NOT: tf.VarHandleOp
+  // CHECK-NOT: tf.ReadVariableOp
+  // CHECK:  "tf._TfrtGetResource"()
+  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+  %1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
+  %2 = "tf.AddV2"(%arg0, %1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
+  %3 = "tf.Identity"(%2) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %3 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]}) 
+  attributes {tf_saved_model.exported_names = ["main"]} {
+  // CHECK-NOT: tf.VarHandleOp
+  // CHECK:  "tf._TfrtGetResource"()
+  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+  // CHECK: "tf.BatchFunction"(%arg0, %0)
+  // CHECK: operandSegmentSizes = array<i32: 1, 1>
+  %1 = "tf.BatchFunction"(%arg0, %0) {allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = "batch/"} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test not hoisting callees in init functions.
+
+"tf_saved_model.session_initializer"() {initializers = [@init]} : () -> ()
+
+func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init"]} {
+  %var0 = "tf.VarHandleOp"() {container = "", shared_name = "var0"} : () -> tensor<!tf_type.resource<tensor<i1>>>
+  %cond = "tf.ReadVariableOp"(%var0) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i1>>>) -> tensor<i1>
+  %x = "tf.StatefulPartitionedCall"(%cond) {device = "/CPU:0", config = "", config_proto = "", executor_type = "", f = @some_func} : (tensor<i1>) -> (tensor<i32>)
+  %var1 = "tf.VarHandleOp"() {container = "", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<i32>>>
+  "tf.AssignVariable"(%var1, %x) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
+  func.return
+}
+
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK-NEXT: return
+
+// CHECK-LABEL: func private @some_func
+func.func private @some_func(%arg: tensor<i1>) -> tensor<i32> {
+  // CHECK-NOT: tf._TfrtGetResource
+  %const = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32> } : () -> tensor<i32>
+  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
+  %0 = "tf.ReadVariableOp"(%handle) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  %r = "tf.SelectV2"(%arg, %const, %0) {device = "/CPU:0"} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %r : tensor<i32>
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test not hoisting callees in xla launch functions.
+
+// CHECK-LABEL: func private @xla_func
+func.func private @xla_func(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32>
+  attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
+  // CHECK-NOT: tf._TfrtGetResource
+  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+  %1 = "tf.ReadVariableOp"(%0) {device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
+  %2 = "tf.AddV2"(%arg0, %1) {device = "/device:CPU:0"} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
+  %3 = "tf.Identity"(%2) {device = "/device:CPU:0"} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %3 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32> {tf_saved_model.index_path = ["input"]}) -> (tensor<*xf32> {tf_saved_model.index_path = ["r"]}) 
+  attributes {tf_saved_model.exported_names = ["main"]} {
+  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+  %1 = "tf.XlaLaunch"(%arg0, %0) {device = "/device:GPU:0", function = @xla_func, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<1x3xf32>, tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<*xf32>
+  func.return  %1 : tensor<*xf32>
+
+}
+
+}
+
+// -----
+
+module attributes {tf_saved_model.semantics} {
+
+// Test not hoisting in TPU functions.
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK-NEXT: return
+
+// CHECK-LABEL: func private @func2
+func.func private @func2(%arg: tensor<i1>) -> tensor<i32> {
+  // CHECK-NOT: tf._TfrtGetResource
+  "tf.TPUReplicateMetadata"() {_tpu_replicate = "0",  allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [], host_compute_core = [], num_cores_per_replica = 4 : i64, num_replicas = 1 : i64, padding_map = [], step_marker_location = "STEP_MARK_AT_ENTRY", topology = "", tpu_compile_options_proto = "", use_spmd_for_xla_partitioning = true, use_tpu = true} : () -> ()
+  %const = "tf.Const"() {device = "/CPU:0", value = dense<1> : tensor<i32> } : () -> tensor<i32>
+  %handle = "tf.VarHandleOp"() {container = "", shared_name = "x"} : () -> tensor<!tf_type.resource<tensor<i32>>>
+  %0 = "tf.ReadVariableOp"(%handle) {device = "/CPU:0"} : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  %r = "tf.SelectV2"(%arg, %const, %0) {device = "/CPU:0"} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %r : tensor<i32>
+}
+
+}

tensorflow/compiler/mlir/tfrt/tests/hoist_invariant_ops_mlrt.mlir

@@ -0,0 +1,45 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tfrt-lower-tf-savedmodel="hoist-invariant-ops=true fuse-get-resource-ops=false" %s | FileCheck %s --dump-input=fail --dump-input-filter=all
+
+module attributes {tf_saved_model.semantics} {
+
+// Test hoisting hash table op.
+
+// CHECK-LABEL: func @_tfrt_resource_init
+// CHECK: [[handle:%.*]] = "tf.HashTableV2"()
+// CHECK-SAME: shared_name = "x"
+// CHECK: "tf._TfrtSetResource"([[handle]]) <{index = [[handle_id:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
+// CHECK: [[x:%.*]] = "tf.LookupTableSizeV2"([[handle]])
+// CHECK: "tf._TfrtSetResource"([[x]]) <{index = [[size_id:.*]] : i64}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i64>) -> ()
+
+// CHECK: func @test_hoist_hash_table
+func.func @hoist_hash_table(%arg: tensor<?x!tf_type.string> {tf_saved_model.index_path = ["input"]}, %default: tensor<i64> {tf_saved_model.index_path = ["default"]}) -> (tensor<i64> {tf_saved_model.index_path = ["r"]}, tensor<*xi64> {tf_saved_model.index_path = ["r1"]})
+  attributes {tf_saved_model.exported_names = ["test_hoist_hash_table"]} {
+  // CHECK-NOT: tf.HashTableV2
+  // CHECK-NOT: tf.LookupTableSizeV2
+  // CHECK-DAG: [[v0:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [[[handle_id]]], shared_name = [{{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
+  // CHECK-DAG: [[v1:%.*]] = "tf._TfrtGetResource"() <{container = [""], indices = [[[size_id]]], shared_name = [{{.*}}]}> {device = "/job:localhost/replica:0/task:0/device:CPU:0"}
+  // CHECK-DAG: [[r:%.*]] = "tf.LookupTableFindV2"([[v0]]
+  // CHECK-DAG: return [[v1]], [[r]]
+  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "x", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
+  %1 = "tf.LookupTableSizeV2"(%0) {device = ""} : (tensor<!tf_type.resource>) -> tensor<i64>
+  %2 = "tf.LookupTableFindV2"(%0, %arg, %default) {device = "/CPU:0"} : (tensor<!tf_type.resource>, tensor<?x!tf_type.string>, tensor<i64>) -> tensor<*xi64>
+  func.return %1, %2 : tensor<i64>, tensor<*xi64>
+}
+
+}
+
+// -----

tensorflow/compiler/mlir/tfrt/tests/ifrt/BUILD

@@ -0,0 +1,21 @@
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+)

tensorflow/compiler/mlir/tfrt/tests/ifrt/lower_to_ifrt_restore_variable.mlir

@@ -0,0 +1,322 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -verify-diagnostics -lower-to-ifrt-restore-variable %s | FileCheck %s
+
+
+// -----
+// single variable
+
+// CHECK-LABEL:   func.func @restore_single() {
+// CHECK-NEXT:     [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:     [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
+// CHECK-SAME:        {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
+// CHECK-NOT:       "tf.RestoreV2"
+// CHECK-NEXT:     return
+
+module {
+  func.func @restore_single() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// single variable: VarHandleOp is before RestoreV2
+
+// CHECK-LABEL:   func.func @varhandle_before_restore() {
+// CHECK-NEXT:     [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:     [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
+// CHECK-SAME:        {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
+// CHECK-NOT:       "tf.RestoreV2"
+// CHECK-NEXT:     return
+
+module {
+  func.func @varhandle_before_restore() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+
+// -----
+// multiple variables
+
+// CHECK-LABEL:   func.func @restore_multiple() {
+// CHECK-NEXT:     [[PREFIX:%.*]] = "tf.Const"()
+// CHECK-NEXT:     [[SLICE:%.*]] = "tf.Const"()
+// CHECK-NEXT:     [[NAME:%.*]] = "tf.Const"()
+// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:     [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "x", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]], [[HANDLEZ]])
+// CHECK-SAME:        {restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}
+// CHECK-NOT:       "tf.RestoreV2"
+// CHECK-NEXT:     return
+
+module {
+  func.func @restore_multiple() {
+    %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<3x1xf32>, tensor<1x3xf32>)
+    %1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    %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
+  }
+}
+
+// -----
+// Restored variable is not assigned with a name is an error.
+
+module {
+  func.func @unassigned_restore_return_error() {
+    %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>
+    //expected-error@below {{'tf.RestoreV2' op expects 2 valid users, but got 1}}
+    %0:2 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<3x1xf32>, tensor<1x3xf32>)
+    %1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+// -----
+//  Restored tensor is consumed by an op other than AssignVariableOp, it is returned as an output by IfrtRestoreVariableOp.
+
+// CHECK-LABEL: func.func @restore_with_consumer() {
+// CHECK-NEXT:   %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:   %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:   %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:   %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:   %1 = "tf.IfrtRestoreVariableOp"(%cst, %cst_1, %cst_0, %0) <{restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: false>}> : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+// CHECK-NEXT:   %2 = "tf.ReluOp"(%1) : (tensor<3x1xf32>) -> tensor<3x1xf32>
+// CHECK-NEXT:   %3 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-NEXT:   "tf.AssignVariableOp"(%3, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+// CHECK-NOT: "tf.RestoreV2"
+// CHECK-NEXT: return
+
+
+module {
+  func.func @restore_with_consumer() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %2 = "tf.ReluOp"(%0) : (tensor<3x1xf32>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "x", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+
+
+// -----
+// variable with cast
+// CHECK-LABEL:   func.func @restore_with_cast() {
+// CHECK-NEXT:     [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:     [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
+// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
+// CHECK-SAME:        {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: false>}
+// CHECK-NOT:       "tf.RestoreV2"
+// CHECK-NEXT:     return
+
+module {
+  func.func @restore_with_cast() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.Cast"(%0) <{Truncate = false}> : (tensor<3x1xf32>) -> tensor<3x1xbf16>
+    %2 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
+    "tf.AssignVariableOp"(%2, %1) : (tensor<!tf_type.resource<tensor<3x1xbf16>>>, tensor<3x1xbf16>) -> ()
+    return
+  }
+}
+
+// -----
+// variable and table lookup
+// CHECK-LABEL: func.func @restore_var_and_table()
+// CHECK-NEXT:  [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"model/foo"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:  [[NAMES:%.*]] = "tf.Const"() <{value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+// CHECK-NEXT:  [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+// CHECK-NEXT:  [[VAR1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "var1"}> : () -> tensor<!tf_type.resource<tensor<*xi32>>>
+// CHECK-NEXT:  [[TABLE1_KEYS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_keys"}> : () -> tensor<!tf_type.resource<tensor<*xi64>>>
+// CHECK-NEXT:  [[TABLE1_VALS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_vals"}> : () -> tensor<!tf_type.resource<tensor<*xf32>>>
+// CHECK-NEXT:  [[RETURNED_VALUES:%.*]]:2 = "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAMES]], [[SLICE]], [[VAR1]], [[TABLE1_KEYS]], [[TABLE1_VALS]])
+// CHECK-SAME:     <{restored_dtypes = [i32, i64, f32], returned_tensor_names = ["table1_keys", "table1_vals"], truncate_in_cast = array<i1: false, false, false>}> : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<!tf_type.resource<tensor<*xi32>>>, tensor<!tf_type.resource<tensor<*xi64>>>, tensor<!tf_type.resource<tensor<*xf32>>>) -> (tensor<*xi64>, tensor<*xf32>)
+// CHECK-NEXT:  [[TABLE1:%.*]] = "tf.HashTableV2"() <{container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32}> : () -> tensor<!tf_type.resource>
+// CHECK-NEXT:  "tf.LookupTableImportV2"([[TABLE1]], [[RETURNED_VALUES]]#0, [[RETURNED_VALUES]]#1) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NEXT:  return
+module {
+  func.func @restore_var_and_table() {
+    %cst = "tf.Const"() {value = dense<"model/foo"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() {value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
+    %cst_1 = "tf.Const"() {value = dense<["", "", ""]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
+    %0:3 = "tf.RestoreV2"(%cst, %cst_0, %cst_1) {dtypes = [i32, i64, f32]} : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi64>, tensor<*xf32>)
+    %1 = "tf.VarHandleOp"() {container = "", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<*xi32>>>
+    "tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<*xi32>>>, tensor<*xi32>) -> ()
+    %2 = "tf.HashTableV2"() {container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32} : () -> tensor<!tf_type.resource>
+    "tf.LookupTableImportV2"(%2, %0#1, %0#2) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// variable and dense table lookup
+// CHECK-LABEL: func.func @restore_var_and_dense_table()
+// CHECK-NEXT:  [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"model/foo"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+// CHECK-NEXT:  [[NAMES:%.*]] = "tf.Const"() <{value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+// CHECK-NEXT:  [[SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<3x!tf_type.string>}> : () -> tensor<3x!tf_type.string>
+// CHECK-NEXT:  [[VAR1:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "var1"}> : () -> tensor<!tf_type.resource<tensor<*xi32>>>
+// CHECK-NEXT:  [[TABLE1_KEYS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_keys"}> : () -> tensor<!tf_type.resource<tensor<*xi64>>>
+// CHECK-NEXT:  [[TABLE1_VALS:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "table1_vals"}> : () -> tensor<!tf_type.resource<tensor<*xf32>>>
+// CHECK-NEXT:  [[RETURNED_VALUES:%.*]]:2 = "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAMES]], [[SLICE]], [[VAR1]], [[TABLE1_KEYS]], [[TABLE1_VALS]])
+// CHECK-SAME:     <{restored_dtypes = [i32, i64, f32], returned_tensor_names = ["table1_keys", "table1_vals"], truncate_in_cast = array<i1: false, false, false>}> : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>, tensor<!tf_type.resource<tensor<*xi32>>>, tensor<!tf_type.resource<tensor<*xi64>>>, tensor<!tf_type.resource<tensor<*xf32>>>) -> (tensor<*xi64>, tensor<*xf32>)
+// CHECK-NEXT:  [[EMPTY_KEY:%.*]] = "tf.Const"() <{value = dense<-1> : tensor<i64>}> : () -> tensor<i64>
+// CHECK-NEXT:  [[DEFAULT_VALUE:%.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<f32>}> : () -> tensor<f32>
+// CHECK-NEXT:  [[TABLE1:%.*]] = "tf.MutableDenseHashTableV2"(%cst_2, %cst_3) <{container = "", shared_name = "table1", value_dtype = f32}> {key_dtype = i64} : (tensor<i64>, tensor<f32>) -> tensor<!tf_type.resource>
+// CHECK-NEXT:  "tf.LookupTableImportV2"([[TABLE1]], [[RETURNED_VALUES]]#0, [[RETURNED_VALUES]]#1) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NEXT: return
+module {
+  func.func @restore_var_and_dense_table() {
+    %cst = "tf.Const"() {value = dense<"model/foo"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() {value = dense<["var1", "table1_keys", "table1_vals"]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
+    %cst_1 = "tf.Const"() {value = dense<["", "", ""]> : tensor<3x!tf_type.string>} : () -> tensor<3x!tf_type.string>
+    %0:3 = "tf.RestoreV2"(%cst, %cst_0, %cst_1) {dtypes = [i32, i64, f32]} : (tensor<!tf_type.string>, tensor<3x!tf_type.string>, tensor<3x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi64>, tensor<*xf32>)
+    %1 = "tf.VarHandleOp"() {container = "", shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<*xi32>>>
+    "tf.AssignVariableOp"(%1, %0#0) : (tensor<!tf_type.resource<tensor<*xi32>>>, tensor<*xi32>) -> ()
+    %empty_key = "tf.Const"() {value = dense<-1> : tensor<i64>} : () -> tensor<i64>
+    %default_value = "tf.Const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
+    %2 = "tf.MutableDenseHashTableV2"(%empty_key, %default_value) {container = "", key_dtype = i64, shared_name = "table1", value_dtype = f32} : (tensor<i64>, tensor<f32>) -> tensor<!tf_type.resource>
+    "tf.LookupTableImportV2"(%2, %0#1, %0#2) : (tensor<!tf_type.resource>, tensor<*xi64>, tensor<*xf32>) -> ()
+    return
+  }
+}
+
+
+// -----
+// restored variable assigned to one var handle and then read and assigned to another
+// CHECK-LABEL: func.func @restore_and_copy() {
+// CHECK-DAG:   [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}>
+// CHECK-DAG:   [[VAR_Y:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[VAR_Z:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
+// CHECK-DAG:   [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
+// CHECK:       "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[VAR_Y]], [[VAR_Z]])
+// CHECK-SAME:  <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}>
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NOT:   "tf.ReadVariableOp"
+// CHECK-NEXT:  return
+
+module {
+  func.func @restore_and_copy() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    %read = "tf.ReadVariableOp"(%1) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %2 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%2, %read) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// restored variable assigned to one var handle and then read and assigned to another in a different block/function
+// CHECK-LABEL: func.func @restore_derived() {
+// CHECK-DAG:   [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_variables"> : tensor<!tf_type.string>}>
+// CHECK-DAG:   [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
+// CHECK-DAG:   [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
+// CHECK:       "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[HANDLEY]], [[HANDLEZ]])
+// CHECK-SAME:  <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, false>}>
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NEXT:  return
+
+module {
+  func.func @restore_derived() {
+    %cst = "tf.Const"() <{value = dense<"restore_variables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+
+  func.func @derived_init() {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %2 = "tf.Identity"(%1) : (tensor<3x1xf32>) -> tensor<3x1xf32>
+    %3 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%3, %2) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    return
+  }
+}
+
+// -----
+// restored variable assigned to one var handle and then read, cast, and assigned to another
+// CHECK-LABEL: func.func @restore_and_cast_copy() {
+// CHECK-DAG:   [[PREFIX:%.*]] = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}>
+// CHECK-DAG:   [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+// CHECK-DAG:   [[HANDLEZ:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
+// CHECK-DAG:   [[NEW_NAME:%.*]] = "tf.Const"() <{value = dense<"y"> : tensor<2x!tf_type.string>}>
+// CHECK-DAG:   [[NEW_SLICE:%.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
+// CHECK:       "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NEW_NAME]], [[NEW_SLICE]], [[HANDLEY]], [[HANDLEZ]])
+// CHECK-SAME:  <{restored_dtypes = [f32, f32], returned_tensor_names = [], truncate_in_cast = array<i1: false, true>}>
+// CHECK-NOT:   "tf.RestoreV2"
+// CHECK-NOT:   "tf.ReadVariableOp"
+// CHECK-NEXT:  return
+module {
+  func.func @restore_and_cast_copy() {
+    %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+    %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+    %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+    %read = "tf.ReadVariableOp"(%1) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %cast = "tf.Cast"(%read) <{Truncate = true}> : (tensor<3x1xf32>) -> tensor<3x1xbf16>
+    %2 = "tf.VarHandleOp"() <{container = "", shared_name = "z"}> : () -> tensor<!tf_type.resource<tensor<3x1xbf16>>>
+    "tf.AssignVariableOp"(%2, %cast) : (tensor<!tf_type.resource<tensor<3x1xbf16>>>, tensor<3x1xbf16>) -> ()
+    return
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/propagate_static_shapes.mlir

@@ -0,0 +1,104 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -propagate-static-shapes %s | FileCheck %s
+
+// -----
+// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 123 : i64}
+// CHECK: return %arg0
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
+// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 123 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: return %[[C1]]
+
+module {
+  func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 123 : i64} {
+    func.return %arg0 : tensor<?x?xi32>
+  }
+  func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
+    %0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
+    %1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %2 = "tf.IfrtCall"(%1) {program_id = 123 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    func.return %2 : tensor<?x?xi32>
+  }
+}
+
+// -----
+// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 3 : i32}, %arg1: tensor<?xi32> {tf._static_shape_arg_idx = 4 : i32}, %arg2: tensor<?x?xf32>, %arg3: tensor<2xi64>, %arg4: tensor<1xi64>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) attributes {tfrt_ifrt_serving.program_id = 456 : i64}
+// CHECK: return %arg0, %arg1, %arg2
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
+// CHECK-NEXT: %[[C1:.*]] = "tf.Const"
+// CHECK-NEXT: %[[R:.*]]:3 = "tf.IfrtCall"(%arg0, %arg1, %arg2, %[[C0]], %[[C1]]) <{operandSegmentSizes = array<i32: 3, 2>, program_id = 456 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>, tensor<2xi64>, tensor<1xi64>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>)
+// CHECK-NEXT: return %[[R]]#0, %[[R]]#1, %[[R]]#2
+
+module {
+  func.func @callee(%arg0: tensor<?x?xi32>, %arg1: tensor<?xi32>, %arg2: tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) attributes {tfrt_ifrt_serving.program_id = 456 : i64} {
+    func.return %arg0, %arg1, %arg2 : tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>
+  }
+  func.func @main(%arg0: tensor<?x?xi32>, %arg1: tensor<?xi32>, %arg2: tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) {
+    %c0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
+    %c1 = "tf.Const"() {value = dense<4> : tensor<1xi64>} : () -> tensor<1xi64>
+    %0 = "tf.SetStaticDimensionBounds"(%arg0, %c0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %1 = "tf.SetStaticDimensionBounds"(%arg1, %c1) : (tensor<?xi32>, tensor<1xi64>) -> tensor<?xi32>
+    %2:3 = "tf.IfrtCall"(%0, %1, %arg2) {program_id = 456 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 3, 0>} : (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>) -> (tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>)
+    func.return %2#0, %2#1, %2#2 : tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xf32>
+  }
+}
+
+// -----
+// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 789 : i64}
+// CHECK: return %arg0
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
+// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 789 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: %[[C2:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 789 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: return %[[C2]]
+
+module {
+  func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 789 : i64} {
+    func.return %arg0 : tensor<?x?xi32>
+  }
+  func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
+    %0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
+    %1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %2 = "tf.IfrtCall"(%1) {program_id = 789 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    %3 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %4 = "tf.IfrtCall"(%3) {program_id = 789 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    func.return %4 : tensor<?x?xi32>
+  }
+}
+
+// -----
+// CHECK-LABEL: func.func @callee(%arg0: tensor<?x?xi32> {tf._static_shape_arg_idx = 1 : i32}, %arg1: tensor<2xi64>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 999 : i64}
+// CHECK: return %arg0
+// CHECK-LABEL: func.func @main
+// CHECK-NEXT: %[[C0:.*]] = "tf.Const"
+// CHECK-NEXT: %[[C1:.*]] = "tf.IfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 999 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: %[[C2:.*]] = "tf.AsyncIfrtCall"(%arg0, %[[C0]]) <{operandSegmentSizes = array<i32: 1, 1>, program_id = 999 : i64, variable_arg_indices = []}> : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+// CHECK-NEXT: return %[[C2]]
+
+module {
+  func.func @callee(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> attributes {tfrt_ifrt_serving.program_id = 999 : i64} {
+    func.return %arg0 : tensor<?x?xi32>
+  }
+  func.func @main(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
+    %0 = "tf.Const"() {value = dense<[1, 2]> : tensor<2xi64>} : () -> tensor<2xi64>
+    %1 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %2 = "tf.IfrtCall"(%1) {program_id = 999 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    %3 = "tf.SetStaticDimensionBounds"(%arg0, %0) : (tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
+    %4 = "tf.AsyncIfrtCall"(%3) {program_id = 999 : i64, variable_arg_indices = [], operandSegmentSizes = array<i32: 1, 0>} : (tensor<?x?xi32>) -> tensor<?x?xi32>
+    func.return %4 : tensor<?x?xi32>
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/rewrite_cluster_to_ifrt_call.mlir

@@ -0,0 +1,140 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -rewrite-cluster-to-ifrt-call %s | FileCheck %s
+// TODO(b/316226111): the printer may not guarantee the same order of fields. Rewrite the checks to be less sensitive to proto serialization formats.
+// -----
+// Non-SPMD: one input and one output
+//
+// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+// CHECK-NEXT:  "tf.IfrtCall"(%arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 1, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}
+// CHECK-SAME:       (tensor<1x3xf32>) -> tensor<1x3xf32>
+// CHECK:    return
+//
+// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>)
+// CHECK-SAME: __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { dim { size: 1 } dim { size: 3 } } kind: PARAMETER sharding { } is_bounded_dynamic_dim: false } retvals { sharding { } } num_replicas: 1 num_cores_per_replica: 1 "
+// CHECK-SAME: device_assignment = []
+// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
+// CHECK:      return
+
+module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
+func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
+  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
+  %0 = "tf_device.cluster_func"(%arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = [""], num_cores_per_replica = 1 : i64, device_assignment = [], topology = "", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = false, use_tpu = true} : (tensor<1x3xf32>) -> (tensor<1x3xf32>)
+  return %0 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: @_func
+func.func private @_func(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
+  return %arg0 : tensor<1x3xf32>
+}
+}
+
+
+// -----
+// SPMD: one input and no return
+//
+// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) {
+// CHECK-NEXT:  "tf.IfrtCall"(%arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 1, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}
+// CHECK-SAME:       (tensor<1x3xf32>) -> ()
+// CHECK:    return
+//
+// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>)
+// CHECK-SAME: __tpu_compile_metadata_text = "args { dtype: DT_FLOAT shape { dim { size: 1 } dim { size: 3 } } 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 } num_replicas: 1 num_cores_per_replica: 2 use_spmd_for_xla_partitioning: true "
+// CHECK-SAME: device_assignment = [0, 0, 0, 0, 0, 0, 0, 1]
+// CHECK-SAME: tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
+// CHECK:      return
+
+module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
+func.func @serving_default(%arg0: tensor<1x3xf32>) -> () {
+  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
+  "tf_device.cluster_func"(%arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}"], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>) -> ()
+  return
+}
+
+// CHECK-LABEL: @_func
+func.func private @_func(%arg0: tensor<1x3xf32>) -> () {
+  return
+}
+}
+
+// -----
+// Multiple ifrt calls and have two sharded arguments
+
+// CHECK-LABEL: func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK-NEXT:  %0 = "tf.IfrtCall"(%arg1, %arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
+// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    %1 = "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+// CHECK-NEXT:    %2 = "tf.IfrtCall"(%1, %arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID]] : i64, variable_arg_indices = []
+// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    %3 = "tf.add"(%0, %2) : (tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<1x1xf32>
+// CHECK:    return
+//
+// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-SAME:      device_assignment = [0, 0, 0, 0, 0, 0, 0, 1]
+// CHECK-SAME:      tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
+// CHECK-NEXT:     %0 = "tf.MatMul"(%arg0, %arg1)
+// CHECK:          return
+
+module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
+func.func @serving_default(%arg0: tensor<3x1xf32>,  %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
+  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
+  %outputs_0 = "tf_device.cluster_func"(%arg1, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  %duplicate_arg =  "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  %outputs_1 = "tf_device.cluster_func"(%duplicate_arg, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], topology = "\0A\04\01\01\01\02\10\01\18\02\22\08\00\00\00\00\00\00\00\01", output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  %outputs_2 = "tf.add"(%outputs_0, %outputs_1): (tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<1x1xf32>
+  return %outputs_2 : tensor<1x1xf32>
+}
+
+// CHECK-LABEL: @_func
+func.func private @_func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> (tensor<1x1xf32>) {
+  %outputs_0 =  "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  return %outputs_0 : tensor<1x1xf32>
+}
+}
+
+
+// -----
+// Missing topology and device assignment attribute in spmd is ok
+
+// CHECK-LABEL: func.func @serving_default(%arg0: tensor<3x1xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK-NEXT:  %0 = "tf.IfrtCall"(%arg1, %arg0)
+// CHECK-SAME:       {operandSegmentSizes = array<i32: 2, 0>, program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
+// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK:    return
+//
+// CHECK:  func.func @_ifrt_program__func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-SAME:       device_assignment = []
+// CHECK-SAME:      tfrt_ifrt_serving.program_id = [[PROGRAM_ID]] : i64
+// CHECK-NEXT:     %0 = "tf.MatMul"(%arg0, %arg1)
+// CHECK:          return
+
+module attributes {tf.devices = ["/job:localhost/replica:0/task:0/device:CPU:0", "/job:localhost/replica:0/task:0/device:TPU_SYSTEM:0", "/job:localhost/replica:0/task:0/device:TPU:0", "/job:localhost/replica:0/task:0/device:TPU:1"], tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1704 : i32}} {
+func.func @serving_default(%arg0: tensor<3x1xf32>,  %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
+  %outputs  =  "tf.TPUCompilationResult"() {_tpu_compilation_status = "cluster", device = ""} : () -> tensor<!tf_type.string>
+  %outputs_0 = "tf_device.cluster_func"(%arg1, %arg0) {_producer_name = "UNKNOWN", func = @_func, input_sharding_configuration = ["{devices=[2,1]0,1}", ""], num_cores_per_replica = 2 : i64, output_sharding_configuration = [""], step_marker_location = "STEP_MARK_AT_ENTRY", use_spmd_for_xla_partitioning = true, use_tpu = true} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  return %outputs_0 : tensor<1x1xf32>
+}
+
+// CHECK-LABEL: @_func
+func.func private @_func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> (tensor<1x1xf32>) {
+  %outputs_0 =  "tf.MatMul"(%arg0, %arg1) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  return %outputs_0 : tensor<1x1xf32>
+}
+}
+

tensorflow/compiler/mlir/tfrt/tests/ifrt/sink_variable_as_named_array.mlir

@@ -0,0 +1,144 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-device-decompose-resource-ops -sink-variable-as-named-array %s | FileCheck %s
+
+// -----
+// Basic test: all variables tensors are for devices and sinked as named ifrt arrays
+//
+//
+// CHECK-LABEL:  func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK-NEXT:   [[HANDLE2:%.*]] = "tf.VarHandleOp"
+// CHECK-NEXT:   [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"([[HANDLE2]])
+// CHECK-SAME:       used_by_host = false 
+// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"([[KEY]], %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}>
+// CHECK-SAME:    : (tensor<!tf_type.string>, tensor<1x3xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    return [[RES]] : tensor<1x1xf32>
+//
+module {
+  func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %result = "tf.IfrtCall"(%2, %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<3x1xf32>, tensor<1x3xf32>) -> (tensor<1x1xf32>)
+    return %result : tensor<1x1xf32>
+  }
+}
+
+// -----
+// Variable tensor for host can still be used.
+//
+// CHECK-LABEL:  func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>) {
+// CHECK:  "tf.VarHandleOp"
+// CHECK-NOT:  [[VARIABLE:%.*]] = "tf.ReadVariableOp"
+// CHECK-NEXT:  [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
+// CHECK-SAME:    used_by_host = true
+// CHECK-NEXT:  [[MATRES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
+// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"(%arg0, [[KEY]]) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}>
+// CHECK-NEXT:    return [[RES]], [[MATRES]] : tensor<1x1xf32>, tensor<1x1xf32>
+//
+module {
+  func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>) {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %3 = "tf.MatMul"(%arg0, %2) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+    %result = "tf.IfrtCall"(%arg0, %2) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<1x3xf32>, tensor<3x1xf32>) -> (tensor<1x1xf32>)
+    return %result, %3 : tensor<1x1xf32>, tensor<1x1xf32>
+  }
+}
+
+// -----
+// Variable tensor is only for host
+//
+// CHECK-LABEL:  func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK:  "tf.VarHandleOp"
+// CHECK-NOT:   tf.ReadVariableOp
+// CHECK-NEXT:  [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
+// CHECK-SAME:    used_by_host = true
+// CHECK-NEXT:  [[RES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
+// CHECK-NEXT:    return [[RES]] : tensor<1x1xf32>
+//
+module {
+  func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %3 = "tf.MatMul"(%arg0, %2) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+    return %3: tensor<1x1xf32>
+  }
+}
+
+// -----
+//  Resources that are created in the same module are not sinked.
+//
+// CHECK-LABEL:  func.func @serving_default
+// CHECK-NOT:  IfrtLoadVariable
+// CHECK:      "tf.VarHandleOp"
+// CHECK-NEXT: "tf.AssignVariableOp"
+// CHECK-NEXT: "tf.ReadVariableOp"
+// CHECK-NEXT: "tf.StatefulPartitionedCall"
+// CHECK-NEXT:  return 
+//
+module {
+  func.func @serving_default() -> tensor<*xi32> {
+    %cst = "tf.Const"() <{value = dense<"some_test.txt"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>
+    "tf.AssignVariableOp"(%0, %cst) <{validate_shape = false}> : (tensor<!tf_type.resource<tensor<!tf_type.string>>>, tensor<!tf_type.string>) -> ()
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<!tf_type.string>>>) -> tensor<*x!tf_type.string>
+    %4 = "tf.StatefulPartitionedCall"(%2) <{config = "", config_proto = "", executor_type = "", f = @__initializer}> : (tensor<*x!tf_type.string>) -> tensor<*xi32>
+    return %4: tensor<*xi32>
+  }
+  func.func @__initializer(%arg0: tensor<*x!tf_type.string>) -> tensor<i32> {
+    %0 = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
+    return %0 : tensor<i32>
+  }
+}
+
+
+// -----
+//  Decomposable Resource Ops usage
+//
+// CHECK-LABEL:  func.func @serving_default
+// CHECK:      "tf.VarHandleOp"
+// CHECK-NEXT: "tf.IfrtLoadVariable"
+// CHECK-NEXT: "tf.GatherV2"
+// CHECK-NEXT:  return 
+//
+module {
+  func.func @serving_default() -> tensor<1x3xbf16> {
+    %cst = "tf.Const"() <{value = dense<[1]> : tensor<1xi32>}> : () -> tensor<1xi32>
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<2x3xbf16>>>
+    %1 = "tf.ResourceGather"(%0, %cst) <{batch_dims = 0 : i64, validate_indices = true}> : (tensor<!tf_type.resource<tensor<2x3xbf16>>>, tensor<1xi32>) -> tensor<1x3xbf16>
+    return %1: tensor<1x3xbf16>
+  }
+}
+
+// -----
+// AsyncIfrtCall test: all variables tensors are for devices and sinked as named ifrt arrays
+//
+//
+// CHECK-LABEL:  func.func @serving_default_async(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+// CHECK-NEXT:   [[HANDLE2:%.*]] = "tf.VarHandleOp"
+// CHECK-NEXT:   [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"([[HANDLE2]])
+// CHECK-SAME:       used_by_host = false
+// CHECK-NEXT:   [[RES:%.*]] = "tf.AsyncIfrtCall"([[KEY]], %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}>
+// CHECK-SAME:    : (tensor<!tf_type.string>, tensor<1x3xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    return [[RES]] : tensor<1x1xf32>
+//
+module {
+  func.func @serving_default_async(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %result = "tf.AsyncIfrtCall"(%2, %arg0) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<3x1xf32>, tensor<1x3xf32>) -> (tensor<1x1xf32>)
+    return %result : tensor<1x1xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_device_cleanup.mlir

@@ -0,0 +1,22 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -tf-device-cleanup | FileCheck %s
+
+// CHECK-LABEL: func @ops_with_device
+func.func @ops_with_device() {
+  %0 = "tf.VarHandleOp"() {container = "", shared_name = "var", device = "/device/..."} : () -> tensor<!tf_type.resource<tensor<1xf32>>>
+  // CHECK-NOT: device = "/device/..."
+  func.return
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_identity_propagation.mlir

@@ -0,0 +1,52 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -tf-identity-propagation -canonicalize | FileCheck %s
+
+// CHECK-LABEL: func @identity
+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>)
+func.func @identity(%arg0: tensor<i32>) -> tensor<i32> {
+  // CHECK-NOT: "tf.Identity"
+  %0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
+  // CHECK: return %[[ARG0]]
+  func.return %0 : tensor<i32>
+}
+
+// CHECK-LABEL: func @identity_terminator
+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>)
+func.func @identity_terminator(%arg0: tensor<i32>) -> (tensor<*xi32>, tensor<i32>) {
+  // CHECK: %[[IDENTITY:.*]] = "tf.Identity"
+  %0 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<*xi32>
+  // CHECK-NOT: "tf.Identity"
+  %1 = "tf.Identity"(%arg0) : (tensor<i32>) -> tensor<i32>
+  // CHECK: return %[[IDENTITY]], %[[ARG0]]
+  func.return %0, %1 : tensor<*xi32>, tensor<i32>
+}
+
+// CHECK-LABEL: func @xla_sharding
+func.func @xla_sharding(%arg0: tensor<i32>) -> tensor<i32> {
+  // CHECK: %[[OUTPUT:.*]] = "tf.Identity"
+  %0 = "tf.Identity"(%arg0) {_XlaSharding = ""} : (tensor<i32>) -> tensor<i32>
+  // CHECK: return %[[OUTPUT]]
+  func.return %0 : tensor<i32>
+}
+
+// CHECK-LABEL: func @identity_n
+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>, %[[ARG1:.*]]: tensor<f32>)
+func.func @identity_n(%arg0: tensor<i32>, %arg1: tensor<f32>) -> (tensor<i32>, tensor<f32>) {
+  // CHECK-NOT: "tf.IdentityN"
+  %0:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>)
+  // CHECK: return %[[ARG0]], %[[ARG1]]
+  func.return %0#0, %0#1 : tensor<i32>, tensor<f32>
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_merging.mlir

@@ -0,0 +1,35 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -tf-restore-merging | FileCheck %s
+
+// CHECK-LABEL: func @single_restore_group
+// CHECK-SAME:    (%[[ARG0:.*]]: {{.*}})
+func.func @single_restore_group(%arg0: tensor<!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>) {
+  %0 = "tf.Const"() {value = dense<"foo"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  %1 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  %2 = "tf.RestoreV2"(%arg0, %0, %1) : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<*xf32>
+
+  %3 = "tf.Const"() {value = dense<"bar"> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  %4 = "tf.Const"() {value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  %5 = "tf.RestoreV2"(%arg0, %3, %4) : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<*xi32>
+
+  // CHECK:      %[[NAMES:.*]] = "tf.Const"() <{value = dense<["foo", "bar"]> : tensor<2x!tf_type.string>}>
+  // CHECK-NEXT:      %[[SHAPES:.*]] = "tf.Const"() <{value = dense<""> : tensor<2x!tf_type.string>}>
+  // CHECK-NEXT:      %[[TENSORS:.*]]:2 = "tf.RestoreV2"(%[[ARG0]], %[[NAMES]], %[[SHAPES]])
+  // CHECK-SAME:   -> (tensor<*xf32>, tensor<*xi32>)
+
+  // CHECK:      return %[[TENSORS]]#0, %[[TENSORS]]#1
+  func.return %2, %5 : tensor<*xf32>, tensor<*xi32>
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_pruning.mlir

@@ -0,0 +1,39 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -tf-restore-pruning %s | FileCheck %s
+
+// CHECK-LABEL:   func.func @prune_unused_restore
+func.func @prune_unused_restore() {
+  %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+  %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+  %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+  // CHECK-NOT: tf.RestoreV2
+  %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+  %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+  return
+}
+
+
+// CHECK-LABEL: func.func @used_restore_remains
+func.func @used_restore_remains() {
+  %cst = "tf.Const"() <{value = dense<"restore_ariables"> : tensor<!tf_type.string>}> : () -> tensor<!tf_type.string>
+  %cst_0 = "tf.Const"() <{value = dense<""> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+  %cst_1 = "tf.Const"() <{value = dense<"y"> : tensor<1x!tf_type.string>}> : () -> tensor<1x!tf_type.string>
+  // CHECK: tf.RestoreV2
+  %0 = "tf.RestoreV2"(%cst, %cst_1, %cst_0): (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>) -> tensor<3x1xf32>
+  %1 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+  "tf.AssignVariableOp"(%1, %0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<3x1xf32>) -> ()
+  return
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt/tf_restore_splitting.mlir

@@ -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.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -tf-restore-splitting | FileCheck %s
+
+// CHECK-LABEL: func @single_restore
+// CHECK-SAME:    (%[[ARG0:.*]]: {{.*}})
+func.func @single_restore(%arg0: tensor<!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>) {
+  %0 = "tf.Const"() {value = dense<["foo", "bar"]> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
+  %1 = "tf.Const"() {value = dense<""> : tensor<2x!tf_type.string>} : () -> tensor<2x!tf_type.string>
+  %2:2 = "tf.RestoreV2"(%arg0, %0, %1) : (tensor<!tf_type.string>, tensor<2x!tf_type.string>, tensor<2x!tf_type.string>) -> (tensor<*xf32>, tensor<*xi32>)
+
+  // CHECK: %[[FOO_NAME:.*]] = "tf.Const"() <{value = dense<"foo"> : tensor<1x!tf_type.string>}>
+  // CHECK: %[[FOO:.*]] = "tf.RestoreV2"(%[[ARG0]], %[[FOO_NAME]], {{.*}})
+
+  // CHECK: %[[BAR_NAME:.*]] = "tf.Const"() <{value = dense<"bar"> : tensor<1x!tf_type.string>}>
+  // CHECK: %[[BAR:.*]] = "tf.RestoreV2"(%[[ARG0]], %[[BAR_NAME]], {{.*}})
+
+  // CHECK: return %[[FOO]], %[[BAR]]
+  func.return %2#0, %2#1 : tensor<*xf32>, tensor<*xi32>
+}

tensorflow/compiler/mlir/tfrt/tests/ifrt_set_tpu_host_allocator.mlir

@@ -0,0 +1,70 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -tf-mlrt-ifrt-set-tpu-host-allocator %s | FileCheck %s --dump-input=fail --dump-input-filter=all
+
+// All arguments are non-variables
+
+// CHECK-LABEL: func.func @serving_default
+// CHECK-NEXT:    "tf.MatMul"
+// CHECK-SAME:      tf_mlrt.custom_device = "tpu_host_device"
+// CHECK-NEXT:    "tf.MatMul"
+// CHECK-SAME:      tf_mlrt.custom_device = "tpu_host_device"
+// CHECK-NEXT:    "tf.IfrtCall"
+func.func @serving_default(%arg0: tensor<3x1xf32>,  %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
+  %producer_0=  "tf.MatMul"(%arg1, %arg0) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  %producer_1=  "tf.MatMul"(%arg1, %arg0) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  %result = "tf.IfrtCall"(%producer_0, %producer_1) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<1x1xf32>)
+  return %result: tensor<1x1xf32>
+}
+
+// Arguments to the IfrtCall are a mix of non-variables and variables
+
+// CHECK-LABEL: func.func @serving_default1
+// CHECK-NEXT:    "tf.VarHandleOp"
+// CHECK-NOT:      tf_mlrt.custom_device
+// CHECK-NEXT:    "tf.ReadVariableOp"
+// CHECK-NOT:      tf_mlrt.custom_device
+// CHECK-NEXT:    "tf.MatMul"
+// CHECK-SAME:      tf_mlrt.custom_device = "tpu_host_device"
+// CHECK-NEXT:    "tf.MatMul"
+// CHECK-NOT:      tf_mlrt.custom_device
+// CHECK-NEXT:    "tf.IfrtCall"
+func.func @serving_default1(%arg0: tensor<3x1xf32>,  %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+    %2 = "tf.MatMul"(%arg1, %arg0) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+    %3 = "tf.MatMul"(%arg1, %arg0) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+    %result = "tf.IfrtCall"(%1, %2) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}> : (tensor<3x1xf32>, tensor<1x1xf32>) -> (tensor<1x1xf32>)
+    return %result: tensor<1x1xf32>
+}
+
+// -----
+// Async test: All arguments are non-variables
+//
+// CHECK-LABEL: func.func @serving_default_async
+// CHECK-NEXT:    "tf.MatMul"
+// CHECK-SAME:      tf_mlrt.custom_device = "tpu_host_device"
+// CHECK-NEXT:    "tf.MatMul"
+// CHECK-SAME:      tf_mlrt.custom_device = "tpu_host_device"
+// CHECK-NEXT:    "tf.AsyncIfrtCall"
+func.func @serving_default_async(%arg0: tensor<3x1xf32>,  %arg1: tensor<1x3xf32>) -> (tensor<1x1xf32>) {
+  %producer_0=  "tf.MatMul"(%arg1, %arg0) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  %producer_1=  "tf.MatMul"(%arg1, %arg0) {transpose_a = false, transpose_b = false} : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+  %result = "tf.AsyncIfrtCall"(%producer_0, %producer_1) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = []}> : (tensor<1x1xf32>, tensor<1x1xf32>) -> (tensor<1x1xf32>)
+  return %result: tensor<1x1xf32>
+}
+
+
+

tensorflow/compiler/mlir/tfrt/tests/ir/BUILD

@@ -0,0 +1,47 @@
+load("//tensorflow:tensorflow.bzl", "if_oss", "tf_cc_test")
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    exclude = ["testdata/**"],
+    features = if_oss(["--path=org_tensorflow/tensorflow/compiler/mlir/tfrt"]),
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//llvm:not",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+)
+
+tf_cc_test(
+    name = "tfrt_fallback_util_test",
+    srcs = ["tfrt_fallback_util_test.cc"],
+    data = [
+        "testdata/test.mlir",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_async_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_sync_opdefs",
+        "//tensorflow/compiler/mlir/tfrt/ir:tfrt_fallback_util",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:resource_loader",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@tf_runtime//:init_tfrt_dialects",
+    ],
+)

tensorflow/compiler/mlir/tfrt/tests/ir/fallback_opt.mlir

@@ -0,0 +1,98 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s | tf-tfrt-opt | FileCheck %s --dump-input=fail
+
+// CHECK-LABEL: func @const_tensor_proto
+func.func @const_tensor_proto() -> !tfrt_fallback.tf_tensor {
+  // CHECK: tfrt_fallback_async.const_tensor_proto "fake serialized proto"
+  %0 = tfrt_fallback_async.const_tensor_proto "fake serialized proto"
+  tfrt.return %0 : !tfrt_fallback.tf_tensor
+}
+
+// CHECK-LABEL: func @const_dense_tensor
+func.func @const_dense_tensor() -> !tfrt_fallback.tf_tensor {
+  // CHECK: tfrt_fallback_async.const_dense_tensor
+  %0 = tfrt_fallback_async.const_dense_tensor dense<0.0> : tensor<f32> {_tfrt_cost = 1 : i64}
+  tfrt.return %0 : !tfrt_fallback.tf_tensor
+}
+
+// CHECK-LABEL: func @const_string_tensor
+func.func @const_string_tensor() -> !tfrt_fallback.tf_tensor {
+  // CHECK: tfrt_fallback_async.const_string_tensor
+  %0 = tfrt_fallback_async.const_string_tensor {shape = [1, 2], value = ["const", "string"], _tfrt_cost = 1 : i64}
+  tfrt.return %0 : !tfrt_fallback.tf_tensor
+}
+
+// CHECK-LABEL: func @convert
+func.func @convert() -> !corert.tensorhandle {
+  %0 = corert.const_dense_tensor dense<0.0> : tensor<f32>
+  // CHECK: tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor
+  %1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %0 {_tfrt_cost = 1 : i64, device = "cpu"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
+  // CHECK: tfrt_fallback_async.executeop key(0) cost(100) device("cpu") "tf.Relu"(%{{.*}}) {T = f32} : 1
+  %2 = tfrt_fallback_async.executeop key(0) cost(100) device("cpu") "tf.Relu"(%1) {T = f32} : 1
+  // CHECK: tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle
+  %3 = tfrt_fallback_async.fallback_tensor_to_corert_tensorhandle %2 {_tfrt_cost = 1 : i64, device="cpu"} : (!tfrt_fallback.tf_tensor) -> (!corert.tensorhandle)
+  tfrt.return %3 : !corert.tensorhandle
+}
+
+// CHECK-LABEL: func @predicate
+func.func @predicate() -> i1 {
+  %0 = tfrt_fallback_async.const_dense_tensor dense<0.0> : tensor<f32> {_tfrt_cost = 1 : i64}
+  // CHECK: tfrt_fallback_async.predicate
+  %1 = tfrt_fallback_async.predicate %0 {_tfrt_cost = 1 : i64, device = "cpu"}
+  tfrt.return %1 : i1
+}
+
+// CHECK-LABEL: func @createop
+func.func @createop(%in_ch: !tfrt.chain) -> !tfrt.chain {
+
+  // CHECK: [[ch:%.*]] = tfrt_fallback_async.createop(%{{.*}}) key(100) device("cpu") "tf.AddV2"() {T = i32} num_args(2)
+  %out_ch = tfrt_fallback_async.createop(%in_ch) key(100) device("cpu") "tf.AddV2"() {T = i32} num_args(2)
+
+  // CHECK: tfrt.return [[ch]]
+  tfrt.return %out_ch: !tfrt.chain
+}
+
+// CHECK-LABEL: func @fallback_resource
+func.func @fallback_resource(%ch0: !tfrt.chain) -> !tfrt.chain {
+
+  %ra = tfrt_fallback_async.const_dense_tensor dense<0.0> : tensor<f32> {_tfrt_cost = 1 : i64}
+  %rb = tfrt_fallback_async.const_dense_tensor dense<0.5> : tensor<f32> {_tfrt_cost = 1 : i64}
+
+  // CHECK: tfrt_fallback_async.set_resource {{%.*}}, {{%.*}} {device = "cpu", index = 0 : i64}
+  // CHECK: tfrt_fallback_async.set_resource {{%.*}}, {{%.*}} {device = "cpu", index = 1 : i64}
+  // CHECK: tfrt_fallback_async.get_resource {{%.*}} {_tfrt_cost = 1 : i64, device = "cpu", indices = [0, 1]}
+  %ch1 = tfrt_fallback_async.set_resource %ch0, %ra {device = "cpu", index = 0 : i64}
+  %ch2 = tfrt_fallback_async.set_resource %ch1, %rb {device = "cpu", index = 1 : i64}
+  %ch3, %a, %b = tfrt_fallback_async.get_resource %ch2 {_tfrt_cost = 1 : i64, device = "cpu", indices = [0 : i64, 1 : i64]} : (!tfrt.chain) -> (!tfrt.chain, !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
+
+  tfrt.return %ch3: !tfrt.chain
+}
+
+// CHECK-LABEL: func @copy_if_small
+func.func @copy_if_small(%arg: !tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor) {
+  // CHECK: tfrt_fallback_async.copy_if_small {{%.*}} {_tfrt_cost = 1 : i64} : (!tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
+  %0:2 = tfrt_fallback_async.copy_if_small %arg {_tfrt_cost = 1 : i64} : (!tfrt_fallback.tf_tensor) -> (!tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor)
+  tfrt.return %0#0, %0#1 : !tfrt_fallback.tf_tensor, !tfrt_fallback.tf_tensor
+}
+
+// CHECK-LABEL: func @custom_allocator
+func.func @custom_allocator(%ch: !tfrt.chain, %arg: !tfrt_fallback.tf_tensor, %allocator: !tfrt_fallback.tf_allocator) -> (!tfrt.chain, !tfrt_fallback.tf_tensor) {
+  // CHECK: tfrt_fallback_async.executeop.allocator(%{{.*}}) key(200) cost(100) device("cpu") "tf.Cast"(%{{.*}}) {Truncate = false, T = f32} : 1
+  %0 = tfrt_fallback_async.executeop.allocator(%allocator) key(200) cost(100) device("cpu") "tf.Cast"(%arg) {Truncate = false, T = f32} : 1
+  // CHECK: tfrt_fallback_async.executeop.seq.allocator(%{{.*}}, %{{.*}}) key(201) cost(100) device("cpu") "tf.Cast"(%{{.*}}) {Truncate = false, T = i32} : 1
+  %out_ch, %1 = tfrt_fallback_async.executeop.seq.allocator(%ch, %allocator) key(201) cost(100) device("cpu") "tf.Cast"(%0) {Truncate = false, T = i32} : 1
+  tfrt.return %out_ch, %1 : !tfrt.chain, !tfrt_fallback.tf_tensor
+}

tensorflow/compiler/mlir/tfrt/tests/ir/testdata/test.mlir

@@ -0,0 +1,22 @@
+// 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.
+// ==============================================================================
+
+func.func @test(%ch: !tfrt.chain, %arg0: !corert.tensorhandle, %arg1_th: !corert.tensorhandle) {
+  %cpu = corert.get_op_handler %ch "cpu"
+  %0 = corert.executeop(%cpu) "tf.Relu"(%arg0) { T = f32 } : 1
+  %arg1 = tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor %arg1_th {_tfrt_cost = 1 : i64, device = "/CPU:0"} : (!corert.tensorhandle) -> (!tfrt_fallback.tf_tensor)
+  %1 = tfrt_fallback_async.executeop key(0) cost(100) device("/CPU:0") "tf.Relu"(%arg1) { T = f32 } : 1
+  tfrt.return
+}

tensorflow/compiler/mlir/tfrt/tests/ir/tfrt_fallback_util_test.cc

@@ -0,0 +1,63 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_util.h"
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#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/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.h"
+#include "tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/test.h"
+#include "tfrt/init_tfrt_dialects.h"  // from @tf_runtime
+
+namespace tfrt {
+namespace fallback_async {
+namespace {
+
+TEST(SavedModelTest, MapFallbackArgs) {
+  std::string saved_model_mlir_path = tensorflow::GetDataDependencyFilepath(
+      "tensorflow/compiler/mlir/tfrt/tests/ir/testdata/test.mlir");
+
+  mlir::DialectRegistry registry;
+  RegisterTFRTDialects(registry);
+  registry.insert<tfrt::fallback_async::FallbackAsyncDialect>();
+  registry.insert<tfrt::fallback_sync::FallbackSyncDialect>();
+
+  mlir::MLIRContext context(registry);
+  auto module =
+      mlir::parseSourceFile<mlir::ModuleOp>(saved_model_mlir_path, &context);
+  ASSERT_TRUE(module);
+
+  std::vector<std::pair<std::string, int>> func_and_index;
+  ForEachArgConsumedByFallback(
+      module.get(),
+      [&func_and_index](llvm::StringRef func_name, int arg_index) {
+        func_and_index.push_back({func_name.str(), arg_index});
+      });
+
+  ASSERT_EQ(func_and_index.size(), 1);
+  EXPECT_EQ(func_and_index[0].first, "test");
+  EXPECT_EQ(func_and_index[0].second, 2);
+}
+
+}  // namespace
+}  // namespace fallback_async
+}  // namespace tfrt

tensorflow/compiler/mlir/tfrt/tests/mlrt/BUILD

@@ -0,0 +1,21 @@
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir/tfrt:tf-tfrt-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+)

tensorflow/compiler/mlir/tfrt/tests/mlrt/assign_op_key.mlir

@@ -0,0 +1,63 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-assign-op-key %s | FileCheck %s
+
+// CHECK-LABEL: func @main
+// CHECK: tf.AddV2
+// CHECK-SAME: {__op_key = 0 : i32}
+
+// CHECK: tf.AddV2
+// CHECK-SAME: {__op_key = 1 : i32}
+
+// CHECK: tf.AddV2
+// CHECK-SAME: {__op_key = 2 : i32}
+
+// CHECK: tf.AddV2
+// CHECK-SAME: {__op_key = 3 : i32}
+
+// CHECK: tf.Sub
+// CHECK-SAME: {__op_key = 4 : i32}
+
+// CHECK: tf.Sub
+// CHECK-SAME: {__op_key = 5 : i32}
+
+// CHECK: tf.Sub
+// CHECK-SAME: {__op_key = 6 : i32}
+
+// CHECK: tf.Sub
+// CHECK-SAME: {__op_key = 7 : i32}
+
+
+// CHECK: [[x:%.*]] = "tf.AddV2"
+// CHECK-SAME: {__op_key = 8 : i32}
+
+// CHECK: return [[x]]
+
+func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
+
+  %a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %b0 = "tf.Sub"(%b, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b3 = "tf.Sub"(%b2, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %c = "tf.AddV2"(%a3, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  func.return %c : tensor<i32>
+}

tensorflow/compiler/mlir/tfrt/tests/mlrt/async_while.mlir

@@ -0,0 +1,244 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-async-while %s | FileCheck %s
+
+// This is a simple case that should be pipelined.
+
+// CHECK-LABEL: func.func private @"map/while_cond"
+func.func private @"map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
+  %0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0 : tensor<i1>
+}
+
+// CHECK-LABEL: func.func private @"map/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i1> {
+// CHECK-NEXT:    %0 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+// CHECK-NEXT:    return %0 : tensor<i1>
+
+// CHECK-LABEL: func.func private @"map/while_body"
+func.func private @"map/while_body"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) {
+  %cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.AddV2"(%loop_count, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %1 = "tf.TensorArrayReadV3"(%handle, %loop_count, %flow_in) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+  %2 = "tf.MatMul"(%1, %matrix)  : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  return %0, %max_iterations, %handle, %flow_in, %2: tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>
+}
+
+// CHECK-LABEL: func.func private @"map/while_body/TfMlrtAsyncWhileBody"(%arg0: !mlrt.promise, %arg1: !mlrt.future, %arg2: !mlrt.promise, %arg3: !mlrt.future, %arg4: !mlrt.promise, %arg5: tensor<i32>, %arg6: tensor<?x!tf_type.resource>, %arg7: tensor<*xf32>) {
+// CHECK-NEXT:    %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
+// CHECK-NEXT:    %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor<i32>
+// CHECK-NEXT:    %1 = "tf.AddV2"(%0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor<i32>) -> ()
+// CHECK-NEXT:    %2 = "tf.PartitionedCall"(%1, %arg5) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor<i1>) -> ()
+// CHECK-NEXT:    %3 = "tf.TensorArrayReadV3"(%arg6, %0, %arg7) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %4 = "tf_mlrt.tf_await"(%arg3) : (!mlrt.future) -> tensor<3x3xf32>
+// CHECK-NEXT:    %5 = "tf.MatMul"(%3, %4) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg4, %5) : (!mlrt.promise, tensor<3x3xf32>) -> ()
+// CHECK-NEXT:    return
+
+//CHECK-LABEL: func.func @serving_default
+func.func @serving_default(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<*xf32>) {
+  %cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  // CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
+  %1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) ->  (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
+  // CHECK-NEXT: %2 = "tf_mlrt.tf_await"(%1#5) : (!mlrt.future) -> tensor<*xf32>
+  // CHECK-NEXT: %3 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32>
+  // CHECK-NEXT:  return %3, %2 : tensor<3x3xf32>, tensor<*xf32>
+  return %1#4, %1#3 : tensor<3x3xf32>, tensor<*xf32>
+}
+
+
+//CHECK-LABEL: func.func @multi_while_test
+func.func @multi_while_test(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>, %array_handle_2: tensor<?x!tf_type.resource>, %array_flow_2: tensor<*xf32>, %matrix_2: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<3x3xf32>) {
+  %cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  // CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
+  %1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) ->  (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
+  // CHECK: %2 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  // CHECK-NEXT: %3:6 = tf_mlrt.tf_async_while @"map/while_body/TfMlrtAsyncWhileBody"(%2, %cst, %arg6, %arg0, %arg4, %arg5) {invariant_size = 3 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
+  %2:5 = "tf.While"(%cst_0, %max_iterations, %array_handle_2, %array_flow_2, %matrix_2) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) ->  (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
+  // CHECK-NEXT: %4 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32>
+  // CHECK-NEXT: %5 = "tf_mlrt.tf_await"(%3#2) : (!mlrt.future) -> tensor<3x3xf32>
+  // CHECK-NEXT:  return %4, %5 : tensor<3x3xf32>, tensor<3x3xf32>
+  return %1#4, %2#4 : tensor<3x3xf32>, tensor<3x3xf32>
+}
+
+// -----
+
+// Test a case in which predicate is updated after mutables and shall not be converted to AsyncWhile.
+
+// CHECK-LABEL: func.func private @"map/while_cond"
+func.func private @"map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
+  %0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0 : tensor<i1>
+}
+
+// CHECK-LABEL: func.func private @"map/while_body"
+func.func private @"map/while_body"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) {
+  %cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %1 = "tf.TensorArrayReadV3"(%handle, %loop_count, %flow_in) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+  %2 = "tf.MatMul"(%1, %matrix)  : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  // Predicate is update at the last stage.
+  %0 = "tf.AddV2"(%loop_count, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  return %0, %max_iterations, %handle, %flow_in, %2: tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>
+}
+
+//CHECK-LABEL: func.func @serving_default
+func.func @serving_default(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<3x3xf32> {
+  %cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: tf.While
+  // CHECK-NOT: AsyncWhile
+  %1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"map/while_body", cond = @"map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) ->  (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
+  return %1#4 : tensor<3x3xf32>
+}
+
+// -----
+
+// The newly create function name may have conflict with existing functions (very rare).
+
+// CHECK-LABEL: func.func private @"random/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
+func.func private @"random/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
+
+  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  return %0: tensor<i32>
+}
+
+// CHECK-LABEL: func.func private @"random/while_body/TfMlrtAsyncWhileBody"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
+func.func private @"random/while_body/TfMlrtAsyncWhileBody"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
+  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  return %0: tensor<i32>
+}
+
+// CHECK-LABEL: func.func private @"random/while_cond"
+func.func private @"random/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
+  %0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0 : tensor<i1>
+}
+
+// CHECK-LABEL: func.func private @"random/while_cond/TfMlrtAsyncWhilePredicate_0"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i1> {
+// CHECK-NEXT:    %0 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+// CHECK-NEXT:    return %0 : tensor<i1>
+
+// CHECK-LABEL: func.func private @"random/while_body"
+func.func private @"random/while_body"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) {
+  %cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.AddV2"(%loop_count, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %1 = "tf.TensorArrayReadV3"(%handle, %loop_count, %flow_in) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+  %2 = "tf.MatMul"(%1, %matrix)  : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  return %0, %max_iterations, %handle, %flow_in, %2: tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>
+}
+
+// CHECK-LABEL: func.func private @"random/while_body/TfMlrtAsyncWhileBody_1"(%arg0: !mlrt.promise, %arg1: !mlrt.future, %arg2: !mlrt.promise, %arg3: !mlrt.future, %arg4: !mlrt.promise, %arg5: tensor<i32>, %arg6: tensor<?x!tf_type.resource>, %arg7: tensor<*xf32>) {
+// CHECK-NEXT:    %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
+// CHECK-NEXT:    %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor<i32>
+// CHECK-NEXT:    %1 = "tf.AddV2"(%0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor<i32>) -> ()
+// CHECK-NEXT:    %2 = "tf.PartitionedCall"(%1, %arg5) <{config = "", config_proto = "", executor_type = "", f = @"random/while_cond/TfMlrtAsyncWhilePredicate_0"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor<i1>) -> ()
+// CHECK-NEXT:    %3 = "tf.TensorArrayReadV3"(%arg6, %0, %arg7) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %4 = "tf_mlrt.tf_await"(%arg3) : (!mlrt.future) -> tensor<3x3xf32>
+// CHECK-NEXT:    %5 = "tf.MatMul"(%3, %4) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg4, %5) : (!mlrt.promise, tensor<3x3xf32>) -> ()
+// CHECK-NEXT:    return
+
+//CHECK-LABEL: func.func @random_serving_default
+func.func @random_serving_default(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> (tensor<3x3xf32>, tensor<*xf32>) {
+  %cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"random/while_cond/TfMlrtAsyncWhilePredicate_0"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  // CHECK-NEXT: %1:6 = tf_mlrt.tf_async_while @"random/while_body/TfMlrtAsyncWhileBody_1"(%0, %cst, %arg3, %arg0, %arg1, %arg2) {invariant_size = 3 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
+  %1:5 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix) {body= @"random/while_body", cond = @"random/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>) ->  (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>)
+  // CHECK-NEXT: %2 = "tf_mlrt.tf_await"(%1#5) : (!mlrt.future) -> tensor<*xf32>
+  // CHECK-NEXT: %3 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32>
+  // CHECK-NEXT:  return %3, %2 : tensor<3x3xf32>, tensor<*xf32>
+  return %1#4, %1#3 : tensor<3x3xf32>, tensor<*xf32>
+}
+
+// -----
+// This case test the re-ordering of the while body function to maximize pipelining between iterations.
+
+// CHECK-LABEL: func.func private @"sort_map/while_cond"
+func.func private @"sort_map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>, %handle_2: tensor<?x!tf_type.resource>, %flow_in_2: tensor<*xf32>, %matrix_2: tensor<3x3xf32>, %bound: tensor<i32>) -> tensor<i1> {
+  %0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  return %0 : tensor<i1>
+}
+
+// CHECK-LABEL: func.func private @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i1> {
+// CHECK-NEXT:    %0 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+// CHECK-NEXT:    return %0 : tensor<i1>
+
+// CHECK-LABEL: func.func private @"sort_map/while_body"
+// CHECK-NEXT:    %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
+// CHECK-NEXT:    %0 = "tf.AddV2"(%arg0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+// CHECK-NEXT:    %1 = "tf.TensorArrayReadV3"(%arg2, %arg0, %arg3) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %2 = "tf.TensorArrayReadV3"(%arg5, %arg0, %arg6) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %3 = "tf.MatMul"(%1, %arg4) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %4 = "tf.GreaterEqual"(%0, %arg8) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+// CHECK-NEXT:    %5 = "tf.MatMul"(%3, %arg7) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %6 = "tf.Select"(%4, %3, %arg4) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:   %7 = "tf.Identity"(%4) : (tensor<i1>) -> tensor<i1>
+// CHECK-NEXT:    %8 = "tf.MatMul"(%5, %arg7) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %9 = "tf.Select"(%7, %8, %arg7) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    return %0, %arg1, %arg2, %arg3, %6, %arg5, %arg6, %9, %arg8
+func.func private @"sort_map/while_body"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>, %handle_2: tensor<?x!tf_type.resource>, %flow_in_2: tensor<*xf32>, %matrix_2: tensor<3x3xf32>, %bound: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<i32>) {
+  %cst_1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %updated_loop_count = "tf.AddV2"(%loop_count, %cst_1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %in_matrix = "tf.TensorArrayReadV3"(%handle, %loop_count, %flow_in) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+  %out_matrix = "tf.MatMul"(%in_matrix, %matrix)  : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  %in_matrix1 = "tf.TensorArrayReadV3"(%handle_2, %loop_count, %flow_in_2) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+  %out_matrix1 = "tf.MatMul"(%out_matrix, %matrix_2)  : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  %out_matrix2 = "tf.MatMul"(%out_matrix1, %matrix_2): (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  %cond = "tf.GreaterEqual"(%updated_loop_count, %bound) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %result = "tf.Select"(%cond, %out_matrix, %matrix) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  %cond2 = "tf.Identity"(%cond) : (tensor<i1>) -> tensor<i1>
+  %result2 = "tf.Select"(%cond2, %out_matrix2, %matrix_2) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  return %updated_loop_count, %max_iterations, %handle, %flow_in, %result, %handle_2, %flow_in_2, %result2, %bound: tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<i32>
+}
+
+// CHECK-LABEL: func.func private @"sort_map/while_body/TfMlrtAsyncWhileBody"
+// CHECK-NEXT:    %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>
+// CHECK-NEXT:    %0 = "tf_mlrt.tf_await"(%arg1) : (!mlrt.future) -> tensor<i32>
+// CHECK-NEXT:    %1 = "tf.AddV2"(%0, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg2, %1) : (!mlrt.promise, tensor<i32>) -> ()
+// CHECK-NEXT:    %2 = "tf.PartitionedCall"(%1, %arg7) <{config = "", config_proto = "", executor_type = "", f = @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg0, %2) : (!mlrt.promise, tensor<i1>) -> ()
+// CHECK-NEXT:    %3 = "tf.TensorArrayReadV3"(%arg8, %0, %arg9) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %4 = "tf.TensorArrayReadV3"(%arg10, %0, %arg11) : (tensor<?x!tf_type.resource>, tensor<i32>, tensor<*xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %5 = "tf_mlrt.tf_await"(%arg3) : (!mlrt.future) -> tensor<3x3xf32>
+// CHECK-NEXT:    %6 = "tf.MatMul"(%3, %5) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %7 = "tf.GreaterEqual"(%1, %arg12) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+// CHECK-NEXT:    %8 = "tf_mlrt.tf_await"(%arg5) : (!mlrt.future) -> tensor<3x3xf32>
+// CHECK-NEXT:    %9 = "tf.MatMul"(%6, %8) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %10 = "tf.Select"(%7, %6, %5) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:   "tf_mlrt.tf_promise"(%arg4, %10) : (!mlrt.promise, tensor<3x3xf32>) -> ()
+// CHECK-NEXT:    %11 = "tf.Identity"(%7) : (tensor<i1>) -> tensor<i1>
+// CHECK-NEXT:    %12 = "tf.MatMul"(%9, %8) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    %13 = "tf.Select"(%11, %12, %8) : (tensor<i1>, tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+// CHECK-NEXT:    "tf_mlrt.tf_promise"(%arg6, %13) : (!mlrt.promise, tensor<3x3xf32>) -> ()
+// CHECK-NEXT:    return
+
+//CHECK-LABEL: func.func @sort_serving_default
+func.func @sort_serving_default(%max_iterations: tensor<i32>, %array_handle: tensor<?x!tf_type.resource>, %array_flow: tensor<*xf32>, %matrix: tensor<3x3xf32>, %bound: tensor<i32>) -> (tensor<3x3xf32>, tensor<3x3xf32>, tensor<*xf32>) {
+  %cst_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: %0 = "tf.PartitionedCall"(%cst, %arg0) <{config = "", config_proto = "", executor_type = "", f = @"sort_map/while_cond/TfMlrtAsyncWhilePredicate"}> : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  // CHECK-NEXT: %1:10 = tf_mlrt.tf_async_while @"sort_map/while_body/TfMlrtAsyncWhileBody"(%0, %cst, %arg3, %arg3, %arg0, %arg1, %arg2, %arg1, %arg2, %arg4) {invariant_size = 6 : i32} : (tensor<i1>, tensor<i32>, tensor<3x3xf32>, tensor<3x3xf32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<i32>) -> (!mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future, !mlrt.future)
+  %1:9 = "tf.While"(%cst_0, %max_iterations, %array_handle, %array_flow, %matrix , %array_handle, %array_flow, %matrix, %bound) {body= @"sort_map/while_body", cond = @"sort_map/while_cond", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>,tensor<i32>) ->  (tensor<i32>, tensor<i32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<?x!tf_type.resource>, tensor<*xf32>, tensor<3x3xf32>, tensor<i32>)
+  // CHECK-NEXT:  %2 = "tf_mlrt.tf_await"(%1#6) : (!mlrt.future) -> tensor<*xf32>
+  // CHECK-NEXT:  %3 = "tf_mlrt.tf_await"(%1#2) : (!mlrt.future) -> tensor<3x3xf32>
+  // CHECK-NEXT:  %4 = "tf_mlrt.tf_await"(%1#3) : (!mlrt.future) -> tensor<3x3xf32>
+  // CHECK-NEXT:  return %3, %4, %2 : 
+  return %1#4, %1#7, %1#3 : tensor<3x3xf32>, tensor<3x3xf32>, tensor<*xf32>
+}
+

tensorflow/compiler/mlir/tfrt/tests/mlrt/fuse_mlrt_ops.mlir

@@ -0,0 +1,92 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-fuse %s | FileCheck %s
+
+// CHECK-LABEL: @main
+// CHECK-SAME: ([[f0:%.*]]: !mlrt.future, [[f1:%.*]]: !mlrt.future, [[f2:%.*]]: !mlrt.future)
+func.func @main(%f0: !mlrt.future, %f1: !mlrt.future, %f2: !mlrt.future) -> (!tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor) {
+  // CHECK-NEXT: [[t:%.*]]:3 = tf_mlrt.await_all [[f0]], [[f1]], [[f2]]
+  // CHECK-NOT: tf_mlrt.await
+  // CHECK-NEXT: return [[t]]#0, [[t]]#1, [[t]]#2
+  %t0 = tf_mlrt.await %f0
+  %t1 = tf_mlrt.await %f1
+  %t2 = tf_mlrt.await %f2
+  func.return %t0, %t1, %t2 : !tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor
+}
+
+// -----
+
+// CHECK-LABEL: @main
+// CHECK-SAME: ([[f0:%.*]]: !mlrt.future, [[f1:%.*]]: !mlrt.future, [[f2:%.*]]: !mlrt.future)
+func.func @main(%f0: !mlrt.future, %f1: !mlrt.future, %f2: !mlrt.future) -> (!tf_mlrt.tensor, !tf_mlrt.tensor) {
+  // CHECK-NEXT: [[t:%.*]]:2 = tf_mlrt.await_all [[f0]], [[f1]]
+  // CHECK-NOT: tf_mlrt.await
+  // CHECK-NEXT: [[t2:%.*]] = tf_mlrt.executeop([[t]]#0, [[t]]#1)
+  // CHECK-NEXT: [[t3:%.*]] = tf_mlrt.await [[f2]]
+  // CHECK-NEXT: return [[t2]], [[t3]]
+  %t0 = tf_mlrt.await %f0
+  %t1 = tf_mlrt.await %f1
+  %t2 = tf_mlrt.executeop(%t0, %t1) {node_def = "AddV2", op_key = 0 : i32} : (!tf_mlrt.tensor, !tf_mlrt.tensor) -> (!tf_mlrt.tensor)
+  %t3 = tf_mlrt.await %f2
+  func.return %t2, %t3 : !tf_mlrt.tensor, !tf_mlrt.tensor
+}
+
+// -----
+
+// CHECK-LABEL: @main
+// CHECK-SAME: ([[f0:%.*]]: !mlrt.async_handle, [[f1:%.*]]: !mlrt.async_handle, [[f2:%.*]]: !mlrt.async_handle)
+func.func @main(%f0: !mlrt.async_handle, %f1: !mlrt.async_handle, %f2: !mlrt.async_handle) -> () {
+  // CHECK-NEXT: mlrt.await_all_handle [[f0]], [[f1]], [[f2]]
+  // CHECK-NOT: mlrt.await_handle
+  // CHECK-NEXT: return
+  mlrt.await_handle %f0
+  mlrt.await_handle %f1
+  mlrt.await_handle %f2
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: @main
+func.func @main() -> (!tf_mlrt.tensor, !tf_mlrt.tensor) {
+  // CHECK-NEXT: [[r:%.*]]:3 = tf_mlrt.get_resource {indices = [2, 0, 1]}
+  // CHECK-NEXT: [[v:%.*]] = tf_mlrt.executeop([[r]]#0, [[r]]#1)
+  // CHECK-NEXT: return [[v]], [[r]]#2
+  %0 = tf_mlrt.get_resource {indices = [2]} : !tf_mlrt.tensor
+  %1 = tf_mlrt.get_resource {indices = [0]} : !tf_mlrt.tensor
+  %r = tf_mlrt.executeop(%0, %1) {node_def = "AddV2", op_key = 0 : i32} : (!tf_mlrt.tensor, !tf_mlrt.tensor) -> (!tf_mlrt.tensor)
+  %2 = tf_mlrt.get_resource {indices = [1]} : !tf_mlrt.tensor
+  func.return %r, %2 : !tf_mlrt.tensor, !tf_mlrt.tensor
+}
+
+// -----
+
+// CHECK-LABEL: @fuse_promise_return
+// CHECK-SAME: ([[p:%.*]]: !mlrt.promise, [[v:%.*]]: !tf_mlrt.tensor)
+func.func @fuse_promise_return(%p: !mlrt.promise, %v: !tf_mlrt.tensor) -> () {
+  // CHECK: tf_mlrt.promise_return [[p]], [[v]]
+  tf_mlrt.promise %p, %v
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: @not_fuse_promise_return
+// CHECK-SAME: ([[p:%.*]]: !mlrt.promise, [[v:%.*]]: !tf_mlrt.tensor)
+func.func @not_fuse_promise_return(%p: !mlrt.promise, %v: !tf_mlrt.tensor) -> (!tf_mlrt.tensor) {
+  // CHECK-NOT: tf_mlrt.promise_return
+  tf_mlrt.promise %p, %v
+  func.return %v : !tf_mlrt.tensor
+}

tensorflow/compiler/mlir/tfrt/tests/mlrt/inline.mlir

@@ -0,0 +1,64 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -pass-pipeline='builtin.module(tf-to-mlrt, inline)' %s | FileCheck %s -dump-input=fail
+
+// Test generated tf_mlrt while body and predicate is inlined.
+
+func.func @then(%x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> tensor<i1> {
+  return %x: tensor<i1>
+}
+
+func.func @else(%x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> tensor<i1> {
+  return %y: tensor<i1>
+}
+
+// CHECK-LABEL: func @while_cond_if
+// CHECK: [[cond:%.*]] = tf_mlrt.predicate
+// CHECK: [[z:%.*]] = mlrt.cond [[cond]] @then @else
+// CHECK: return [[z]]
+func.func @while_cond_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> (tensor<i1>) {
+  %r = "tf.If"(%cond, %x, %y, %z) {then_branch = @then, else_branch = @else, is_stateless = true} : (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) -> tensor<i1>
+  return %r : tensor<i1>
+}
+
+// CHECK-LABEL: func @while_body_if
+func.func @while_body_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) {
+  %0 = "tf.Const"() {__op_key = 0: i32, device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
+  %1 = "tf.Add"(%z, %0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %cond, %x, %y, %1 : tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>
+}
+
+// CHECK-LABEL: func @while_test_if
+// CHECK-SAME: -> !tf_mlrt.tensor
+func.func @while_test_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>) -> (tensor<i32>) {
+  // CHECK: [[CONST:%.*]] = tf_mlrt.constop {tensor_proto = "\08\03\12\00"}
+  %cst = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // Predicate should be inlined.
+  // CHECK-NEXT: tf_mlrt.predicate
+  // CHECK-NEXT: mlrt.cond
+  // CHECK-NEXT: tf_mlrt.predicate
+
+  // CHECK-NEXT: mlrt.while
+  %0:4 = "tf.While"(%cond, %x, %y, %cst) { cond = @while_cond_if, body = @while_body_if, is_stateless = false, parallel_iterations = 1} : (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) -> (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>)
+  // CHECK: return
+  // CHECK-SAME: !tf_mlrt.tensor
+  func.return %0#3 : tensor<i32>
+}
+
+// CHECK-LABEL: func @"while_body_if/tf_mlrt_body"
+// CHECK-NOT: call
+
+// CHECK-LABEL: func @"while_cond_if/tf_mlrt_predicate"
+// CHECK-NOT: call

tensorflow/compiler/mlir/tfrt/tests/mlrt/parallelization.mlir

@@ -0,0 +1,392 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-parallelization %s | FileCheck %s --dump-input=fail --dump-input-filter=all
+
+// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
+// CHECK-SAME: ({{%.*}}: tensor<i32>, [[PROMISE:%.*]]: !mlrt.promise)
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: [[RES:%.*]] = "tf.Sub"
+// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
+// CHECK: return
+
+// CHECK-LABEL: func @main
+// CHECK: [[PROMISE:%.*]], [[FUTURE:%.*]] = "tf_mlrt.allocate_futures"
+// CHECK: [[HANDLE:%.*]] = mlrt.async({{%.*}}, [[PROMISE]])
+// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: [[x:%.*]] = "tf.AddV2"
+// CHECK: [[y:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
+// CHECK: [[RES:%.*]] = "tf.AddV2"([[x]], [[y]])
+// CHECK: mlrt.await_handle [[HANDLE]]
+// CHECK: return [[RES]]
+
+func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
+
+  %a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %b0 = "tf.Sub"(%b, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b3 = "tf.Sub"(%b2, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %c = "tf.AddV2"(%a3, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  func.return %c : tensor<i32>
+}
+
+// -----
+
+// Test merging child streams
+
+// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
+// CHECK-SAME: ({{%.*}}: tensor<i32>, {{%.*}}: tensor<i32>, [[PROMISE:%.*]]: !mlrt.promise)
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: [[RES:%.*]] = "tf.Sub"
+// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
+// CHECK: return
+
+// CHECK-LABEL: func @main
+// CHECK: [[PROMISE:%.*]], [[FUTURE:%.*]] = "tf_mlrt.allocate_futures"
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: [[VALUE:%.*]] = "tf.AddV2"
+// CHECK: [[HANDLE:%.*]] = mlrt.async([[VALUE]], {{%.*}}, [[PROMISE]])
+// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: [[x:%.*]] = "tf.AddV2"
+// CHECK: [[y:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
+// CHECK: [[RES:%.*]] = "tf.AddV2"([[x]], [[y]])
+// CHECK: mlrt.await_handle [[HANDLE]]
+// CHECK: return [[RES]]
+
+func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
+
+  %a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a4 = "tf.AddV2"(%a3, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a5 = "tf.AddV2"(%a4, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a6 = "tf.AddV2"(%a5, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a7 = "tf.AddV2"(%a6, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %b0 = "tf.Sub"(%a3, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b3 = "tf.Sub"(%b2, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %c = "tf.AddV2"(%a7, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  func.return %c : tensor<i32>
+}
+
+// -----
+
+// Test side-effecting ops
+
+// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
+// CHECK-SAME: ([[ARG:%.*]]: tensor<i32>, [[FUTURE:%.*]]: !mlrt.future, [[CONTROL_PROMISE:%.*]]: !mlrt.promise)
+// CHECK: [[HANDLE:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
+// CHECK: "tf.AssignVariableOp"([[HANDLE]], [[ARG]])
+// CHECK-NEXT: mlrt.promise_control [[CONTROL_PROMISE]]
+
+// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
+// CHECK-SAME: ({{%.*}}: tensor<i32>, {{%.*}}: tensor<i32>, [[FUTURE:%.*]]: !mlrt.future, [[PROMISE:%.*]]: !mlrt.promise)
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: [[V:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
+// CHECK-NEXT: [[RES:%.*]] = "tf.Sub"({{%.*}}, [[V]])
+// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
+// CHECK: return
+
+// CHECK-LABEL: func private @main_stream_{{[0-9]*}}
+// CHECK-SAME: ([[CONTROL_FUTURE:%.*]]: !mlrt.future, [[PROMISE:%.*]]: !mlrt.promise, [[PROMISE_HANDLE:%.*]]: !mlrt.promise)
+// CHECK: [[HANDLE:%.*]] = "tf.VarHandleOp"
+// CHECK-NEXT: "tf_mlrt.tf_promise"([[PROMISE_HANDLE]], [[HANDLE]])
+// CHECK: mlrt.await_control [[CONTROL_FUTURE]]
+// CHECK-NEXT: [[V:%.*]] = "tf.ReadVariableOp"([[HANDLE]])
+// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[V]])
+
+// CHECK-LABEL: func @main
+// CHECK: [[PROMISE:%.*]]:3, [[FUTURE:%.*]]:3 = "tf_mlrt.allocate_futures"
+// CHECK: [[CONTROL_PROMISE:%.*]], [[CONTROL_FUTURE:%.*]] = "mlrt.allocate_control_futures"
+// CHECK: [[ASYNC_HANDLE_0:%.*]] = mlrt.async([[CONTROL_FUTURE]], [[PROMISE]]#0, [[PROMISE]]#1)
+// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
+// CHECK: [[ASYNC_HANDLE_1:%.*]] = mlrt.async({{%.*}}, {{%.*}}, [[FUTURE]]#0, [[PROMISE]]#2)
+// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: [[x:%.*]] = "tf.AddV2"
+// CHECK: [[ASYNC_HANDLE_2:%.*]] = mlrt.async([[x]], [[FUTURE]]#1, [[CONTROL_PROMISE]])
+// CHECK-SAME: callee = @main_stream_{{[0-9]*}}
+// CHECK: [[y:%.*]] = "tf_mlrt.tf_await"([[FUTURE]]#2)
+// CHECK: [[RES:%.*]] = "tf.AddV2"([[x]], [[y]])
+// CHECK: mlrt.await_handle [[ASYNC_HANDLE_0]]
+// CHECK-NEXT: mlrt.await_handle [[ASYNC_HANDLE_1]]
+// CHECK-NEXT: mlrt.await_handle [[ASYNC_HANDLE_2]]
+// CHECK-NEXT: return [[RES]]
+
+func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
+  %handle = "tf.VarHandleOp"() {container = "", shared_name = "var"} : () -> tensor<!tf_type.resource_handle<tensor<i32>>>
+
+  %a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  "tf.AssignVariableOp"(%handle, %a3) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
+
+  %b0 = "tf.Sub"(%a, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %var = "tf.ReadVariableOp"(%handle) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
+  %b3 = "tf.Sub"(%b2, %var) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %c = "tf.AddV2"(%a3, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  func.return %c : tensor<i32>
+}
+
+// -----
+
+// Test multiple promises and futures
+
+// CHECK-LABEL: func private @main_stream_1
+// CHECK: mlrt.await_control
+// CHECK: "tf.DummySideEffecting"() {id = 4
+// CHECK: return
+
+// CHECK-LABEL: func private @main_stream_2
+// CHECK: mlrt.await_control
+// CHECK: "tf.DummySideEffecting"() {id = 3
+// CHECK: mlrt.promise_control
+// CHECK: return
+
+// CHECK-LABEL: func private @main_stream_3
+// CHECK: mlrt.await_control
+// CHECK: "tf.DummySideEffecting"() {id = 2
+// CHECK: mlrt.promise_control
+// CHECK: return
+
+// CHECK-LABEL: func private @main_stream_4
+// CHECK: "tf.DummySideEffecting"() {id = 1
+// CHECK: mlrt.promise_control
+// CHECK: return
+
+// CHECK-LABEL: func @main()
+// CHECK: [[PROMISES:%.*]]:3, [[FUTURES:%.*]]:3 = "mlrt.allocate_control_futures"
+// CHECK: mlrt.async([[PROMISES]]#2) {callee = @main_stream_4
+// CHECK: mlrt.async([[FUTURES]]#2, [[PROMISES]]#1) {callee = @main_stream_3
+// CHECK: mlrt.async([[FUTURES]]#1, [[PROMISES]]#0) {callee = @main_stream_2
+// CHECK: mlrt.async([[FUTURES]]#0) {callee = @main_stream_1
+// CHECK: mlrt.await_handle
+// CHECK: mlrt.await_handle
+// CHECK: mlrt.await_handle
+// CHECK: mlrt.await_handle
+
+func.func @main() {
+  "tf.DummySideEffecting"() {id = 1} : () -> ()
+  "tf.DummySideEffecting"() {id = 2} : () -> ()
+  "tf.DummySideEffecting"() {id = 3} : () -> ()
+  "tf.DummySideEffecting"() {id = 4} : () -> ()
+  func.return
+}
+
+// -----
+
+// Test correctness when there are both data and control promises in a stream function.
+
+// CHECK-LABEL: func private @main_stream_1
+// CHECK-SAME: ([[PROMISE:%.*]]: !mlrt.promise, [[CONTROL_PROMISE:%.*]]: !mlrt.promise)
+// CHECK: tf.DummySideEffecting
+// CHECK: "tf_mlrt.tf_promise"([[PROMISE]]
+// CHECK: mlrt.promise_control [[CONTROL_PROMISE]]
+
+func.func @main() -> tensor<i32> {
+  %v = "tf.DummySideEffecting"() {id = 1} : () -> tensor<i32>
+
+  %w = "tf.DummySideEffecting"() {id = 2} : () -> tensor<i32>
+  %r = "tf.AddV2"(%w, %v) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %r : tensor<i32>
+}
+
+// -----
+
+// Test inputs to the child streams are merged to the parent streams
+
+// CHECK-LABEL: func private @main_stream_1
+// CHECK-SAME: ([[INPUT0:%.*]]: tensor<i32>, [[INPUT1:%.*]]: tensor<i32>
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: mlrt.async({{%.*}}, [[INPUT1]]
+
+// CHECK-LABEL: func @main
+func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
+
+  %a0 = "tf.AddV2"(%a, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a1 = "tf.AddV2"(%a0, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a2 = "tf.AddV2"(%a1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a3 = "tf.AddV2"(%a2, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %b0 = "tf.Sub"(%b, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b1 = "tf.Sub"(%b0, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %c = "tf.AddV2"(%b1, %a) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %b2 = "tf.Sub"(%b1, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b3 = "tf.Sub"(%b2, %b) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %d = "tf.AddN"(%a3, %b3, %c) : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %d : tensor<i32>
+}
+
+// -----
+
+// Test that constants are copied instead of using promise/await.
+
+// CHECK-LABEL: func private @main_stream_1
+// CHECK-SAME: ({{%.*}}: tensor<i32>, [[PROMISE:%.*]]: !mlrt.promise)
+// CHECK: tf._TfrtGetResource
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: tf.Sub
+// CHECK: [[RES:%.*]] = "tf.Sub"
+// CHECK: "tf_mlrt.tf_promise"([[PROMISE]], [[RES]])
+// CHECK: return
+
+// CHECK-NOT: func private @main_stream
+
+// CHECK-LABEL: func @main
+// CHECK: [[PROMISE:%.*]], [[FUTURE:%.*]] = "tf_mlrt.allocate_futures"
+// CHECK-NEXT: [[HANDLE:%.*]] = mlrt.async({{%.*}}, [[PROMISE]])
+// CHECK-SAME: callee = @main_stream_1
+// CHECK: tf._TfrtGetResource
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: tf.AddV2
+// CHECK: [[x:%.*]] = "tf.AddV2"
+// CHECK: [[y:%.*]] = "tf_mlrt.tf_await"([[FUTURE]])
+// CHECK: [[RES:%.*]] = "tf.AddV2"([[x]], [[y]])
+// CHECK: mlrt.await_handle [[HANDLE]]
+// CHECK: return [[RES]]
+
+func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
+
+  %c0 = "tf._TfrtGetResource"() {indices = [0], shared_name = [""], container = [""]} : () -> (tensor<i32>)
+
+  %a0 = "tf.AddV2"(%a, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a1 = "tf.AddV2"(%a0, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a2 = "tf.AddV2"(%a1, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %a3 = "tf.AddV2"(%a2, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %b0 = "tf.Sub"(%b, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b1 = "tf.Sub"(%b0, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b2 = "tf.Sub"(%b1, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %b3 = "tf.Sub"(%b2, %c0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  %c = "tf.AddV2"(%a3, %b3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  func.return %c : tensor<i32>
+}
+
+// -----
+
+// Test that constants private to a stream are still handled properly when we are copying shared constants.
+
+// CHECK-LABEL: func private @main_stream_1
+// CHECK: [[r:%.*]] = "tf._TfrtGetResource"
+// CHECK-SAME: indices = [1]
+// CHECK: "tf.DummySideEffecting"([[r]])
+
+// CHECK-LABEL: func private @main_stream_2
+// CHECK: [[r:%.*]] = "tf._TfrtGetResource"
+// CHECK-SAME: indices = [0]
+// CHECK: "tf.DummySideEffecting"([[r]])
+
+// CHECK-LABEL: func @main
+
+func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> () {
+
+  %c0 = "tf._TfrtGetResource"() {indices = [0], shared_name = [""], container = [""]} : () -> (tensor<i32>)
+  "tf.DummySideEffecting"(%c0) : (tensor<i32>) -> ()
+
+  %c1 = "tf._TfrtGetResource"() {indices = [1], shared_name = [""], container = [""]} : () -> (tensor<i32>)
+  "tf.DummySideEffecting"(%c1) : (tensor<i32>) -> ()
+
+  func.return
+}
+
+// -----
+
+// Test that streams with no args but side-effecting ops are still created properly
+
+// CHECK-LABEL: func private @main_stream_1()
+// CHECK: [[r:%.*]] = "tf._TfrtGetResource"
+// CHECK-SAME: indices = [0]
+// CHECK: "tf.DummySideEffecting"([[r]])
+
+// CHECK-LABEL: func @main
+
+func.func @main(%a: tensor<i32>, %b: tensor<i32>) -> () {
+  %c0 = "tf._TfrtGetResource"() {indices = [0], shared_name = [""], container = [""]} : () -> (tensor<i32>)
+  "tf.DummySideEffecting"(%c0) : (tensor<i32>) -> ()
+  func.return
+}
+
+// -----
+
+// Test control deps of tf.Assert is skipped.
+
+// CHECK-LABEL: func.func private @skip_assert_stream_3(
+// CHECK-NOT: mlrt.await_control
+// CHECK: tf.Assert
+// CHECK-NOT: mlrt.promise_control
+// CHECK: return
+
+// CHECK-LABEL: func.func private @skip_assert_stream_2(
+// CHECK-NOT: mlrt.await_control
+// CHECK: tf.Assert
+// CHECK-NOT: mlrt.promise_control
+// CHECK: return
+
+func.func @skip_assert(%key: tensor<!tf_type.string>) -> (tensor<i64>, tensor<i64>) {
+  %error_message = "tf.Const"() {value = dense<"error"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+  %default = "tf.Const"() {value = dense<-1> : tensor<i64>} : () -> tensor<i64>
+  %handle = "tf.HashTableV2"() {container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", key_dtype = !tf_type.string, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>
+
+
+  %keys = "tf.Const"() {value = dense<["a", "b", "c", "d"]> : tensor<4x!tf_type.string>} : () -> tensor<4x!tf_type.string>
+  %values = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi64>} : () -> tensor<4xi64>
+  "tf.LookupTableImportV2"(%handle, %keys, %values) {device = ""} : (tensor<!tf_type.resource>, tensor<4x!tf_type.string>, tensor<4xi64>) -> ()
+  %value0 = "tf.LookupTableFindV2"(%handle, %key, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>
+  %cond = "tf.Equal"(%value0, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0", incompatible_shape_error = true} : (tensor<i64>, tensor<i64>) -> tensor<i1>
+  "tf.Assert"(%cond, %error_message) {device = "/job:localhost/replica:0/task:0/device:CPU:0", summarize = 3 : i64} : (tensor<i1>, tensor<!tf_type.string>) -> ()
+  "tf.Assert"(%cond, %error_message) {device = "/job:localhost/replica:0/task:0/device:CPU:0", summarize = 3 : i64} : (tensor<i1>, tensor<!tf_type.string>) -> ()
+  %value1 = "tf.LookupTableFindV2"(%handle, %key, %default) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource>, tensor<!tf_type.string>, tensor<i64>) -> tensor<i64>
+  func.return %value0, %value1 : tensor<i64>, tensor<i64>
+}

tensorflow/compiler/mlir/tfrt/tests/mlrt/rewrite_ifrt_load_variable.mlir

@@ -0,0 +1,57 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-rewrite-ifrt-load-variable %s | FileCheck %s
+
+// Variable is used by both CPU and TPU
+//
+// CHECK-LABEL: func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    [[HANDLE:%.*]] = "tf.VarHandleOp"()
+// CHECK-NEXT:    [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])
+// CHECK-SAME:       <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)
+// CHECK-NEXT:    [[TENSOR:%.*]] = "tf_mlrt.tf_await"([[FURTURE]]) : (!mlrt.future) -> tensor<3x1xf32>
+// CHECK-NEXT:    "tf.MatMul"(%arg0, [[TENSOR]]) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+// CHECK-NEXT:    "tf.IfrtCall"(%arg0, [[ARRAYKEY]]) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}> {__tpu_compile_metadata_text = "retvals { sharding { } }"} : (tensor<1x3xf32>, tensor<!tf_type.string>) -> tensor<1x1xf32>
+// CHECK-NEXT:    return
+//
+ func.func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %array_key, %tensor = "tf.IfrtLoadVariable"(%0) <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, tensor<3x1xf32>)
+    %1 = "tf.MatMul"(%arg0, %tensor) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
+    %2 = "tf.IfrtCall"(%arg0, %array_key) <{operandSegmentSizes = array<i32: 2, 0>, program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}> {__tpu_compile_metadata_text = "retvals { sharding { } }"} : (tensor<1x3xf32>, tensor<!tf_type.string>) -> tensor<1x1xf32>
+    return %2 : tensor<1x1xf32>
+  }
+
+
+// -----
+// Variable is used by two CPU ops
+//
+// CHECK-LABEL: func @serving_default
+// CHECK-NEXT:    [[HANDLE:%.*]] = "tf.VarHandleOp"()
+// CHECK-NEXT:    [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])
+// CHECK-SAME:       <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)
+// CHECK:         [[TENSOR:%.*]] = "tf_mlrt.tf_await"([[FURTURE]]) : (!mlrt.future) -> tensor<3x1xf32>
+// CHECK-NEXT:    "tf.AddV2"([[TENSOR]], %cst) : (tensor<3x1xf32>, tensor<3x1xf32>) -> tensor<3x1xf32>
+// CHECK-NEXT:    "tf.Sub"([[TENSOR]], %cst) : (tensor<3x1xf32>, tensor<3x1xf32>) -> tensor<3x1xf32>
+// CHECK-NEXT:    return
+//
+ func.func @serving_default() {
+    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+    %array_key, %tensor = "tf.IfrtLoadVariable"(%0) <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, tensor<3x1xf32>)
+    %cst_24 = "tf.Const"() <{value = dense<[[0.0], [1.0], [2.0]]> : tensor<3x1xf32>}> : () -> tensor<3x1xf32>
+    %1 = "tf.AddV2"(%tensor, %cst_24) : (tensor<3x1xf32>, tensor<3x1xf32>) -> tensor<3x1xf32>
+    %2 = "tf.Sub"(%tensor, %cst_24) : (tensor<3x1xf32>, tensor<3x1xf32>) -> tensor<3x1xf32>
+
+    return 
+  }

tensorflow/compiler/mlir/tfrt/tests/mlrt/tf_to_mlrt.mlir

@@ -0,0 +1,616 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-to-mlrt %s | FileCheck %s
+
+// CHECK-LABEL: @main_stream_0
+// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[promise_b:%.*]]: !mlrt.promise)
+func.func @main_stream_0(%input0: tensor<i32>, %promise_b: !mlrt.promise) {
+  %const = "tf.Const"() {__op_key = 0 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[a:%.*]] = tf_mlrt.executeop([[input0]],
+  // CHECK-SAME: AddV2
+  %a = "tf.AddV2"(%input0, %const) {__op_key = 1: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: [[b:%.*]] = tf_mlrt.executeop([[a]])
+  // CHECK-SAME: Abs
+  %b = "tf.Abs"(%a) {__op_key = 2 : i32}: (tensor<i32>) -> tensor<i32>
+  // CHECK: tf_mlrt.promise [[promise_b]], [[b]]
+  "tf_mlrt.tf_promise"(%promise_b, %b) : (!mlrt.promise, tensor<i32>) -> ()
+  // CHECK: return
+  return
+}
+
+// CHECK-LABEL: @main_stream_1
+// CHECK-SAME: ([[input1:%.*]]: !tf_mlrt.tensor, [[promise_c:%.*]]: !mlrt.promise, [[promise_d:%.*]]: !mlrt.promise)
+func.func @main_stream_1(%input1: tensor<i32>, %promise_c: !mlrt.promise, %promise_d: !mlrt.promise) {
+  %const = "tf.Const"() {__op_key = 3 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[c:%.*]] = tf_mlrt.executeop([[input1]],
+  // CHECK-SAME: Sub
+  %c = "tf.Sub"(%input1, %const) {__op_key = 4: i32} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: tf_mlrt.promise [[promise_c]], [[c]]
+  "tf_mlrt.tf_promise"(%promise_c, %c) : (!mlrt.promise, tensor<i32>) -> ()
+  // CHECK: [[d:%.*]] = tf_mlrt.executeop([[c]])
+  // CHECK-SAME: Abs
+  %d = "tf.Abs"(%c) {__op_key = 5: i32}: (tensor<i32>) -> tensor<i32>
+  // CHECK: tf_mlrt.promise [[promise_d]], [[d]]
+  "tf_mlrt.tf_promise"(%promise_d, %d) : (!mlrt.promise, tensor<i32>) -> ()
+  // CHECK: return
+  return
+}
+
+// CHECK-LABEL: @main
+// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[input1:%.*]]: !tf_mlrt.tensor)
+func.func @main(%input0: tensor<i32>, %input1: tensor<i32>) -> tensor<i32> {
+  // CHECK: [[promises:%.*]]:3, [[futures:%.*]]:3 = "tf_mlrt.allocate_futures"
+  // CHECK-SAME: num_futures = 3
+  %promise_b, %promise_c, %promise_d, %future_b, %future_c, %future_d =
+    "tf_mlrt.allocate_futures"()
+    {num_futures = 3 : i32, resultSegmentSizes = array<i32: 3, 3>} : () ->
+    (!mlrt.promise, !mlrt.promise, !mlrt.promise,
+     !mlrt.future, !mlrt.future, !mlrt.future)
+
+  // CHECK: [[handle_0:%.*]] = mlrt.async([[input0]], [[promises]]#0)
+  // CHECK-SAME: callee = @main_stream_0
+  %handle_0 = mlrt.async(%input0, %promise_b)
+    {callee = @main_stream_0} :
+    (tensor<i32>, !mlrt.promise) -> !mlrt.async_handle
+  // CHECK: [[handle_1:%.*]] = mlrt.async([[input1]], [[promises]]#1, [[promises]]#2)
+  // CHECK-SAME: callee = @main_stream_1
+  %handle_1 = mlrt.async(%input1, %promise_c, %promise_d)
+    {callee = @main_stream_1} :
+    (tensor<i32>, !mlrt.promise, !mlrt.promise) -> !mlrt.async_handle
+
+  %const = "tf.Const"() {__op_key = 6: i32, value = dense<2> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[e:%.*]] = tf_mlrt.executeop([[input1]],
+  // CHECK-SAME: Mul
+  %e = "tf.Mul"(%input1, %const) {__op_key = 7: i32} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: [[c:%.*]] = tf_mlrt.await [[futures]]#1
+  %c = "tf_mlrt.tf_await"(%future_c) : (!mlrt.future) ->tensor<i32>
+  // CHECK: [[f:%.*]] = tf_mlrt.executeop([[e]], [[c]])
+  // CHECK-SAME: Div
+  %f = "tf.Div"(%e, %c) {__op_key = 8: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  // CHECK: [[b:%.*]] = tf_mlrt.await [[futures]]#0
+  %b = "tf_mlrt.tf_await"(%future_b) : (!mlrt.future) ->tensor<i32>
+  // CHECK: [[d:%.*]] = tf_mlrt.await [[futures]]#2
+  %d = "tf_mlrt.tf_await"(%future_d) : (!mlrt.future) ->tensor<i32>
+
+  // CHECK: [[result:%.*]] = tf_mlrt.executeop([[b]], [[d]], [[f]])
+  // CHECK-SAME: AddN
+  %result = "tf.AddN"(%b, %d, %f) {__op_key = 9: i32}: (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  // CHECK: mlrt.await_handle [[handle_0]]
+  // CHECK: mlrt.await_handle [[handle_1]]
+  mlrt.await_handle %handle_0
+  mlrt.await_handle %handle_1
+
+  // CHECK: return [[result]]
+  return %result : tensor<i32>
+}
+
+// -----
+
+// Test lowering tf.If
+
+func.func @then(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
+  return %x: tensor<i32>
+}
+
+func.func @else(%x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
+  return %y: tensor<i32>
+}
+
+// CHECK-LABEL: func @main
+// CHECK-SAME: ([[cond_tensor:%.*]]: !tf_mlrt.tensor, [[x:%.*]]: !tf_mlrt.tensor, [[y:%.*]]: !tf_mlrt.tensor)
+// CHECK: [[cond:%.*]] = tf_mlrt.predicate [[cond_tensor]]
+// CHECK: [[z:%.*]] = mlrt.cond [[cond]] @then @else([[x]], [[y]])
+// CHECK: return [[z]]
+func.func @main(%cond: tensor<i1>, %x: tensor<i32>, %y: tensor<i32>) -> tensor<i32> {
+  %z = "tf.If"(%cond, %x, %y) {then_branch = @then, else_branch = @else, is_stateless = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> tensor<i32>
+  return %z: tensor<i32>
+}
+
+// -----
+
+// Test lowering AsyncOpKernel
+
+// CHECK-LABEL: func @main
+func.func @main(%x: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
+  // CHECK: [[y_future:%.*]] = tf_mlrt.async_executeop
+  %y = "tf.TestAsyncIdentity"(%x) {__op_key = 0: i32, T = i32} : (tensor<i32>) -> tensor<i32>
+  // CHECK: [[z:%.*]] = tf_mlrt.executeop
+  %z = "tf.Identity"(%x) {__op_key = 1: i32}: (tensor<i32>) -> tensor<i32>
+  // CHECK: [[y:%.*]] = tf_mlrt.await [[y_future]]
+  // CHECK-NEXT: tf_mlrt.executeop([[y]]
+  %w = "tf.AddV2"(%y, %z) {__op_key = 2: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK-NEXT: tf_mlrt.executeop([[y]]
+  %u = "tf.AddV2"(%y, %z) {__op_key = 3: i32} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK-NEXT: tf_mlrt.executeop([[y]]
+  %v = "tf.AddV2"(%y, %z) {__op_key = 4: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  return %w, %u, %v : tensor<i32>, tensor<i32>, tensor<i32>
+}
+
+// -----
+
+// Test lowering BatchFunction op.
+
+func.func @batched_function(%x: tensor<?xi32>) -> tensor<?xi32> {
+  return %x : tensor<?xi32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%x: tensor<1xi32>) -> (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) {
+  // CHECK: [[y_future:%.*]] = tf_mlrt.batch_function
+  // CHECK-SAME: f = @batched_function
+  // CHECK-SAME: \22batch_function\22
+  %y = "tf.BatchFunction"(%x) {
+    allowed_batch_sizes = [6], batch_timeout_micros = 100000 : i64,
+    batching_queue = "", container = "", device = "/device:CPU:0",
+    enable_large_batch_splitting = false, f = @batched_function,
+    max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64,
+    num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>,
+    shared_name = "batch_function"
+  } : (tensor<1xi32>) -> tensor<1xi32>
+
+  // CHECK: [[z:%.*]] = tf_mlrt.executeop
+  %z = "tf.Identity"(%x) {__op_key = 0: i32} : (tensor<1xi32>) -> tensor<1xi32>
+  // CHECK: [[y:%.*]] = tf_mlrt.await [[y_future]]
+  // CHECK-NEXT: tf_mlrt.executeop([[y]]
+  %w = "tf.AddV2"(%y, %z) {__op_key = 1: i32}: (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+  // CHECK-NEXT: tf_mlrt.executeop([[y]]
+  %u = "tf.AddV2"(%y, %z) {__op_key = 2: i32}: (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+  // CHECK-NEXT: tf_mlrt.executeop([[y]]
+  %v = "tf.AddV2"(%y, %z) {__op_key = 3: i32}: (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
+  return %w, %u, %v : tensor<1xi32>, tensor<1xi32>, tensor<1xi32>
+}
+
+// -----
+
+// Test node names are preserved.
+
+// CHECK-LABEL: func @main
+func.func @main(%x: tensor<i32>) -> tensor<i32> {
+  // CHECK: tf_mlrt.executeop
+  // CHECK-SAME: name: \22name_loc/AddV2_0\22
+  %y = "tf.AddV2"(%x, %x) {__op_key = 0: i32} : (tensor<i32>, tensor<i32>) -> tensor<i32> loc("name_loc:AddV2")
+  // CHECK: tf_mlrt.executeop
+  // CHECK-SAME: name: \22fused_loc/AddV2_1\22
+  %z = "tf.AddV2"(%y, %x) {__op_key = 1: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32> loc(fused["fused_loc:", "AddV2"])
+  // CHECK: tf_mlrt.executeop
+  // CHECK-SAME: name: \22AddV2_2\22
+  %w = "tf.AddV2"(%z, %x) {__op_key = 2: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  return %z : tensor<i32>
+}
+
+// -----
+
+// Test function name canonicalization
+
+// CHECK-LABEL: func @__inference_pruned_35
+func.func @__inference_pruned_35() -> tensor<!tf_type.variant> attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "flatmapdataset__4_RetVal"}} {
+  %0 = "tf.Const"() {__op_key = 0: i32, device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
+  %1 = "tf.Const"() {__op_key = 1: i32, device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
+  %2 = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
+  %3 = "tf.RangeDataset"(%0, %1, %2) {__op_key = 3: i32, device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
+  // CHECK: tf_mlrt.executeop{{.*}}op: \22FlatMapDataset\22
+  // CHECK-SAME: \22__inference_Dataset_flat_map_lambda_19\22
+  %4 = "tf.FlatMapDataset"(%3) {__op_key = 4: i32, Targuments = [], device = "/device:CPU:0", f = @__inference_Dataset_flat_map_lambda_190, output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
+  func.return %4 : tensor<!tf_type.variant>
+}
+// CHECK-LABEL: __inference_Dataset_flat_map_lambda_190
+func.func private @__inference_Dataset_flat_map_lambda_190(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}) -> tensor<!tf_type.variant> attributes {tf._original_func_name = "__inference_Dataset_flat_map_lambda_19", tf._tf_data_function = true, tf.signature.is_stateful} {
+  %0 = "tf.Const"() {__op_key = 5: i32, device = "/device:CPU:0", value = dense<0> : tensor<i64>} : () -> tensor<i64>
+  %1 = "tf.Const"() {__op_key = 6: i32,device = "/device:CPU:0", value = dense<1> : tensor<i64>} : () -> tensor<i64>
+  %2 = "tf.Const"() {__op_key = 7: i32,device = "/device:CPU:0", value = dense<5> : tensor<i64>} : () -> tensor<i64>
+  %3 = "tf.RangeDataset"(%0, %2, %1) {__op_key = 8: i32, device = "/device:CPU:0", output_shapes = [#tf_type.shape<>], output_types = [i64], metadata = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<!tf_type.variant>
+  // CHECK: tf_mlrt.executeop{{.*}}op: \22MapDataset\22
+  // CHECK-SAME: \22__inference_Dataset_map_lambda_16\22
+  %4 = "tf.MapDataset"(%3) {__op_key = 9: i32, device = "/device:CPU:0", f = @__inference_Dataset_map_lambda_160, f._tf_data_function = true, output_shapes = [#tf_type.shape<>], output_types = [i64], preserve_cardinality = true, use_inter_op_parallelism = true, metadata = ""} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
+  %5 = "tf.Identity"(%4) {__op_key = 10: i32, device = "/device:CPU:0"} : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>
+  func.return %5 : tensor<!tf_type.variant>
+}
+// CHECK-LABEL: __inference_Dataset_map_lambda_160
+func.func private @__inference_Dataset_map_lambda_160(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}) -> tensor<i64> attributes {tf._tf_data_function = true} {
+  %0 = "tf.Const"() {__op_key = 11: i32, device = "/device:CPU:0", value = dense<2> : tensor<i64>} : () -> tensor<i64>
+  %1 = "tf.Mul"(%arg0, %0) {__op_key = 12: i32, device = "/device:CPU:0"} : (tensor<i64>, tensor<i64>) -> tensor<i64>
+  %2 = "tf.Identity"(%1) {__op_key = 13: i32, device = "/device:CPU:0"} : (tensor<i64>) -> tensor<i64>
+  func.return %2 : tensor<i64>
+}
+
+// -----
+
+// Test while conversion
+
+// CHECK-LABEL: func @while_cond_lt9
+// CHECK-SAME: ([[arg0:%.*]]: !tf_mlrt.tensor) -> !tf_mlrt.tensor
+func.func @while_cond_lt9(%arg0: tensor<i32>) -> tensor<i1> {
+  %0 = "tf.Const"() {__op_key = 0: i32, device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
+  %1 = "tf.Less"(%arg0, %0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  func.return %1 : tensor<i1>
+}
+
+// CHECK-LABEL: func @while_body_add2
+// CHECK-SAME: ([[arg0:%.*]]: !tf_mlrt.tensor) -> !tf_mlrt.tensor
+func.func @while_body_add2(%arg0: tensor<i32>) -> tensor<i32> {
+  %0 = "tf.Const"() {__op_key = 2: i32, device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
+  %1 = "tf.Add"(%arg0, %0) {__op_key = 3: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %1 : tensor<i32>
+}
+
+// CHECK-LABEL: func @while_test()
+// CHECK-SAME: -> !tf_mlrt.tensor
+func.func @while_test() -> (tensor<i32>) {
+  // CHECK: [[CONST:%.*]] = tf_mlrt.constop
+  %0 = "tf.Const"() {__op_key = 4: i32, device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[pred_res:%.*]] = call @"while_cond_lt9/tf_mlrt_predicate"([[CONST]]) : (!tf_mlrt.tensor) -> i1
+  // CHECK: [[while_res:%.*]]:2 = mlrt.while
+  // CHECK-SAME: @"while_body_add2/tf_mlrt_body"([[CONST]])
+  // CHECK-SAME: (!tf_mlrt.tensor) -> (!tf_mlrt.tensor, i1)
+  %1 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
+  // CHECK: return [[while_res]]#0 : !tf_mlrt.tensor
+  func.return %1 : tensor<i32>
+}
+// CHECK: func @"while_body_add2/tf_mlrt_body"([[arg:%.*]]: !tf_mlrt.tensor) -> (!tf_mlrt.tensor, i1)
+// CHECK: [[body_res:%.*]] = call @while_body_add2([[arg]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
+// CHECK: [[pred_res:%.*]] = call @"while_cond_lt9/tf_mlrt_predicate"([[body_res]]) : (!tf_mlrt.tensor) -> i1
+// CHECK: return [[body_res]], [[pred_res]] : !tf_mlrt.tensor, i1
+
+// CHECK: func @"while_cond_lt9/tf_mlrt_predicate"([[arg:%.*]]: !tf_mlrt.tensor) -> i1
+// CHECK: [[cond_res:%.*]] = call @while_cond_lt9([[arg]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
+// CHECK: [[bool_res:%.*]] = tf_mlrt.predicate [[cond_res]]
+// CHECK: return [[bool_res]] : i1
+
+// CHECK-LABEL: func @multi_while_test
+func.func @multi_while_test() -> (tensor<i32>, tensor<i32>) {
+  %0 = "tf.Const"() {__op_key = 5: i32, device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %1 = "tf.Const"() {__op_key = 6: i32, device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[pred_0:%.*]] = call @"while_cond_lt9/tf_mlrt_predicate"
+  // CHECK: mlrt.while [[pred_0]] @"while_body_add2/tf_mlrt_body"
+  // CHECK: [[pred_1:%.*]] = call @"while_cond_lt9/tf_mlrt_predicate"
+  // CHECK: mlrt.while [[pred_1]] @"while_body_add2/tf_mlrt_body"
+  %2 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
+  %3 = "tf.While"(%1) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
+  func.return %2, %3 : tensor<i32>, tensor<i32>
+}
+
+// -----
+
+// Test async output to function is converted
+
+// CHECK-LABEL: @serving_default_stream_1
+// CHECK-SAME: !mlrt.future
+func.func private @serving_default_stream_1(%arg0: tensor<i32>) {
+  // CHECK: [[tensor:%.*]] = tf_mlrt.await
+  // CHECK: tf_mlrt.executeop([[tensor]])
+  %0 = "tf.StringFormat"(%arg0) {__op_key = 0: i32, device = "/job:localhost/replica:0/task:0/device:CPU:0", placeholder = "{}", strtemplate = "%s", summarize = 3 : i64, template = "Outside compiled {}"} : (tensor<i32>) -> tensor<!tf_type.string>
+  "tf.PrintV2"(%0) {__op_key = 1: i32, device = "/job:localhost/replica:0/task:0/device:CPU:0", end = "\0A", output_stream = "stderr"} : (tensor<!tf_type.string>) -> ()
+  return
+}
+
+func.func @callee(%arg: tensor<i32>) -> (tensor<i32>) {
+  func.return %arg: tensor<i32>
+}
+
+// CHECK-LABEL: @executeop_input
+func.func @executeop_input(%arg0: tensor<i32>) -> (tensor<i32>) {
+  // CHECK: [[async_out:%.*]] = tf_mlrt.batch_function
+  %2 = "tf.BatchFunction"(%arg0) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = ""} : (tensor<i32>) -> tensor<i32>
+  // CHECK-NEXT: mlrt.async([[async_out]]) {{.*}} : (!mlrt.future)
+  %3 = mlrt.async(%2) {callee = @serving_default_stream_1} : (tensor<i32>) -> !mlrt.async_handle
+  // CHECK: mlrt.await_handle
+  mlrt.await_handle %3
+  // CHECK: return
+  // CHECK-SAME: !tf_mlrt.tensor
+  func.return %2 : tensor<i32>
+}
+
+// -----
+
+// Support pre-assigned op_key
+
+// CHECK-LABEL: @main
+// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[promise_b:%.*]]: !mlrt.promise)
+func.func @main(%input0: tensor<i32>, %promise_b: !mlrt.promise) {
+  %const = "tf.Const"() {__op_key = 0 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[a:%.*]] = tf_mlrt.executeop([[input0]],
+  // CHECK-SAME: AddV2
+  // CHECK-SAME: op_key = 1
+  // CHECK-NOT: __op_key
+  %a = "tf.AddV2"(%input0, %const) {__op_key = 1: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: [[b:%.*]] = tf_mlrt.executeop([[a]])
+  // CHECK-SAME: Abs
+  // CHECK-SAME: op_key = 2
+  // CHECK-NOT: __op_key
+  %b = "tf.Abs"(%a) {__op_key = 2: i32 }: (tensor<i32>) -> tensor<i32>
+  // CHECK: tf_mlrt.promise [[promise_b]], [[b]]
+  "tf_mlrt.tf_promise"(%promise_b, %b) : (!mlrt.promise, tensor<i32>) -> ()
+  // CHECK: return
+  return
+}
+
+// -----
+
+// Test future as input to promise
+
+// CHECK-LABEL: func @main_stream_0
+func.func @main_stream_0(%x: tensor<i32>, %p: !mlrt.promise) -> () {
+  // CHECK: [[y_future:%.*]] = tf_mlrt.async_executeop
+  %y = "tf.TestAsyncIdentity"(%x) {__op_key = 0: i32, T = i32} : (tensor<i32>) -> tensor<i32>
+  // CHECK: tf_mlrt.promise_future
+  // CHECK-SAME: [[y_future]]
+  "tf_mlrt.tf_promise"(%p, %y): (!mlrt.promise, tensor<i32>) -> ()
+  return
+}
+
+// CHECK-LABEL: @main
+// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor)
+func.func @main(%input0: tensor<i32>) -> tensor<i32> {
+  // CHECK: [[promises:%.*]], [[futures:%.*]] = "tf_mlrt.allocate_futures"
+  // CHECK-SAME: num_futures = 1
+  %promise_b, %future_b = "tf_mlrt.allocate_futures"()
+    {num_futures = 1 : i32, resultSegmentSizes = array<i32: 1, 1>} : () ->
+    (!mlrt.promise, !mlrt.future)
+
+  // CHECK: [[handle_0:%.*]] = mlrt.async([[input0]], [[promises]])
+  // CHECK-SAME: callee = @main_stream_0
+  %handle_0 = mlrt.async(%input0, %promise_b)
+    {callee = @main_stream_0} :
+    (tensor<i32>, !mlrt.promise) -> !mlrt.async_handle
+
+  // CHECK: [[const:%.*]]  = tf_mlrt.const
+  %const = "tf.Const"() {__op_key = 1: i32, value = dense<2> : tensor<i32>} : () -> tensor<i32>
+
+  // CHECK: [[b:%.*]] = tf_mlrt.await [[futures]]
+  %b = "tf_mlrt.tf_await"(%future_b) : (!mlrt.future) ->tensor<i32>
+
+  // CHECK: [[result:%.*]] = tf_mlrt.executeop([[b]], [[const]])
+  // CHECK-SAME: AddV2
+  %result = "tf.AddV2"(%b, %const) {__op_key = 2: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+  // CHECK: mlrt.await_handle [[handle_0]]
+  mlrt.await_handle %handle_0
+
+  // CHECK: return [[result]]
+  return %result : tensor<i32>
+}
+
+// -----
+
+// Test lowering of tf call ops
+
+// CHECK-LABEL: @callee
+func.func @callee(%arg0: tensor<i32>) -> (tensor<i32>) {
+  func.return %arg0: tensor<i32>
+}
+
+// CHECK-LABEL: func @call_test
+func.func @call_test(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<i32>) {
+  %0 = "tf.Add"(%arg0, %arg0) {__op_key = 0, device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  // CHECK: [[results_0:%.*]] = call @callee(
+  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
+  %1 = "tf.StatefulPartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
+  // CHECK-NEXT: [[results_1:%.*]] = call @callee(
+  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
+  %2 = "tf.PartitionedCall"(%0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> (tensor<i32>)
+  // CHECK-NEXT: [[results_2:%.*]] = call @callee(
+  // CHECK-SAME: (!tf_mlrt.tensor) -> !tf_mlrt.tensor
+  %3 = "tf.LegacyCall"(%0) {f = @callee} : (tensor<i32>) -> (tensor<i32>)
+  // CHECK: [[results_0]], [[results_1]], [[results_2]]
+  func.return %1, %2, %3 : tensor<i32>, tensor<i32>, tensor<i32>
+}
+
+// CHECK-LABEL: @branch0
+func.func @branch0(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
+  %0 = "tf.Add" (%arg0, %arg1) {__op_key = 1, device = "/device:CPU:0"}  : (tensor<f32>, tensor<f32>) -> tensor<f32>
+  func.return %0 : tensor<f32>
+}
+
+// CHECK-LABEL: @branch1
+func.func @branch1(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
+  %0 = "tf.Add" (%arg0, %arg1) {__op_key = 2, device = "/device:CPU:0"}  : (tensor<f32>, tensor<f32>) -> tensor<f32>
+  %1 = "tf.Add" (%arg0, %0) {__op_key = 3, device = "/device:CPU:0"}  : (tensor<f32>, tensor<f32>) -> tensor<f32>
+  func.return %1 : tensor<f32>
+}
+
+// CHECK-LABEL: func @case_test
+// CHECK-SAME: ([[tf_idx:%.*]]: !tf_mlrt.tensor, [[branch_arg0:%.*]]: !tf_mlrt.tensor, [[branch_arg1:%.*]]: !tf_mlrt.tensor)
+func.func @case_test(%arg0: tensor<i32>, %arg1: tensor<f32>,  %arg2: tensor<f32>) -> tensor<f32> {
+  // CHECK: [[idx:%.*]] = tf_mlrt.tensor_to_int32 [[tf_idx]]
+  // CHECK-NEXT: [[out:%.*]] = mlrt.case [[idx]] [@branch0, @branch1]([[branch_arg0]], [[branch_arg1]])
+  %0 = "tf.Case"(%arg0, %arg1, %arg2) {_lower_using_switch_merge = true, branches = [@branch0, @branch1], is_stateless = true} : (tensor<i32>, tensor<f32>, tensor<f32>) -> tensor<f32>
+  func.return %0 : tensor<f32>
+}
+
+// -----
+
+// Test await is added for unused futures
+
+// CHECK-LABEL: func @unused_future_arg
+// CHECK-SAME: ({{%.*}}: !tf_mlrt.tensor, [[unused:%.*]]: !mlrt.future)
+func.func @unused_future_arg(%x: tensor<i32>, %unused: !mlrt.future) -> tensor<i32> {
+  // CHECK: mlrt.await_all_control [[unused]]
+  return %x : tensor<i32>
+}
+
+// CHECK-LABEL: func @unused_future
+func.func @unused_future(%x: tensor<i32>) -> tensor<i32> {
+  // CHECK: [[unused:%.*]] = tf_mlrt.async_executeop
+  %unused = "tf.TestAsyncIdentity"(%x) {__op_key = 0: i32, T = i32} : (tensor<i32>) -> tensor<i32>
+  // CHECK: mlrt.await_all_control [[unused]]
+  return %x : tensor<i32>
+}
+
+// -----
+
+// Test for XlaLaunch
+
+func.func private @xla_func_0(%arg0: tensor<1x3xf32>, %arg1: tensor<1x3xf32>) -> tensor<1x3xf32> attributes {tf._XlaMustCompile = true, tf._noinline = true, tf._original_func_name = "should_not_be_used"} {
+  %1 = "tf.AddV2"(%arg0, %arg1) {__op_key = 0: i32} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %1 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @xla_func
+func.func @xla_func(%arg0: tensor<1x3xf32>) -> tensor<*xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "output:0"}} {
+  %0 = "tf.VarHandleOp"() {__op_key = 1: i32, device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+  %1 = "tf.ReadVariableOp"(%0) {__op_key = 2: i32, device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> tensor<1x3xf32>
+  // CHECK: tf_mlrt.executeop
+  // CHECK: tf_mlrt.async_executeop{{.*}}op: \22XlaLaunch\22\0A
+  // CHECK: tf_mlrt.await
+  // CHECK: return
+  // CHECK-SAME: !tf_mlrt.tensor
+  %2 = "tf.XlaLaunch"(%arg0, %1) {__op_key = 3: i32, _noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %2 : tensor<*xf32>
+}
+
+// -----
+
+// Test lowering of IfrtLoadVariableOp
+
+// CHECK-LABEL: func @ifrt_load_variable_test
+func.func @ifrt_load_variable_test() -> () {
+  // CHECK: [[HANDLE:%.*]] = tf_mlrt.executeop()
+  // CHECK-SAME:  VarHandleOp
+  %0 = "tf.VarHandleOp"() {__op_key = 1: i32, device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>
+  // CHECK-NEXT: "tf_mlrt.ifrt_load_variable"([[HANDLE]])
+  // CHECK-SAME: used_by_host = true
+  %1, %2 = "tf_mlrt.tf_ifrt_load_variable"(%0) {used_by_host = true, __op_key = 2: i32, device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<1x3xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)
+  // CHECK-NEXT: mlrt.await_all_control
+  // CHECK-NEXT: return
+  func.return
+}
+
+// -----
+
+// Test lowering of IfrtRestoreVariableOp
+
+// CHECK-LABEL: func @ifrt_restore_variable_test
+func.func @ifrt_restore_variable_test() -> () {
+  // CHECK-NEXT: [[PREFIX:%.*]] = tf_mlrt.constop
+  %cst = "tf.Const"() {__op_key = 0: i32, value = dense<"restore_ariables"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+  // CHECK-NEXT: [[SLICE:%.*]] = tf_mlrt.constop
+  %cst_0 = "tf.Const"()  {__op_key = 1: i32, value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  // CHECK-NEXT: [[NAME:%.*]] = tf_mlrt.constop
+  %cst_1 = "tf.Const"()  {__op_key = 2: i32, value = dense<["y"]> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  // CHECK-NEXT: [[HANDLE:%.*]] = tf_mlrt.executeop
+  %handle = "tf.VarHandleOp"() {__op_key = 3: i32, container = "x", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+  // CHECK-NEXT: "tf_mlrt.ifrt_restore_variable"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLE]]) <{restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: true>}>
+  "tf.IfrtRestoreVariableOp"(%cst, %cst_1, %cst_0, %handle) {restored_dtypes = [f32], returned_tensor_names = [], truncate_in_cast = array<i1: true>} : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> ()
+  // CHECK-NEXT: return
+  func.return
+}
+
+// -----
+
+// Test lowering of tf.IfrtResourceDeserializeOp to tf_mlrt.ifrt_resource_deserialize
+
+// CHECK-LABEL: func @ifrt_resource_deserialize_test
+func.func @ifrt_resource_deserialize_test(%arg0: tensor<!tf_type.resource<tensor<f32>>>) {
+  %input_dir = "tf.Const"() { value = dense<"some/path"> : tensor<!tf_type.string> } : () -> tensor<!tf_type.string>
+  // CHECK: "tf_mlrt.ifrt_resource_deserialize"(%arg0, %{{.*}}) <{tensor_name = "my_tensor"}>
+  "tf.IfrtResourceDeserialize"(%arg0, %input_dir) {require_matching_crc = false, tensor_name = "my_tensor"} : (tensor<!tf_type.resource<tensor<f32>>>, tensor<!tf_type.string>) -> ()
+  func.return
+}
+
+
+// -----
+
+// Test lowering of tf.IfrtRestoreVariableOp with outputs to tf_mlrt.ifrt_restore_variable
+
+// CHECK-LABEL: func @ifrt_restore_variable_with_output_test
+func.func @ifrt_restore_variable_with_output_test() -> (tensor<3x1xf32>) {
+  // CHECK-NEXT: [[PREFIX:%.*]] = tf_mlrt.constop
+  %cst = "tf.Const"() {__op_key = 0: i32, value = dense<"restore_ariables"> : tensor<!tf_type.string>} : () -> tensor<!tf_type.string>
+  // CHECK-NEXT: [[SLICE:%.*]] = tf_mlrt.constop
+  %cst_0 = "tf.Const"()  {__op_key = 1: i32, value = dense<""> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  // CHECK-NEXT: [[NAME:%.*]] = tf_mlrt.constop
+  %cst_1 = "tf.Const"()  {__op_key = 2: i32, value = dense<["y"]> : tensor<1x!tf_type.string>} : () -> tensor<1x!tf_type.string>
+  // CHECK-NEXT: [[HANDLE:%.*]] = tf_mlrt.executeop
+  %handle = "tf.VarHandleOp"() {__op_key = 3: i32, container = "x", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+  // CHECK-NEXT: [[RESFUTURE:%.*]] = "tf_mlrt.ifrt_restore_variable"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLE]]) <{restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: true>}> : (!tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor, !tf_mlrt.tensor) -> !mlrt.future
+  %result = "tf.IfrtRestoreVariableOp"(%cst, %cst_1, %cst_0, %handle) {restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: true>} : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+  // CHECK-NEXT: [[RESULT:%.*]] = tf_mlrt.await [[RESFUTURE]]
+  // CHECK-NEXT: return [[RESULT]] : !tf_mlrt.tensor
+  func.return %result : tensor<3x1xf32>
+}
+
+// -----
+
+// Test lowering of tf.IfrtCall (should use fallback)
+
+// CHECK-LABEL: func @ifrt_call_fallback_test
+func.func @ifrt_call_fallback_test(%arg0: tensor<i32>) -> tensor<i32> {
+  // CHECK: tf_mlrt.executeop
+  // CHECK-SAME: IfrtCall
+  %0 = "tf.IfrtCall"(%arg0) {program_id = 123 : i64, variable_arg_indices = [], __op_key = 0 : i32, operandSegmentSizes = array<i32: 1, 0>} : (tensor<i32>) -> tensor<i32>
+  return %0 : tensor<i32>
+}
+
+// -----
+
+// Test lowering of tf.AsyncIfrtCall (should split into async call and await)
+
+// CHECK-LABEL: func @async_ifrt_call_test
+func.func @async_ifrt_call_test(%arg0: tensor<i32>) -> tensor<i32> {
+  // CHECK: [[FUTURE:%.*]] = tf_mlrt.async_ifrt_call
+  // CHECK: tf_mlrt.await [[FUTURE]]
+  %0 = "tf.AsyncIfrtCall"(%arg0) {program_id = 123 : i64, variable_arg_indices = [], __op_key = 0 : i32, operandSegmentSizes = array<i32: 1, 0>} : (tensor<i32>) -> tensor<i32>
+  return %0 : tensor<i32>
+}
+
+// -----
+
+// Test lowering of tf.PartitionedCall and tf.StatefulPartitionedCall to func.call
+
+// CHECK-LABEL: func private @callee
+// CHECK-SAME: (%arg0: !tf_mlrt.tensor) -> !tf_mlrt.tensor
+func.func private @callee(%arg0: tensor<i32>) -> tensor<i32> {
+  return %arg0 : tensor<i32>
+}
+
+// CHECK-LABEL: func @test_call_ops
+// CHECK-SAME: ([[ARG:%.*]]: !tf_mlrt.tensor) -> (!tf_mlrt.tensor, !tf_mlrt.tensor)
+func.func @test_call_ops(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
+  // CHECK: [[RES1:%.*]] = call @callee([[ARG]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
+  %0 = "tf.PartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> tensor<i32>
+  
+  // CHECK: [[RES2:%.*]] = call @callee([[ARG]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
+  %1 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @callee} : (tensor<i32>) -> tensor<i32>
+  
+  // CHECK: return [[RES1]], [[RES2]] : !tf_mlrt.tensor, !tf_mlrt.tensor
+  return %0, %1 : tensor<i32>, tensor<i32>
+}
+
+// -----
+
+// Test lowering of tf.PartitionedCall with a future input (should insert await)
+
+// CHECK-LABEL: func private @callee_future
+// CHECK-SAME: (%arg0: !tf_mlrt.tensor) -> !tf_mlrt.tensor
+func.func private @callee_future(%arg0: tensor<3x1xf32>) -> tensor<3x1xf32> {
+  return %arg0 : tensor<3x1xf32>
+}
+
+// CHECK-LABEL: func @test_call_with_future
+// CHECK-SAME: ([[ARG0:%.*]]: !tf_mlrt.tensor, [[ARG1:%.*]]: !tf_mlrt.tensor, [[ARG2:%.*]]: !tf_mlrt.tensor, [[ARG3:%.*]]: !tf_mlrt.tensor) -> !tf_mlrt.tensor
+func.func @test_call_with_future(%arg0: tensor<!tf_type.string>, %arg1: tensor<1x!tf_type.string>, %arg2: tensor<1x!tf_type.string>, %arg3: tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32> {
+  // CHECK: [[FUTURE:%.*]] = "tf_mlrt.ifrt_restore_variable"([[ARG0]], [[ARG1]], [[ARG2]], [[ARG3]])
+  %result = "tf.IfrtRestoreVariableOp"(%arg0, %arg1, %arg2, %arg3) {restored_dtypes = [f32], returned_tensor_names = ["y"], truncate_in_cast = array<i1: true>} : (tensor<!tf_type.string>, tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+  // CHECK-NEXT: [[AWAITED:%.*]] = tf_mlrt.await [[FUTURE]]
+  // CHECK-NEXT: [[RESULT:%.*]] = call @callee_future([[AWAITED]]) : (!tf_mlrt.tensor) -> !tf_mlrt.tensor
+  %0 = "tf.PartitionedCall"(%result) {config = "", config_proto = "", executor_type = "", f = @callee_future} : (tensor<3x1xf32>) -> tensor<3x1xf32>
+  // CHECK: return [[RESULT]] : !tf_mlrt.tensor
+  return %0 : tensor<3x1xf32>
+}
+
+
+

tensorflow/compiler/mlir/tfrt/tests/mlrt/tpu_conversions.mlir

@@ -0,0 +1,243 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt --split-input-file -pass-pipeline='builtin.module(pre-parallel-tf-to-mlrt{use-tpu-host-allocator-for-inputs=true},tf-mlrt-parallelization{tfrt-cost-threshold=4},tf-to-mlrt)'  %s | FileCheck %s --dump-input=fail --dump-input-filter=all
+
+func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
+  func.return %arg0: tensor<i32>
+}
+
+// CHECK-LABEL: func @batch_function
+func.func @batch_function(%arg0: tensor<i32>) -> (tensor<i32>) {
+  // CHECK: [[batch_result_future:%.*]] = tf_mlrt.batch_function
+  // CHECK: [[batch_result:%.*]] = tf_mlrt.await [[batch_result_future]]
+  // CHECK-NEXT: [[rendezvous_key_base:%.*]] = tf_mlrt_tpu.compile_and_execute([[batch_result]])
+  // CHECK-NEXT: return [[rendezvous_key_base]]
+  %0 = "tf.BatchFunction"(%arg0, %arg0) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %1 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
+  func.return %1 : tensor<i32>
+}
+
+// -----
+
+func.func @executeop_input(%arg0: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
+  // CHECK-NOT: tf_mlrt.executeop(
+  // CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
+  // CHECK: [[cast:%.*]] = tf_mlrt.executeop.device([[device]]){{.*}}op: \22Cast\22
+  // CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute([[cast]])
+  // CHECK: tf_mlrt.await [[result_future]]
+  %0 = "tf.Cast"(%arg0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  %1, %2 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> (tensor<i32>, tensor<i32>)
+  func.return %1, %2 : tensor<i32>, tensor<i32>
+}
+
+// -----
+
+func.func @executeop_side_effecting_input(%arg0: tensor<!tf_type.resource<tensor<4xf32>>>, %indices: tensor<i32>) -> (tensor<i32>) {
+  // CHECK-NOT: tf_mlrt.executeop(
+  // CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
+  // CHECK: [[var:%.*]] = tf_mlrt.executeop.device([[device]]){{.*}}op: \22ResourceGather\22
+  // CHECK: [[rendezvous_key_base:%.*]] = tf_mlrt_tpu.compile_and_execute([[var]])
+  %0 = "tf.ResourceGather"(%arg0, %indices) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<!tf_type.resource<tensor<4xf32>>>, tensor<i32>) -> tensor<f32>
+  %1 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> tensor<i32>
+  func.return %1 : tensor<i32>
+}
+
+// -----
+
+func.func @executeop_input_same_execute_op(%arg0: tensor<i32>, %arg1: tensor<2xf32>) -> (tensor<i32>) {
+  // CHECK-NOT: tf_mlrt.executeop(
+  // CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
+  // CHECK: [[split:%.*]]:2 = tf_mlrt.executeop.device([[device]])
+  // CHECK: tf_mlrt_tpu.compile_and_execute([[split]]#0, [[split]]#1)
+  %0, %1 = "tf.Split"(%arg0, %arg1) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>, tensor<2xf32>) -> (tensor<f32>, tensor<f32>)
+  %2 = "tf.TPUCompileMlirAndExecute"(%0, %1) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, producer_name = "producer_name"} : (tensor<f32>, tensor<f32>) -> tensor<i32>
+  func.return %2 : tensor<i32>
+}
+
+// -----
+
+// Test that inputs are lowered correctly when they form a DAG.
+
+// CHECK-LABEL: executeop_dag
+func.func @executeop_dag(%arg0: tensor<i32>) -> (tensor<i32>) {
+  // CHECK-NEXT: tf_mlrt_tpu.get_tpu_host_device
+  // CHECK-NEXT: tf_mlrt.executeop.device{{.*}}op: \22Cast\22
+  // CHECK-NEXT: tf_mlrt_tpu.get_tpu_host_device
+  // CHECK-NEXT: tf_mlrt.executeop.device{{.*}}op: \22Relu\22
+  // CHECK-NEXT: tf_mlrt_tpu.compile_and_execute
+  %0 = "tf.Cast"(%arg0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  %1 = "tf.Relu"(%0) {__op_key = 1: i32, device = "/device:CPU:0"} : (tensor<f32>) -> (tensor<f32>)
+  %2 = "tf.TPUCompileMlirAndExecute"(%1, %0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, producer_name = "producer_name"} : (tensor<f32>, tensor<f32>) -> tensor<i32>
+  func.return %2 : tensor<i32>
+}
+
+// -----
+
+func.func @test_fuse_dynamic_dimension_ops(%arg0: tensor<*xi32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<*xi32>, %arg3: tensor<*xi32>, %arg4: tensor<*xi32>, %arg5: tensor<?xi64>, %arg6: tensor<?xi64>, %arg7: tensor<?xi64>) -> tensor<*xi32> {
+  %0 = "tf.ReadVariableOp"(%arg1) {__op_key = 0: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
+  %1 = "tf.Shape"(%arg0) {__op_key = 1: i32, device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
+  %2 = "tf.Shape"(%0) {__op_key = 2: i32, device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
+  // CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute
+  // CHECK-SAME: constant_operand_indices = array<i32: 2>
+  // CHECK-SAME: num_operands = 4
+  // CHECK-SAME: operands_with_static_shape = array<i32: 0, 1, 3>
+  %rendezvous_key_base, %results = "tf.TPUCompileMlirAndExecute"(%arg0, %2, %0, %1, %arg5, %arg6, %arg7) {operands_with_static_shape = [0 : i32, 1 : i32, 3 : i32], metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 4, 3>, producer_name = "producer_name"} : (tensor<*xi32>, tensor<?xi64>, tensor<*xi32>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>, tensor<?xi64>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
+  func.return %results : tensor<*xi32>
+}
+
+// -----
+
+// Test async output of tf.TPUCompileMlirAndExecute to function is converted
+
+// CHECK-LABEL: @executeop_input_stream_1
+// CHECK-SAME: ([[future:%.*]]: !mlrt.future
+// CHECK: [[tensor:%.*]] = tf_mlrt.await [[future]]
+// CHECK: tf_mlrt.executeop([[tensor]])
+// CHECK-SAME: StringFormat
+
+// CHECK-LABEL: @executeop_input
+func.func @executeop_input(%arg0: tensor<i32>) -> (tensor<i32>) {
+  // CHECK: tf_mlrt.executeop
+  %0 = "tf.Cast"(%arg0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  // CHECK: [[rendezvous_key_base:%.*]], [[result:%.*]] = tf_mlrt_tpu.compile_and_execute
+  %1, %2 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> (tensor<i32>, tensor<i32>)
+  %3 = "tf.StringFormat"(%2) {__op_key = 1: i32, device = "/job:localhost/replica:0/task:0/device:CPU:0", placeholder = "{}", strtemplate = "%s", summarize = 3 : i64, template = "Outside compiled {}"} : (tensor<i32>) -> tensor<!tf_type.string>
+  "tf.PrintV2"(%3) {__op_key = 2: i32, device = "/job:localhost/replica:0/task:0/device:CPU:0", end = "\0A", output_stream = "stderr"} : (tensor<!tf_type.string>) -> ()
+  // CHECK: [[handle:%.*]] = mlrt.async([[result]])
+  // CHECK-SAME: (!mlrt.future)
+  // CHECK: mlrt.await_handle [[handle]]
+  // CHECK: return [[rendezvous_key_base]]
+  // CHECK-SAME: !tf_mlrt.tensor
+  func.return %1 : tensor<i32>
+}
+
+// -----
+
+// Test constant arguments to tf.TPUCompileMlirAndExecute are preserved during parallelization.
+
+// CHECK-LABEL: @preserve_constant_args(
+func.func @preserve_constant_args(%arg0: tensor<i32>, %arg1: tensor<*x!tf_type.resource>, %arg2: tensor<*x!tf_type.resource>, %arg3: tensor<*x!tf_type.resource>) -> (tensor<i32>) {
+  // CHECK-NOT: ReadVariableOp
+  // CHECK: mlrt.async(
+  %v0 = "tf.ReadVariableOp"(%arg1) {__op_key = 0: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<i32>
+  %v1 = "tf.ReadVariableOp"(%arg2) {__op_key = 1: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<i32>
+  // CHECK: [[cast:%.*]] = tf_mlrt.executeop(
+  // CHECK-SAME: ReadVariableOp
+  %v2 = "tf.ReadVariableOp"(%arg3) {__op_key = 2: i32, device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<i32>
+  // CHECK: [[cast:%.*]] = tf_mlrt.executeop.device
+  // CHECK-SAME: Cast
+  %0 = "tf.Cast"(%arg0) {__op_key = 3: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  // CHECK: tf_mlrt_tpu.compile_and_execute({{%.*}}, [[cast]]
+  // CHECK-SAME: constant_operand_indices = array<i32: 1, 3, 4>
+  %1, %2 = "tf.TPUCompileMlirAndExecute"(%0, %v1, %0, %v2, %v0, %arg0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 6, 0>, producer_name = "producer_name"} : (tensor<f32>, tensor<i32>, tensor<f32>, tensor<i32>, tensor<i32>, tensor<i32>) -> (tensor<i32>, tensor<i32>)
+  func.return %2 : tensor<i32>
+}
+
+// -----
+
+func.func @executeop_input_async() -> (tensor<i32>, tensor<i32>) {
+  // CHECK-NOT: tf_mlrt.executeop(
+  // CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
+  // CHECK: [[recv_future:%.*]] = tf_mlrt.async_executeop.device([[device]]){{.*}}op: \22Recv\22
+  // CHECK: [[recv:%.*]] = tf_mlrt.await [[recv_future]]
+  // CHECK: [[rendezvous_key_base:%.*]], [[result_future:%.*]] = tf_mlrt_tpu.compile_and_execute([[recv]])
+  // CHECK: tf_mlrt.await [[result_future]]
+  %0 = "tf.Recv"() {__op_key = 0: i32, device = "/device:CPU:0", tensor_name = "tensor", send_device = "/device:CPU:0", send_device_incarnation = 0, recv_device = "/device:CPU:0"} : () -> tensor<f32>
+  %1, %2 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> (tensor<i32>, tensor<i32>)
+  func.return %1, %2 : tensor<i32>, tensor<i32>
+}
+
+// -----
+
+// Test the output from TPU op is properly awaited before its use by map_fn.
+// CHECK-LABEL: @main
+// CHECK-SAME: ([[input0:%.*]]: !tf_mlrt.tensor, [[input1:%.*]]: !tf_mlrt.tensor)
+func.func @main(%input0: tensor<i32>, %input1: tensor<i32>, %input2: tensor<!tf_type.variant<tensor<*xf32>>> ) -> tensor<i32> {
+  %0 = "tf.Cast"(%input0) {__op_key = 0: i32, device = "/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  // CHECK: tf_mlrt_tpu.compile_and_execute
+  %1, %2 = "tf.TPUCompileMlirAndExecute"(%0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<f32>) -> (tensor<i32>, tensor<i32>)
+  %max_iter = "tf.Const"() {__op_key = 1, value = dense<2> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: tf_mlrt.map_fn
+  %result = "tf_mlrt.tf_map_fn"(%max_iter, %input2, %2) { operandSegmentSizes = array<i32: 1, 1, 1>, body_fn = @NopMapFnBody, num_tensor_list_or_flow_in = 1 : i32} : (tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) -> tensor<i32>
+  return %result : tensor<i32>
+}
+
+// CHECK-LABEL: @NopMapFnBody
+func.func private @NopMapFnBody(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant<tensor<*xf32>>>) -> () {
+  %const = "tf.Const"() {__op_key = 2 : i32, value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %a = "tf.AddV2"(%arg2, %const) {__op_key = 3: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  return
+}
+
+// -----
+
+// If not all in_tensors of the batched op are used by TPUCompileMlirAndExecute,
+// should not lower batch_function to tf_mlrt.batch_function.device.
+
+func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
+  %1 = "tf.TPUCompileMlirAndExecute"(%arg0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
+  %const = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %a = "tf.AddV2"(%arg1, %const) {__op_key = 3: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %1: tensor<i32>
+}
+
+// CHECK-LABEL: func @batch_function
+func.func @batch_function(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
+  // CHECK: tf_mlrt.batch_function(%arg0, %arg1)
+  // CHECK-NOT:  batch_function.device
+  %0 = "tf.BatchFunction"(%arg0, %arg1) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 2, 0>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %0 : tensor<i32>
+}
+
+// -----
+
+// captured tensors not used by TPUCompileMlirAndExecute is ok.
+
+func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
+  %1 = "tf.TPUCompileMlirAndExecute"(%arg0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
+  %const = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %a = "tf.AddV2"(%arg1, %const) {__op_key = 3: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %1: tensor<i32>
+}
+
+// CHECK-LABEL: func @batch_function
+func.func @batch_function(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
+  // CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
+  // CHECK: [[batch_result_future:%.*]] = tf_mlrt.batch_function.device([[device]]) (%arg0, %arg1)
+  // CHECK: [[batch_result:%.*]] = tf_mlrt.await [[batch_result_future]]
+  // CHECK: return [[batch_result]]
+  %0 = "tf.BatchFunction"(%arg0, %arg1) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @callee, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 1>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %0 : tensor<i32>
+}
+
+// -----
+func.func @batched_func(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
+  %0 = "tf.TPUCompileMlirAndExecute"(%arg0) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
+  %2 = "tf.TPUCompileMlirAndExecute"(%arg1) {metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 1, 0>, producer_name = "producer_name"} : (tensor<i32>) -> tensor<i32>
+  func.return %2: tensor<i32>
+}
+
+// CHECK-LABEL: func @batch_function
+func.func @batch_function(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
+  // CHECK: [[device:%.*]] = tf_mlrt_tpu.get_tpu_host_device
+  // CHECK: [[batch_result_future:%.*]] = tf_mlrt.batch_function.device([[device]]) (%arg0, %arg1)
+  // CHECK: [[batch_result:%.*]] = tf_mlrt.await [[batch_result_future]]
+  // CHECK: return [[batch_result]]
+  %0 = "tf.BatchFunction"(%arg0, %arg1) {device = "/device:CPU:0", allowed_batch_sizes = [64], batch_timeout_micros = 1 : i64, batching_queue = "", container = "", f = @batched_func, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 2, 0>, shared_name = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %0 : tensor<i32>
+}
+
+
+

tensorflow/compiler/mlir/tfrt/tests/mlrt/while_to_map_fn.mlir

@@ -0,0 +1,829 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-while-to-map-fn %s | FileCheck %s
+
+// Test a while to map_fn conversion in which the max iteration is hard coded inside the predicate body.
+
+// CHECK-LABEL: map/while_cond
+func.func private @"map/while_cond"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
+  return %2 : tensor<i1>
+}
+
+// CHECK-LABEL: map/while_body
+func.func private @"map/while_body"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00, 9.000000e+00]> : tensor<9xf32>} : () -> tensor<9xf32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32>
+  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<2xi32>} : () -> tensor<2xi32>
+  %cst_3 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00]> : tensor<9xf32>} : () -> tensor<9xf32>
+  %cst_4 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.AddV2"(%arg0, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %1 = "tf.Mul"(%arg3, %cst_3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<9xf32>) -> tensor<9xf32>
+  %2 = "tf.Reshape"(%1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
+  %3 = "tf.AddV2"(%arg1, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %4 = "tf.GatherV2"(%cst_1, %arg1, %cst_0) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>
+  %5 = "tf.Cast"(%4) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  %6 = "tf.Mul"(%5, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<9xf32>) -> tensor<9xf32>
+  %7 = "tf.Reshape"(%6, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
+  %8 = "tf.MatMul"(%2, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  %9 = "tf.MatrixDeterminant"(%8) {T = f32, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3x3xf32>) -> tensor<f32>
+  %10 = "tf.TensorListSetItem"(%arg2, %arg1, %9) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<f32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  return %0, %3, %10, %arg3 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
+}
+
+// CHECK-LABEL: map/while_body/MapFnBody
+// CHECK-SAME: (%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<?xf32>)
+// CHECK: [[det:%.*]] = "tf.MatrixDeterminant"
+// CHECK-NEXT: [[ta_0:%.*]] = "tf_mlrt.tf_await"(%arg0) : (!mlrt.future) -> tensor<!tf_type.variant<tensor<*xf32>>>
+// CHECK-NEXT: [[ta_1:%.*]] = "tf.TensorListSetItem"([[ta_0]], %arg3, [[det]]) <{
+// CHECK-NEXT:  "tf_mlrt.tf_promise"(%arg1, [[ta_1]]) : (!mlrt.promise, tensor<!tf_type.variant<tensor<*xf32>>>) -> ()
+// CHECK-NEXT: return
+
+//CHECK-LABEL: @serving_default
+func.func @serving_default(%arg0: tensor<?xf32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> tensor<3xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input:0", outputs = "PartitionedCall:0"}} {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[tensor_list:%.*]] = "tf.TensorListReserve"([[shape:%.*]], [[reserve_size:%.*]]) {
+  %0 = "tf.TensorListReserve"(%cst_1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  // CHECK: [[map_fn_result:%.*]] = tf_mlrt.tf_map_fn([[reserve_size]], [[tensor_list]], %arg0)
+  // CHECK-SAME: {body_fn = @"map/while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
+  // CHECK-NOT: tf.While
+  %1:4 = "tf.While"(%cst, %cst, %0, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @"map/while_body", cond = @"map/while_cond", device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
+  // CHECK-NEXT: "tf.TensorListStack"([[map_fn_result]], %cst_0) <{
+  %2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 3 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<0xi32>) -> tensor<3xf32>
+  return %2 : tensor<3xf32>
+}
+
+// -----
+
+// Test a while to map_fn conversion in which max_iterations are passed
+// into the predicate function.
+
+// CHECK-LABEL: @"map/while_cond"
+func.func private @"map/while_cond"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>, %arg4: tensor<!tf_type.resource<tensor<3x1xf32>>>, %arg5: tensor<?x3xf32>, %arg6: tensor<?x4xf32>) -> tensor<i1> {
+  %outputs =  "tf.Less"(%arg0, %arg3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %outputs_0  =  "tf.Less"(%arg1, %arg3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %outputs_2  =  "tf.LogicalAnd"(%outputs_0, %outputs) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
+
+  return %outputs_2 : tensor<i1>
+}
+
+// CHECK-LABEL: @"map/while_body"
+func.func private @"map/while_body"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>, %arg4: tensor<!tf_type.resource<tensor<3x1xf32>>>, %arg5: tensor<?x3xf32>, %arg6: tensor<?x4xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>) {
+  %outputs =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_0  =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_2  =  "tf.AddV2"(%arg0, %outputs_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %outputs_4  =  "tf.ReadVariableOp"(%arg4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>
+  %outputs_6  =  "tf.Identity"(%outputs_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<i32>
+  %outputs_8  =  "tf.MatMul"(%arg5, %outputs_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<?x3xf32>, tensor<3x1xf32>) -> tensor<?x1xf32>
+  %outputs_10  =  "tf.AddV2"(%arg1, %outputs_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %outputs_12  =  "tf.Identity"(%outputs_10) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<i32>
+  %outputs_14  =  "tf.GatherV2"(%arg6, %arg1, %outputs) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x4xf32>, tensor<i32>, tensor<i32>) -> tensor<4xf32>
+  %outputs_16  =  "tf.AddV2"(%outputs_8, %outputs_14) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x1xf32>, tensor<4xf32>) -> tensor<?x4xf32>
+  %outputs_18  =  "tf.TensorListSetItem"(%arg2, %arg1, %outputs_16) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<?x4xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  return  %outputs_6, %outputs_12, %outputs_18, %arg3, %arg4, %arg5, %arg6 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>
+}
+
+// CHECK-LABEL: @"map/while_body/MapFnBody"
+// CHECK-SAME (%arg0: !mlrt.Future, %arg1: !mlrt.Promise, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<i32>, %arg5: tensor<!tf_type.resource<tensor<3x1xf32>>>, %arg6: tensor<?x3xf32>, %arg7: tensor<?x4xf32>)
+// CHECK-NEXT: [[cst_0:%.*]] = "tf.Const"
+// CHECK-NEXT: [[cst_1:%.*]] = "tf.Const"
+// CHECK-NEXT: [[loop_counter:%.*]] = "tf.AddV2"(%arg2, [[cst_1]])
+// CHECK-NEXT: [[weight:%.*]] = "tf.ReadVariableOp"(%arg5)
+// CHECK-NEXT: [[mpy:%.*]] = "tf.MatMul"(%arg6, [[weight]])
+// CHECK-NEXT: [[element_index:%.*]] = "tf.AddV2"(%arg3, [[cst_1]])
+// CHECK-NEXT: [[bias:%.*]] = "tf.GatherV2"(%arg7, %arg3, [[cst_0]])
+// CHECK-NEXT: [[res:%.*]] = "tf.AddV2"([[mpy]], [[bias]])
+// CHECK-NEXT: [[ta_0:%.*]] = "tf_mlrt.tf_await"(%arg0)
+// CHECK-NEXT: [[ta_1:%.*]] = "tf.TensorListSetItem"([[ta_0]], %arg3, [[res]])
+// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[ta_1]])
+// CHECK-NEXT: return
+
+// CHECK-LABEL: func @main_while
+func.func @main_while(%arg0: tensor<?x3xf32>, %arg1: tensor<?x4xf32>) -> tensor<?x?x4xf32> {
+  %outputs =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[-1, 4]> : tensor<2xi32>} : () -> tensor<2xi32>
+  %outputs_0  =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_2  =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
+  %outputs_4  =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
+  %outputs_6  =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_8  =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[elems:%.*]] = "tf.VarHandleOp"
+  %outputs_10  =  "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", shared_name = "w"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
+  %outputs_12  =  "tf.Shape"(%arg1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x4xf32>) -> tensor<2xi32>
+  // CHECK: [[max_iter:%.*]] = "tf.StridedSlice"
+  %outputs_14  =  "tf.StridedSlice"(%outputs_12, %outputs_2, %outputs_4, %outputs_4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<2xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
+  // CHECK: [[tensor_list:%.*]] = "tf.TensorListReserve"
+  %outputs_16  =  "tf.TensorListReserve"(%outputs_0, %outputs_14) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  // CHECK: tf_mlrt.tf_map_fn
+  // CHECK-SAME: ([[max_iter]], [[tensor_list]], [[max_iter]], [[elems]], %arg0, %arg1)
+  // CHECK-SAME: {body_fn = @"map/while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
+  // CHECK-NOT: tf.while
+  %outputs_18:7  =  "tf.While"(%outputs_6, %outputs_6, %outputs_16, %outputs_14, %outputs_10, %arg0, %arg1) {_lower_using_switch_merge = true, _num_original_outputs = 8 : i64, _read_only_resource_inputs = [6], _xla_propagate_compile_time_consts = true, body = @"map/while_body", cond = @"map/while_cond", device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<!tf_type.resource<tensor<3x1xf32>>>, tensor<?x3xf32>, tensor<?x4xf32>)
+  %outputs_20  =  "tf.TensorListStack"(%outputs_18#2, %outputs) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = -1 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<2xi32>) -> tensor<?x?x4xf32>
+  return %outputs_20 : tensor<?x?x4xf32>
+}
+
+// -----
+
+// Test a while to map_fn conversion in which the passed in max_iterations 
+// is not in typical location of %arg3 and there are identify chains in function bodies.
+
+// CHECK-LABEL: @map_while_cond_170
+func.func private @map_while_cond_170(%arg0: tensor<i32> {tf._user_specified_name = "map/while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "map/while/maximum_iterations"}, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>, %arg4: tensor<*x!tf_type.variant>, %arg5: tensor<*xf32>) -> tensor<*xi1> attributes {tf._construction_context = "kEagerRuntime", tf._original_func_name = "map_while_cond_17"} {
+  %outputs =  "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  %outputs_0 =  "tf.Less"(%arg0, %arg1) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
+  %outputs_2 =  "tf.Less"(%arg2, %outputs) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
+  %outputs_4 =  "tf.LogicalAnd"(%outputs_0, %outputs_2) {device = ""} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
+  %outputs_6 =  "tf.Identity"(%outputs_4) {device = ""} : (tensor<*xi1>) -> tensor<*xi1>
+  return %outputs_6 : tensor<*xi1>
+}
+
+// Original input argument list (loop_counter, max_iterations, element_index, tensor_list, read_only_tensor_list, scale)
+// CHECK-LABEL: @map_while_body_180
+func.func private @map_while_body_180(%arg0: tensor<i32> {tf._user_specified_name = "map/while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "map/while/maximum_iterations"}, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant>, %arg4: tensor<!tf_type.variant> {tf._user_specified_name = "map/TensorArrayUnstack/TensorListFromTensor"}, %arg5: tensor<?xf32> {tf._user_specified_name = "input"}) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>, tensor<!tf_type.variant>, tensor<?xf32>) attributes {tf._construction_context = "kEagerRuntime", tf._original_func_name = "map_while_body_18"} {
+  %outputs =  "tf.Const"() {device = "", value = dense<16> : tensor<2xi32>} : () -> tensor<2xi32>
+  %outputs_0 =  "tf.Const"() {device = "", value = dense<16> : tensor<2xi32>} : () -> tensor<2xi32>
+  %outputs_2 =  "tf.Const"() {device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
+  %outputs_4 =  "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_6 =  "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_8 =  "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_10 =  "tf.Const"() {device = "", value = dense<256> : tensor<i32>} : () -> tensor<i32>
+  %outputs_12 =  "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_14 =  "tf.Range"(%outputs_12, %outputs_10, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
+  %outputs_16 =  "tf.Cast"(%outputs_14) {Truncate = false, device = ""} : (tensor<*xi32>) -> tensor<*xf32>
+  %outputs_18 =  "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_20 =  "tf.Const"() {device = "", value = dense<257> : tensor<i32>} : () -> tensor<i32>
+  %outputs_22 =  "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_24 =  "tf.Range"(%outputs_22, %outputs_20, %outputs_18) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
+  %outputs_26 =  "tf.Cast"(%outputs_24) {Truncate = false, device = ""} : (tensor<*xi32>) -> tensor<*xf32>
+  %outputs_28 =  "tf.Const"() {device = "", value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
+  %outputs_30 =  "tf.Transpose"(%outputs_26, %outputs_28) {device = ""} : (tensor<*xf32>, tensor<1xi32>) -> tensor<*xf32>
+  %outputs_32 =  "tf.AddV2"(%arg0, %outputs_6) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
+  %outputs_34 =  "tf.Identity"(%outputs_32) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
+  %outputs_36 =  "tf.Identity"(%arg1) {device = ""} : (tensor<i32>) -> tensor<*xi32>
+  %outputs_38 =  "tf.Mul"(%outputs_16, %arg5) {device = ""} : (tensor<*xf32>, tensor<?xf32>) -> tensor<*xf32>
+  %outputs_40 =  "tf.Reshape"(%outputs_38, %outputs) {device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>
+  %outputs_42 =  "tf.AddV2"(%arg2, %outputs_4) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
+  %outputs_44 =  "tf.Identity"(%outputs_42) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
+  %outputs_46 =  "tf.TensorListGetItem"(%arg4, %arg2, %outputs_2) {device = ""} : (tensor<!tf_type.variant>, tensor<i32>, tensor<0xi32>) -> tensor<*xi32>
+  %outputs_48 =  "tf.Cast"(%outputs_46) {Truncate = false, device = ""} : (tensor<*xi32>) -> tensor<*xf32>
+  %outputs_50 =  "tf.Mul"(%outputs_30, %outputs_48) {device = ""} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+  %outputs_52 =  "tf.Reshape"(%outputs_50, %outputs_0) {device = ""} : (tensor<*xf32>, tensor<2xi32>) -> tensor<*xf32>
+  %outputs_54 =  "tf.MatMul"(%outputs_40, %outputs_52) {device = "", transpose_a = false, transpose_b = false} : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+  %outputs_56 =  "tf.MatrixDeterminant"(%outputs_54) {T = f32, device = ""} : (tensor<*xf32>) -> tensor<*xf32>
+  %outputs_58 =  "tf.TensorListSetItem"(%arg3, %arg2, %outputs_56) {device = "", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant>, tensor<i32>, tensor<*xf32>) -> tensor<*x!tf_type.variant>
+  %outputs_60 =  "tf.Identity"(%outputs_58) {device = ""} : (tensor<*x!tf_type.variant>) -> tensor<*x!tf_type.variant>
+  return %outputs_34, %outputs_36, %outputs_44, %outputs_60, %arg4, %arg5 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*x!tf_type.variant>, tensor<!tf_type.variant>, tensor<?xf32>
+}
+
+// Converted input argument list (loop_counter, element_index, max_iterations, tensor_list, read_only_tensor_list, scale)
+// CHECK-LABEL: @"map_while_body_180/MapFnBody"
+// CHECK-SAME: (%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: tensor<i32> {tf._user_specified_name = "map/while/loop_counter"}, %arg3: tensor<i32>, %arg4: tensor<i32> {tf._user_specified_name = "map/while/maximum_iterations"}, %arg5: tensor<!tf_type.variant> {tf._user_specified_name = "map/TensorArrayUnstack/TensorListFromTensor"}, %arg6: tensor<?xf32> {tf._user_specified_name = "input"})
+// CHECK: [[res:%.*]] = "tf.MatrixDeterminant"
+// CHECK-NEXT: [[ta_0:%.*]] = "tf_mlrt.tf_await"(%arg0)
+// CHECK-NEXT: [[ta_1:%.*]] = "tf.TensorListSetItem"([[ta_0]], %arg3, [[res]])
+// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[ta_1]])
+// CHECK-NEXT: return
+
+
+// CHECK-LABEL: __inference_while_from_map_fn_810
+// CHECK-SAME: ([[scale:%.*]]: tensor<?xf32>
+func.func private @__inference_while_from_map_fn_810(%arg0: tensor<?xf32> {tf._user_specified_name = "input"}) -> tensor<*xf32> attributes {tf._construction_context = "kEagerRuntime", tf._original_func_name = "__inference_while_from_map_fn_81"} {
+  // CHECK: [[element_index:%.*]] = "tf.Const"
+  %outputs =  "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_0 =  "tf.Const"() {device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
+  %outputs_2=  "tf.Const"() {device = "", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_4 =  "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: tf.TensorListReserve
+  %outputs_6 =  "tf.TensorListReserve"(%outputs_2, %outputs_4) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi32>>>
+  %outputs_8 =  "tf.Const"() {device = "", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
+  %outputs_10 =  "tf.Const"() {device = "", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_12 =  "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: [[tensor_list:%.*]] = "tf.TensorListReserve"([[shape:%.*]], [[reserve_size:%.*]]) {
+  %outputs_14 =  "tf.TensorListReserve"(%outputs_10, %outputs_12) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  // CHECK-NEXT: [[loop_counter:%.*]] = "tf.Const"
+  %outputs_16 =  "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK-NEXT: [[max_iterations:%.*]] = "tf.Const"
+  %outputs_18 =  "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  %outputs_20 =  "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_22 =  "tf.Const"() {device = "", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  %outputs_24 =  "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_26 =  "tf.Range"(%outputs_24, %outputs_22, %outputs_20) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<*xi32>
+  // CHECK: [[read_only_tensor_list:%.*]] = "tf.TensorListFromTensor"
+  %outputs_28 =  "tf.TensorListFromTensor"(%outputs_26, %outputs_0) {device = ""} : (tensor<*xi32>, tensor<0xi32>) -> tensor<*x!tf_type.variant>
+// CHECK: [[map_fn_out:%.*]] = tf_mlrt.tf_map_fn
+  // CHECK-SAME: ([[reserve_size]], [[tensor_list]], [[max_iterations]], [[read_only_tensor_list]], [[scale]])
+  // CHECK-SAME: {body_fn = @"map_while_body_180/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
+  // CHECK-NOT: tf.While
+  %outputs_30:6 =  "tf.While"(%outputs_16, %outputs_18, %outputs, %outputs_14, %outputs_28, %arg0) {T = [i32, i32, i32, !tf_type.variant, !tf_type.variant, f32], _lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], body = @map_while_body_180, cond = @map_while_cond_170, device = "", is_stateless = true, output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<?>], parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<*x!tf_type.variant>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant>, tensor<!tf_type.variant>, tensor<?xf32>)
+  // CHECK-NEXT: "tf.TensorListStack"
+  // CHECK-SAME: ([[map_fn_out]],
+  %outputs_32 =  "tf.TensorListStack"(%outputs_30#3, %outputs_8) {device = "", num_elements = 16 : i64} : (tensor<!tf_type.variant>, tensor<0xi32>) -> tensor<*xf32>
+  %outputs_34 =  "tf.Identity"(%outputs_32) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>
+  return %outputs_34 : tensor<*xf32>
+}
+
+// -----
+
+// Test a while to map_fn conversion in which tensor array is used instead of
+// tensor list.
+
+// CHECK-LABEL: map/while/LoopCond_cond
+func.func private @"map/while/LoopCond_cond"(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, %arg5: tensor<f32>, %arg6: tensor<2x!tf_type.resource<tensor<*xui8>>>) -> tensor<i1> {
+  %outputs = "tf.Less"(%arg0, %arg3) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
+  %outputs_0 = "tf.Less"(%arg1, %arg3) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
+  %outputs_2 = "tf.LogicalAnd"(%outputs, %outputs_0) {device = ""} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
+  %outputs_4 = "tf.ToBool"(%outputs_2) : (tensor<*xi1>) -> tensor<i1>
+  return %outputs_4 : tensor<i1>
+}
+
+// CHECK-LABEL: map/while/LoopCond_body
+func.func private @"map/while/LoopCond_body"(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, %arg5: tensor<f32>, %arg6: tensor<2x!tf_type.resource<tensor<*xui8>>>) -> (tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>) {
+  %outputs = "tf.Const"() {value = dense<224> : tensor<2xi32>} : () -> tensor<2xi32>
+  %outputs_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_4 = "tf.Identity"(%arg0) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
+  %outputs_6 = "tf.AddV2"(%outputs_4, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
+  %outputs_8 = "tf.Identity"(%arg1) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
+  %outputs_10 = "tf.AddV2"(%outputs_8, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
+  %outputs_12 = "tf.Identity"(%arg2) {device = ""} : (tensor<f32>) -> tensor<f32>
+  %outputs_14 = "tf.TensorArrayReadV3"(%arg4, %outputs_8, %arg5) {device = ""} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<*xi32>, tensor<f32>) -> tensor<*x!tf_type.string>
+  %outputs_16 = "tf.DecodeJpeg"(%outputs_14) {acceptable_fraction = 1.000000e+00 : f32, channels = 3 : i64, dct_method = "INTEGER_FAST", device = "", fancy_upscaling = true, ratio = 1 : i64, try_recover_truncated = false} : (tensor<*x!tf_type.string>) -> tensor<?x?x3xui8>
+  %outputs_18 = "tf.ExpandDims"(%outputs_16, %outputs_0) {device = ""} : (tensor<?x?x3xui8>, tensor<i32>) -> tensor<1x?x?x3xui8>
+  %outputs_20 = "tf.ResizeBilinear"(%outputs_18, %outputs) {align_corners = false, device = "", half_pixel_centers = false} : (tensor<1x?x?x3xui8>, tensor<2xi32>) -> tensor<1x224x224x3xf32>
+  %outputs_22 = "tf.Squeeze"(%outputs_20) {device = "", squeeze_dims = [0]} : (tensor<1x224x224x3xf32>) -> tensor<224x224x3xf32>
+  %outputs_24 = "tf.Cast"(%outputs_22) {Truncate = false, device = ""} : (tensor<224x224x3xf32>) -> tensor<224x224x3xui8>
+  %outputs_26 = "tf.TensorArrayWriteV3"(%arg6, %outputs_8, %outputs_24, %outputs_12) {device = ""} : (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<*xi32>, tensor<224x224x3xui8>, tensor<f32>) -> tensor<f32>
+  return %outputs_6, %outputs_10, %outputs_26, %arg3, %arg4, %arg5, %arg6: tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>
+}
+
+// CHECK-LABEL: @"map/while/LoopCond_body/MapFnBody"
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.TensorArrayReadV3
+// CHECK-NEXT: tf.DecodeJpeg
+// CHECK-NEXT: tf.ExpandDims
+// CHECK-NEXT: tf.ResizeBilinear
+// CHECK-NEXT: tf.Squeeze
+// CHECK-NEXT: tf.Cast
+// CHECK-NEXT: tf_mlrt.tf_await
+// CHECK-NEXT: tf.TensorArrayWriteV3
+// CHECK-NEXT: tf_mlrt.tf_promise
+// CHECK-NEXT: return
+
+//CHECK-LABEL: map_while_test
+func.func @map_while_test(%arg0: tensor<?x!tf_type.string>) -> tensor<?x224x224x3xui8> {
+  %outputs = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<1xi32>
+  %outputs_0 = "tf.Const"() {value = dense<224> : tensor<2xi32>} : () -> tensor<2xi32>
+  %outputs_2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_4 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
+  %outputs_6 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
+  %outputs_8 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_10 = "tf.Shape"(%arg0) {device = ""} : (tensor<?x!tf_type.string>) -> tensor<1xi32>
+  // CHECK: [[max_iter:%.*]] = "tf.StridedSlice"
+  %outputs_12 = "tf.StridedSlice"(%outputs_10, %outputs_6, %outputs_4, %outputs_4) {begin_mask = 0 : i64, device = "", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
+  // CHECK-NEXT: tf.Range
+  %outputs_14 = "tf.Range"(%outputs_2, %outputs_12, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
+  // CHECK-NEXT: [[handle_1:%.*]], [[flow_in_1:%.*]] = "tf.TensorArrayV3"
+  %outputs_16:2 = "tf.TensorArrayV3"(%outputs_12) {clear_after_read = true, device = "", dtype = !tf_type.string, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>)
+  // CHECK-NEXT: [[handle_2:%.*]] = "tf.TensorArrayScatterV3"
+  %outputs_18 = "tf.TensorArrayScatterV3"(%outputs_16#0, %outputs_14, %arg0, %outputs_16#1) {device = ""} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<?xi32>, tensor<?x!tf_type.string>, tensor<f32>) -> tensor<f32>
+  // CHECK-NEXT: tf.Range
+  %outputs_20 = "tf.Range"(%outputs_2, %outputs_12, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
+  // CHECK-NEXT: [[tensor_array:%.*]], [[flow_in:%.*]] = "tf.TensorArrayV3"
+  %outputs_22:2 = "tf.TensorArrayV3"(%outputs_12) {clear_after_read = true, device = "", dtype = ui8, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<f32>)
+  // CHECK-NEXT: tf_mlrt.tf_map_fn
+  // CHECK-SAME: ([[max_iter]], [[flow_in]], [[max_iter]], [[handle_1]], [[handle_2]], [[tensor_array]])
+  // CHECK-SAME: {body_fn = @"map/while/LoopCond_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32} 
+  // CHECK-NOT: tf.While
+  %outputs_24:7 = "tf.While"(%outputs, %outputs, %outputs_22#1, %outputs_12, %outputs_16#0, %outputs_18, %outputs_22#0) {_xla_propagate_compile_time_consts = true, body = @"map/while/LoopCond_body", cond = @"map/while/LoopCond_cond", device = "", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<1xi32>, tensor<1xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>) -> (tensor<1xi32>, tensor<1xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>)
+  // CHECK-NEXT: tf.TensorArrayGatherV3
+  %outputs_26 = "tf.TensorArrayGatherV3"(%outputs_22#0, %outputs_20, %outputs_24#2) {device = "", element_shape = #tf_type.shape<224x224x3>} : (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<?xi32>, tensor<f32>) -> tensor<?x224x224x3xui8>
+  return %outputs_26 : tensor<?x224x224x3xui8>
+}
+
+// -----
+// Test non-applicable while is NOT converted to map_fn.
+
+// CHECK-LABEL: func @while_cond_lt9
+func.func @while_cond_lt9(%arg0: tensor<i32>) -> tensor<i1> {
+  %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<9> : tensor<i32>} : () -> tensor<i32>
+  %1 = "tf.Less"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  func.return %1 : tensor<i1>
+}
+
+// CHECK-LABEL: func @while_body_add2
+func.func @while_body_add2(%arg0: tensor<i32>) -> tensor<i32> {
+  %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
+  %1 = "tf.Add"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  func.return %1 : tensor<i32>
+}
+
+// CHECK-LABEL: func @while_test()
+func.func @while_test() -> (tensor<i32>) {
+  %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: tf.While
+  %1 = "tf.While"(%0) { cond = @while_cond_lt9, body = @while_body_add2, is_stateless = false, parallel_iterations = 1} : (tensor<i32>) -> (tensor<i32>)
+  func.return %1 : tensor<i32>
+}
+
+// -----
+
+// Test a case that the while body has multiple tensor lists.
+
+// CHECK-LABEL: tf.MultiListWhileRegion_body
+func.func private @tf.MultiListWhileRegion_body(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<!tf_type.variant<tensor<*xf32>>>, %arg4: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00], [8.000000e+00, 9.000000e+00, 1.000000e+01, 1.100000e+01, 1.200000e+01, 1.300000e+01, 1.400000e+01, 1.500000e+01]]> : tensor<2x8xf32>} : () -> tensor<2x8xf32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[[1.600000e+01, 1.700000e+01, 1.800000e+01, 1.900000e+01, 2.000000e+01, 2.100000e+01, 2.200000e+01, 2.300000e+01], [2.400000e+01, 2.500000e+01, 2.600000e+01, 2.700000e+01, 2.800000e+01, 2.900000e+01, 3.000000e+01, 3.100000e+01]]> : tensor<2x8xf32>} : () -> tensor<2x8xf32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.GatherV2"(%arg4, %cst_2, %cst_2) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<i32>, tensor<i32>) -> tensor<f32>
+  %1 = "tf.AddV2"(%arg0, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %2 = "tf.AddV2"(%arg1, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %3 = "tf.GatherV2"(%cst_0, %arg1, %cst_2) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x8xf32>, tensor<i32>, tensor<i32>) -> tensor<8xf32>
+  %4 = "tf.Mul"(%0, %3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<8xf32>) -> tensor<8xf32>
+  %5 = "tf.TensorListSetItem"(%arg2, %arg1, %4) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<8xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  %6 = "tf.GatherV2"(%cst, %arg1, %cst_2) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x8xf32>, tensor<i32>, tensor<i32>) -> tensor<8xf32>
+  %7 = "tf.Mul"(%0, %6) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<8xf32>) -> tensor<8xf32>
+  %8 = "tf.TensorListSetItem"(%arg3, %arg1, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<8xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  return %1, %2, %5, %8, %arg4 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
+}
+
+// CHECK-LABEL: tf.MultiListWhileRegion_body/MapFnBody
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.GatherV2
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.GatherV2
+// CHECK-NEXT: tf.Mul
+// CHECK-NEXT: tf.GatherV2
+// CHECK-NEXT: tf.Mul
+// CHECK-NEXT: tf_mlrt.tf_await
+// CHECK-NEXT: tf_mlrt.tf_await
+// CHECK-NEXT: tf.TensorListSetItem
+// CHECK-NEXT: tf.TensorListSetItem
+// CHECK-NEXT: tf_mlrt.tf_promise
+// CHECK-NEXT: tf_mlrt.tf_promise
+// CHECK-NEXT: return
+
+// CHECK-LABEL: tf.MultiListWhileRegion_cond
+func.func private @tf.MultiListWhileRegion_cond(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<!tf_type.variant<tensor<*xf32>>>, %arg4: tensor<?xf32>) -> tensor<i1> {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
+  return %2 : tensor<i1>
+}
+
+// CHECK-LABEL: multilist_serving
+func.func private @multilist_serving(%arg0: tensor<?xf32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<2x8xf32>, tensor<2x8xf32>) {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<8> : tensor<1xi32>} : () -> tensor<1xi32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: TensorListReserve
+  %0 = "tf.TensorListReserve"(%cst_1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  // CHECK-NEXT: tf_mlrt.tf_map_fn
+  %1:5 = "tf.While"(%cst, %cst, %0, %0, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 8 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @tf.MultiListWhileRegion_body, cond = @tf.MultiListWhileRegion_cond, device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
+  // CHECK-NEXT: TensorListStack
+  %2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<2x8xf32>
+  %3 = "tf.TensorListStack"(%1#3, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<2x8xf32>
+  return %3, %2 : tensor<2x8xf32>, tensor<2x8xf32>
+}
+
+
+// -----
+
+// Convert a while with multiple tensor array to map_fn
+
+// CHECK-LABEL: tf.WhileRegion1_body(
+func.func private @tf.WhileRegion1_body(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<f32>, %arg3: tensor<f32>, %arg4: tensor<i32>, %arg5: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg6: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg7: tensor<*xi32>) -> (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<*xi32>) {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.AddV2"(%arg0, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %1 = "tf.AddV2"(%arg1, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %3 = "tf.RaggedTensorToVariant"(%arg7) {RAGGED_RANK = 0 : i64, Tsplits = i64, Tvalues = i32, batched_input = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>) -> tensor<!tf_type.variant>
+  %4 = "tf.TensorArrayWriteV3"(%arg5, %arg1, %3, %arg2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<i32>, tensor<!tf_type.variant>, tensor<f32>) -> tensor<f32>
+  %5 = "tf.RaggedTensorToVariant"(%arg7) {RAGGED_RANK = 0 : i64, Tsplits = i64, Tvalues = f32, batched_input = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi32>) -> tensor<!tf_type.variant>
+  %6 = "tf.TensorArrayWriteV3"(%arg6, %arg1, %5, %arg3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<i32>, tensor<!tf_type.variant>, tensor<f32>) -> tensor<f32>
+  return %0, %1, %4, %6, %arg4, %arg5, %arg6, %arg7 : tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<*xi32>
+}
+
+// CHECK-LABEL: func.func private @"tf.WhileRegion1_body/MapFnBody"(%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: !mlrt.future, %arg3: !mlrt.promise, %arg4: tensor<i32>, %arg5: tensor<i32>, %arg6: tensor<i32>, %arg7: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg8: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg9: tensor<*xi32>) attributes {tfrt.cost_threshold = 4294967295 : i64} 
+// CHECK: [[result_0:%.*]] = "tf.RaggedTensorToVariant"
+// CHECK: [[result_1:%.*]] = "tf.RaggedTensorToVariant"
+// CHECK-NEXT: [[flow_in_0:%.*]] = "tf_mlrt.tf_await"(%arg0) : (!mlrt.future) -> tensor<f32>
+// CHECK-NEXT: [[flow_in_1:%.*]] = "tf_mlrt.tf_await"(%arg2) : (!mlrt.future) -> tensor<f32>
+// CHECK-NEXT: [[flow_out_0:%.*]] = "tf.TensorArrayWriteV3"(%arg7, %arg5, [[result_0]], [[flow_in_0]])
+// CHECK-NEXT: [[flow_out_1:%.*]] = "tf.TensorArrayWriteV3"(%arg8, %arg5, [[result_1]], [[flow_in_1]])
+// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[flow_out_0]]) : (!mlrt.promise, tensor<f32>) -> ()
+// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg3, [[flow_out_1]]) : (!mlrt.promise, tensor<f32>) -> ()
+// CHECK-NEXT: return
+
+// CHECK-LABEL: tf.WhileRegion1_cond
+func.func private @tf.WhileRegion1_cond(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<f32>, %arg3: tensor<f32>, %arg4: tensor<i32>, %arg5: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg6: tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, %arg7: tensor<*xi32>) -> (tensor<i1>) {
+  %0 = "tf.Less"(%arg0, %arg4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
+  %1 = "tf.Less"(%arg1, %arg4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<*xi1>
+  %2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
+  %3 = "tf.ToBool"(%2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<*xi1>) -> tensor<i1>
+  return %3 : tensor<i1>
+}
+
+// CHECK-LABEL: func.func private @tf.WhileRegion2_body(
+func.func private @tf.WhileRegion2_body(%arg0: tensor<*xi32>) -> (tensor<?x!tf_type.variant>, tensor<?x!tf_type.variant>) {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[-1, 4]> : tensor<2xi32>} : () -> tensor<2xi32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %max_iter = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<4> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: "tf.TensorArrayV3"
+  %handle_12, %flow_13 = "tf.TensorArrayV3"(%max_iter) {device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = !tf_type.variant, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<f32>)
+  // CHECK: "tf.TensorArrayV3"
+  %handle_14, %flow_15 = "tf.TensorArrayV3"(%max_iter) {device = "/job:localhost/replica:0/task:0/device:CPU:0", dtype = !tf_type.variant, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<f32>)
+  // CHECK: tf_mlrt.tf_map_fn
+  // CHECK-SAME: {body_fn = @"tf.WhileRegion1_body/MapFnBody", num_tensor_list_or_flow_in = 2 : i32}
+  %4:8 = "tf.While"(%cst_0, %cst_0, %flow_13, %flow_15, %max_iter, %handle_12, %handle_14, %arg0) {body = @tf.WhileRegion1_body, cond = @tf.WhileRegion1_cond, device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<*xi32>)
+  // CHECK: TensorArrayGatherV3
+  %5 = "tf.TensorArrayGatherV3"(%handle_12, %1, %4#2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<i32>, tensor<f32>) -> tensor<?x!tf_type.variant>
+  // CHECK: TensorArrayGatherV3
+  %6 = "tf.TensorArrayGatherV3"(%handle_14, %2, %4#3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.variant>>>, tensor<i32>, tensor<f32>) -> tensor<?x!tf_type.variant>
+  return %5, %6 : tensor<?x!tf_type.variant>, tensor<?x!tf_type.variant>
+}
+
+// -----
+
+// Test a while to map_fn conversion in which tensor array is used instead of
+// tensor list and the tensor array size and the number of iterations are bounded
+// by separate constants of the same value.
+
+// CHECK-LABEL: map2/while/LoopCond_body
+func.func private @"map2/while/LoopCond_body"(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, %arg5: tensor<f32>, %arg6: tensor<2x!tf_type.resource<tensor<*xui8>>>) -> (tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>) {
+  %outputs = "tf.Const"() {value = dense<224> : tensor<2xi32>} : () -> tensor<2xi32>
+  %outputs_0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_4 = "tf.Identity"(%arg0) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
+  %outputs_6 = "tf.AddV2"(%outputs_4, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
+  %outputs_8 = "tf.Identity"(%arg1) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
+  %outputs_10 = "tf.AddV2"(%outputs_8, %outputs_2) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
+  %outputs_12 = "tf.Identity"(%arg2) {device = ""} : (tensor<f32>) -> tensor<f32>
+  %outputs_14 = "tf.TensorArrayReadV3"(%arg4, %outputs_8, %arg5) {device = ""} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<*xi32>, tensor<f32>) -> tensor<*x!tf_type.string>
+  %outputs_16 = "tf.DecodeJpeg"(%outputs_14) {acceptable_fraction = 1.000000e+00 : f32, channels = 3 : i64, dct_method = "INTEGER_FAST", device = "", fancy_upscaling = true, ratio = 1 : i64, try_recover_truncated = false} : (tensor<*x!tf_type.string>) -> tensor<?x?x3xui8>
+  %outputs_18 = "tf.ExpandDims"(%outputs_16, %outputs_0) {device = ""} : (tensor<?x?x3xui8>, tensor<i32>) -> tensor<1x?x?x3xui8>
+  %outputs_20 = "tf.ResizeBilinear"(%outputs_18, %outputs) {align_corners = false, device = "", half_pixel_centers = false} : (tensor<1x?x?x3xui8>, tensor<2xi32>) -> tensor<1x224x224x3xf32>
+  %outputs_22 = "tf.Squeeze"(%outputs_20) {device = "", squeeze_dims = [0]} : (tensor<1x224x224x3xf32>) -> tensor<224x224x3xf32>
+  %outputs_24 = "tf.Cast"(%outputs_22) {Truncate = false, device = ""} : (tensor<224x224x3xf32>) -> tensor<224x224x3xui8>
+  %outputs_26 = "tf.TensorArrayWriteV3"(%arg6, %outputs_8, %outputs_24, %outputs_12) {device = ""} : (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<*xi32>, tensor<224x224x3xui8>, tensor<f32>) -> tensor<f32>
+  return %outputs_6, %outputs_10, %outputs_26, %arg3, %arg4, %arg5, %arg6: tensor<*xi32>, tensor<*xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>
+}
+
+// CHECK-LABEL: @"map2/while/LoopCond_body/MapFnBody"
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.Const
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.TensorArrayReadV3
+// CHECK-NEXT: tf.DecodeJpeg
+// CHECK-NEXT: tf.ExpandDims
+// CHECK-NEXT: tf.ResizeBilinear
+// CHECK-NEXT: tf.Squeeze
+// CHECK-NEXT: tf.Cast
+// CHECK-NEXT: tf_mlrt.tf_await
+// CHECK-NEXT: tf.TensorArrayWriteV3
+// CHECK-NEXT: tf_mlrt.tf_promise
+// CHECK-NEXT: return
+
+// CHECK-LABEL: map2/while/LoopCond_cond
+func.func private @"map2/while/LoopCond_cond"(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<f32>, %arg3: tensor<i32>, %arg4: tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, %arg5: tensor<f32>, %arg6: tensor<2x!tf_type.resource<tensor<*xui8>>>) -> tensor<i1> {
+  %cst = "tf.Const"() {value = dense<224> : tensor<i32>} : () -> tensor<i32>
+  %outputs = "tf.Less"(%arg0, %cst) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
+  %outputs_0 = "tf.Less"(%arg1, %cst) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
+  %outputs_2 = "tf.LogicalAnd"(%outputs, %outputs_0) {device = ""} : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
+  %outputs_4 = "tf.ToBool"(%outputs_2) : (tensor<*xi1>) -> tensor<i1>
+  return %outputs_4 : tensor<i1>
+}
+
+//CHECK-LABEL: map2_while_test
+func.func private @map2_while_test(%arg0: tensor<?x!tf_type.string>) -> tensor<?x224x224x3xui8> {
+  // CHECK-NEXT: tf.Const
+  %outputs = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<1xi32>
+  // CHECK-NEXT: [[max_iter:%.*]] = "tf.Const"
+  %cst_0 = "tf.Const"() {value = dense<224> : tensor<i32>} : () -> tensor<i32>
+  %cst_1 = "tf.Const"() {value = dense<256> : tensor<i32>} : () -> tensor<i32>
+  %outputs_2 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %outputs_4 = "tf.Const"() {value = dense<1> : tensor<1xi32>} : () -> tensor<1xi32>
+  %outputs_6 = "tf.Const"() {value = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
+  %outputs_8 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %outputs_10 = "tf.Shape"(%arg0) {device = ""} : (tensor<?x!tf_type.string>) -> tensor<1xi32>
+  // CHECK: tf.Range
+  %outputs_14 = "tf.Range"(%outputs_2, %cst_0, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
+  // CHECK-NEXT: tf.TensorArrayV3
+  %outputs_16:2 = "tf.TensorArrayV3"(%cst_0) {clear_after_read = true, device = "", dtype = !tf_type.string, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>)
+  // CHECK-NEXT: tf.TensorArrayScatterV3
+  %outputs_18 = "tf.TensorArrayScatterV3"(%outputs_16#0, %outputs_14, %arg0, %outputs_16#1) {device = ""} : (tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<?xi32>, tensor<?x!tf_type.string>, tensor<f32>) -> tensor<f32>
+  // CHECK-NEXT: tf.Range
+  %outputs_20 = "tf.Range"(%outputs_2, %cst_0, %outputs_8) {device = ""} : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<?xi32>
+  // CHECK-NEXT: [[tensor_array:%.*]], [[flow_in:%.*]] = "tf.TensorArrayV3"
+  %outputs_22:2 = "tf.TensorArrayV3"(%cst_0) {clear_after_read = true, device = "", dtype = ui8, dynamic_size = false, element_shape = #tf_type.shape<*>, identical_element_shapes = true, tensor_array_name = ""} : (tensor<i32>) -> (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<f32>)
+  // CHECK-NEXT: tf_mlrt.tf_map_fn
+  // CHECK-SAME: ([[max_iter]], [[flow_in]], %cst_1
+  // CHECK-SAME: {body_fn = @"map2/while/LoopCond_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
+  // CHECK-NOT: tf.While
+  %outputs_24:7 = "tf.While"(%outputs, %outputs, %outputs_22#1, %cst_1, %outputs_16#0, %outputs_18, %outputs_22#0) {_xla_propagate_compile_time_consts = true, body = @"map2/while/LoopCond_body", cond = @"map2/while/LoopCond_cond", device = "", is_stateless = false, parallel_iterations = 10 : i64, shape_invariant} : (tensor<1xi32>, tensor<1xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>) -> (tensor<1xi32>, tensor<1xi32>, tensor<f32>, tensor<i32>, tensor<2x!tf_type.resource<tensor<*x!tf_type.string>>>, tensor<f32>, tensor<2x!tf_type.resource<tensor<*xui8>>>)
+  // CHECK-NEXT: tf.TensorArrayGatherV3
+  %outputs_26 = "tf.TensorArrayGatherV3"(%outputs_22#0, %outputs_20, %outputs_24#2) {device = "", element_shape = #tf_type.shape<224x224x3>} : (tensor<2x!tf_type.resource<tensor<*xui8>>>, tensor<?xi32>, tensor<f32>) -> tensor<?x224x224x3xui8>
+  return %outputs_26 : tensor<?x224x224x3xui8>
+}
+
+// -----
+// Test a nest while in which the while body is after the usage.
+
+// CHECK-LABEL: nested_while
+func.func @nested_while(%arg0: tensor<?xf32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"}) -> (tensor<16x16x?xf32>)  {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[16, -1]> : tensor<2xi32>} : () -> tensor<2xi32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  // CHECK: tf.TensorListReserve
+  %0 = "tf.TensorListReserve"(%cst_1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  // CHECK-NEXT: tf_mlrt.tf_map_fn
+  %1:4 = "tf.While"(%cst, %cst, %0, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @tf.NestedWhileRegion1_body, cond = @tf.NestedWhileRegion1_cond, device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
+  %2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 16 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<2xi32>) -> tensor<16x16x?xf32>
+  return %2 : tensor<16x16x?xf32>
+}
+// CHECK-LABEL: tf.NestedWhileRegion1_body
+func.func private @tf.NestedWhileRegion1_body(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<1xi32>} : () -> tensor<1xi32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi32>} : () -> tensor<16xi32>
+  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %cst_3 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  %cst_4 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.TensorListReserve"(%cst_4, %cst_3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  %1 = "tf.AddV2"(%arg0, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %2 = "tf.AddV2"(%arg1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %3 = "tf.GatherV2"(%cst_1, %arg1, %cst_0) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<16xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>
+  %4 = "tf.Cast"(%3) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  %5 = "tf.Mul"(%arg3, %4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<f32>) -> tensor<?xf32>
+  %6:4 = "tf.While"(%cst_0, %cst_0, %0, %5) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @tf.NestedWhileRegion_body, cond = @tf.NestedWhileRegion_cond, device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
+  %7 = "tf.TensorListStack"(%6#2, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 16 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<16x?xf32>
+  %8 = "tf.TensorListSetItem"(%arg2, %arg1, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<16x?xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  return %1, %2, %8, %arg3 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
+}
+
+//CHECK-LABEL: @"tf.NestedWhileRegion1_body/MapFnBody"(%arg0: !mlrt.future, %arg1: !mlrt.promise, %arg2: tensor<i32>, %arg3: tensor<i32>, %arg4: tensor<?xf32>) 
+// CHECK: tf.TensorListReserve
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.GatherV2
+// CHECK-NEXT: tf.Cast
+// CHECK-NEXT: tf.Mul
+// CHECK-NEXT: tf_mlrt.tf_map_fn
+// CHECK-NEXT: tf.TensorListStack
+// CHECK-NEXT: tf_mlrt.tf_await
+// CHECK-NEXT: tf.TensorListSetItem
+// CHECK-NEXT: tf_mlrt.tf_promise
+// CHECK-NEXT: return
+
+// CHECK-LABEL: tf.NestedWhileRegion1_cond
+func.func private @tf.NestedWhileRegion1_cond(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
+  return %2 : tensor<i1>
+}
+// CHECK-LABEL: tf.NestedWhileRegion_body
+func.func private @tf.NestedWhileRegion_body(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi32>} : () -> tensor<16xi32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.AddV2"(%arg0, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %1 = "tf.AddV2"(%arg1, %cst_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %2 = "tf.GatherV2"(%cst_0, %arg1, %cst) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<16xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>
+  %3 = "tf.Cast"(%2) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  %4 = "tf.Mul"(%arg3, %3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<f32>) -> tensor<?xf32>
+  %5 = "tf.TensorListSetItem"(%arg2, %arg1, %4) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<?xf32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  return %0, %1, %5, %arg3 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
+}
+
+// CHECK-LABEL: tf.NestedWhileRegion_body/MapFnBody
+// CHECK: tf.AddV2
+// CHECK-NEXT: tf.AddV2
+// CHECK-NEXT: tf.GatherV2
+// CHECK-NEXT: tf.Cast
+// CHECK-NEXT: tf.Mul
+// CHECK-NEXT: tf_mlrt.tf_await
+// CHECK-NEXT: tf.TensorListSetItem
+// CHECK-NEXT: "tf_mlrt.tf_promise
+// CHECK-NEXT: return
+
+// CHECK-LABEL: tf.NestedWhileRegion_cond
+func.func private @tf.NestedWhileRegion_cond(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<16> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
+  return %2 : tensor<i1>
+}
+
+// -----
+
+// Test a while to map_fn conversion is skipped if the tensor list cannot be found in the current function body.
+
+// CHECK-LABEL: map/while_cond
+func.func private @"map/while_cond"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.Less"(%arg0, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = "tf.Less"(%arg1, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = "tf.LogicalAnd"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i1>, tensor<i1>) -> tensor<i1>
+  return %2 : tensor<i1>
+}
+
+// CHECK-LABEL: map/while_body
+func.func private @"map/while_body"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00, 9.000000e+00]> : tensor<9xf32>} : () -> tensor<9xf32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0, 1, 2]> : tensor<3xi32>} : () -> tensor<3xi32>
+  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<2xi32>} : () -> tensor<2xi32>
+  %cst_3 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00, 7.000000e+00, 8.000000e+00]> : tensor<9xf32>} : () -> tensor<9xf32>
+  %cst_4 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  %0 = "tf.AddV2"(%arg0, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %1 = "tf.Mul"(%arg3, %cst_3) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?xf32>, tensor<9xf32>) -> tensor<9xf32>
+  %2 = "tf.Reshape"(%1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
+  %3 = "tf.AddV2"(%arg1, %cst_4) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %4 = "tf.GatherV2"(%cst_1, %arg1, %cst_0) {batch_dims = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3xi32>, tensor<i32>, tensor<i32>) -> tensor<i32>
+  %5 = "tf.Cast"(%4) {Truncate = false, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>) -> tensor<f32>
+  %6 = "tf.Mul"(%5, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<f32>, tensor<9xf32>) -> tensor<9xf32>
+  %7 = "tf.Reshape"(%6, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<9xf32>, tensor<2xi32>) -> tensor<3x3xf32>
+  %8 = "tf.MatMul"(%2, %7) {device = "/job:localhost/replica:0/task:0/device:CPU:0", transpose_a = false, transpose_b = false} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
+  %9 = "tf.MatrixDeterminant"(%8) {T = f32, device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<3x3xf32>) -> tensor<f32>
+  %10 = "tf.TensorListSetItem"(%arg2, %arg1, %9) {device = "/job:localhost/replica:0/task:0/device:CPU:0", resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<f32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  return %0, %3, %10, %arg3 : tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>
+}
+
+//CHECK-LABEL: @func
+func.func @func(%arg0: tensor<?xf32>, %arg1: tensor<!tf_type.variant<tensor<*xf32>>>) -> tensor<3xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "serving_default_input:0", outputs = "PartitionedCall:0"}} {
+  %cst = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  %cst_0 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
+  %cst_1 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %cst_2 = "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<3> : tensor<i32>} : () -> tensor<i32>
+  // CHECK-NOT: tf_map_fn
+  %1:4 = "tf.While"(%cst, %cst, %arg1, %arg0) {_lower_using_switch_merge = true, _num_original_outputs = 6 : i64, _read_only_resource_inputs = [], _xla_propagate_compile_time_consts = true, body = @"map/while_body", cond = @"map/while_cond", device = "/job:localhost/replica:0/task:0/device:CPU:0", is_stateless = true, parallel_iterations = 4 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<?xf32>)
+  %2 = "tf.TensorListStack"(%1#2, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 3 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<0xi32>) -> tensor<3xf32>
+  return %2 : tensor<3xf32>
+}
+
+// -----
+
+// Test a while to map_fn conversion in which a tf.StopGradient is inserted to consume the while result.
+// CHECK-LABEL: @while_map_while_body_884030
+func.func private @while_map_while_body_884030(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, %arg4: tensor<!tf_type.variant<tensor<?x?x1xui8>>> {tf._user_specified_name = "while/map/TensorArrayUnstack/TensorListFromTensor"}) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>) {
+  %cst = "tf.Const"() <{value = dense<[-1, -1, 1]> : tensor<3xi32>}>  : () -> tensor<3xi32>
+  %cst_0 = "tf.Const"() <{value = dense<1> : tensor<i32>}>  : () -> tensor<i32>
+  %0 = "tf.AddV2"(%arg2, %cst_0)  : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %1 = "tf.Identity"(%0)  : (tensor<i32>) -> tensor<i32>
+  %2 = "tf.TensorListGetItem"(%arg4, %arg2, %cst)  : (tensor<!tf_type.variant<tensor<?x?x1xui8>>>, tensor<i32>, tensor<3xi32>) -> tensor<?x?x1xui8>
+  %3 = "tf.EncodePng"(%2) <{compression = -1 : i64}>  : (tensor<?x?x1xui8>) -> tensor<!tf_type.string>
+  %4 = "tf.TensorListSetItem"(%arg3, %arg2, %3) <{resize_if_index_out_of_bounds = false}>  : (tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<i32>, tensor<!tf_type.string>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
+  %5 = "tf.Identity"(%4)  : (tensor<!tf_type.variant<tensor<*x!tf_type.string>>>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
+  %6 = "tf.AddV2"(%arg0, %cst_0)  : (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %7 = "tf.Identity"(%6)  : (tensor<i32>) -> tensor<i32>
+  %8 = "tf.Identity"(%arg1)  : (tensor<i32>) -> tensor<i32>
+  return %7, %8, %1, %5, %arg4 : tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>
+}
+
+// CHECK-LABEL:  while_map_while_body_884030/MapFnBody
+// CHECK:         tf.AddV2
+// CHECK-NEXT:    tf.TensorListGetItem
+// CHECK-NEXT:    tf.EncodePng
+// CHECK-NEXT:    tf.AddV2
+// CHECK-NEXT:    tf_await
+// CHECK-NEXT:    tf.TensorListSetItem
+// CHECK-NEXT:    tf_promise
+ 
+// CHECK-LABEL: @while_map_while_cond_884020
+func.func private @while_map_while_cond_884020(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<i32>, %arg3: tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, %arg4: tensor<!tf_type.variant<tensor<?x?x1xui8>>>) -> tensor<i1> {
+  %cst = "tf.Const"() <{value = dense<11> : tensor<i32>}>  : () -> tensor<i32>
+  %0 = "tf.Less"(%arg2, %cst)  : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %1 = "tf.Less"(%arg0, %arg1)  : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  %2 = "tf.LogicalAnd"(%1, %0)  : (tensor<i1>, tensor<i1>) -> tensor<i1>
+  %3 = "tf.Identity"(%2)  : (tensor<i1>) -> tensor<i1>
+  return %3 : tensor<i1>
+}
+
+// CHECK-LABEL: @main
+// CHECK:       tf.Cast
+// CHECK-NEXT:  tf.TensorListReserve
+// CHECK-NEXT:  tf.Transpose
+// CHECK-NEXT:  tf.TensorListFromTensor
+// CHECK-NEXT:  tf_mlrt.tf_map_fn
+// CHECK-SAME:   {body_fn = @"while_map_while_body_884030/MapFnBody", num_tensor_list_or_flow_in = 1 : i32} : (tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<i32>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
+// CHECK-NEXT:  tf.StopGradient
+// CHECK-NEXT:  tf.TensorListStack
+func.func @main(%arg0: tensor<1x?x?x11xf32>) -> tensor<11x!tf_type.string> {
+  %cst_0 = "tf.Const"() <{value = dense<[3, 1, 2, 0]> : tensor<4xi32>}>  : () -> tensor<4xi32>
+  %cst_10 = "tf.Const"() <{value = dense<0> : tensor<i32>}>  : () -> tensor<i32>
+  %cst_11 = "tf.Const"() <{value = dense<2> : tensor<i32>}>  : () -> tensor<i32>
+  %cst_12 = "tf.Const"() <{value = dense<1> : tensor<i32>}>  : () -> tensor<i32>
+  %cst_13 = "tf.Const"() <{value = dense<[-1, -1, 1]> : tensor<3xi32>}>  : () -> tensor<3xi32>
+  %cst_14 = "tf.Const"() <{value = dense<> : tensor<0xi32>}>  : () -> tensor<0xi32>
+  %cst_15 = "tf.Const"() <{value = dense<-1> : tensor<i32>}>  : () -> tensor<i32>
+  %cst_16 = "tf.Const"() <{value = dense<11> : tensor<i32>}>  : () -> tensor<i32>
+  %92 = "tf.Cast"(%arg0) <{Truncate = false}>  : (tensor<1x?x?x11xf32>) -> tensor<1x?x?x11xui8>
+  %0 = "tf.TensorListReserve"(%cst_15, %cst_16)  : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
+  %93 = "tf.Transpose"(%92, %cst_0)  : (tensor<1x?x?x11xui8>, tensor<4xi32>) -> tensor<11x?x?x1xui8>
+  %94 = "tf.TensorListFromTensor"(%93, %cst_13)  : (tensor<11x?x?x1xui8>, tensor<3xi32>) -> tensor<!tf_type.variant<tensor<?x?x1xui8>>>
+  %95:5 = "tf.While"(%cst_10, %cst_16, %cst_10, %0, %94) <{body = @while_map_while_body_884030, cond = @while_map_while_cond_884020, is_stateless = true, parallel_iterations = 16 : i64, shape_invariant}> {T = [i32, i32, i32, !tf_type.variant, !tf_type.variant], _lower_using_switch_merge = true, _num_original_outputs = 5 : i64, _read_only_resource_inputs = [], device = "", output_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>]} : (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>) -> (tensor<i32>, tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<!tf_type.variant<tensor<?x?x1xui8>>>)
+  %96 = "tf.StopGradient"(%95#3)  : (tensor<!tf_type.variant<tensor<*x!tf_type.string>>>) -> tensor<!tf_type.variant<tensor<*x!tf_type.string>>>
+  %97 = "tf.TensorListStack"(%96, %cst_14) <{num_elements = 11 : i64}>  : (tensor<!tf_type.variant<tensor<*x!tf_type.string>>>, tensor<0xi32>) -> tensor<11x!tf_type.string>
+  return %97 : tensor<11x!tf_type.string>
+}
+
+// -----
+
+// Test a while to map_fn conversion where constants are hoisted to resources.
+
+// CHECK-LABEL: hoisted_constants_map_while_cond
+func.func private @hoisted_constants_map_while_cond(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>) -> tensor<i1> {
+  %0 = "tf.Less"(%arg0, %arg3) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
+  %1 = "tf.Less"(%arg1, %arg3) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi1>
+  %2 = "tf.LogicalAnd"(%0, %1) : (tensor<*xi1>, tensor<*xi1>) -> tensor<*xi1>
+  %3 = "tf.ToBool"(%2) : (tensor<*xi1>) -> tensor<i1>
+  return %3 : tensor<i1>
+}
+
+// CHECK-LABEL: hoisted_constants_map_while_body
+func.func private @hoisted_constants_map_while_body(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<i32>) -> (tensor<*xi32>, tensor<*xi32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) {
+  // Resource 1 is constant 1
+  %c1 = "tf._TfrtGetResource"() <{container = [""], indices = [1 : i64], shared_name = [""]}> : () -> tensor<i32>
+  %0 = "tf.AddV2"(%arg0, %c1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
+  %1 = "tf.AddV2"(%arg1, %c1) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
+  
+  // Dummy work
+  %dummy = "tf.Const"() {value = dense<1.0> : tensor<f32>} : () -> tensor<f32>
+  %2 = "tf.TensorListSetItem"(%arg2, %arg1, %dummy) {resize_if_index_out_of_bounds = false} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<*xi32>, tensor<f32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+  
+  return %0, %1, %2, %arg3 : tensor<*xi32>, tensor<*xi32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>
+}
+
+// CHECK-LABEL: @hoisted_constants_main
+// CHECK:       tf.TensorListReserve
+// CHECK-NEXT:  tf_mlrt.tf_map_fn
+// CHECK-SAME:   {body_fn = @"hoisted_constants_map_while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
+// CHECK-NEXT:  tf.Const
+// CHECK-NEXT:  tf.TensorListStack
+// CHECK-NOT:   tf.While
+func.func @hoisted_constants_main() -> tensor<3xf32> {
+  // Resource 0 is constant 0
+  %c0 = "tf._TfrtGetResource"() <{container = [""], indices = [0 : i64], shared_name = [""]}> : () -> tensor<i32>
+  // Resource 2 is constant 3 (max iterations)
+  %c3 = "tf._TfrtGetResource"() <{container = [""], indices = [2 : i64], shared_name = [""]}> : () -> tensor<i32>
+
+  %c_neg1 = "tf.Const"() {value = dense<-1> : tensor<i32>} : () -> tensor<i32>
+  %tensor_list = "tf.TensorListReserve"(%c_neg1, %c3) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
+
+  %0:4 = "tf.While"(%c0, %c0, %tensor_list, %c3) {body = @hoisted_constants_map_while_body, cond = @hoisted_constants_map_while_cond, is_stateless = true, parallel_iterations = 10 : i64, shape_invariant} : (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) -> (tensor<i32>, tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>)
+  %cst_shape = "tf.Const"() {value = dense<[3]> : tensor<1xi32>} : () -> tensor<1xi32>
+  %stacked = "tf.TensorListStack"(%0#2, %cst_shape) {num_elements = 3 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<3xf32>
+  return %stacked : tensor<3xf32>
+}
+
+// This function defines the resources.
+func.func private @_tfrt_resource_init() {
+  %c0 = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
+  "tf._TfrtSetResource"(%c0) {index = 0 : i64} : (tensor<i32>) -> ()
+
+  %c1 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+  "tf._TfrtSetResource"(%c1) {index = 1 : i64} : (tensor<i32>) -> ()
+
+  %c3 = "tf.Const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
+  "tf._TfrtSetResource"(%c3) {index = 2 : i64} : (tensor<i32>) -> ()
+
+  return
+}

tensorflow/compiler/mlir/tfrt/tests/optimize.mlir

@@ -0,0 +1,63 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -optimize-tf-for-tfrt -split-input-file -verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: @fold_device_index
+func.func @fold_device_index() -> tensor<i32> {
+  // CHECK-NOT: tf.DeviceIndex
+  // CHECK: tf.Const
+  // CHECK-SAME: value = dense<1> : tensor<i32>
+  %0 = "tf.DeviceIndex"() {device = "/device:CPU:0", device_names = ["GPU", "CPU"]} : () -> tensor<i32>
+  func.return %0 : tensor<i32>
+}
+
+// -----
+
+// CHECK-LABEL: @not_fold_device_index
+func.func @not_fold_device_index() -> tensor<i32> {
+  // CHECK-NOT: tf.Const
+  // CHECK: tf.DeviceIndex
+  %0 = "tf.DeviceIndex"() {device = "", device_names = ["CPU", "GPU"]} : () -> tensor<i32>
+  func.return %0 : tensor<i32>
+}
+
+// -----
+
+// CHECK-LABEL: @eliminate_multinomial
+func.func @eliminate_multinomial(%0: tensor<*xf32>, %1: tensor<*xi32>) -> (tensor<*xi64>, tensor<*xi64>) {
+  // CHECK-NEXT: tf.Multinomial
+  // CHECK-NEXT: return
+  %2 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
+  %3 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
+  func.return %2, %3 : tensor<*xi64>, tensor<*xi64>
+}
+
+// -----
+
+// CHECK-LABEL: @not_eliminate_multinomial
+func.func @not_eliminate_multinomial(%0: tensor<*xf32>, %1: tensor<*xi32>) -> (tensor<*xi64>, tensor<*xi64>) {
+  // CHECK-NEXT: tf.Multinomial
+  // CHECK-SAME: seed = 0
+  // CHECK-NEXT: tf.Multinomial
+  // CHECK-SAME: seed = 1
+  // CHECK-NEXT: tf.Multinomial
+  // CHECK-SAME: seed = 0
+  // CHECK-NEXT: return
+  %2 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
+  %3 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 1 : i64, seed2 = 1 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
+  %4 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
+  %5 = "tf.Multinomial"(%0, %1) {device = "/job:localhost/replica:0/task:0/device:CPU:0", seed = 0 : i64, seed2 = 0 : i64} : (tensor<*xf32>, tensor<*xi32>) -> tensor<*xi64>
+  func.return %2, %3 : tensor<*xi64>, tensor<*xi64>
+}

tensorflow/compiler/mlir/tfrt/tests/pack_inputs.mlir

@@ -0,0 +1,34 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -pack-inputs="slices=1,0,24,2,24,24" | FileCheck %s
+
+// Tests basic contiguous slice packing:
+// - Input 0 is preserved (not in slice list).
+// - Inputs 1 and 2 are packed into a single contiguous i8 buffer with explicit slices.
+
+module {
+  // CHECK-LABEL: func.func @main
+  // CHECK-SAME: (%arg0: tensor<10x10xf32>, %arg1: tensor<48xi8>) -> (tensor<10x10xf32>, tensor<3x2xf32>, tensor<3x2xf32>)
+  func.func @main(%arg0: tensor<10x10xf32>, %arg1: tensor<3x2xf32>, %arg2: tensor<3x2xf32>) -> (tensor<10x10xf32>, tensor<3x2xf32>, tensor<3x2xf32>) attributes {ifrt.function} {
+    // CHECK: %[[SLICE0:.*]] = stablehlo.slice %arg1 [0:24] : (tensor<48xi8>) -> tensor<24xi8>
+    // CHECK: %[[RESHAPE0:.*]] = stablehlo.reshape %[[SLICE0]] : (tensor<24xi8>) -> tensor<3x2x4xi8>
+    // CHECK: %[[BITCAST0:.*]] = stablehlo.bitcast_convert %[[RESHAPE0]] : (tensor<3x2x4xi8>) -> tensor<3x2xf32>
+    // CHECK: %[[SLICE1:.*]] = stablehlo.slice %arg1 [24:48] : (tensor<48xi8>) -> tensor<24xi8>
+    // CHECK: %[[RESHAPE1:.*]] = stablehlo.reshape %[[SLICE1]] : (tensor<24xi8>) -> tensor<3x2x4xi8>
+    // CHECK: %[[BITCAST1:.*]] = stablehlo.bitcast_convert %[[RESHAPE1]] : (tensor<3x2x4xi8>) -> tensor<3x2xf32>
+    // CHECK: return %arg0, %[[BITCAST0]], %[[BITCAST1]] : tensor<10x10xf32>, tensor<3x2xf32>, tensor<3x2xf32>
+    return %arg0, %arg1, %arg2 : tensor<10x10xf32>, tensor<3x2xf32>, tensor<3x2xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/pack_inputs_exclude.mlir

@@ -0,0 +1,35 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt %s -pack-inputs="slices=1,0,24,2,28,16" | FileCheck %s
+
+// Tests advanced slice packing with gaps/padding and unmerged variables:
+// - Input 0 is preserved.
+// - Inputs 1 and 2 are packed into a single i8 buffer with a 4-byte gap (start 0 and start 28).
+// - Inputs 3 and 4 are skipped and remain unchanged.
+
+module {
+  // CHECK-LABEL: func.func @main
+  // CHECK-SAME: (%arg0: tensor<10x10xf32>, %arg1: tensor<5x2xf32>, %arg2: tensor<6x1xf32>, %arg3: tensor<44xi8>) -> (tensor<10x10xf32>, tensor<2x3xf32>, tensor<4x1xf32>, tensor<5x2xf32>, tensor<6x1xf32>)
+  func.func @main(%arg0: tensor<10x10xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor<4x1xf32>, %arg3: tensor<5x2xf32>, %arg4: tensor<6x1xf32>) -> (tensor<10x10xf32>, tensor<2x3xf32>, tensor<4x1xf32>, tensor<5x2xf32>, tensor<6x1xf32>) attributes {ifrt.function} {
+    // CHECK: %[[SLICE0:.*]] = stablehlo.slice %arg3 [0:24] : (tensor<44xi8>) -> tensor<24xi8>
+    // CHECK: %[[RESHAPE0:.*]] = stablehlo.reshape %[[SLICE0]] : (tensor<24xi8>) -> tensor<2x3x4xi8>
+    // CHECK: %[[BITCAST0:.*]] = stablehlo.bitcast_convert %[[RESHAPE0]] : (tensor<2x3x4xi8>) -> tensor<2x3xf32>
+    // CHECK: %[[SLICE1:.*]] = stablehlo.slice %arg3 [28:44] : (tensor<44xi8>) -> tensor<16xi8>
+    // CHECK: %[[RESHAPE1:.*]] = stablehlo.reshape %[[SLICE1]] : (tensor<16xi8>) -> tensor<4x1x4xi8>
+    // CHECK: %[[BITCAST1:.*]] = stablehlo.bitcast_convert %[[RESHAPE1]] : (tensor<4x1x4xi8>) -> tensor<4x1xf32>
+    // CHECK: return %arg0, %[[BITCAST0]], %[[BITCAST1]], %arg1, %arg2 : tensor<10x10xf32>, tensor<2x3xf32>, tensor<4x1xf32>, tensor<5x2xf32>, tensor<6x1xf32>
+    return %arg0, %arg1, %arg2, %arg3, %arg4 : tensor<10x10xf32>, tensor<2x3xf32>, tensor<4x1xf32>, tensor<5x2xf32>, tensor<6x1xf32>
+  }
+}

tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op.mlir

@@ -0,0 +1,123 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tfrt-reconfig-batch-op="tfrt-min-num-batch-threads=2 tfrt-min-max-enqueued-batches=3 tfrt-batch-padding-policy=PAD_UP" %s | FileCheck %s --dump-input=always
+
+// -----
+
+// The num_batch_threads is lowered bound to 2 from the original attribute of 1
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
+  // CHECK:  "tf.BatchFunction"
+  // CHECK-SAME: allowed_batch_sizes = [6]
+  // CHECK-SAME: batch_padding_policy = "PAD_UP"
+  // CHECK-SAME: batch_timeout_micros = 100000 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = false
+  // CHECK-SAME: max_batch_size = 6 : i64
+  // CHECK-SAME: max_enqueued_batches = 10 : i64
+  // CHECK-SAME: num_batch_threads = 2 : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}
+
+// -----
+
+// The num_batch_threads remains 3 (the same as the original attribute)
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
+  // CHECK:  "tf.BatchFunction"
+  // CHECK-SAME: allowed_batch_sizes = [6]
+  // CHECK-SAME: batch_padding_policy = "PAD_UP"
+  // CHECK-SAME: batch_timeout_micros = 100000 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = false
+  // CHECK-SAME: max_batch_size = 6 : i64
+  // CHECK-SAME: max_enqueued_batches = 10 : i64
+  // CHECK-SAME: num_batch_threads = 3 : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 3 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}
+
+// -----
+
+// The max_enqueued_batches is changed to 3 from the original attribute of 2
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
+  // CHECK:  "tf.BatchFunction"
+  // CHECK-SAME: allowed_batch_sizes = [6]
+  // CHECK-SAME: batch_padding_policy = "PAD_UP"
+  // CHECK-SAME: batch_timeout_micros = 100000 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = false
+  // CHECK-SAME: max_batch_size = 6 : i64
+  // CHECK-SAME: max_enqueued_batches = 3 : i64
+  // CHECK-SAME: num_batch_threads = 2 : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 2 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}
+
+// -----
+
+// The max_enqueued_batches remains 10 (the same as the original attribute)
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
+  // CHECK:  "tf.BatchFunction"
+  // CHECK-SAME: allowed_batch_sizes = [6]
+  // CHECK-SAME: batch_padding_policy = "PAD_UP"
+  // CHECK-SAME: batch_timeout_micros = 100000 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = false
+  // CHECK-SAME: max_batch_size = 6 : i64
+  // CHECK-SAME: max_enqueued_batches = 10 : i64
+  // CHECK-SAME: num_batch_threads = 3 : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 3 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}

tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op_global_prioritization.mlir

@@ -0,0 +1,81 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tfrt-reconfig-batch-op="tfrt-batch-queue-global-prioritization-num-threads=4" %s | FileCheck %s --dump-input=always
+
+// -----
+
+// The num_batch_threads is updated to 4 from the original attribute of 2,
+// a mixed_priority_policy is set along with all low priority batching params
+// being copied from the high priority batching params since no low priority
+// settings are provided.
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
+  // CHECK:  "tf.BatchFunction"
+  // CHECK-SAME: allowed_batch_sizes = [1, 2]
+  // CHECK-SAME: batch_padding_policy = "PAD_UP"
+  // CHECK-SAME: batch_timeout_micros = 2 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = false
+  // CHECK-SAME: enable_priority_aware_batch_scheduler = true
+  // CHECK-SAME: max_batch_size = 2 : i64
+  // CHECK-SAME: max_enqueued_batches = 2 : i64
+  // CHECK-SAME: num_batch_threads = 4 : i64
+  // CHECK-SAME: num_warmup_batch_threads = {{[0-9]+}} : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [1, 2], batch_padding_policy = "PAD_UP", batch_timeout_micros = 2 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 2 : i64, max_enqueued_batches = 2 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}
+
+// -----
+
+// Same as first test, but low_priority_* parameters are already provided
+// so not overriden.
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
+  // CHECK:  "tf.BatchFunction"
+  // CHECK-SAME: allowed_batch_sizes = [1, 2]
+  // CHECK-SAME: batch_padding_policy = "PAD_UP"
+  // CHECK-SAME: batch_timeout_micros = 2 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = false
+  // CHECK-SAME: enable_priority_aware_batch_scheduler = true
+  // CHECK-SAME: low_priority_allowed_batch_sizes = [1, 10]
+  // CHECK-SAME: low_priority_batch_timeout_micros = 7 : i64
+  // CHECK-SAME: low_priority_max_batch_size = 8 : i64
+  // CHECK-SAME: low_priority_max_enqueued_batches = 9 : i64
+  // CHECK-SAME: max_batch_size = 2 : i64
+  // CHECK-SAME: max_enqueued_batches = 2 : i64
+  // CHECK-SAME: num_batch_threads = 4 : i64
+  // CHECK-SAME: num_warmup_batch_threads = {{[0-9]+}} : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [1, 2], batch_padding_policy = "PAD_UP", batch_timeout_micros = 2 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 2 : i64, max_enqueued_batches = 2 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/", low_priority_batch_timeout_micros = 7 : i64, low_priority_max_batch_size = 8 : i64, low_priority_max_enqueued_batches = 9 : i64, low_priority_allowed_batch_sizes = [1, 10]} : (tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}

tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op_params.mlir

@@ -0,0 +1,64 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tfrt-reconfig-batch-op="tfrt-min-num-batch-threads=4 tfrt-min-max-enqueued-batches=4 tfrt-num-batch-threads=3 tfrt-max-batch-size=3 tfrt-batch-timeout-micros=3 tfrt-allowed-batch-sizes=3,4 tfrt-max-enqueued-batches=3 tfrt-enable-large-batch-splitting=true tfrt-mixed-priority-batching-policy=priority_merge tfrt-low-priority-max-batch-size=5 tfrt-low-priority-batch-timeout-micros=5 tfrt-low-priority-allowed-batch-sizes=5,6 tfrt-low-priority-max-enqueued-batches=5 tfrt-num-warmup-batch-threads=2" %s | FileCheck %s --dump-input=always
+
+// -----
+
+// The num_batch_threads is updated to 3 from the original attribute of 2,
+// overriding the min_num_batch_threads of 4.
+// The max_batch_size is updated to 3 from the original attribute of 2.
+// The batch_timeout_micros is updated to 3 from the original attribute of 2.
+// The allowed_batch_sizes is updated to [3, 4] from the original attribute of
+// [1, 2].
+// The max_enqueued_batches is updated to 3 from the original attribute of 2,
+// overriding the min_max_enqueued_batches of 4.
+// The enable_large_batch_splitting is updated to true from the original
+// attribute of false.
+// The mixed_priority_policy is updated to "priority_merge" from the original
+// attribute of "".
+// The low_priority_max_batch_size is updated to 5 from the original attribute of 2.
+// The low_priority_batch_timeout_micros is updated to 5 from the original attribute of 2.
+// The low_priority_allowed_batch_sizes is updated to [5, 6] from the original attribute of [1, 2].
+// The low_priority_max_enqueued_batches is updated to 5 from the original attribute of 2.
+// The num_warmup_batch_threads is updated to 2 from the original attribute of 0.
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
+  // CHECK:  "tf.BatchFunction"
+  // CHECK-SAME: allowed_batch_sizes = [3, 4]
+  // CHECK-SAME: batch_padding_policy = "PAD_UP"
+  // CHECK-SAME: batch_timeout_micros = 3 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = true
+  // CHECK-SAME: low_priority_allowed_batch_sizes = [5, 6]
+  // CHECK-SAME: low_priority_batch_timeout_micros = 5 : i64
+  // CHECK-SAME: low_priority_max_batch_size = 5 : i64
+  // CHECK-SAME: low_priority_max_enqueued_batches = 5 : i64
+  // CHECK-SAME: max_batch_size = 3 : i64
+  // CHECK-SAME: max_enqueued_batches = 3 : i64
+  // CHECK-SAME: mixed_priority_policy = "priority_merge"
+  // CHECK-SAME: num_batch_threads = 3 : i64
+  // CHECK-SAME: num_warmup_batch_threads = 2 : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [1, 2], batch_padding_policy = "PAD_UP", batch_timeout_micros = 2 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, low_priority_allowed_batch_sizes = [1, 2], low_priority_batch_timeout_micros = 2 : i64, low_priority_max_batch_size = 2 : i64, low_priority_max_enqueued_batches = 2 : i64, max_batch_size = 2 : i64, max_enqueued_batches = 2 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}

tensorflow/compiler/mlir/tfrt/tests/reconfig_batch_op_skip_pass.mlir

@@ -0,0 +1,43 @@
+// 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.
+// ==============================================================================
+// RUN: tf-tfrt-opt -split-input-file -tfrt-reconfig-batch-op="tfrt-min-num-batch-threads=0 tfrt-min-max-enqueued-batches=0" %s | FileCheck %s --dump-input=always
+
+// -----
+
+// Confirm no batch options change when the tfrt-reconfig-batch-op flag is
+// empty.
+
+// CHECK-LABEL: func private @batched_function
+func.func private @batched_function(%arg0: tensor<1x3xf32>) -> tensor<1x3xf32> {
+  %2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
+  func.return %2 : tensor<1x3xf32>
+}
+
+// CHECK-LABEL: func @main
+func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
+  // CHECK:  "tf.BatchFunction"
+  // CHECK-SAME: allowed_batch_sizes = [6]
+  // CHECK-SAME: batch_padding_policy = "PAD_UP"
+  // CHECK-SAME: batch_timeout_micros = 100000 : i64
+  // CHECK-SAME: batching_queue = ""
+  // CHECK-SAME: container = ""
+  // CHECK-SAME: enable_large_batch_splitting = false
+  // CHECK-SAME: max_batch_size = 6 : i64
+  // CHECK-SAME: max_enqueued_batches = 10 : i64
+  // CHECK-SAME: num_batch_threads = 1 : i64
+  // CHECK-SAME: shared_name = "batch/"
+  %1 = "tf.BatchFunction"(%arg0) {allowed_batch_sizes = [6], batch_padding_policy = "PAD_UP", batch_timeout_micros = 100000 : i64, batching_queue = "", container = "", device = "/device:CPU:0", enable_large_batch_splitting = false, f = @batched_function, max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64, num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>, shared_name = "batch/"} : (tensor<1x3xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}