chore: import upstream snapshot with attribution

tensorflow/core/tfrt/BUILD

@@ -0,0 +1,23 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:license"])
+
+cc_library(
+    name = "stream_ops_op_lib",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core/tfrt/kernels:stream_ops",
+        "//tensorflow/core/tfrt/ops:stream_ops_op_lib",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "ifrt_program_ops_op_lib",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core/tfrt/kernels:ifrt_program_ops",
+        "//tensorflow/core/tfrt/ops:ifrt_program_ops_op_lib",
+    ],
+    alwayslink = 1,
+)

tensorflow/core/tfrt/common/BUILD

@@ -0,0 +1,329 @@
+load(
+    "@xla//xla/tsl/platform:build_config_root.bzl",
+    "tf_cuda_tests_tags",
+)
+load("//tensorflow:tensorflow.bzl", "tf_cc_test", "tf_cuda_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+# Authorized users go here.
+package_group(
+    name = "friends",
+    packages = [
+        # copybara:uncomment "//learning/brain/experimental/dtensor/...",
+        # copybara:uncomment "//learning/brain/experimental/tfrt/...",
+        # copybara:uncomment "//learning/brain/google/monitoring/...",
+        # copybara:uncomment "//learning/brain/google/xla/...",
+        # copybara:uncomment "//learning/brain/tfrc/...",
+        # copybara:uncomment "//learning/brain/tfrt/...",
+        # copybara:uncomment "//learning/serving/model_servers/...",
+        "//tensorflow/c/...",
+        "//tensorflow/compiler/jit/...",
+        "//tensorflow/compiler/tf2xla/...",
+        "//tensorflow/core/common_runtime/...",
+        "//tensorflow/core/common_runtime/next_pluggable_device/...",
+        "//tensorflow/core/tfrt/...",
+        "//tensorflow/core/tpu/...",
+        "//tensorflow/dtensor/...",
+        "//third_party/tf_runtime_google/...",
+    ],
+)
+
+cc_library(
+    name = "global_state",
+    srcs = [
+        "global_state.cc",
+    ],
+    hdrs = [
+        "global_state.h",
+    ],
+    visibility = [":friends"],
+    deps = [
+        "//tensorflow/core:framework",
+        "@tf_runtime//:hostcontext",
+        "@xla//xla/pjrt:utils",
+    ],
+)
+
+cc_library(
+    name = "async_value_tensor",
+    srcs = [
+        "async_value_tensor.cc",
+    ],
+    hdrs = [
+        "async_value_tensor.h",
+    ],
+    visibility = [":friends"],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core/common_runtime:dma_helper",
+        "@com_google_absl//absl/log:check",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:support",
+        "@xla//xla/pjrt:pjrt_client",
+    ],
+)
+
+tf_cc_test(
+    name = "async_value_tensor_test",
+    srcs = ["async_value_tensor_test.cc"],
+    deps = [
+        ":async_value_tensor",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/tsl/concurrency:async_value",
+    ],
+)
+
+cc_library(
+    name = "pjrt_state",
+    srcs = [
+        "pjrt_state.cc",
+    ],
+    hdrs = [
+        "pjrt_state.h",
+    ],
+    visibility = [":friends"],
+    deps = [
+        ":pjrt_client_factory_options",
+        ":pjrt_client_factory_registry",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core/framework:resource_base",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/synchronization",
+        "@tsl//tsl/platform:statusor",
+        "@xla//xla/client:local_client",
+        "@xla//xla/pjrt:host_memory_allocator",
+        "@xla//xla/pjrt:local_device_state",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt:tf_pjrt_client",
+        "@xla//xla/stream_executor/integrations:tf_allocator_adapter",
+    ],
+)
+
+# This target does not have depend on GPU and XLA compilation to have a small size.
+# Uses this target if you do not need to create a PJRT client.
+cc_library(
+    name = "pjrt_util",
+    srcs = [
+        "pjrt_util.cc",
+    ],
+    hdrs = [
+        "pjrt_util.h",
+    ],
+    visibility = [":friends"],
+    deps = [
+        ":global_state",
+        ":pjrt_state",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_types_hdr",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:refcount",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@tsl//tsl/platform:errors",
+        "@xla//xla/pjrt:pjrt_client",
+    ],
+)
+
+# Uses this target if you need to create a PJRT client.
+# TODO(b/280671896) Combines pjrt_util and create_pjrt_client_util
+cc_library(
+    name = "create_pjrt_client_util",
+    srcs = [
+        "create_pjrt_client_util.cc",
+    ],
+    hdrs = [
+        "create_pjrt_client_util.h",
+    ],
+    visibility = [":friends"],
+    deps = [
+        ":global_state",
+        ":pjrt_state",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_types_hdr",
+        "//tensorflow/core/platform:refcount",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@tsl//tsl/platform:errors",
+        "@xla//xla/pjrt:pjrt_client",
+    ],
+)
+
+tf_cc_test(
+    name = "pjrt_state_test",
+    srcs = ["pjrt_state_test.cc"],
+    deps = [
+        ":pjrt_cpu_client_registration",
+        ":pjrt_state",
+        "//tensorflow/core:framework_types_hdr",
+        "//tensorflow/core:test",
+        "//tensorflow/core/platform:refcount",
+        "//tensorflow/core/protobuf:error_codes_proto_impl_cc",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:status_matchers",
+        "@tsl//tsl/platform:statusor",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options",
+        "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)
+
+tf_cc_test(
+    name = "pjrt_util_test",
+    srcs = ["pjrt_util_test.cc"],
+    deps = [
+        ":pjrt_state",
+        ":pjrt_util",
+        "//tensorflow/core:framework",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:status_matchers",
+        "@tsl//tsl/platform:statusor",
+        "@tsl//tsl/platform:test",
+        "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options",
+        "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
+        "@xla//xla/tsl/lib/core:status_test_util",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)
+
+tf_cuda_cc_test(
+    name = "create_pjrt_client_util_test",
+    srcs = ["create_pjrt_client_util_test.cc"],
+    tags = tf_cuda_tests_tags() + [
+        "noasan",
+        "nomsan",
+        "notsan",
+    ],
+    deps = [
+        ":create_pjrt_client_util",
+        ":global_state",
+        ":pjrt_gpu_client_registration",
+        ":pjrt_state",
+        "//tensorflow/compiler/tf2xla:xla_op_registry",
+        "//tensorflow/core:framework",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:status_matchers",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/service:gpu_plugin",
+        "@xla//xla/tsl/lib/core:status_test_util",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)
+
+cc_library(
+    name = "pjrt_client_factory_options",
+    hdrs = ["pjrt_client_factory_options.h"],
+    visibility = [":friends"],
+)
+
+cc_library(
+    name = "pjrt_client_factory_registry",
+    srcs = ["pjrt_client_factory_registry.cc"],
+    hdrs = ["pjrt_client_factory_registry.h"],
+    deps = [
+        ":pjrt_client_factory_options",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_lite",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@tsl//tsl/platform:errors",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/tsl/framework:device_type",
+    ],
+)
+
+cc_library(
+    name = "pjrt_cpu_client_registration",
+    srcs = ["pjrt_cpu_client_registration.cc"],
+    deps = [
+        ":pjrt_client_factory_options",
+        ":pjrt_client_factory_registry",
+        "//tensorflow/core:framework_types_hdr",
+        "@com_google_absl//absl/status:statusor",
+        "@tsl//tsl/platform:statusor",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options",
+        "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
+    ],
+    alwayslink = True,
+)
+
+tf_cc_test(
+    name = "pjrt_cpu_client_registration_test",
+    srcs = ["pjrt_cpu_client_registration_test.cc"],
+    deps = [
+        ":pjrt_client_factory_options",
+        ":pjrt_client_factory_registry",
+        ":pjrt_cpu_client_registration",
+        "//tensorflow/core:framework_types_hdr",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "pjrt_gpu_client_registration",
+    srcs = ["pjrt_gpu_client_registration.cc"],
+    deps = [
+        ":pjrt_client_factory_options",
+        ":pjrt_client_factory_registry",
+        "//tensorflow/compiler/tf2xla:xla_op_registry",
+        "//tensorflow/core:framework_types_hdr",
+        "@com_google_absl//absl/status:statusor",
+        "@tsl//tsl/platform:statusor",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
+        "@xla//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client",
+    ],
+    alwayslink = True,
+)
+
+cc_library(
+    name = "metrics",
+    srcs = ["metrics.cc"],
+    hdrs = ["metrics.h"],
+    deps = [
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+tf_cuda_cc_test(
+    name = "pjrt_gpu_client_registration_test",
+    size = "small",
+    srcs = ["pjrt_gpu_client_registration_test.cc"],
+    tags = tf_cuda_tests_tags() + [
+        "noasan",
+        "nomsan",
+        "notsan",
+    ],
+    deps = [
+        ":pjrt_client_factory_options",
+        ":pjrt_client_factory_registry",
+        ":pjrt_gpu_client_registration",
+        "//tensorflow/core:framework_types_hdr",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:statusor",
+        "@xla//xla/service:gpu_plugin",
+    ],
+)

tensorflow/core/tfrt/common/async_value_tensor.cc

@@ -0,0 +1,86 @@
+/* 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/core/tfrt/common/async_value_tensor.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include "absl/log/check.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "tensorflow/core/common_runtime/dma_helper.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tfrt/host_context/async_value.h"  // from @tf_runtime
+#include "tfrt/support/forward_decls.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+namespace {
+constexpr uintptr_t kTag = 0x1ULL;
+}
+
+/*static*/ AsyncValueTensor* AsyncValueTensor::FromTensor(
+    const Tensor* tensor) {
+  if (tensor == nullptr) return nullptr;
+  const TensorBuffer* buf = DMAHelper::buffer(tensor);
+  if (buf == nullptr || !buf->AllocatesOpaqueHandle()) {
+    return nullptr;
+  }
+  AsyncValueTensor* av_tensor = FromOpaquePointer(buf->data());
+  return av_tensor;
+}
+
+const tfrt::RCReference<tfrt::AsyncValue>& AsyncValueTensor::GetAsyncRef() {
+  return av_ref_;
+}
+
+void AsyncValueTensor::SetAsyncRef(tfrt::RCReference<tfrt::AsyncValue> av_ref) {
+  av_ref_ = std::move(av_ref);
+}
+
+std::shared_ptr<xla::PjRtBuffer> AsyncValueTensor::GetBuffer() {
+  return buffer_;
+}
+
+void AsyncValueTensor::SetBuffer(std::shared_ptr<xla::PjRtBuffer> buffer) {
+  buffer_ = std::move(buffer);
+}
+
+/*static*/ AsyncValueTensor* AsyncValueTensor::FromOpaquePointer(void* ptr) {
+  uintptr_t value = reinterpret_cast<uintptr_t>(ptr);
+  if (value & kTag) {
+    return reinterpret_cast<AsyncValueTensor*>(value & ~kTag);
+  } else {
+    return nullptr;
+  }
+}
+
+/*static*/ void* AsyncValueTensor::ToOpaquePointer(AsyncValueTensor* tensor) {
+  uintptr_t value = reinterpret_cast<uintptr_t>(tensor);
+  CHECK_EQ(value & kTag, 0);  // Crash OK
+  value |= kTag;
+  return reinterpret_cast<AsyncValueTensor*>(value);
+}
+
+void* AsyncValueAllocator::AllocateRaw(size_t alignment, size_t num_bytes) {
+  return AsyncValueTensor::ToOpaquePointer(new AsyncValueTensor);
+}
+
+void AsyncValueAllocator::DeallocateRaw(void* ptr) {
+  delete AsyncValueTensor::FromOpaquePointer(ptr);
+}
+
+}  // namespace tensorflow

tensorflow/core/tfrt/common/async_value_tensor.h

@@ -0,0 +1,72 @@
+/* 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_CORE_TFRT_COMMON_ASYNC_VALUE_TENSOR_H_
+#define TENSORFLOW_CORE_TFRT_COMMON_ASYNC_VALUE_TENSOR_H_
+
+#include <cstddef>
+#include <memory>
+
+#include "xla/pjrt/pjrt_client.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/types.h"
+#include "tfrt/support/forward_decls.h"  // from @tf_runtime
+#include "tfrt/support/ref_count.h"  // from @tf_runtime
+
+namespace tensorflow {
+
+// The implementation of a Tensor for an AsyncValue and PjRtBuffer. We used it
+// to integrate TF with TFRT.
+// TODO(b/243983834) After the migration of using PjRt for data transfer is
+// completed, GetAsyncRef and SetAsyncRef will be removed and this class will be
+// renamed to PjRtBufferTensor.
+class AsyncValueTensor {
+ public:
+  // Downcast from a Tensor to an AsyncValueTensor. Return nullptr if the
+  // downcast fails.
+  static AsyncValueTensor* FromTensor(const Tensor* tensor);
+
+  const tfrt::RCReference<tfrt::AsyncValue>& GetAsyncRef();
+
+  void SetAsyncRef(tfrt::RCReference<tfrt::AsyncValue> av_ref);
+
+  std::shared_ptr<xla::PjRtBuffer> GetBuffer();
+
+  void SetBuffer(std::shared_ptr<xla::PjRtBuffer> buffer);
+
+  // Convert from a raw pointer to an AsyncValueTensor, removing the pointer
+  // tag.
+  static AsyncValueTensor* FromOpaquePointer(void* ptr);
+
+  // Convert to a raw pointer from an AsyncValueTensor, adding the pointer tag.
+  static void* ToOpaquePointer(AsyncValueTensor* tensor);
+
+ private:
+  tfrt::RCReference<tfrt::AsyncValue> av_ref_;
+  std::shared_ptr<xla::PjRtBuffer> buffer_;
+};
+
+class AsyncValueAllocator : public Allocator {
+ public:
+  void* AllocateRaw(size_t alignment, size_t num_bytes) override;
+  void DeallocateRaw(void* ptr) override;
+
+  bool AllocatesOpaqueHandle() const override { return true; }
+  std::string Name() override { return "async-value"; }
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_COMMON_ASYNC_VALUE_TENSOR_H_

tensorflow/core/tfrt/common/async_value_tensor_test.cc

@@ -0,0 +1,76 @@
+/* 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/core/tfrt/common/async_value_tensor.h"
+
+#include <cstdint>
+#include <memory>
+
+#include <gtest/gtest.h>
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/tsl/concurrency/async_value_ref.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace {
+
+TEST(AsyncValueTensorTest, InvalidTensor) {
+  tensorflow::Tensor tensor(tensorflow::DT_INT64, tensorflow::TensorShape({1}));
+
+  AsyncValueTensor* avt = AsyncValueTensor::FromTensor(&tensor);
+
+  ASSERT_EQ(avt, nullptr);
+}
+
+TEST(AsyncValueTensorTest, SetAndGetAsyncValue) {
+  AsyncValueAllocator allocator;
+  tensorflow::Tensor tensor(&allocator, tensorflow::DT_INT64,
+                            tensorflow::TensorShape({1}));
+
+  AsyncValueTensor* avt = AsyncValueTensor::FromTensor(&tensor);
+
+  ASSERT_NE(avt, nullptr);
+
+  tsl::AsyncValueRef<int32_t> value =
+      tsl::MakeConstructedAsyncValueRef<int32_t>(123);
+
+  avt->SetAsyncRef(value.CopyRCRef());
+
+  auto ret_value = avt->GetAsyncRef();
+  ASSERT_EQ(ret_value, value.CopyRCRef());
+}
+
+TEST(AsyncValueTensorTest, SetAndGetBuffer) {
+  AsyncValueAllocator allocator;
+  tensorflow::Tensor tensor(&allocator, tensorflow::DT_INT64,
+                            tensorflow::TensorShape({1}));
+
+  AsyncValueTensor* avt = AsyncValueTensor::FromTensor(&tensor);
+
+  ASSERT_NE(avt, nullptr);
+
+  std::shared_ptr<xla::PjRtBuffer> buffer;
+
+  avt->SetBuffer(buffer);
+
+  auto ret_buffer = avt->GetBuffer();
+
+  ASSERT_EQ(ret_buffer, buffer);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/tfrt/common/create_pjrt_client_util.cc

@@ -0,0 +1,47 @@
+/* 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/core/tfrt/common/create_pjrt_client_util.h"
+
+#include <optional>
+#include <set>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/tfrt/common/global_state.h"
+#include "tensorflow/core/tfrt/common/pjrt_state.h"
+#include "tsl/platform/errors.h"
+
+namespace tensorflow {
+
+absl::StatusOr<xla::PjRtClient*> GetOrCreatePjRtClient(
+    const DeviceType& device_type,
+    std::optional<std::set<int>> allowed_devices) {
+  ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr();
+  PjRtState* pjrt_state;
+  TF_RETURN_IF_ERROR(rmgr->LookupOrCreate<PjRtState>(
+      rmgr->default_container(), kPjRtStateResourceName, &pjrt_state,
+      [&](PjRtState** ret) {
+        *ret = PjRtState::Create();
+        return absl::OkStatus();
+      }));
+  core::ScopedUnref pjrt_state_ref(pjrt_state);
+  return pjrt_state->GetOrCreatePjRtClient(device_type);
+}
+
+}  // namespace tensorflow

tensorflow/core/tfrt/common/create_pjrt_client_util.h

@@ -0,0 +1,45 @@
+/* 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_CORE_TFRT_COMMON_CREATE_PJRT_CLIENT_UTIL_H_
+#define TENSORFLOW_CORE_TFRT_COMMON_CREATE_PJRT_CLIENT_UTIL_H_
+
+#include <optional>
+#include <set>
+
+#include "absl/status/statusor.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "tensorflow/core/framework/types.h"
+
+namespace tensorflow {
+
+// Gets PJRT client from TFGlobalResourceManager. If it is not found, creates a
+// PJRT client and adds it to TFGlobalResourceManager. Different `DeviceType`
+// can choose to create the PJRT client explicitly (e.g. in ops) and add it to
+// TFGlobalResourceManager, or create a PJRT client on the first use implicitly
+// in this method.
+// The inputs are the device_type of the caller, and an optional
+// set of device IDs `allowed_devices` for which the stream executor will be
+// created. `allowed_devices` is only used for GPU.
+// TODO(b/260802979): consider passing `XlaPlatformInfo` for the options to
+// create a client, or creating a class similar to `LocalClientOptions`.
+// TODO(b/280111106): make PjrtClientFactoryOptions an input of
+// GetOrCreatePjRtClient.
+absl::StatusOr<xla::PjRtClient*> GetOrCreatePjRtClient(
+    const DeviceType& device_type,
+    std::optional<std::set<int>> allowed_devices = std::nullopt);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_COMMON_CREATE_PJRT_CLIENT_UTIL_H_

tensorflow/core/tfrt/common/create_pjrt_client_util_test.cc

@@ -0,0 +1,47 @@
+/* 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/core/tfrt/common/create_pjrt_client_util.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"  // IWYU pragma: keep
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tsl/platform/status_matchers.h"
+
+namespace tensorflow {
+namespace {
+
+using ::testing::HasSubstr;
+
+TEST(CreatePjRtClientTest, GetNotExistPjRtClientNotImplemented) {
+  EXPECT_THAT(GetOrCreatePjRtClient(DEVICE_CPU),
+              absl_testing::StatusIs(
+                  error::NOT_FOUND,
+                  HasSubstr(absl::StrCat("The PJRT client factory of `",
+                                         DEVICE_CPU, "` is not registered"))));
+}
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+TEST(CreatePjRtClientTest, GetNotExistGpuPjRtClient) {
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client,
+                          GetOrCreatePjRtClient(DEVICE_XLA_GPU));
+  EXPECT_THAT(pjrt_client, ::testing::NotNull());
+}
+#endif
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/tfrt/common/global_state.cc

@@ -0,0 +1,71 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/tfrt/common/global_state.h"
+
+#include <memory>
+#include <utility>
+
+#include "xla/pjrt/utils.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+#include "tfrt/host_context/host_allocator.h"  // from @tf_runtime
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_global {
+namespace {
+
+tfrt::HostContext* GetStaticHostContext() {
+  static ::tfrt::HostContext* host_context = [] {
+    // Create host context.
+    auto decoded_diagnostic_handler =
+        [&](const ::tfrt::DecodedDiagnostic& diag) { abort(); };
+    std::unique_ptr<::tfrt::ConcurrentWorkQueue> work_queue =
+        ::tfrt::CreateMultiThreadedWorkQueue(
+            /*num_threads=*/xla::DefaultThreadPoolSize(),
+            /*num_blocking_threads=*/64);
+    std::unique_ptr<::tfrt::HostAllocator> host_allocator =
+        ::tfrt::CreateMallocAllocator();
+    return new ::tfrt::HostContext(decoded_diagnostic_handler,
+                                   std::move(host_allocator),
+                                   std::move(work_queue));
+  }();
+  return host_context;
+}
+
+}  // namespace
+
+/*static*/ ::tfrt::HostContext* GlobalHostContext::host_ctx_ = nullptr;
+
+/*static*/ void GlobalHostContext::Set(::tfrt::HostContext* host_ctx) {
+  host_ctx_ = host_ctx;
+}
+
+/*static*/ ::tfrt::HostContext* GlobalHostContext::Get() {
+  // If HostContext is explicitly injected at context creation, use it here.
+  if (host_ctx_) return host_ctx_;
+
+  // Otherwise we assume it is running TFRT TF OpKernels, and currently it is
+  // implicitly created.
+  return GetStaticHostContext();
+}
+
+ResourceMgr* GetTFGlobalResourceMgr() {
+  static ResourceMgr* const rmgr = new ResourceMgr();
+  return rmgr;
+}
+
+}  // namespace tfrt_global
+}  // namespace tensorflow

tensorflow/core/tfrt/common/global_state.h

@@ -0,0 +1,40 @@
+/* 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_CORE_TFRT_COMMON_GLOBAL_STATE_H_
+#define TENSORFLOW_CORE_TFRT_COMMON_GLOBAL_STATE_H_
+
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_global {
+
+class GlobalHostContext {
+ public:
+  static void Set(::tfrt::HostContext* host_ctx);
+  static ::tfrt::HostContext* Get();
+
+ private:
+  static ::tfrt::HostContext* host_ctx_;
+};
+
+// A global resource manager in TF core framework. It can be used to store
+// resources that are per host.
+ResourceMgr* GetTFGlobalResourceMgr();
+
+}  // namespace tfrt_global
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_COMMON_GLOBAL_STATE_H_

tensorflow/core/tfrt/common/metrics.cc

@@ -0,0 +1,49 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/common/metrics.h"
+
+#include <cstdint>
+#include <string>
+
+#include "absl/strings/str_cat.h"
+#include "xla/tsl/lib/monitoring/sampler.h"
+
+namespace tensorflow {
+namespace tfrt_metrics {
+
+tsl::monitoring::SamplerCell* GetTfrtGraphExecutorLatencySampler(
+    const std::string& model_name, int64_t model_version,
+    const std::string& graph_name) {
+  static auto* cell = tsl::monitoring::Sampler<3>::New(
+      {"/tfrt/graph_executor/latency",
+       "Tracks the latency of GraphExecutor (in microseconds) of a graph.",
+       "model_name", "model_version", "graph_name"},
+      tsl::monitoring::Buckets::Exponential(10, 1.5, 33));
+  return cell->GetCell(model_name, absl::StrCat(model_version), graph_name);
+}
+
+tsl::monitoring::SamplerCell* GetTfrtDeviceExecutionLatency(
+    const std::string& model_name, int64_t model_version) {
+  static auto* cell = tsl::monitoring::Sampler<2>::New(
+      {"/tfrt/device_execution/latency",
+       "Tracks the latency of device execution (in microseconds).",
+       "model_name", "model_version"},
+      tsl::monitoring::Buckets::Exponential(10, 1.5, 33));
+  return cell->GetCell(model_name, absl::StrCat(model_version));
+}
+
+}  // namespace tfrt_metrics
+}  // namespace tensorflow

tensorflow/core/tfrt/common/metrics.h

@@ -0,0 +1,37 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TFRT_COMMON_METRICS_H_
+#define TENSORFLOW_CORE_TFRT_COMMON_METRICS_H_
+
+#include <cstdint>
+#include <string>
+
+#include "xla/tsl/lib/monitoring/sampler.h"
+
+namespace tensorflow {
+namespace tfrt_metrics {
+
+tsl::monitoring::SamplerCell* GetTfrtGraphExecutorLatencySampler(
+    const std::string& model_name, int64_t model_version,
+    const std::string& graph_name);
+
+tsl::monitoring::SamplerCell* GetTfrtDeviceExecutionLatency(
+    const std::string& model_name, int64_t model_version);
+
+}  // namespace tfrt_metrics
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_COMMON_METRICS_H_

tensorflow/core/tfrt/common/pjrt_client_factory_options.h

@@ -0,0 +1,42 @@
+/* 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_CORE_TFRT_COMMON_PJRT_CLIENT_FACTORY_OPTIONS_H_
+#define TENSORFLOW_CORE_TFRT_COMMON_PJRT_CLIENT_FACTORY_OPTIONS_H_
+
+#include <optional>
+#include <set>
+#include <string>
+
+namespace xla {
+// PjrtClientFactoryOptions store arguments to create PJRT client.
+// Caller is responsible to set option value for corresponding PJRT client
+// factory.
+struct PjrtClientFactoryOptions {
+  struct GpuClientCreateOptions {
+    bool asynchronous = false;
+    int node_id = 0;
+    std::optional<std::set<int>> allowed_devices = std::nullopt;
+    std::optional<std::string> platform_name = std::nullopt;
+  };
+
+  struct CpuClientCreateOptions {
+    bool asynchronous = false;
+  };
+  GpuClientCreateOptions gpu_options;
+  CpuClientCreateOptions cpu_options;
+};
+}  // namespace xla
+
+#endif  // TENSORFLOW_CORE_TFRT_COMMON_PJRT_CLIENT_FACTORY_OPTIONS_H_

tensorflow/core/tfrt/common/pjrt_client_factory_registry.cc

@@ -0,0 +1,72 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_registry.h"
+
+#include <memory>
+#include <string>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/tsl/framework/device_type.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_options.h"
+#include "tsl/platform/errors.h"
+
+namespace xla {
+PjrtClientFactoryRegistry& PjrtClientFactoryRegistry::Get() {
+  static PjrtClientFactoryRegistry* kInstance = new PjrtClientFactoryRegistry();
+  return *kInstance;
+}
+
+tensorflow::InitOnStartupMarker
+PjrtClientFactoryRegistry::RegisterPjrtClientFactory(
+    const tsl::DeviceType& device_type,
+    const PjrtClientFactory& client_factory) {
+  tensorflow::mutex_lock l(mu_);
+  const std::string& device_type_str = device_type.type_string();
+  if (registry_.find(device_type_str) != registry_.end()) {
+    LOG(ERROR) << "Duplicate device type " << device_type_str;
+  }
+  registry_.emplace(device_type_str, client_factory);
+  return {};
+}
+
+absl::StatusOr<std::unique_ptr<PjRtClient>>
+PjrtClientFactoryRegistry::GetPjrtClient(
+    const tsl::DeviceType& device_type,
+    const PjrtClientFactoryOptions& options) {
+  tensorflow::tf_shared_lock l(mu_);
+  const std::string& device_type_str = device_type.type_string();
+  const auto client_factory_it = registry_.find(device_type_str);
+  if (client_factory_it == registry_.end()) {
+    std::string error_msg;
+    absl::StrAppend(&error_msg, " The PJRT client factory of `",
+                    device_type_str,
+                    "` is not registered, available client factory: [");
+    for (const auto& [device_name, ignored_func] : registry_) {
+      absl::StrAppend(&error_msg, "`", device_name, "`, ");
+    }
+    absl::StrAppend(&error_msg,
+                    "]. Did you forget to link with the appropriate "
+                    "`pjrt_*_client_registration` library?");
+    return absl::NotFoundError(error_msg);
+  }
+  return client_factory_it->second(options);
+}
+}  // namespace xla

tensorflow/core/tfrt/common/pjrt_client_factory_registry.h

@@ -0,0 +1,73 @@
+/* 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_CORE_TFRT_COMMON_PJRT_CLIENT_FACTORY_REGISTRY_H_
+#define TENSORFLOW_CORE_TFRT_COMMON_PJRT_CLIENT_FACTORY_REGISTRY_H_
+
+#include <functional>
+#include <memory>
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/statusor.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/tsl/framework/device_type.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_options.h"
+#include "tsl/platform/thread_annotations.h"
+
+namespace xla {
+
+using PjrtClientFactory =
+    std::function<absl::StatusOr<std::unique_ptr<PjRtClient>>(
+        const PjrtClientFactoryOptions&)>;
+
+// The Pjrt client factory registry holds all the registered client factories.
+class PjrtClientFactoryRegistry {
+ public:
+  explicit PjrtClientFactoryRegistry() = default;
+
+  // Registers PjrtClientFactory with DeviceType as key.
+  tensorflow::InitOnStartupMarker RegisterPjrtClientFactory(
+      const tsl::DeviceType& device_type,
+      const PjrtClientFactory& client_factory);
+
+  // Given the device type, finds related PjrtClientFactory function which takes
+  // factory option and returns PjrtClient if succeeds.
+  absl::StatusOr<std::unique_ptr<PjRtClient>> GetPjrtClient(
+      const tsl::DeviceType& device_type,
+      const PjrtClientFactoryOptions& options);
+
+  // Returns singleton instance of PjrtClientFactoryRegistry class.
+  static PjrtClientFactoryRegistry& Get();
+
+ private:
+  absl::flat_hash_map<std::string, PjrtClientFactory> registry_
+      TF_GUARDED_BY(mu_);
+
+  mutable ::tensorflow::mutex mu_;
+};
+
+// The `REGISTER_PJRT_CLIENT_FACTORY()` calls RegisterPjrtClientFactory on
+// program startup.
+#define REGISTER_PJRT_CLIENT_FACTORY(pjrt_client, device_type, client_factory) \
+  static ::tensorflow::InitOnStartupMarker const register_##pjrt_client =      \
+      ::tensorflow::InitOnStartupMarker{}                                      \
+      << ::xla::PjrtClientFactoryRegistry::Get().RegisterPjrtClientFactory(    \
+             tsl::DeviceType(device_type), client_factory)
+
+}  // namespace xla
+
+#endif  // TENSORFLOW_CORE_TFRT_COMMON_PJRT_CLIENT_FACTORY_REGISTRY_H_

tensorflow/core/tfrt/common/pjrt_cpu_client_registration.cc

@@ -0,0 +1,43 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_options.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_registry.h"
+#include "tsl/platform/statusor.h"
+
+namespace xla {
+
+absl::StatusOr<std::unique_ptr<xla::PjRtClient>> GetCpuClient(
+    const PjrtClientFactoryOptions& option) {
+  xla::CpuClientOptions cpu_options;
+  cpu_options.asynchronous = option.cpu_options.asynchronous;
+
+  TF_ASSIGN_OR_RETURN(std::unique_ptr<PjRtClient> client,
+                      xla::GetXlaPjrtCpuClient(cpu_options));
+
+  return std::move(client);
+}
+
+REGISTER_PJRT_CLIENT_FACTORY(cpu_client, tensorflow::DEVICE_CPU, GetCpuClient);
+
+}  // namespace xla

tensorflow/core/tfrt/common/pjrt_cpu_client_registration_test.cc

@@ -0,0 +1,35 @@
+/* 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 <gtest/gtest.h>
+#include "xla/tsl/framework/device_type.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_options.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_registry.h"
+#include "tsl/platform/statusor.h"
+
+namespace xla {
+namespace {
+
+TEST(PjrtCpuClientCreateTest, TestCpuCreateoption) {
+  PjrtClientFactoryOptions options = PjrtClientFactoryOptions();
+  options.cpu_options.asynchronous = true;
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto client, xla::PjrtClientFactoryRegistry::Get().GetPjrtClient(
+                       tsl::DeviceType(tensorflow::DEVICE_CPU), options));
+}
+}  // namespace
+
+}  // namespace xla

tensorflow/core/tfrt/common/pjrt_gpu_client_registration.cc

@@ -0,0 +1,47 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
+#include "xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_options.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_registry.h"
+#include "tsl/platform/statusor.h"
+namespace xla {
+
+absl::StatusOr<std::unique_ptr<xla::PjRtClient>> GetGpuClient(
+    const PjrtClientFactoryOptions& option) {
+  xla::GpuClientOptions gpu_client_options;
+  gpu_client_options.node_id = option.gpu_options.node_id;
+  gpu_client_options.num_nodes = 1;
+  gpu_client_options.allowed_devices = option.gpu_options.allowed_devices;
+  gpu_client_options.platform_name = option.gpu_options.platform_name;
+  TF_ASSIGN_OR_RETURN(std::unique_ptr<PjRtClient> client,
+                      xla::GetXlaPjrtGpuClient(gpu_client_options));
+  return std::move(client);
+}
+
+REGISTER_PJRT_CLIENT_FACTORY(gpu_client, tensorflow::DEVICE_GPU, GetGpuClient);
+
+REGISTER_PJRT_CLIENT_FACTORY(xla_gpu_client, tensorflow::DEVICE_XLA_GPU,
+                             GetGpuClient);
+
+}  // namespace xla

tensorflow/core/tfrt/common/pjrt_gpu_client_registration_test.cc

@@ -0,0 +1,35 @@
+/* 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 <gtest/gtest.h>
+#include "xla/tsl/framework/device_type.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_options.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_registry.h"
+#include "tsl/platform/statusor.h"
+
+namespace xla {
+namespace {
+
+TEST(PjrtGpuClientCreateTest, TestGpuCreateOption) {
+  PjrtClientFactoryOptions options = PjrtClientFactoryOptions();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto client, xla::PjrtClientFactoryRegistry::Get().GetPjrtClient(
+                       tsl::DeviceType(tensorflow::DEVICE_GPU), options));
+}
+
+}  // namespace
+
+}  // namespace xla

tensorflow/core/tfrt/common/pjrt_state.cc

@@ -0,0 +1,105 @@
+/* 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/core/tfrt/common/pjrt_state.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/synchronization/mutex.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/tf_pjrt_client.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_options.h"
+#include "tensorflow/core/tfrt/common/pjrt_client_factory_registry.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow {
+
+PjRtState* PjRtState::Create() { return new PjRtState(); }
+
+absl::StatusOr<xla::PjRtClient*> PjRtState::GetPjRtClient(
+    const DeviceType& device_type) {
+  absl::MutexLock lock(mu_);
+  if (auto it = clients_.find(device_type); it != clients_.end()) {
+    return it->second.get();
+  }
+  return errors::NotFound("PjRt client not found for device type ",
+                          device_type);
+}
+
+absl::StatusOr<xla::PjRtClient*> PjRtState::GetOrCreatePjRtClient(
+    const DeviceType& device_type) {
+  absl::MutexLock lock(mu_);
+  if (auto it = clients_.find(device_type); it != clients_.end()) {
+    return it->second.get();
+  }
+  std::unique_ptr<xla::PjRtClient> pjrt_client;
+  // TODO(b/260799193): use XlaPlatformInfo to pass device-specific options.
+  // This info should be set in the plugin init for next pluggable device.
+
+  // TODO(b/280111106): make PjrtClientFactoryOptions an input of
+  // GetOrCreatePjRtClient.
+  xla::PjrtClientFactoryOptions options = xla::PjrtClientFactoryOptions();
+  TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::PjRtClient> client,
+                      xla::PjrtClientFactoryRegistry::Get().GetPjrtClient(
+                          device_type, options));
+  pjrt_client = xla::TfPjRtClient::CreateTfPjRtClient(std::move(client));
+
+  clients_[device_type] = std::move(pjrt_client);
+  return clients_[device_type].get();
+}
+
+absl::Status PjRtState::SetPjRtClient(const DeviceType& device_type,
+                                      std::unique_ptr<xla::PjRtClient> client) {
+  absl::MutexLock lock(mu_);
+  if (auto it = clients_.find(device_type); it != clients_.end()) {
+    unused_.push_back(std::move(it->second));
+  }
+  clients_[device_type] = std::move(client);
+  return absl::OkStatus();
+}
+
+absl::Status PjRtState::MovePjRtClientToUnused(const DeviceType& device_type) {
+  absl::MutexLock lock(mu_);
+  if (auto it = clients_.find(device_type); it != clients_.end()) {
+    unused_.push_back(std::move(it->second));
+    clients_.erase(it);
+    return absl::OkStatus();
+  }
+  return errors::NotFound("PjRt client not found for device type ",
+                          device_type);
+}
+
+absl::Status PjRtState::SetPjRtGpuClientCreationInfo(
+    std::unique_ptr<PjRtGpuClientCreationInfo> info) {
+  absl::MutexLock lock(mu_);
+  pjrt_gpu_client_creation_info_ = std::move(info);
+  return absl::OkStatus();
+}
+
+PjRtGpuClientCreationInfo* PjRtState::GetPjRtGpuClientCreationInfo() {
+  absl::MutexLock lock(mu_);
+  return pjrt_gpu_client_creation_info_.get();
+}
+
+std::string PjRtState::DebugString() const { return "PjRtState"; }
+
+}  // namespace tensorflow

tensorflow/core/tfrt/common/pjrt_state.h

@@ -0,0 +1,91 @@
+/* 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_CORE_TFRT_COMMON_PJRT_STATE_H_
+#define TENSORFLOW_CORE_TFRT_COMMON_PJRT_STATE_H_
+
+#include <map>
+#include <memory>
+#include <set>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/synchronization/mutex.h"
+#include "xla/client/local_client.h"
+#include "xla/pjrt/host_memory_allocator.h"
+#include "xla/pjrt/local_device_state.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/stream_executor/integrations/tf_allocator_adapter.h"
+#include "xla/tsl/framework/allocator.h"
+#include "tensorflow/core/framework/resource_base.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+const char kPjRtStateResourceName[] = "pjrt_state";
+using PjRtClientsMap = std::map<DeviceType, std::unique_ptr<xla::PjRtClient>>;
+
+// Information needed to create a PjRt GPU Client which is used when creating
+// a client after after information about remote devices is available.
+struct PjRtGpuClientCreationInfo {
+  std::set<int> allowed_devices;
+  std::unique_ptr<se::MultiDeviceAdapter> allocator;
+  std::unique_ptr<xla::HostMemoryAllocator> host_memory_allocator;
+  std::map<int, std::unique_ptr<xla::LocalDeviceState>> local_device_states;
+  xla::LocalClient* local_client;
+};
+
+// The class for the state related to PjRt. It contains a map from `DeviceType`
+// to `PjRtClient`. It will be stored in the global `ResourceManager`.
+class PjRtState : public ResourceBase {
+ public:
+  static PjRtState* Create();
+  absl::StatusOr<xla::PjRtClient*> GetPjRtClient(const DeviceType& device_type);
+  absl::StatusOr<xla::PjRtClient*> GetOrCreatePjRtClient(
+      const DeviceType& device_type);
+  absl::Status SetPjRtClient(const DeviceType& device_type,
+                             std::unique_ptr<xla::PjRtClient> client);
+  // Moves PJRT client to `unused_`. The PJRT client moved to `unused_` will not
+  // be returned by `GetPjRtClient`.
+  absl::Status MovePjRtClientToUnused(const DeviceType& device_type);
+  std::string DebugString() const override;
+
+  // Saves information needed to create a PJRT client (to enable creating a
+  // client with remote devices).
+  absl::Status SetPjRtGpuClientCreationInfo(
+      std::unique_ptr<PjRtGpuClientCreationInfo> info);
+
+  // Retrieves information needed to create a PJRT client (for creating a
+  // client with remote devices).
+  PjRtGpuClientCreationInfo* GetPjRtGpuClientCreationInfo();
+
+ private:
+  explicit PjRtState() = default;
+  absl::Mutex mu_;
+  PjRtClientsMap clients_ ABSL_GUARDED_BY(mu_);
+  // Store the PJRT clients that are no longer used to guarantee that PJRT
+  // clients outlive PJRT buffers.
+  std::vector<std::unique_ptr<xla::PjRtClient>> unused_ ABSL_GUARDED_BY(mu_);
+
+  std::unique_ptr<PjRtGpuClientCreationInfo> pjrt_gpu_client_creation_info_
+      ABSL_GUARDED_BY(mu_);
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_COMMON_PJRT_STATE_H_

tensorflow/core/tfrt/common/pjrt_state_test.cc

@@ -0,0 +1,132 @@
+/* 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/core/tfrt/common/pjrt_state.h"
+
+#include <memory>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
+
+namespace {
+
+using tensorflow::PjRtState;
+using ::testing::HasSubstr;
+using ::tsl::testing::StatusIs;
+
+class PjRtStateTestFixture : public testing::Test {
+ protected:
+  PjRtStateTestFixture() { pjrt_state_ = PjRtState::Create(); }
+  ~PjRtStateTestFixture() override {
+    tensorflow::core::ScopedUnref pjrt_state_ref(pjrt_state_);
+  }
+  PjRtState* pjrt_state_;
+};
+
+TEST_F(PjRtStateTestFixture, SetAndGetPjRtClient) {
+  xla::CpuClientOptions options;
+  options.asynchronous = true;
+  options.cpu_device_count = 1;
+  TF_ASSERT_OK(pjrt_state_->SetPjRtClient(
+      tensorflow::DEVICE_CPU, xla::GetXlaPjrtCpuClient(options).value()));
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client,
+                          pjrt_state_->GetPjRtClient(tensorflow::DEVICE_CPU));
+  EXPECT_THAT(pjrt_client, testing::NotNull());
+}
+
+TEST_F(PjRtStateTestFixture, AddAlreadyExistsPjRtClient) {
+  xla::CpuClientOptions options;
+  options.asynchronous = true;
+  options.cpu_device_count = 1;
+
+  TF_ASSERT_OK(pjrt_state_->SetPjRtClient(
+      tensorflow::DEVICE_CPU, xla::GetXlaPjrtCpuClient(options).value()));
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client_1,
+                          pjrt_state_->GetPjRtClient(tensorflow::DEVICE_CPU));
+
+  TF_ASSERT_OK(pjrt_state_->SetPjRtClient(
+      tensorflow::DEVICE_CPU, xla::GetXlaPjrtCpuClient(options).value()));
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client_2,
+                          pjrt_state_->GetPjRtClient(tensorflow::DEVICE_CPU));
+
+  EXPECT_NE(pjrt_client_1, pjrt_client_2);
+}
+
+TEST_F(PjRtStateTestFixture, GetNotExistPjRtClient) {
+  EXPECT_THAT(pjrt_state_->GetPjRtClient(tensorflow::DEVICE_CPU),
+              absl_testing::StatusIs(
+                  tensorflow::error::NOT_FOUND,
+                  HasSubstr("PjRt client not found for device type")));
+}
+
+TEST_F(PjRtStateTestFixture, DeletePjRtClient) {
+  xla::CpuClientOptions options;
+  options.asynchronous = true;
+  options.cpu_device_count = 1;
+
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetXlaPjrtCpuClient(options));
+  xla::PjRtClient* pjrt_client_ptr = pjrt_client.get();
+  TF_ASSERT_OK(pjrt_state_->SetPjRtClient(tensorflow::DEVICE_CPU,
+                                          std::move(pjrt_client)));
+
+  TF_ASSERT_OK(pjrt_state_->MovePjRtClientToUnused(tensorflow::DEVICE_CPU));
+
+  EXPECT_THAT(pjrt_state_->GetPjRtClient(tensorflow::DEVICE_CPU),
+              absl_testing::StatusIs(
+                  tensorflow::error::NOT_FOUND,
+                  HasSubstr("PjRt client not found for device type")));
+  // Verifies that the PJRT client is still alive.
+  EXPECT_EQ(pjrt_client_ptr->platform_name(), "cpu");
+}
+
+TEST_F(PjRtStateTestFixture, DeleteNotExistPjRtClient) {
+  EXPECT_THAT(pjrt_state_->MovePjRtClientToUnused(tensorflow::DEVICE_CPU),
+              absl_testing::StatusIs(
+                  tensorflow::error::NOT_FOUND,
+                  HasSubstr("PjRt client not found for device type")));
+}
+
+TEST_F(PjRtStateTestFixture, GetOrCreatePjRtClientExist) {
+  xla::CpuClientOptions options;
+  options.asynchronous = true;
+  options.cpu_device_count = 1;
+
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetXlaPjrtCpuClient(options));
+  auto pjrt_client_ptr = pjrt_client.get();
+  TF_ASSERT_OK(pjrt_state_->SetPjRtClient(tensorflow::DEVICE_CPU,
+                                          std::move(pjrt_client)));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto pjrt_client_get,
+      pjrt_state_->GetOrCreatePjRtClient(tensorflow::DEVICE_CPU));
+  EXPECT_THAT(pjrt_client_get, pjrt_client_ptr);
+}
+
+TEST_F(PjRtStateTestFixture, GetOrCreatePjRtClientNotExist) {
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, pjrt_state_->GetOrCreatePjRtClient(
+                                                tensorflow::DEVICE_CPU));
+  EXPECT_THAT(pjrt_client, testing::NotNull());
+}
+
+}  // namespace

tensorflow/core/tfrt/common/pjrt_util.cc

@@ -0,0 +1,95 @@
+/* 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/core/tfrt/common/pjrt_util.h"
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/common/global_state.h"
+#include "tensorflow/core/tfrt/common/pjrt_state.h"
+#include "tsl/platform/errors.h"
+
+namespace tensorflow {
+
+absl::Status SetPjRtClientInTFGlobalResourceManager(
+    const DeviceType& device_type, std::unique_ptr<xla::PjRtClient> client) {
+  ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr();
+  PjRtState* pjrt_state;
+  TF_RETURN_IF_ERROR(rmgr->LookupOrCreate<PjRtState>(
+      rmgr->default_container(), kPjRtStateResourceName, &pjrt_state,
+      [&](PjRtState** ret) {
+        *ret = PjRtState::Create();
+        return absl::OkStatus();
+      }));
+  core::ScopedUnref pjrt_state_ref(pjrt_state);
+  if (client == nullptr) {
+    return absl::InvalidArgumentError("PJRT client is nullptr.");
+  }
+  TF_RETURN_IF_ERROR(pjrt_state->SetPjRtClient(device_type, std::move(client)));
+  return absl::OkStatus();
+}
+
+absl::StatusOr<xla::PjRtClient*> GetPjRtClient(const DeviceType& device_type) {
+  ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr();
+  PjRtState* pjrt_state;
+  TF_RETURN_IF_ERROR(rmgr->LookupOrCreate<PjRtState>(
+      rmgr->default_container(), kPjRtStateResourceName, &pjrt_state,
+      [&](PjRtState** ret) {
+        *ret = PjRtState::Create();
+        return absl::OkStatus();
+      }));
+  core::ScopedUnref pjrt_state_ref(pjrt_state);
+  return pjrt_state->GetPjRtClient(device_type);
+}
+
+absl::Status SetPjRtGpuClientCreationInfoInTFGlobalResourceManager(
+    std::unique_ptr<PjRtGpuClientCreationInfo> info) {
+  ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr();
+  PjRtState* pjrt_state;
+  TF_RETURN_IF_ERROR(rmgr->LookupOrCreate<PjRtState>(
+      rmgr->default_container(), kPjRtStateResourceName, &pjrt_state,
+      [&](PjRtState** ret) {
+        *ret = PjRtState::Create();
+        return absl::OkStatus();
+      }));
+  core::ScopedUnref pjrt_state_ref(pjrt_state);
+  if (info == nullptr) {
+    return absl::InvalidArgumentError("PJRT client creation info is nullptr.");
+  }
+  TF_RETURN_IF_ERROR(pjrt_state->SetPjRtGpuClientCreationInfo(std::move(info)));
+  return absl::OkStatus();
+}
+
+absl::StatusOr<PjRtGpuClientCreationInfo*> GetPjRtGpuClientCreationInfo() {
+  ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr();
+  PjRtState* pjrt_state;
+  TF_RETURN_IF_ERROR(rmgr->LookupOrCreate<PjRtState>(
+      rmgr->default_container(), kPjRtStateResourceName, &pjrt_state,
+      [&](PjRtState** ret) {
+        *ret = PjRtState::Create();
+        return absl::OkStatus();
+      }));
+  core::ScopedUnref pjrt_state_ref(pjrt_state);
+  return pjrt_state->GetPjRtGpuClientCreationInfo();
+}
+}  // namespace tensorflow

tensorflow/core/tfrt/common/pjrt_util.h

@@ -0,0 +1,46 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TFRT_COMMON_PJRT_UTIL_H_
+#define TENSORFLOW_CORE_TFRT_COMMON_PJRT_UTIL_H_
+
+#include <memory>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/common/pjrt_state.h"
+
+namespace tensorflow {
+
+// Sets PJRT client for device_type in TFGlobalResourceManager. If a PJRT client
+// for this device_type already exists, the existing PJRT client will not be
+// destroyed, and will be kept alive in an "unused client" vector. PJRT API
+// semantics require the PJRT client to outlive PJRT buffers.
+absl::Status SetPjRtClientInTFGlobalResourceManager(
+    const DeviceType& device_type, std::unique_ptr<xla::PjRtClient> client);
+
+// Gets (the most recent) PJRT client for device_type from
+// TFGlobalResourceManager.
+absl::StatusOr<xla::PjRtClient*> GetPjRtClient(const DeviceType& device_type);
+
+absl::Status SetPjRtGpuClientCreationInfoInTFGlobalResourceManager(
+    std::unique_ptr<PjRtGpuClientCreationInfo> info);
+absl::StatusOr<PjRtGpuClientCreationInfo*> GetPjRtGpuClientCreationInfo();
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_COMMON_PJRT_UTIL_H_

tensorflow/core/tfrt/common/pjrt_util_test.cc

@@ -0,0 +1,68 @@
+/* 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/core/tfrt/common/pjrt_util.h"
+
+#include <memory>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/tfrt/common/pjrt_state.h"
+#include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
+#include "tsl/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::HasSubstr;
+using ::tsl::testing::StatusIs;
+
+TEST(PjRtUtilTest, SetGetAndDeletePjRtClient) {
+  xla::CpuClientOptions options;
+  options.asynchronous = true;
+  options.cpu_device_count = 1;
+  TF_ASSERT_OK(SetPjRtClientInTFGlobalResourceManager(
+      DEVICE_CPU, xla::GetXlaPjrtCpuClient(options).value()));
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, GetPjRtClient(DEVICE_CPU));
+  EXPECT_THAT(pjrt_client, ::testing::NotNull());
+}
+
+TEST(PjRtStateResourceManagerTest, SetNullPjRtClient) {
+  EXPECT_THAT(SetPjRtClientInTFGlobalResourceManager(DEVICE_CPU, nullptr),
+              absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                                     HasSubstr("PJRT client is nullptr")));
+}
+
+TEST(PjRtGpuClientCreationInfoTest, SetAndGet) {
+  auto info = std::make_unique<PjRtGpuClientCreationInfo>();
+  info->allowed_devices.insert(123);
+  TF_ASSERT_OK(
+      SetPjRtGpuClientCreationInfoInTFGlobalResourceManager(std::move(info)));
+
+  TF_ASSERT_OK_AND_ASSIGN(PjRtGpuClientCreationInfo * retrieved_info,
+                          GetPjRtGpuClientCreationInfo());
+
+  EXPECT_THAT(retrieved_info->allowed_devices, ElementsAre(123));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/tfrt/fallback/BUILD

@@ -0,0 +1,195 @@
+load(
+    "//tensorflow:tensorflow.bzl",
+    "if_mobile",
+    "if_not_mobile",
+    "tf_cc_test",
+)
+load("//tensorflow:tensorflow.default.bzl", "tf_cuda_cc_test")
+load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
+load(
+    "//tensorflow/core/platform:build_config_root.bzl",
+    "if_static",
+    "tf_cuda_tests_tags",
+)
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    packages = [
+        # Authorized users go here.
+        "//tensorflow/compiler/mlir/...",
+        "//tensorflow/core/...",
+        # copybara:uncomment "//third_party/neptune/...",
+        # copybara:uncomment "//learning/brain/experimental/tfrt/...",
+        # copybara:uncomment "//learning/brain/tfrt/...",
+        # copybara:uncomment "//learning/brain/mobile/lite/delegates/tfmrt/...",
+        # TODO(b/476497239): Remove the old JSV codebase once the migration is done.
+        # copybara:uncomment "//learning/infra/mira/experimental/orbax_model/...",
+    ],
+)
+
+cc_library(
+    name = "fallback_state",
+    srcs = ["fallback_state.cc"],
+    hdrs = ["fallback_state.h"],
+    deps = [
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:portable_gif_internal",
+        "//tensorflow/core:session_options",
+        "//tensorflow/core/common_runtime:core_cpu_internal",
+        "//tensorflow/core/common_runtime:device_set",
+        "//tensorflow/core/framework:device_attributes_proto_cc",
+        "//tensorflow/core/framework:function_proto_cc",
+        "//tensorflow/core/framework:graph_proto_cc",
+        "//tensorflow/core/tpu:virtual_device",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@tsl//tsl/platform:refcount",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)
+
+tf_cc_test(
+    name = "fallback_state_test",
+    srcs = ["fallback_state_test.cc"],
+    deps = [
+        ":fallback_state",
+        "//tensorflow/cc:cc_ops",
+        "//tensorflow/cc:const_op",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/core:all_kernels",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/framework:function_proto_cc",
+        "//tensorflow/core/platform:status_matchers",
+        "//tensorflow/core/protobuf:error_codes_proto_impl_cc",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_googletest//:gtest",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)
+
+cc_library(
+    name = "op_kernel_runner",
+    srcs = ["op_kernel_runner.cc"],
+    hdrs = ["op_kernel_runner.h"],
+    visibility = [
+        # copybara:uncomment "//tensorflow/core/runtime_fallback:internal",
+        "//visibility:public",
+    ],
+    deps = [
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+    ] + if_mobile([
+        "//tensorflow/core:portable_tensorflow_lib_lite",
+    ]) + if_not_mobile([
+        "//tensorflow/core:framework",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core/framework:node_def_proto_cc",
+        "//tensorflow/core/framework:op_def_proto_cc",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+    ]),
+)
+
+cc_library(
+    name = "op_kernel_runner_cache",
+    srcs = ["op_kernel_runner_cache.cc"],
+    hdrs = ["op_kernel_runner_cache.h"],
+    deps = [
+        ":op_kernel_runner",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@tf_runtime//:hostcontext",
+    ],
+)
+
+cc_library(
+    name = "cost_recorder",
+    srcs = ["cost_recorder.cc"],
+    hdrs = ["cost_recorder.h"],
+    deps = [
+        ":op_cost_map_proto_cc",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/util:env_var",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+cc_library(
+    name = "device_with_custom_allocator",
+    hdrs = ["device_with_custom_allocator.h"],
+    deps = [
+        "//tensorflow/core:framework",
+        "@xla//xla/tsl/framework:allocator",
+    ],
+)
+
+tf_cc_test(
+    name = "cost_recorder_test",
+    srcs = ["cost_recorder_test.cc"],
+    deps = [
+        ":cost_recorder",
+        ":op_cost_map_proto_cc",
+        "//tensorflow/core:lib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+tf_cuda_cc_test(
+    name = "op_kernel_runner_test",
+    size = "small",
+    srcs = ["op_kernel_runner_test.cc"],
+    tags = tf_cuda_tests_tags(),
+    deps = [
+        ":fallback_state",
+        ":op_kernel_runner",
+        ":op_kernel_runner_cache",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:session_options",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ] + if_static(
+        [
+            "//tensorflow/core/common_runtime:function",
+        ],
+        [],
+    ),
+)
+
+tf_proto_library(
+    name = "op_cost_map_proto",
+    srcs = ["op_cost_map.proto"],
+)

tensorflow/core/tfrt/fallback/cost_recorder.cc

@@ -0,0 +1,82 @@
+/* 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/core/tfrt/fallback/cost_recorder.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/fallback/op_cost_map.pb.h"
+#include "tensorflow/core/util/env_var.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+void CostRecorder::RecordCost(int64_t op_key, uint64_t execution_time) {
+  mutex_lock l(op_cost_map_mutex_);
+  op_cost_map_[op_key].first += execution_time;
+  op_cost_map_[op_key].second += 1;
+}
+
+uint64_t CostRecorder::GetCost(int64_t op_key) const {
+  tf_shared_lock l(op_cost_map_mutex_);
+
+  const auto iter = op_cost_map_.find(op_key);
+  if (iter == op_cost_map_.end()) return std::numeric_limits<uint32_t>::max();
+
+  const auto total_cost = iter->second.first;
+  const auto num_ops = iter->second.second;
+
+  auto r =
+      std::max(static_cast<uint64_t>(1),
+               static_cast<uint64_t>(total_cost / num_ops));
+
+  VLOG(2) << "Get cost for op_key=" << op_key << ", cost=" << r;
+
+  return r;
+}
+
+absl::Status CostRecorder::WriteToFile() const {
+  OpCostMapProto op_cost_map_proto;
+  {
+    tf_shared_lock l(op_cost_map_mutex_);
+    for (const auto& [op_key, op_cost] : op_cost_map_) {
+      const uint64_t avg_op_cost = op_cost.first / op_cost.second;
+      (*op_cost_map_proto.mutable_op_cost_map())[op_key] = avg_op_cost;
+    }
+  }
+
+  std::string measured_cost_path;
+  TF_RETURN_IF_ERROR(ReadStringFromEnvVar(MesuredCostPathEnvVarName(), "",
+                                          &measured_cost_path));
+  return tensorflow::WriteTextProto(tensorflow::Env::Default(),
+                                    measured_cost_path, op_cost_map_proto);
+}
+
+size_t CostRecorder::size() const {
+  tf_shared_lock l(op_cost_map_mutex_);
+  return op_cost_map_.size();
+}
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/fallback/cost_recorder.h

@@ -0,0 +1,69 @@
+/* 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.
+==============================================================================*/
+
+// This file defines a recorder for op cost measurement.
+
+#ifndef TENSORFLOW_CORE_TFRT_FALLBACK_COST_RECORDER_H_
+#define TENSORFLOW_CORE_TFRT_FALLBACK_COST_RECORDER_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/thread_annotations.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+// Thread-safe.
+// Maintains the execution durations by `op_key`. Note that `op_key` is only
+// unique within a model.
+class CostRecorder {
+ public:
+  // Records an execution duration for the op keyed by `op_key`.
+  void RecordCost(int64_t op_key, uint64_t execution_time);
+
+  // Returns the normalized average execution duration of the op keyed by
+  // `op_key`. If there is no record for `op_key`, returns the uint32_t::max to
+  // avoid stream merging. Note that we don't use uint64_t::max because
+  // otherwise adding op costs would cause overflow.
+  uint64_t GetCost(int64_t op_key) const;
+
+  // Writes the op cost map (in format of `OpCostMapProto`) to a file specified
+  // by the env var name `MesuredCostPathEnvVarName()`.
+  // TODO(b/263837451): Fix the op_key unstableness during serialization.
+  absl::Status WriteToFile() const;
+
+  size_t size() const;
+
+  static const char* MesuredCostPathEnvVarName() {
+    return "TF_TFRT_MEASURED_COST_PATH";
+  }
+
+ private:
+  mutable tensorflow::mutex op_cost_map_mutex_;
+  // Map op key to {sum of op execution duration, #occurences of the op}.
+  absl::flat_hash_map<int64_t, std::pair<uint64_t, uint64_t>> op_cost_map_
+      TF_GUARDED_BY(op_cost_map_mutex_);
+};
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_FALLBACK_COST_RECORDER_H_

tensorflow/core/tfrt/fallback/cost_recorder_test.cc

@@ -0,0 +1,104 @@
+/* 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/core/tfrt/fallback/cost_recorder.h"
+
+#include <cstdint>
+#include <limits>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/tfrt/fallback/op_cost_map.pb.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+namespace {
+
+constexpr int64_t kTestOpKey = 1;
+constexpr uint64_t kTestCost = 1234;
+constexpr uint64_t kTestAvgCost = 1851;
+
+TEST(CostRecorderTest, RecordCostTest) {
+  CostRecorder recorder;
+
+  recorder.RecordCost(kTestOpKey, kTestCost);
+  recorder.RecordCost(kTestOpKey, kTestCost);
+
+  EXPECT_EQ(recorder.size(), 1);
+}
+
+TEST(CostRecorderTest, GetCostTest) {
+  CostRecorder recorder;
+
+  recorder.RecordCost(kTestOpKey, kTestCost);
+  recorder.RecordCost(kTestOpKey, 2 * kTestCost);
+
+  EXPECT_EQ(recorder.size(), 1);
+  EXPECT_EQ(recorder.GetCost(kTestOpKey), kTestAvgCost);
+}
+
+TEST(CostRecorderTest, GetCostDefaultValueTest) {
+  CostRecorder recorder;
+  ASSERT_EQ(recorder.size(), 0);
+
+  EXPECT_EQ(recorder.GetCost(kTestOpKey),
+            std::numeric_limits<uint32_t>::max());
+}
+
+TEST(CostRecorderTest, WriteToFileTest) {
+  CostRecorder recorder;
+  ASSERT_EQ(recorder.size(), 0);
+
+  std::string measured_cost_path;
+  tensorflow::Env::Default()->LocalTempFilename(&measured_cost_path);
+  ASSERT_EQ(setenv("TF_TFRT_MEASURED_COST_PATH", measured_cost_path.c_str(), 1),
+            0);
+  TF_CHECK_OK(recorder.WriteToFile());
+
+  OpCostMapProto op_cost_map_proto;
+  TF_CHECK_OK(tensorflow::ReadTextProto(
+      tensorflow::Env::Default(), measured_cost_path, &op_cost_map_proto));
+
+  EXPECT_EQ(op_cost_map_proto.op_cost_map_size(), 0);
+}
+
+TEST(CostRecorderTest, ProtoRecordsTest) {
+  CostRecorder recorder;
+
+  // Records the cost of op.
+  recorder.RecordCost(kTestOpKey, kTestCost);
+  recorder.RecordCost(kTestOpKey, 2 * kTestCost);
+  ASSERT_EQ(recorder.size(), 1);
+
+  // Writes op's cost to the disk.
+  std::string measured_cost_path;
+  tensorflow::Env::Default()->LocalTempFilename(&measured_cost_path);
+  ASSERT_EQ(setenv(CostRecorder::MesuredCostPathEnvVarName(),
+                   measured_cost_path.c_str(), 1),
+            0);
+  TF_CHECK_OK(recorder.WriteToFile());
+
+  // Reads op's cost from the disk.
+  OpCostMapProto op_cost_map_proto;
+  TF_CHECK_OK(tensorflow::ReadTextProto(
+      tensorflow::Env::Default(), measured_cost_path, &op_cost_map_proto));
+
+  EXPECT_EQ(op_cost_map_proto.op_cost_map().find(kTestOpKey)->second,
+            kTestAvgCost);
+}
+
+}  // namespace
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/fallback/device_with_custom_allocator.h

@@ -0,0 +1,101 @@
+/* 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_CORE_TFRT_FALLBACK_DEVICE_WITH_CUSTOM_ALLOCATOR_H_
+#define TENSORFLOW_CORE_TFRT_FALLBACK_DEVICE_WITH_CUSTOM_ALLOCATOR_H_
+
+#include <utility>
+
+#include "xla/tsl/framework/allocator.h"
+#include "tensorflow/core/framework/device.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+class DeviceWithCustomAllocator : public tensorflow::Device {
+ public:
+  DeviceWithCustomAllocator(tensorflow::Device* device,
+                            tensorflow::Allocator* allocator)
+      : Device(device->env(), device->attributes()),
+        device_(device),
+        allocator_(allocator) {
+    DCHECK(device_);
+    DCHECK(allocator_);
+  }
+
+  Allocator* GetAllocator(AllocatorAttributes attr) override {
+    return allocator_;
+  }
+
+  const DeviceBase* UnderlyingDevice() const override {
+    return device_->UnderlyingDevice();
+  }
+  DeviceBase* UnderlyingDevice() override {
+    return device_->UnderlyingDevice();
+  }
+
+  const CpuWorkerThreads* tensorflow_cpu_worker_threads() const override {
+    return device_->tensorflow_cpu_worker_threads();
+  }
+
+  Allocator* GetScopedAllocator(AllocatorAttributes attr,
+                                int64_t step_id) override {
+    return device_->GetScopedAllocator(attr, step_id);
+  }
+
+  ScopedAllocatorMgr* GetScopedAllocatorMgr() const override {
+    return device_->GetScopedAllocatorMgr();
+  }
+
+  const Eigen::ThreadPoolDevice* eigen_cpu_device() override {
+    return device_->eigen_cpu_device();
+  }
+
+  thread::ThreadPool* tensorflow_device_thread_pool() override {
+    return device_->tensorflow_device_thread_pool();
+  }
+
+  bool has_eigen_cpu_device() const override {
+    return device_->has_eigen_cpu_device();
+  }
+
+  absl::Status MakeTensorFromProto(const TensorProto& tensor_proto,
+                                   const AllocatorAttributes alloc_attrs,
+                                   Tensor* tensor) override {
+    return device_->MakeTensorFromProto(tensor_proto, alloc_attrs, tensor);
+  }
+
+  void CopyTensorInSameDevice(const Tensor* input_tensor, Tensor* output_tensor,
+                              const DeviceContext* device_context,
+                              StatusCallback done) override {
+    device_->CopyTensorInSameDevice(input_tensor, output_tensor, device_context,
+                                    std::move(done));
+  }
+
+  absl::Status Sync() override { return device_->Sync(); }
+
+  // Returns the resource manager associated w/ this device.
+  ResourceMgr* resource_manager() override {
+    return device_->resource_manager();
+  }
+
+ private:
+  tensorflow::Device* device_ = nullptr;
+  tensorflow::Allocator* allocator_ = nullptr;
+};
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_FALLBACK_DEVICE_WITH_CUSTOM_ALLOCATOR_H_

tensorflow/core/tfrt/fallback/fallback_state.cc

@@ -0,0 +1,186 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "absl/base/nullability.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/strings/strip.h"
+#include "xla/tsl/platform/errors.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/graph_execution_state.h"
+#include "tensorflow/core/common_runtime/rendezvous_mgr.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/framework/device_factory.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/rendezvous.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/graph/types.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/public/version.h"
+#include "tensorflow/core/tpu/virtual_device.h"
+#include "tsl/platform/refcount.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+namespace {
+
+std::string DeviceName(absl::string_view name_prefix,
+                       absl::string_view device_type, int32_t task_id,
+                       size_t device_id) {
+  return strings::StrCat(absl::StripSuffix(name_prefix, "0"), task_id,
+                         "/device:", device_type, ":", device_id);
+}
+
+DeviceAttributes BuildDeviceAttributes(absl::string_view name_prefix,
+                                       const char *device_type, int32_t task_id,
+                                       size_t device_id) {
+  const DeviceAttributes attrs = Device::BuildDeviceAttributes(
+      DeviceName(name_prefix, device_type, task_id, device_id),
+      DeviceType(device_type), Bytes(16ULL << 30), DeviceLocality(),
+      absl::StrCat("device: ", device_type, " device"));
+  return attrs;
+}
+
+}  // namespace
+
+absl::StatusOr<std::unique_ptr<FallbackState>> FallbackState::Create(
+    const SessionOptions &session_options,
+    const tensorflow::FunctionDefLibrary &fdef_lib) {
+  // Create devices.
+  std::vector<std::unique_ptr<Device>> devices;
+  TF_RETURN_IF_ERROR(DeviceFactory::AddDevices(
+      session_options, "/job:localhost/replica:0/task:0", &devices));
+
+  return std::make_unique<FallbackState>(session_options, std::move(devices),
+                                         fdef_lib);
+}
+
+absl::StatusOr<std::unique_ptr<FallbackState>>
+FallbackState::CreateWithCpuDevice(
+    const SessionOptions &session_options,
+    const tensorflow::FunctionDefLibrary &fdef_lib) {
+  // Create devices.
+  std::vector<std::unique_ptr<Device>> devices;
+  TF_RETURN_IF_ERROR(DeviceFactory::AddCpuDevices(
+      session_options, "/job:localhost/replica:0/task:0", &devices));
+
+  return std::make_unique<FallbackState>(session_options, std::move(devices),
+                                         fdef_lib);
+}
+
+absl::StatusOr<std::unique_ptr<FallbackState>>
+FallbackState::CreateWithMockGpuDevice(
+    const SessionOptions &session_options,
+    const tensorflow::FunctionDefLibrary &fdef_lib) {
+  // Create devices.
+  std::vector<std::unique_ptr<Device>> devices;
+  TF_RETURN_IF_ERROR(DeviceFactory::AddCpuDevices(
+      session_options, "/job:localhost/replica:0/task:0", &devices));
+
+  auto device_attrs =
+      BuildDeviceAttributes("/job:localhost/replica:0/task:0", "GPU", 0, 0);
+  devices.push_back(
+      std::make_unique<VirtualDevice>(session_options.env, device_attrs));
+
+  return std::make_unique<FallbackState>(session_options, std::move(devices),
+                                         fdef_lib);
+}
+
+absl::StatusOr<std::unique_ptr<FallbackState>>
+FallbackState::CreateWithDeviceMgr(
+    const SessionOptions &session_options,
+    const tensorflow::FunctionDefLibrary &fdef_lib,
+    DynamicDeviceMgr *absl_nonnull device_mgr) {
+  return std::make_unique<FallbackState>(session_options, device_mgr, fdef_lib);
+}
+
+FallbackState::FallbackState(const SessionOptions &session_options,
+                             std::variant<std::vector<std::unique_ptr<Device>>,
+                                          DynamicDeviceMgr *absl_nonnull>
+                                 device_mgr,
+                             const tensorflow::FunctionDefLibrary &fdef_lib)
+    : session_options_(session_options),
+      device_manager_(
+          std::holds_alternative<std::vector<std::unique_ptr<Device>>>(
+              device_mgr)
+              ? std::move(
+                    std::get<std::vector<std::unique_ptr<Device>>>(device_mgr))
+              : std::vector<std::unique_ptr<Device>>()),
+      device_manager_ptr_(
+          std::holds_alternative<DynamicDeviceMgr *absl_nonnull>(device_mgr)
+              ? std::get<DynamicDeviceMgr *absl_nonnull>(device_mgr)
+              : &device_manager_),
+      func_lib_def_(OpRegistry::Global(), fdef_lib),
+      pflr_(device_manager_ptr_, session_options.env, &session_options.config,
+            TF_GRAPH_DEF_VERSION, &func_lib_def_,
+            session_options.config.graph_options().optimizer_options(),
+            /*thread_pool=*/nullptr, /*parent=*/nullptr,
+            /*session_metadata=*/nullptr,
+            Rendezvous::Factory{[](const int64_t, const DeviceMgr *device_mgr,
+                                   tsl::core::RefCountPtr<Rendezvous> *r) {
+              *r = tsl::core::RefCountPtr<Rendezvous>(
+                  new IntraProcessRendezvous(device_mgr));
+              return absl::OkStatus();
+            }}) {
+  for (auto *d : device_manager_ptr_->ListDevices()) {
+    device_set_.AddDevice(d);
+  }
+
+  // client_device is the device for feed and fetch tensors.
+  device_set_.set_client_device(device_manager().HostCPU());
+}
+
+absl::StatusOr<std::unique_ptr<GraphExecutionState>>
+FallbackState::CreateGraphExecutionState(GraphDef graph_def, bool run_placer,
+                                         bool enable_tf2xla_mlir_bridge) const {
+  // Create GraphExecutionState which contains the preprocessed graph including
+  // device information. The following code is adapted from
+  // http://cs?q=tensorflow/core/common_runtime/direct_session.cc:427%20at_cl:352783230
+
+  GraphExecutionStateOptions options;
+  options.device_set = &device_set_;
+  options.session_options = &session_options_;
+  options.session_handle = "tfrt_fallback_handle";
+  options.run_placer = run_placer;
+  options.enable_tf2xla_mlir_bridge = enable_tf2xla_mlir_bridge;
+
+  std::unique_ptr<GraphExecutionState> execution_state;
+  TF_RETURN_IF_ERROR(GraphExecutionState::MakeForBaseGraph(
+      std::move(graph_def), options, &execution_state));
+  return execution_state;
+}
+
+absl::Status FallbackState::AddFunctionDef(const FunctionDef &func_def) {
+  return func_lib_def_.AddFunctionDef(func_def);
+}
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/fallback/fallback_state.h

@@ -0,0 +1,105 @@
+/* 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_CORE_TFRT_FALLBACK_FALLBACK_STATE_H_
+#define TENSORFLOW_CORE_TFRT_FALLBACK_FALLBACK_STATE_H_
+
+#include <memory>
+#include <variant>
+#include <vector>
+
+#include "absl/base/nullability.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/device_set.h"
+#include "tensorflow/core/common_runtime/graph_execution_state.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+// FallbackState contains the necessary runtime states (eg. Devices) used in
+// current tensorflow. It also provides methods used in current tensorflow.
+class FallbackState {
+ public:
+  // The FunctionDefLibrary is passed in to initialize the
+  // ProcessFunctionLibraryRuntime member of this class
+  static absl::StatusOr<std::unique_ptr<FallbackState>> Create(
+      const SessionOptions &session_options,
+      const tensorflow::FunctionDefLibrary &fdef_lib);
+
+  static absl::StatusOr<std::unique_ptr<FallbackState>> CreateWithCpuDevice(
+      const SessionOptions &session_options,
+      const tensorflow::FunctionDefLibrary &fdef_lib);
+
+  static absl::StatusOr<std::unique_ptr<FallbackState>> CreateWithMockGpuDevice(
+      const SessionOptions &session_options,
+      const tensorflow::FunctionDefLibrary &fdef_lib);
+
+  static absl::StatusOr<std::unique_ptr<FallbackState>> CreateWithDeviceMgr(
+      const SessionOptions &session_options,
+      const tensorflow::FunctionDefLibrary &fdef_lib,
+      DynamicDeviceMgr *absl_nonnull device_mgr);
+
+  FallbackState(const SessionOptions &session_options,
+                std::variant<std::vector<std::unique_ptr<Device>>,
+                             DynamicDeviceMgr *absl_nonnull>
+                    device_mgr,
+                const tensorflow::FunctionDefLibrary &fdef_lib);
+
+  // Create GraphExecutionState from the `graph_def`. The result will contain a
+  // preprocessed graph with runtime information such as devices.
+  absl::StatusOr<std::unique_ptr<GraphExecutionState>>
+  CreateGraphExecutionState(GraphDef graph_def, bool run_placer = true,
+                            bool enable_tf2xla_mlir_bridge = true) const;
+
+  // Adds `func_def` to the function library.
+  absl::Status AddFunctionDef(const FunctionDef &func_def);
+
+  const SessionOptions &session_options() const { return session_options_; }
+
+  const DeviceMgr &device_manager() const { return *device_manager_ptr_; }
+
+  DeviceMgr &device_manager() { return *device_manager_ptr_; }
+
+  const DeviceSet &device_set() const { return device_set_; }
+
+  const ProcessFunctionLibraryRuntime &process_function_library_runtime()
+      const {
+    return pflr_;
+  }
+
+  const FunctionLibraryDefinition &func_lib_def() const {
+    return func_lib_def_;
+  }
+
+ private:
+  SessionOptions session_options_;
+  DynamicDeviceMgr device_manager_;
+  DynamicDeviceMgr *absl_nonnull device_manager_ptr_;
+  DeviceSet device_set_;
+  FunctionLibraryDefinition func_lib_def_;
+  ProcessFunctionLibraryRuntime pflr_;
+};
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_FALLBACK_FALLBACK_STATE_H_

tensorflow/core/tfrt/fallback/fallback_state_test.cc

@@ -0,0 +1,144 @@
+/* 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/tfrt/fallback/fallback_state.h"
+
+#include <memory>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include "absl/base/nullability.h"
+#include "absl/status/status_matchers.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/framework/device_factory.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+
+namespace tensorflow {
+namespace {
+
+using ::testing::HasSubstr;
+using ::testing::Not;
+
+TEST(FallbackStateTest, CreateWithCpuDeviceVector) {
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+
+  std::vector<std::unique_ptr<Device>> devices;
+  TF_ASSERT_OK(DeviceFactory::AddCpuDevices(
+      session_options, "/job:localhost/replica:0/task:0", &devices));
+
+  std::variant<std::vector<std::unique_ptr<Device>>,
+               DynamicDeviceMgr* absl_nonnull>
+      device_variant = std::move(devices);
+
+  auto fallback_state = std::make_unique<tfrt_stub::FallbackState>(
+      session_options, std::move(device_variant), fdef_lib);
+
+  const auto& device_manager = fallback_state->device_manager();
+  EXPECT_GT(device_manager.NumDevices(), 0);
+  EXPECT_EQ(device_manager.NumDeviceType("CPU"), 1);
+}
+
+TEST(FallbackStateTest, CreateWithDynamicDeviceMgr) {
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+
+  std::vector<std::unique_ptr<Device>> devices;
+  TF_ASSERT_OK(DeviceFactory::AddCpuDevices(
+      session_options, "/job:localhost/replica:0/task:0", &devices));
+  auto static_device_mgr =
+      std::make_unique<DynamicDeviceMgr>(std::move(devices));
+
+  DynamicDeviceMgr* absl_nonnull device_mgr_ptr(static_device_mgr.get());
+
+  auto fallback_state = std::make_unique<tfrt_stub::FallbackState>(
+      session_options, device_mgr_ptr, fdef_lib);
+  const auto& device_manager = fallback_state->device_manager();
+
+  EXPECT_GT(device_manager.NumDevices(), 0);
+  EXPECT_EQ(device_manager.NumDeviceType("CPU"), 1);
+}
+
+TEST(FallbackStateTest, CreateRendezvous) {
+  // Given a FallbackState, when a function is launched by function library
+  // runtime without an explicit rendezvous, it should be able to create one
+  // from the rendezvous factory.
+  FunctionDefLibrary flib;
+  *flib.add_function() = FunctionDefHelper::Define(
+      /*function_name=*/"dummy_fn",
+      /*arg_def=*/{},
+      /*return values=*/{},
+      /*attr def=*/{},
+      /*node_def=*/{});
+
+  TF_ASSERT_OK_AND_ASSIGN(auto fallback_state,
+                          tfrt_stub::FallbackState::Create({}, flib));
+
+  const ProcessFunctionLibraryRuntime& pflr =
+      fallback_state->process_function_library_runtime();
+  FunctionLibraryRuntime::Options opts;
+  opts.source_device = "/job:localhost/replica:0/task:0";
+  opts.remote_execution = true;
+
+  auto status = pflr.RunSync(opts, pflr.GetHandle("dummy_fn"), {}, nullptr);
+
+  EXPECT_THAT(status, Not(absl_testing::StatusIs(error::FAILED_PRECONDITION,
+                                                 HasSubstr("rendezvous"))));
+}
+
+TEST(FallbackStateTest, CreateGraphExecutionState) {
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tfrt_stub::FallbackState::CreateWithCpuDevice(session_options, fdef_lib));
+
+  GraphDef graphdef;
+  {
+    auto scope = tensorflow::Scope::NewRootScope().WithDevice(
+        "/job:localhost/replica:0/task:0/device:CPU:0");
+
+    Output a = ops::Const(scope.WithOpName("a"), 2.0, {1, 1});
+
+    TF_ASSERT_OK(scope.ToGraphDef(&graphdef));
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(auto graph_execution_state,
+                          fallback_state->CreateGraphExecutionState(
+                              std::move(graphdef), /*run_placer=*/true,
+                              /*enable_tf2xla_mlir_bridge=*/false));
+}
+
+TEST(FallbackStateTest, CreateWithMockGpuDevice) {
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+  TF_ASSERT_OK_AND_ASSIGN(auto fallback_state,
+                          tfrt_stub::FallbackState::CreateWithMockGpuDevice(
+                              session_options, fdef_lib));
+  const auto& device_manager = fallback_state->device_manager();
+  EXPECT_GT(device_manager.NumDeviceType("GPU"), 0);
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/core/tfrt/fallback/op_cost_map.proto

@@ -0,0 +1,26 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tfrt_stub;
+
+// For serializing and restoring the cost of op, see cost_recorder.h for
+// details.
+// NEXT_ID: 2
+message OpCostMapProto {
+  // Maps an op_key to a cost measured in nanoseconds.
+  map<int64, uint64> op_cost_map = 1;
+}

tensorflow/core/tfrt/fallback/op_kernel_runner.cc

@@ -0,0 +1,204 @@
+/* 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/tfrt/fallback/op_kernel_runner.h"
+
+#include <cstddef>
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+namespace {
+
+absl::Status CheckOpDefCompatibility(const tensorflow::OpDef& op_def) {
+  auto check_arg_def = [&](const auto& arg_def) {
+    if (arg_def.is_ref())
+      return absl::InternalError(
+          absl::StrCat("TFRT kernel fallback error: Unsupported ref args in ",
+                       op_def.name()));
+    return absl::OkStatus();
+  };
+
+  for (const auto& arg_def : op_def.input_arg())
+    TF_RETURN_IF_ERROR(check_arg_def(arg_def));
+  for (const auto& arg_def : op_def.output_arg())
+    TF_RETURN_IF_ERROR(check_arg_def(arg_def));
+
+  return absl::OkStatus();
+}
+
+// Create a tensorflow::NodeDef from the tensorflow::OpDef and the attributes.
+absl::StatusOr<tensorflow::NodeDef> BuildNodeDef(
+    const tensorflow::OpDef& op_def, absl::string_view node_name, int num_args,
+    const std::function<absl::Status(tensorflow::AttrValueMap*)>&
+        attr_builder) {
+  tensorflow::NodeDef node_def;
+  node_def.set_name(node_name);
+  node_def.set_op(op_def.name());
+  for (int i = 0; i < num_args; ++i) {
+    node_def.add_input("dummy_input");
+  }
+
+  auto* attr_value_map = node_def.mutable_attr();
+  TF_RETURN_IF_ERROR(attr_builder(attr_value_map));
+
+  // For any attr-value pairs that exist in the op def (from op registry)
+  // but not in `attr_value_map`, fill them into `attr_value_map`, so that we
+  // can run a TFE_Op without having to specify all the default attr values
+  // (e.g. for matmul, the `transpose_a` attr defaults to false).
+  for (const auto& attr_def : op_def.attr()) {
+    if (attr_def.has_default_value()) {
+      // Insertion will fail if this attribute already has a value.
+      attr_value_map->insert({attr_def.name(), attr_def.default_value()});
+    }
+  }
+  return node_def;
+}
+
+absl::Status CreateOpKernel(tensorflow::FunctionLibraryRuntime* flr,
+                            tensorflow::NodeDef ndef,
+                            std::unique_ptr<tensorflow::OpKernel>* result) {
+  std::shared_ptr<const tensorflow::NodeProperties> props;
+  TF_RETURN_IF_ERROR(tensorflow::NodeProperties::CreateFromNodeDef(
+      std::move(ndef), flr->GetFunctionLibraryDefinition(), &props));
+  tensorflow::OpKernel* k = nullptr;
+  TF_RETURN_IF_ERROR(flr->CreateKernel(props, &k));
+  result->reset(k);
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::StatusOr<OpKernelRunner> OpKernelRunner::Create(
+    absl::string_view op_name, absl::string_view node_name,
+    absl::string_view device_name, int num_args,
+    const std::function<absl::Status(tensorflow::AttrValueMap*)>& attr_builder,
+    const tensorflow::DeviceMgr& device_manager,
+    const tensorflow::ProcessFunctionLibraryRuntime&
+        process_function_library_runtime) {
+  tensorflow::Device* device = nullptr;
+  absl::Status s = device_manager.LookupDevice(device_name, &device);
+
+  // Fall back to host device if it fails to find the specified device.
+  if (!s.ok()) {
+    LOG_EVERY_N_SEC(WARNING, 30)
+        << "Failed to find device " << device_name
+        << " when creating OpKernel: " << op_name << ". Error: " << s
+        << ", fallback to host device instead";
+    device = device_manager.HostCPU();
+  }
+
+  return Create(op_name, node_name, num_args, attr_builder,
+                process_function_library_runtime, device);
+}
+
+absl::StatusOr<OpKernelRunner> OpKernelRunner::Create(
+    absl::string_view op_name, absl::string_view node_name, int num_args,
+    const std::function<absl::Status(tensorflow::AttrValueMap*)>& attr_builder,
+    const tensorflow::ProcessFunctionLibraryRuntime&
+        process_function_library_runtime,
+    tensorflow::Device* device) {
+  const OpDef* op_def = nullptr;
+  TF_RETURN_IF_ERROR(tensorflow::OpRegistry::Global()->LookUpOpDef(
+      std::string(op_name), &op_def));
+  TF_RETURN_IF_ERROR(CheckOpDefCompatibility(*op_def));
+  VLOG(1) << "KernelFallbackExecuteCompat creating op " << op_name
+          << " from OpDef: " << op_def->DebugString();
+
+  TF_ASSIGN_OR_RETURN(auto node_def,
+                      BuildNodeDef(*op_def, node_name, num_args, attr_builder));
+
+  VLOG(1) << "KernelFallbackExecuteCompat created NodeDef: "
+          << node_def.DebugString();
+
+  tensorflow::FunctionLibraryRuntime* function_library_runtime = nullptr;
+
+  function_library_runtime =
+      process_function_library_runtime.GetFLR(device->name());
+
+  std::unique_ptr<OpKernel> op_kernel;
+  TF_RETURN_IF_ERROR(CreateOpKernel(function_library_runtime,
+                                    std::move(node_def), &op_kernel));
+
+  if (!op_kernel) {
+    return absl::InternalError(
+        absl::StrCat("Failed to create OpKernel for op: ", op_name));
+  }
+  return OpKernelRunner(device, function_library_runtime, std::move(op_kernel));
+}
+
+OpKernelRunner::OpKernelRunner(
+    tensorflow::Device* device,
+    tensorflow::FunctionLibraryRuntime* function_library_runtime,
+    std::unique_ptr<tensorflow::OpKernel> op_kernel)
+    : op_kernel_(std::move(op_kernel)), info_(std::make_unique<Info>()) {
+  DCHECK(device);
+  DCHECK(function_library_runtime);
+
+  info_->device = device;
+  info_->function_library_runtime = function_library_runtime;
+  info_->resource_manager = device->resource_manager();
+  info_->is_async = (op_kernel_->AsAsync() != nullptr);
+
+  const auto& input_memory_types = op_kernel_->input_memory_types();
+
+  auto& input_alloc_attrs = info_->input_alloc_attrs;
+  auto& output_alloc_attrs = info_->output_alloc_attrs;
+
+  input_alloc_attrs.resize(op_kernel_->num_inputs());
+  for (size_t i = 0, e = op_kernel_->num_inputs(); i < e; ++i) {
+    input_alloc_attrs[i].set_on_host(input_memory_types[i] ==
+                                     tensorflow::HOST_MEMORY);
+  }
+  const auto& output_memory_types = op_kernel_->output_memory_types();
+  output_alloc_attrs.resize(op_kernel_->num_outputs());
+  for (size_t i = 0, e = output_alloc_attrs.size(); i < e; ++i) {
+    output_alloc_attrs[i].set_on_host(output_memory_types[i] ==
+                                      tensorflow::HOST_MEMORY);
+  }
+
+  input_alloc_attrs_ = input_alloc_attrs;
+  output_alloc_attrs_ = output_alloc_attrs;
+}
+
+void OpKernelRunner::RunAsync(OpKernelContext* context,
+                              AsyncOpKernel::DoneCallback done_callback) const {
+  DVLOG(1) << "KernelFallbackExecuteCompat Running Async Op: "
+           << op_kernel_->def().DebugString()
+           << ", on Device: " << context->device()->name();
+
+  AsyncOpKernel* async = op_kernel_->AsAsync();
+  DCHECK(async);
+
+  // For TFRT GPU or TPU, we currently only run xla clusters on GPU or TPU, and
+  // all other ops are run on CPU.
+
+  async->ComputeAsync(context, std::move(done_callback));
+}
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/fallback/op_kernel_runner.h

@@ -0,0 +1,236 @@
+/* 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_CORE_TFRT_FALLBACK_OP_KERNEL_RUNNER_H_
+#define TENSORFLOW_CORE_TFRT_FALLBACK_OP_KERNEL_RUNNER_H_
+
+#include <assert.h>
+#include <stddef.h>
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/base/attributes.h"
+#include "absl/container/inlined_vector.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+class OpKernelRunner {
+ public:
+  static absl::StatusOr<OpKernelRunner> Create(
+      absl::string_view op_name, absl::string_view node_name,
+      absl::string_view device_name, int num_args,
+      const std::function<absl::Status(tensorflow::AttrValueMap*)>&
+          attr_builder,
+      const tensorflow::DeviceMgr& device_manager,
+      const tensorflow::ProcessFunctionLibraryRuntime&
+          process_function_library_runtime);
+
+  ABSL_DEPRECATED("Please use the Create() method that takes node_name.")
+  static absl::StatusOr<OpKernelRunner> Create(
+      absl::string_view op_name, absl::string_view device_name, int num_args,
+      const std::function<absl::Status(tensorflow::AttrValueMap*)>&
+          attr_builder,
+      const tensorflow::DeviceMgr& device_manager,
+      const tensorflow::ProcessFunctionLibraryRuntime&
+          process_function_library_runtime) {
+    return Create(op_name, /*node_name=*/op_name, device_name, num_args,
+                  attr_builder, device_manager,
+                  process_function_library_runtime);
+  }
+
+  static absl::StatusOr<OpKernelRunner> Create(
+      absl::string_view op_name, absl::string_view node_name, int num_args,
+      const std::function<absl::Status(tensorflow::AttrValueMap*)>&
+          attr_builder,
+      const tensorflow::ProcessFunctionLibraryRuntime&
+          process_function_library_runtime,
+      tensorflow::Device* device);
+
+  ABSL_DEPRECATED("Please use the Create() method that takes node_name.")
+  static absl::StatusOr<OpKernelRunner> Create(
+      absl::string_view op_name, int num_args,
+      const std::function<absl::Status(tensorflow::AttrValueMap*)>&
+          attr_builder,
+      const tensorflow::ProcessFunctionLibraryRuntime&
+          process_function_library_runtime,
+      tensorflow::Device* device) {
+    return Create(op_name, /*node_name=*/op_name, num_args, attr_builder,
+                  process_function_library_runtime, device);
+  }
+
+  OpKernelRunner() = default;
+
+  explicit operator bool() const { return op_kernel_ != nullptr; }
+
+  void Run(OpKernelContext* context) const {
+    DVLOG(1) << "KernelFallbackExecuteCompat Running Op: "
+             << op_kernel_->def().DebugString()
+             << ", on Device: " << context->device()->name();
+
+    // For TFRT GPU or TPU, we currently only run xla clusters on GPU or TPU,
+    // and all other ops are run on CPU.
+
+    op_kernel_->Compute(context);
+  }
+
+  void RunAsync(OpKernelContext* context,
+                AsyncOpKernel::DoneCallback done_callback) const;
+
+  bool IsAsync() const { return info_->is_async; }
+
+  tensorflow::OpKernel* op_kernel() const { return op_kernel_.get(); }
+  tensorflow::Device* device() const { return info_->device; }
+  tensorflow::FunctionLibraryRuntime* function_library_runtime() const {
+    return info_->function_library_runtime;
+  }
+  tensorflow::ResourceMgr* resource_manager() const {
+    return info_->resource_manager;
+  }
+
+  absl::Span<const AllocatorAttributes> input_alloc_attrs() const {
+    return input_alloc_attrs_;
+  }
+  absl::Span<const AllocatorAttributes> output_alloc_attrs() const {
+    return output_alloc_attrs_;
+  }
+
+ private:
+  explicit OpKernelRunner(
+      tensorflow::Device* device,
+      tensorflow::FunctionLibraryRuntime* function_library_runtime,
+      std::unique_ptr<OpKernel> op_kernel);
+
+  std::unique_ptr<OpKernel> op_kernel_;
+  absl::Span<const AllocatorAttributes> input_alloc_attrs_;
+  absl::Span<const AllocatorAttributes> output_alloc_attrs_;
+
+  struct Info {
+    tensorflow::Device* device = nullptr;
+    tensorflow::FunctionLibraryRuntime* function_library_runtime = nullptr;
+    tensorflow::ResourceMgr* resource_manager = nullptr;
+    bool is_async = false;
+    absl::InlinedVector<AllocatorAttributes, 4UL> input_alloc_attrs;
+    absl::InlinedVector<AllocatorAttributes, 1UL> output_alloc_attrs;
+  };
+  std::unique_ptr<Info> info_;
+};
+
+// OpKernelRunState keeps the states needed for per-kernel execution.
+struct OpKernelRunState {
+  std::vector<const tensorflow::TensorBuffer*> tensor_buffers;
+  std::vector<tensorflow::TensorValue> input_tf_tensor_values;
+  OpKernelContext::Params params;
+  absl::InlinedVector<tensorflow::Tensor, 4UL> input_tf_tensors;
+
+  OpKernelRunState() = default;
+  OpKernelRunState(absl::Span<const tensorflow::TensorValue> tensor_values,
+                   const OpKernelContext::Params& p,
+                   tensorflow::DeviceBase* device = nullptr) {
+    // `input_tf_tensor_values` contains the reference to all tensor used,
+    // while `input_tf_tensors` only contains those needs ownership so their
+    // sizes may not match. For this copy assignment, we conservatively copy all
+    // tensors.
+    input_tf_tensors.reserve(tensor_values.size());
+    for (const auto& tensor_value : tensor_values) {
+      input_tf_tensors.push_back(*tensor_value.tensor);
+    }
+    for (auto& tensor : input_tf_tensors) {
+      input_tf_tensor_values.emplace_back(&tensor);
+    }
+
+    // Since `input_tf_tensor_values` and `params` contains pointers to
+    // `input_tf_tensors`, we need to change those pointers to the correct ones
+    // after copying.
+    params = p;
+    params.inputs = input_tf_tensor_values;
+    // Clear eigen_gpu_device to ensure OpKernelContext constructor will make a
+    // new eigen GPU device.
+    params.eigen_gpu_device = nullptr;
+    if (device != nullptr) params.device = device;
+  }
+
+  OpKernelRunState(const OpKernelRunState& other) = delete;
+  OpKernelRunState& operator=(const OpKernelRunState& other) = delete;
+
+  ~OpKernelRunState() = default;
+};
+
+// OpKernelRunnerTable for keeping OpKernelRunner instances to avoid expensive
+// reinstantiation of OpKernel and other repeated setup per kernel execution.
+// OpKernelRunnerTable is thread-compatible.
+class OpKernelRunnerTable {
+ public:
+  OpKernelRunnerTable() = default;
+
+  // Return true if it successfully inserts `runner`. `index` is supposed to be
+  // dense.
+  bool Insert(int64_t index, OpKernelRunner runner) {
+    if (runners_.size() <= index) runners_.resize(index + 1);
+    if (runners_[index]) return false;
+    runners_[index] = std::move(runner);
+    return true;
+  }
+
+  // Return the OpKernelRunner at the corresponding `index` in the table. The
+  // result can never be nullptr. It is a fatal error to use an index that is
+  // not in the table. Note that the returned pointer will be invalidated if
+  // Insert() is called.
+  const OpKernelRunner* Get(int64_t index) const {
+    // Out of bounds vector access will throw an exception and anyway will crash
+    // the binary, prefer a more readable error message.
+    CHECK_GT(runners_.size(), index)  // Crash OK
+        << "runner index is out of bounds: index=" << index
+        << " size=" << runners_.size();
+    CHECK(runners_[index])  // Crash OK
+        << "runner is not available: index=" << index;
+    return GetUnsafe(index);
+  }
+
+  const OpKernelRunner* GetUnsafe(int64_t index) const {
+    DCHECK_GT(runners_.size(), index);
+    auto& result = runners_[index];
+    DCHECK(result);
+    return &result;
+  }
+
+ private:
+  std::vector<OpKernelRunner> runners_;
+};
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_FALLBACK_OP_KERNEL_RUNNER_H_

tensorflow/core/tfrt/fallback/op_kernel_runner_cache.cc

@@ -0,0 +1,80 @@
+/* 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/core/tfrt/fallback/op_kernel_runner_cache.h"
+
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/base/casts.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+absl::StatusOr<OpKernelRunner*> OpKernelRunnerCache::GetOrCreate(
+    tfrt::Location loc, absl::string_view op_name,
+    absl::string_view device_name, int num_args,
+    const std::function<absl::Status(tensorflow::AttrValueMap*)>& attr_builder,
+    const tensorflow::DeviceMgr& device_manager,
+    const tensorflow::ProcessFunctionLibraryRuntime&
+        process_function_library_runtime) {
+  OpLocationKey key(loc);
+  {
+    tf_shared_lock lock(mu_);
+    auto it = map_.find(key);
+    if (it != map_.end()) {
+      DCHECK_EQ(it->second->op_kernel()->def().op(), op_name);
+      return it->second.get();
+    }
+  }
+
+  mutex_lock lock(mu_);
+
+  auto it = map_.find(key);
+  if (it != map_.end()) {
+    DCHECK_EQ(it->second->op_kernel()->def().op(), op_name);
+    return it->second.get();
+  }
+
+  VLOG(1) << "KernelFallbackExecuteCompat creating op " << op_name
+          << " at location " << loc.data << " on device " << device_name;
+
+  std::string node_name = absl::StrCat(
+      op_name, "_", loc.data, "_", absl::bit_cast<uintptr_t>(loc.GetHandler()));
+
+  TF_ASSIGN_OR_RETURN(
+      auto runner, OpKernelRunner::Create(
+                       op_name, node_name, device_name, num_args, attr_builder,
+                       device_manager, process_function_library_runtime));
+
+  auto runner_uptr = std::make_unique<OpKernelRunner>(std::move(runner));
+
+  auto* runner_ptr = runner_uptr.get();
+  auto r = map_.emplace(key, std::move(runner_uptr)).second;
+  DCHECK(r);
+
+  return runner_ptr;
+}
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/fallback/op_kernel_runner_cache.h

@@ -0,0 +1,74 @@
+/* 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_CORE_TFRT_FALLBACK_OP_KERNEL_RUNNER_CACHE_H_
+#define TENSORFLOW_CORE_TFRT_FALLBACK_OP_KERNEL_RUNNER_CACHE_H_
+
+#include <functional>
+#include <memory>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
+#include "tfrt/host_context/location.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+class OpLocationKey {
+ public:
+  explicit OpLocationKey(tfrt::Location loc) : loc_(loc) {}
+
+  template <typename H>
+  friend H AbslHashValue(H h, const OpLocationKey& key) {
+    // NOTE: Each BEF file has its own LocationHandler. Using LocationHandler
+    // as part of cache key here can avoid cache collision between different
+    // BEF file.
+    return H::combine(std::move(h), key.loc_.data, key.loc_.GetHandler());
+  }
+
+  friend bool operator==(const OpLocationKey& x, const OpLocationKey& y) {
+    return x.loc_.data == y.loc_.data &&
+           x.loc_.GetHandler() == y.loc_.GetHandler();
+  }
+
+ private:
+  tfrt::Location loc_;
+};
+
+// OpKernelRunnerCache is similar to OpKernelRunnerTable but thread-safe.
+class OpKernelRunnerCache {
+ public:
+  OpKernelRunnerCache() = default;
+
+  absl::StatusOr<OpKernelRunner*> GetOrCreate(
+      tfrt::Location loc, absl::string_view op_name,
+      absl::string_view device_name, int num_args,
+      const std::function<absl::Status(tensorflow::AttrValueMap*)>&
+          attr_builder,
+      const tensorflow::DeviceMgr& device_manager,
+      const tensorflow::ProcessFunctionLibraryRuntime&
+          process_function_library_runtime);
+
+ private:
+  mutable mutex mu_;
+  absl::flat_hash_map<OpLocationKey, std::unique_ptr<OpKernelRunner>> map_
+      TF_GUARDED_BY(mu_);
+};
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+#endif  // TENSORFLOW_CORE_TFRT_FALLBACK_OP_KERNEL_RUNNER_CACHE_H_

tensorflow/core/tfrt/fallback/op_kernel_runner_test.cc

@@ -0,0 +1,167 @@
+/* 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/core/tfrt/fallback/op_kernel_runner.h"
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/framework/device_factory.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tensorflow/core/tfrt/fallback/op_kernel_runner_cache.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+namespace {
+
+using ::testing::IsNull;
+using ::testing::SizeIs;
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+constexpr const char* kDeviceType = "GPU";
+#else
+constexpr const char* kDeviceType = "CPU";
+#endif
+
+class TestOpKernel : public OpKernel {
+ public:
+  using OpKernel::OpKernel;
+
+  ~TestOpKernel() override = default;
+
+  void Compute(OpKernelContext* context) override {
+    context->set_output(0, context->input(0));
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("TestOp").Device(DEVICE_CPU), TestOpKernel);
+
+// Identical to BatchFunction except it has 2 extra TFRT attributes and it does
+// not have `f` attribute. Users will not invoke this op directly.
+REGISTER_OP("TestOp").Input("x: int32").Output("y: int32");
+
+TEST(OpKernelRunnerTest, Create) {
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+  TF_ASSERT_OK_AND_ASSIGN(auto fallback_state,
+                          FallbackState::Create(session_options, fdef_lib));
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto runner,
+      OpKernelRunner::Create(
+          /*op_name=*/
+          "TestOp", /*node_name=*/"TestOp_node_name",
+          /*device_name=*/"/job:localhost/replica:0/task:0/device:CPU:0",
+          /*num_args=*/1,
+          /*attr_builder=*/
+          [](tensorflow::AttrValueMap*) { return absl::OkStatus(); },
+          fallback_state->device_manager(),
+          fallback_state->process_function_library_runtime()));
+
+  ASSERT_TRUE(runner);
+
+  EXPECT_EQ(runner.op_kernel()->name(), "TestOp_node_name");
+}
+
+TEST(OpKernelRunnerTest, OpKernelRunnerCache) {
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+  TF_ASSERT_OK_AND_ASSIGN(auto fallback_state,
+                          FallbackState::Create(session_options, fdef_lib));
+
+  OpKernelRunnerCache cache;
+
+  tfrt::Location loc(/*handler=*/nullptr, /*data=*/100);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto* runner,
+      cache.GetOrCreate(
+          loc,
+          /*op_name=*/"TestOp",
+          /*device_name=*/"/job:localhost/replica:0/task:0/device:CPU:0",
+          /*num_args=*/1,
+          /*attr_builder=*/
+          [](tensorflow::AttrValueMap*) { return absl::OkStatus(); },
+          fallback_state->device_manager(),
+          fallback_state->process_function_library_runtime()));
+
+  ASSERT_TRUE(runner);
+
+  EXPECT_EQ(runner->op_kernel()->name(), "TestOp_100_0");
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      runner,
+      cache.GetOrCreate(
+          loc,
+          /*op_name=*/"TestOp",
+          /*device_name=*/"/job:localhost/replica:0/task:0/device:CPU:0",
+          /*num_args=*/1,
+          /*attr_builder=*/
+          [](tensorflow::AttrValueMap*) { return absl::OkStatus(); },
+          fallback_state->device_manager(),
+          fallback_state->process_function_library_runtime()));
+
+  ASSERT_TRUE(runner);
+
+  EXPECT_EQ(runner->op_kernel()->name(), "TestOp_100_0");
+}
+
+TEST(OpKernelRunnerTest, OpKernelRunState) {
+  SessionOptions options;
+  auto* device_count = options.config.mutable_device_count();
+  device_count->insert({kDeviceType, 1});
+  std::vector<std::unique_ptr<Device>> devices;
+  TF_ASSERT_OK(DeviceFactory::GetFactory(kDeviceType)
+                   ->CreateDevices(options,
+                                   /*name_prefix=*/"/job:a/replica:0/task:0",
+                                   &devices));
+  ASSERT_EQ(devices.size(), 1);
+
+  OpKernelContext::Params params;
+  params.device = devices[0].get();
+  params.ensure_eigen_gpu_device();
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+  ASSERT_THAT(params.eigen_gpu_device, ::testing::NotNull());
+#endif
+
+  Tensor a(DT_FLOAT, TensorShape({}));
+  Tensor b(DT_INT32, TensorShape({}));
+  absl::InlinedVector<TensorValue, 4UL> inputs{TensorValue(&a),
+                                               TensorValue(&b)};
+  params.inputs = inputs;
+
+  Tensor c(DT_UINT8, TensorShape({}));
+  absl::InlinedVector<TensorValue, 4UL> new_inputs{TensorValue(&c)};
+
+  OpKernelRunState run_state(new_inputs, params);
+
+  EXPECT_THAT(run_state.input_tf_tensors, SizeIs(1));
+  EXPECT_THAT(run_state.input_tf_tensor_values, SizeIs(1));
+  EXPECT_EQ(run_state.params.inputs.data(),
+            run_state.input_tf_tensor_values.data());
+  EXPECT_THAT(run_state.params.eigen_gpu_device, IsNull());
+}
+
+}  // namespace
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/gpu/kernel/BUILD

@@ -0,0 +1,129 @@
+load("//tensorflow:tensorflow.bzl", "tf_cuda_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/core/tfrt/saved_model:__pkg__"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "gpurt_kernels",
+    srcs = ["gpurt_kernels.cc"],
+    deps = [
+        ":gpu_runner",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_compat_request_state",
+        "//tensorflow/core/runtime_fallback/kernel:tensor_util",
+        "//tensorflow/core/tfrt/utils:fallback_tensor",
+        "//tensorflow/core/tfrt/utils:gpu_variables_table",
+        "@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",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:support",
+        "@tf_runtime//:tensor_alwayslink",
+    ],
+    alwayslink = True,
+)
+
+cc_library(
+    name = "gpu_runner",
+    srcs = ["gpu_runner.cc"],
+    hdrs = ["gpu_runner.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/jit:pjrt_compile_util",
+        "//tensorflow/compiler/jit:pjrt_tensor_buffer_util",
+        "//tensorflow/compiler/jit:xla_compile_util",
+        "//tensorflow/compiler/jit:xla_device_no_jit_rewrite_registration",
+        "//tensorflow/compiler/jit:xla_launch_util",
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/framework:attr_value_proto_cc",
+        "//tensorflow/core/framework:function_proto_cc",
+        "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_compat_request_state",
+        "//tensorflow/core/tfrt/common:global_state",
+        "//tensorflow/core/tfrt/utils:fallback_tensor",
+        "//tensorflow/core/tfrt/utils:gpu_variables_table",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Support",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:support",
+        "@tf_runtime//:tensor_alwayslink",
+        "@tsl//tsl/platform:fingerprint",
+        "@tsl//tsl/platform:protobuf",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt:pjrt_common",
+        "@xla//xla/tsl/framework:device_id",
+        "@xla//xla/tsl/framework:serving_device_selector",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cuda_cc_test(
+    name = "gpu_runner_test",
+    srcs = ["gpu_runner_test.cc"],
+    tags = [
+        "gpu",  # Only enables test on GPU.
+        "no_oss",  # This test only runs with GPU.
+        "noasan",
+        "nomsan",
+        "noopt",
+        "notsan",
+    ],
+    deps = [
+        ":gpu_runner",
+        "//tensorflow/cc:function_ops",
+        "//tensorflow/cc:math_ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/compiler/jit:xla_gpu_jit",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:test",
+        "//tensorflow/core/common_runtime/gpu:gpu_serving_device_selector",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/kernels:ops_testutil",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:status_matchers",
+        "//tensorflow/core/platform:statusor",
+        "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_compat_request_state",
+        "//tensorflow/core/tfrt/common:pjrt_util",
+        "//tensorflow/core/tfrt/fallback:fallback_state",
+        "@com_google_googletest//:gtest_main",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:tensor",
+        "@xla//xla/tsl/framework:serving_device_selector_policies",
+    ],
+)
+
+cc_library(
+    name = "tfrt_gpu_init",
+    srcs = ["tfrt_gpu_init.cc"],
+    hdrs = ["tfrt_gpu_init.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":gpu_runner",
+        "//tensorflow/core/common_runtime/gpu:gpu_serving_device_selector",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/tfrt/runtime",
+        "@com_google_absl//absl/status",
+        "@tf_runtime//:hostcontext",
+        "@xla//xla/tsl/framework:serving_device_selector_policies",
+    ],
+)

tensorflow/core/tfrt/gpu/kernel/gpu_runner.cc

@@ -0,0 +1,491 @@
+/* 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/core/tfrt/gpu/kernel/gpu_runner.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/synchronization/notification.h"
+#include "absl/types/span.h"
+#include "llvm/ADT/SmallVector.h"
+#include "tensorflow/compiler/jit/pjrt_compile_util.h"
+#include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h"
+#include "tensorflow/compiler/jit/xla_compile_util.h"
+#include "tensorflow/compiler/jit/xla_launch_util.h"
+#include "tensorflow/compiler/jit/xla_platform_info.h"
+#include "tensorflow/compiler/tf2xla/xla_compiler.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/pjrt_common.h"
+#include "xla/tsl/framework/device_id.h"
+#include "xla/tsl/framework/device_id_manager.h"
+#include "xla/tsl/framework/serving_device_selector.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h"
+#include "tensorflow/core/tfrt/common/global_state.h"
+#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
+#include "tensorflow/core/tfrt/utils/gpu_variables_table.h"
+#include "tsl/platform/fingerprint.h"
+#include "tsl/platform/protobuf.h"
+#include "tfrt/host_context/async_dispatch.h"  // from @tf_runtime
+#include "tfrt/host_context/async_value_ref.h"  // from @tf_runtime
+#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
+#include "tfrt/host_context/kernel_registry.h"  // from @tf_runtime
+#include "tfrt/support/forward_decls.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace gpu {
+
+namespace {
+
+// TODO(b/298478068): Consider to integrate this into
+// tfd::TransferTensorToDevice().
+tfrt::AsyncValueRef<tfrt_stub::FallbackTensor> TransferTensorToDevice(
+    const tfrt_stub::FallbackTensor& tensor, tfrt::HostContext* host_ctx,
+    Device* gpu_device) {
+  const tensorflow::Tensor& src = tensor.tensor();
+  tensorflow::AllocatorAttributes attr;
+  attr.set_use_pjrt_allocator(true);
+
+  tensorflow::Tensor dst(gpu_device->GetAllocator(attr), src.dtype(),
+                         src.shape());
+  if (src.shape().num_elements() == 0) {
+    return tfrt::MakeAvailableAsyncValueRef<tfrt_stub::FallbackTensor>(dst);
+  }
+  auto result =
+      tfrt::MakeUnconstructedAsyncValueRef<tfrt_stub::FallbackTensor>();
+
+  DeviceContext* pjrt_device_context =
+      gpu_device->tensorflow_accelerator_device_info()->pjrt_context;
+  bool enqueued = tfrt::EnqueueBlockingWork(
+      host_ctx, [result = result.CopyRef(), gpu_device, pjrt_device_context,
+                 src, dst = std::move(dst)]() mutable {
+        absl::Notification n;
+        absl::Status status;
+        pjrt_device_context->CopyCPUTensorToDevice(
+            &src, gpu_device, &dst, [&status, &n](absl::Status s) mutable {
+              status = s;
+              n.Notify();
+            });
+        n.WaitForNotification();
+        if (!status.ok()) {
+          result.SetError(absl::InternalError(status.message()));
+        } else {
+          result.emplace(std::move(dst));
+        }
+      });
+
+  if (!enqueued) {
+    return tfrt::MakeErrorAsyncValueRef(absl::InternalError(
+        "Failed to enqueue blocking task to transfer tensor."));
+  }
+  return result;
+}
+
+tfrt::AsyncValueRef<tfrt_stub::FallbackTensor> TransferTensorFromDevice(
+    const tfrt_stub::FallbackTensor& tensor, tfrt::HostContext* host_ctx,
+    Device* cpu_device, Device* gpu_device) {
+  const tensorflow::Tensor& src = tensor.tensor();
+
+  tensorflow::AllocatorAttributes attr;
+  tensorflow::Tensor dst(cpu_device->GetAllocator(attr), src.dtype(),
+                         src.shape());
+  if (src.shape().num_elements() == 0) {
+    return tfrt::MakeAvailableAsyncValueRef<tfrt_stub::FallbackTensor>(dst);
+  }
+  auto result =
+      tfrt::MakeUnconstructedAsyncValueRef<tfrt_stub::FallbackTensor>();
+
+  DeviceContext* pjrt_device_context =
+      gpu_device->tensorflow_accelerator_device_info()->pjrt_context;
+  bool enqueued = tfrt::EnqueueBlockingWork(
+      host_ctx, [result = result.CopyRef(), gpu_device, pjrt_device_context,
+                 src, dst = std::move(dst)]() mutable {
+        absl::Notification n;
+        absl::Status status;
+        pjrt_device_context->CopyDeviceTensorToCPU(
+            &src, "tensor_name", gpu_device, &dst,
+            [&status, &n](absl::Status s) mutable {
+              status = s;
+              n.Notify();
+            });
+        n.WaitForNotification();
+        if (!status.ok()) {
+          result.SetError(absl::InternalError(status.message()));
+        } else {
+          result.emplace(std::move(dst));
+        }
+      });
+
+  if (!enqueued) {
+    return tfrt::MakeErrorAsyncValueRef(absl::InternalError(
+        "Failed to enqueue blocking task to transfer tensor."));
+  }
+  return result;
+}
+
+absl::StatusOr<
+    llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>>
+PopulateResultsFromPjRtExecutableOutputs(
+    const XlaCompiler::CompilationResult& compilation_result,
+    std::vector<std::unique_ptr<xla::PjRtBuffer>>& executable_outputs,
+    Device* device, int num_outputs) {
+  llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>
+      fallback_tensor_results;
+
+  for (int i = 0; i < num_outputs; ++i) {
+    const DataType& dtype = compilation_result.outputs[i].type;
+    CHECK(!compilation_result.outputs[i].is_constant);  // Crash OK
+    CHECK(dtype != DT_RESOURCE);                        // Crash OK
+
+    xla::PjRtBuffer* output_buffer = executable_outputs[i].get();
+    if (output_buffer->IsTuple()) {
+      return absl::InvalidArgumentError(
+          "Tuple PJRT buffer output is not supported.");
+    }
+    absl::Span<const int64_t> dims;
+    std::optional<std::vector<int64_t>> logical_dims_storage;
+    if (output_buffer->has_dynamic_dimensions()) {
+      TF_ASSIGN_OR_RETURN(std::vector<int64_t> logical_dims,
+                          output_buffer->logical_dimensions());
+      logical_dims_storage.emplace(std::move(logical_dims));
+      dims = *logical_dims_storage;
+    } else {
+      dims = output_buffer->dimensions();
+    }
+    TensorShape tensor_shape;
+    for (int i = 0; i < dims.size(); ++i) {
+      TF_RETURN_IF_ERROR(tensor_shape.AddDimWithStatus(dims[i]));
+    }
+
+    TF_ASSIGN_OR_RETURN(
+        Tensor output_tensor,
+        MakeTensorFromPjRtBuffer(dtype, tensor_shape,
+                                 std::move(executable_outputs[i])));
+    auto result = tfrt::MakeAvailableAsyncValueRef<tfrt_stub::FallbackTensor>(
+        output_tensor);
+    fallback_tensor_results.emplace_back(std::move(result));
+  }
+  return fallback_tensor_results;
+}
+
+absl::StatusOr<
+    llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>>
+TransferOutputsToHostIfNeeded(
+    llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>> outputs,
+    absl::Span<const int64_t> used_output_indices, Device* cpu_device,
+    Device* gpu_device, tfrt::HostContext* host_ctx) {
+  llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>> results;
+  for (int i = 0, j = 0; i < outputs.size(); ++i) {
+    if (j < used_output_indices.size() && i == used_output_indices[j]) {
+      CHECK(outputs[i].IsAvailable());  // Crash OK
+      tfrt::AsyncValueRef<tfrt_stub::FallbackTensor> output_on_cpu =
+          TransferTensorFromDevice(outputs[i].get(), host_ctx, cpu_device,
+                                   gpu_device);
+      results.push_back(std::move(output_on_cpu));
+      ++j;
+    } else {
+      results.push_back(std::move(outputs[i]));
+    }
+  }
+  return results;
+}
+
+absl::StatusOr<
+    llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>>
+TransferVariablesAndInputs(int device_idx,
+                           absl::Span<const tfrt_stub::FallbackTensor> args,
+                           absl::Span<const int64_t> resource_indices,
+                           Device* cpu_device,
+                           const absl::flat_hash_map<int, Device*>& gpu_devices,
+                           tfrt::gpu::GpuVariablesTable& vars_table,
+                           bool variables_are_shared,
+                           tfrt::HostContext* host_ctx) {
+  llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>> results;
+
+  // Find all devices that are on the same platform (physical GPU). Variables
+  // will be distributed to the memory of virtual devices on the same GPU.
+  tsl::PlatformDeviceId platform_device_id;
+  DeviceType device_type(DEVICE_GPU);
+  TF_RETURN_IF_ERROR(tsl::DeviceIdManager::TfToPlatformDeviceId(
+      device_type, tsl::TfDeviceId(device_idx), &platform_device_id));
+  TF_ASSIGN_OR_RETURN(const std::vector<tsl::TfDeviceId> devices_on_platform,
+                      tsl::DeviceIdManager::GetTfDevicesOnPlatform(
+                          device_type, platform_device_id));
+  absl::flat_hash_set<int64_t> resource_indices_set(resource_indices.begin(),
+                                                    resource_indices.end());
+
+  // If variables are shared, there is only one copy of variables for all
+  // logical devices on the same physical GPU device; otherwise, each logical
+  // device has its own copy of variables.
+  const int cache_copy_idx =
+      variables_are_shared ? platform_device_id.value() : device_idx;
+
+  for (int i = 0, resource_idx = 0; i < args.size(); ++i) {
+    if (resource_indices_set.contains(i)) {
+      // Transfer resources.
+      VLOG(2) << "Transfer resource arg[" << i << "].";
+      tfrt::AsyncValueRef<tfrt_stub::FallbackTensor> device_tensor;
+      auto cached_device_variable =
+          vars_table.GetDeviceVariable(args[i], cache_copy_idx);
+      if (cached_device_variable) {
+        VLOG(2) << "Cache hit for resource arg[" << i << "].";
+        device_tensor = cached_device_variable.CopyRef();
+      } else {
+        VLOG(2) << "Cache miss for resource arg[" << i << "].";
+
+        int gpu_device_idx;
+        if (variables_are_shared) {
+          // Distribute variables on logical devices on the same GPU.
+          const int idx = resource_idx % devices_on_platform.size();
+          gpu_device_idx = devices_on_platform[idx].value();
+        } else {
+          gpu_device_idx = device_idx;
+        }
+
+        VLOG(2) << "Transfer the resource arg[" << i << "] to device "
+                << gpu_device_idx << ".";
+        device_tensor = TransferTensorToDevice(args[i], host_ctx,
+                                               gpu_devices.at(gpu_device_idx));
+        vars_table.AddOrUpdateDeviceVariable(args[i], cache_copy_idx,
+                                             std::move(device_tensor));
+        device_tensor =
+            vars_table.GetDeviceVariable(args[i], cache_copy_idx).CopyRef();
+      }
+      results.push_back(device_tensor);
+      ++resource_idx;
+    } else {
+      // Transfer inputs.
+      VLOG(2) << "Transfer input arg[" << i << "].";
+      tfrt::AsyncValueRef<tfrt_stub::FallbackTensor> device_tensor =
+          TransferTensorToDevice(args[i], host_ctx, gpu_devices.at(device_idx));
+      results.push_back(device_tensor);
+    }
+  }
+  return results;
+}
+
+absl::StatusOr<uint64_t> GenerateFingerprint(
+    const std::string& function_name,
+    const tfd::KernelFallbackCompatRequestState* fallback_request_state) {
+  const FunctionLibraryDefinition* flib_def =
+      fallback_request_state->cpu_function_library_runtime()
+          ->GetFunctionLibraryDefinition();
+  const FunctionDef* fdef = flib_def->Find(function_name);
+  if (!fdef) {
+    return absl::InternalError(
+        absl::StrCat("Failed to find the function ", function_name));
+  }
+  return tsl::Fingerprint64(
+      absl::StrCat(fallback_request_state->session_metadata().name(),
+                   fallback_request_state->session_metadata().version(),
+                   tsl::LegacyUnredactedDebugString(fdef->signature())));
+}
+
+std::vector<XlaCompiler::Argument> BuildXlaCompilerArguments(
+    absl::Span<const tfrt_stub::FallbackTensor> inputs) {
+  std::vector<XlaCompiler::Argument> out;
+  out.resize(inputs.size());
+
+  for (int input_num = 0; input_num < inputs.size(); ++input_num) {
+    const tensorflow::Tensor& input = inputs[input_num].tensor();
+    CHECK_GT(input.NumElements(), 0);     // Crash OK
+    CHECK(input.dtype() != DT_RESOURCE);  // Crash OK
+
+    XlaCompiler::Argument& arg = out[input_num];
+    arg.kind = XlaCompiler::Argument::kParameter;
+    arg.type = input.dtype();
+    arg.shape = input.shape();
+  }
+  return out;
+}
+
+absl::Status CompileProgram(
+    const GpuRunInputs& run_inputs, int device_idx,
+    const XlaCompiler::CompilationResult** compilation_result,
+    xla::PjRtClient** pjrt_client,
+    xla::PjRtLoadedExecutable** pjrt_executable) {
+  std::vector<XlaCompiler::Argument> xla_compiler_args =
+      BuildXlaCompilerArguments(run_inputs.args);
+
+  DeviceBase* device = run_inputs.gpu_devices.at(device_idx);
+  FunctionLibraryRuntime* flr =
+      run_inputs.fallback_request_state->process_function_library_runtime()
+          .GetFLR(run_inputs.gpu_devices.at(device_idx)->name());
+  XlaPlatformInfo platform_info =
+      XlaPlatformInfoFromDevice(run_inputs.gpu_devices.at(device_idx));
+  NameAttrList function;
+  function.set_name(run_inputs.func_name);
+
+  // We store information about the JIT-compiled XLA computation in the
+  // ResourceMgr.
+  ResourceMgr* rm = tfrt_global::GetTFGlobalResourceMgr();
+
+  return CompileToPjRtLoadedExecutable(
+      device, platform_info, function, xla_compiler_args,
+      /*compile_mode=*/DeviceCompileMode::kStrict,
+      /*has_ref_vars=*/false,
+      /*may_alias_resource_update=*/false, flr, rm, compilation_result,
+      pjrt_client, pjrt_executable);
+}
+
+// Execute the program and transfer the results to the host.
+absl::StatusOr<
+    llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>>
+ExecuteProgram(
+    const GpuRunInputs& run_inputs,
+    const llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>&
+        transferred_args,
+    const XlaCompiler::CompilationResult* compilation_result,
+    xla::PjRtClient* pjrt_client, xla::PjRtLoadedExecutable* pjrt_executable,
+    int device_idx) {
+  // Execute the program.
+  std::vector<const Tensor*> inputs;
+  for (const auto& arg : transferred_args) {
+    if (arg.IsError()) {
+      return absl::InternalError(
+          absl::StrCat("Data transfer failed: ", arg.GetError().message()));
+    }
+    inputs.push_back(&arg->tensor());
+  }
+
+  if (compilation_result->collective_info.has_value()) {
+    return absl::UnimplementedError(
+        "Execution with collectives is not supported.");
+  }
+
+  TF_ASSIGN_OR_RETURN(
+      xla::PjRtDevice * pjrt_device,
+      pjrt_client->LookupAddressableDevice(xla::LocalDeviceId(device_idx)));
+  TF_ASSIGN_OR_RETURN(
+      std::vector<std::unique_ptr<xla::PjRtBuffer>> executable_outputs,
+      RunPjRtExecutable(/*num_missing_prefix_ctx_inputs=*/0, inputs,
+                        /*variable_snapshots=*/{}, /*updated_variables=*/{},
+                        DeviceType(DEVICE_GPU),
+                        /*use_pjrt_tensor_buffer=*/true, *compilation_result,
+                        pjrt_device, pjrt_client, pjrt_executable));
+
+  // Populate the results and transfer the results to host.
+  TF_ASSIGN_OR_RETURN(
+      llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>> results,
+      PopulateResultsFromPjRtExecutableOutputs(
+          *compilation_result, executable_outputs,
+          run_inputs.gpu_devices.at(device_idx), run_inputs.num_outputs));
+
+  return TransferOutputsToHostIfNeeded(
+      results, run_inputs.used_output_indices, run_inputs.cpu_device,
+      run_inputs.gpu_devices.at(device_idx), run_inputs.host_ctx);
+}
+
+}  // namespace
+
+absl::StatusOr<
+    llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>>
+GpuRunner::Run(GpuRunInputs run_inputs) {
+  // Select a device to run this input.
+  TF_ASSIGN_OR_RETURN(uint64_t fingerprint,
+                      GenerateFingerprint(run_inputs.func_name,
+                                          run_inputs.fallback_request_state));
+  tsl::DeviceReservation device_reservation =
+      serving_device_selector_->ReserveDevice(absl::StrCat(fingerprint));
+  const int device_idx = device_reservation.device_index();
+  VLOG(1) << "GpuRunner selected device " << device_idx << ".";
+
+  // Compile the program.
+  const XlaCompiler::CompilationResult* compilation_result;
+  xla::PjRtClient* pjrt_client;                // Not owned.
+  xla::PjRtLoadedExecutable* pjrt_executable;  // Not owned.
+  TF_RETURN_IF_ERROR(CompileProgram(run_inputs, device_idx, &compilation_result,
+                                    &pjrt_client, &pjrt_executable));
+
+  // Transfer variables and inputs.
+  TF_ASSIGN_OR_RETURN(
+      llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>
+          transferred_args,
+      TransferVariablesAndInputs(
+          device_idx, run_inputs.args, run_inputs.resource_indices,
+          run_inputs.cpu_device, run_inputs.gpu_devices, vars_table_,
+          /*variables_are_shared=*/false, run_inputs.host_ctx));
+
+  llvm::SmallVector<tfrt::RCReference<tfrt::AsyncValue>, 4>
+      transferred_args_to_wait;
+  for (const auto& arg : transferred_args) {
+    if (!arg.IsAvailable()) {
+      transferred_args_to_wait.push_back(arg.CopyRCRef());
+    }
+  }
+
+  llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>> results;
+  results.reserve(run_inputs.num_outputs);
+  for (size_t i = 0; i < run_inputs.num_outputs; ++i) {
+    results.emplace_back(
+        tfrt::MakeUnconstructedAsyncValueRef<tfrt_stub::FallbackTensor>());
+  }
+
+  tfrt::RunWhenReady(
+      transferred_args_to_wait,
+      [run_inputs = std::move(run_inputs),
+       transferred_args = std::move(transferred_args), results = results,
+       compilation_result, pjrt_client, pjrt_executable, device_idx]() mutable {
+        auto execution_outputs =
+            ExecuteProgram(run_inputs, transferred_args, compilation_result,
+                           pjrt_client, pjrt_executable, device_idx);
+        CHECK_EQ(results.size(), execution_outputs->size());  // Crash OK.
+
+        if (!execution_outputs.ok()) {
+          // Set all outputs as the error returned by the execution.
+          for (size_t i = 0; i < results.size(); ++i) {
+            results[i].SetError(
+                absl::InternalError(execution_outputs.status().message()));
+          }
+          return;
+        }
+        // Populate each output once it is available.
+        for (int i = 0; i < results.size(); ++i) {
+          auto& result = results[i];
+          auto& output_av = (*execution_outputs)[i];
+          output_av.AndThen([result = result, output_av = output_av] {
+            result.emplace(std::move(output_av.get().tensor()));
+          });
+        }
+      });
+  return results;
+}
+
+}  // namespace gpu
+}  // namespace tensorflow

tensorflow/core/tfrt/gpu/kernel/gpu_runner.h

@@ -0,0 +1,73 @@
+/* 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_CORE_TFRT_GPU_KERNEL_GPU_RUNNER_H_
+#define TENSORFLOW_CORE_TFRT_GPU_KERNEL_GPU_RUNNER_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/statusor.h"
+#include "llvm/ADT/SmallVector.h"
+#include "xla/tsl/framework/serving_device_selector.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h"
+#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
+#include "tensorflow/core/tfrt/utils/gpu_variables_table.h"
+#include "tfrt/host_context/async_value_ref.h"  // from @tf_runtime
+#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
+#include "tfrt/support/forward_decls.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace gpu {
+
+constexpr char kGpuRunnerResourceName[] = "GpuRunnerResource";
+
+struct GpuRunInputs {
+  std::vector<tfrt_stub::FallbackTensor> args;
+  int num_outputs;
+  std::vector<int64_t> resource_indices;
+  std::vector<int64_t> used_output_indices;
+  std::string func_name;
+  Device* cpu_device;
+  absl::flat_hash_map<int, Device*> gpu_devices;
+  const tfd::KernelFallbackCompatRequestState* fallback_request_state;
+  tfrt::HostContext* host_ctx;
+};
+
+class GpuRunner {
+ public:
+  explicit GpuRunner(tsl::ServingDeviceSelector* serving_device_selector)
+      : serving_device_selector_(serving_device_selector) {}
+
+  // This compiles the given program and runs the given input tensors in
+  // `run_inputs`, and returns the output tensor AsyncValues.
+  absl::StatusOr<
+      llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>>>
+  Run(GpuRunInputs run_inputs);
+
+ private:
+  tsl::ServingDeviceSelector* serving_device_selector_;
+  tfrt::gpu::GpuVariablesTable vars_table_;
+};
+
+}  // namespace gpu
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_GPU_KERNEL_GPU_RUNNER_H_

tensorflow/core/tfrt/gpu/kernel/gpu_runner_test.cc

@@ -0,0 +1,209 @@
+/* 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.
+==============================================================================*/
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#include "tensorflow/core/tfrt/gpu/kernel/gpu_runner.h"
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/function_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "xla/tsl/framework/serving_device_selector_policies.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_serving_device_selector.h"
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/graph_to_functiondef.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+#include "tfrt/host_context/diagnostic.h"  // from @tf_runtime
+#include "tfrt/host_context/function.h"  // from @tf_runtime
+#include "tfrt/host_context/host_allocator.h"  // from @tf_runtime
+#include "tfrt/host_context/host_context.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace gpu {
+namespace {
+// TODO(b/297939691): Change this to greater than 1 after PJRT device supports
+// both logical and physical device IDs.
+constexpr int kNumVirtualGpuDevices = 1;
+constexpr char kFunctionName[] = "foo";
+
+StatusOr<std::unique_ptr<Graph>> SampleGraphAddXY() {
+  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto a = ops::_Arg(scope.WithOpName("A"), DT_INT32, 0);
+  auto b = ops::_Arg(scope.WithOpName("B"), DT_INT32, 1);
+  auto c = ops::Add(scope.WithOpName("C"), a, b);
+  auto d = ops::_Retval(scope.WithOpName("D"), c, 0);
+  TF_RETURN_IF_ERROR(scope.ToGraph(graph.get()));
+  return graph;
+}
+
+StatusOr<FunctionDef> SampleFunctionAddXY(const std::string& name) {
+  TF_ASSIGN_OR_RETURN(auto graph, SampleGraphAddXY());
+  FunctionDef fdef;
+  TF_RETURN_IF_ERROR(GraphToFunctionDef(*graph, name, &fdef));
+  return fdef;
+}
+
+Status GetDevices(const tensorflow::tfd::KernelFallbackCompatRequestState*
+                      fallback_request_state,
+                  Device** cpu_device,
+                  absl::flat_hash_map<int, Device*>& gpu_devices) {
+  *cpu_device = fallback_request_state->device_manager().HostCPU();
+  if (!*cpu_device) {
+    return absl::InternalError(
+        "Fallback request state must have a valid host cpu device.");
+  }
+  for (Device* device :
+       fallback_request_state->device_manager().ListDevices()) {
+    if (device->device_type() != DEVICE_GPU) continue;
+    if (!gpu_devices.try_emplace(device->parsed_name().id, device).second) {
+      return absl::InternalError(absl::StrCat(
+          "A device with the same device ID already exists when adding ",
+          device->name()));
+    }
+  }
+  if (gpu_devices.empty()) {
+    return absl::InternalError("No GPU device is found.");
+  }
+  for (const auto& [id, device] : gpu_devices) {
+    if (id >= gpu_devices.size()) {
+      return absl::InternalError("Device IDs are not consecutive.");
+    }
+  }
+  return OkStatus();
+}
+
+template <typename T>
+Tensor CreateTensor(const TensorShape& input_shape,
+                    gtl::ArraySlice<T> input_data,
+                    Allocator* allocator = nullptr) {
+  Tensor tensor(DataTypeToEnum<T>::value, input_shape);
+  test::FillValues<T>(&tensor, input_data);
+  return tensor;
+}
+
+class GpuRunnerTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Create FallbackState.
+    tensorflow::SessionOptions session_options;
+    TF_ASSERT_OK_AND_ASSIGN(FunctionDef fdef,
+                            SampleFunctionAddXY(kFunctionName));
+    tensorflow::FunctionDefLibrary fdef_lib;
+    *fdef_lib.add_function() = fdef;
+    TF_ASSERT_OK_AND_ASSIGN(fallback_state_, tfrt_stub::FallbackState::Create(
+                                                 session_options, fdef_lib));
+
+    // Create KernelFallbackCompatRequestState.
+    std::function<void(std::function<void()>)> runner =
+        [](const std::function<void()>& f) { f(); };
+    tfrt_stub::OpKernelRunnerTable runner_table;
+    tfd::FallbackResourceArray resource_array;
+    fallback_request_state_ =
+        std::make_unique<tfd::KernelFallbackCompatRequestState>(
+            &runner, &fallback_state_->device_manager(), /*step_id=*/0,
+            &runner_table, &resource_array,
+            /*user_intra_op_threadpool=*/nullptr,
+            /*model_metadata=*/std::nullopt,
+            &fallback_state_->process_function_library_runtime());
+
+    // Create execution context.
+    auto host_allocator = tfrt::CreateMallocAllocator();
+    auto work_queue = tfrt::CreateMultiThreadedWorkQueue(
+        /*num_threads=*/2, /*num_blocking_threads=*/2);
+
+    host_context_ = std::make_unique<tfrt::HostContext>(
+        [&](const tfrt::DecodedDiagnostic& diag) {}, std::move(host_allocator),
+        std::move(work_queue));
+    tfrt::RequestContextBuilder req_ctx_builder =
+        tfrt::RequestContextBuilder(host_context_.get(), nullptr);
+    tfrt::Expected<tfrt::RCReference<tfrt::RequestContext>> req_ctx(
+        std::move(req_ctx_builder).build());
+    ASSERT_TRUE(!!req_ctx);
+    exec_ctx_ = std::make_unique<tfrt::ExecutionContext>(std::move(*req_ctx));
+
+    // Create a gpu runner.
+    auto policy = std::make_unique<tsl::RoundRobinPolicy>();
+    serving_device_selector_ = std::make_unique<GpuServingDeviceSelector>(
+        kNumVirtualGpuDevices, std::move(policy));
+    gpu_runner_ = std::make_unique<GpuRunner>(serving_device_selector_.get());
+  }
+
+  std::unique_ptr<tfrt_stub::FallbackState> fallback_state_;
+  std::unique_ptr<tfd::KernelFallbackCompatRequestState>
+      fallback_request_state_;
+  std::unique_ptr<tfrt::HostContext> host_context_;
+  std::unique_ptr<tfrt::ExecutionContext> exec_ctx_;
+  std::unique_ptr<GpuServingDeviceSelector> serving_device_selector_;
+  std::unique_ptr<GpuRunner> gpu_runner_;
+};
+
+TEST_F(GpuRunnerTest, Basic) {
+  // Construct GpuRunInputs.
+  GpuRunInputs run_inputs;
+
+  llvm::SmallVector<tfrt_stub::FallbackTensor> args;
+  Tensor tensor1 = CreateTensor<int32>(TensorShape({1, 2}), {1, 2});
+  Tensor tensor2 = CreateTensor<int32>(TensorShape({1, 2}), {3, 4});
+  args.push_back(tfrt_stub::FallbackTensor(tensor1));
+  args.push_back(tfrt_stub::FallbackTensor(tensor2));
+  run_inputs.args = &args;
+
+  run_inputs.num_outputs = 1;
+  run_inputs.resource_indices = tfrt::ArrayRef<int64_t>(0);
+  run_inputs.used_output_indices = tfrt::ArrayRef<int64_t>(0);
+  run_inputs.func_name = kFunctionName;
+
+  absl::flat_hash_map<int, Device*> gpu_devices;
+  ASSERT_OK(GetDevices(fallback_request_state_.get(), &run_inputs.cpu_device,
+                       gpu_devices));
+  run_inputs.gpu_devices = &gpu_devices;
+  run_inputs.fallback_request_state = fallback_request_state_.get();
+  run_inputs.exec_ctx = exec_ctx_.get();
+
+  // Run the input.
+  TF_ASSERT_OK_AND_ASSIGN(
+      llvm::SmallVector<tfrt::AsyncValueRef<tfrt_stub::FallbackTensor>> outputs,
+      gpu_runner_->Run(run_inputs));
+
+  llvm::SmallVector<tfrt::RCReference<tfrt::AsyncValue>, 4> outputs_to_wait;
+  for (const auto& output : outputs) {
+    if (!output.IsAvailable()) {
+      outputs_to_wait.push_back(output.CopyRCRef());
+    }
+  }
+  exec_ctx_->host()->Await(outputs_to_wait);
+
+  ASSERT_EQ(outputs.size(), 1);
+  auto expected = CreateTensor<int32>(TensorShape({1, 2}), {4, 6});
+  test::ExpectTensorEqual<int32>(expected, outputs[0].get().tensor());
+}
+
+}  // namespace
+}  // namespace gpu
+}  // namespace tensorflow
+#endif

tensorflow/core/tfrt/gpu/kernel/gpurt_kernels.cc

@@ -0,0 +1,243 @@
+/* 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 <cstdint>
+#include <optional>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/STLExtras.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h"
+#include "tensorflow/core/runtime_fallback/kernel/tensor_util.h"
+#include "tensorflow/core/tfrt/gpu/kernel/gpu_runner.h"
+#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
+#include "tensorflow/core/tfrt/utils/gpu_variables_table.h"
+#include "tfrt/host_context/async_value_ref.h"  // from @tf_runtime
+#include "tfrt/host_context/attribute_utils.h"  // from @tf_runtime
+#include "tfrt/host_context/execution_context.h"  // from @tf_runtime
+#include "tfrt/host_context/kernel_registry.h"  // from @tf_runtime
+#include "tfrt/host_context/kernel_utils.h"  // from @tf_runtime
+#include "tfrt/support/forward_decls.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace gpu {
+
+namespace {
+using tfrt_stub::FallbackTensor;
+
+// Transfers `tensor` from `src_device` to `dst_device`.
+tfrt::AsyncValueRef<FallbackTensor> TransferTensor(
+    const tfrt::ExecutionContext& exec_ctx, const FallbackTensor& tensor,
+    Device* src_device, Device* dst_device) {
+  const tensorflow::Tensor& src = tensor.tensor();
+  return tfd::TransferTensorToDevice<FallbackTensor>(exec_ctx, src, src_device,
+                                                     dst_device);
+}
+
+struct Devices {
+  Device* cpu_device = nullptr;
+  // A map from Device ID to the Device, which should contain IDs [0, num_gpus).
+  absl::flat_hash_map<int, Device*> gpu_devices;
+};
+
+// Gets CPU and GPU devices from the fallback state. Here we assume the devices
+// saved in the fallback state remain the same for all requests.
+const absl::StatusOr<Devices>& GetDevices(
+    const tfrt::ExecutionContext& exec_ctx) {
+  static const auto* const result = [&]() -> absl::StatusOr<Devices>* {
+    Devices devices;
+    tfrt::RequestContext* req_ctx = exec_ctx.request_ctx();
+    const auto* fallback_request_state =
+        req_ctx->GetDataIfExists<tfd::KernelFallbackCompatRequestState>();
+    if (!fallback_request_state) {
+      return new absl::StatusOr<Devices>(
+          absl::InternalError("Fallback request state is not found."));
+    }
+
+    devices.cpu_device = fallback_request_state->device_manager().HostCPU();
+    if (!devices.cpu_device) {
+      return new absl::StatusOr<Devices>(absl::InternalError(
+          "Fallback request state must have a valid host cpu device."));
+    }
+    for (Device* device :
+         fallback_request_state->device_manager().ListDevices()) {
+      if (device->device_type() == DEVICE_GPU) {
+        if (!devices.gpu_devices.try_emplace(device->parsed_name().id, device)
+                 .second) {
+          return new absl::StatusOr<Devices>(absl::InternalError(absl::StrCat(
+              "A device with the same device ID already exists when adding ",
+              device->name())));
+        }
+      }
+    }
+    if (devices.gpu_devices.empty()) {
+      return new absl::StatusOr<Devices>(
+          absl::InternalError("No GPU device is found."));
+    }
+    for (const auto& [id, device] : devices.gpu_devices) {
+      if (id >= devices.gpu_devices.size()) {
+        return new absl::StatusOr<Devices>(
+            absl::InternalError("Device IDs are not consecutive."));
+      }
+    }
+    return new absl::StatusOr<Devices>(std::move(devices));
+  }();
+  return *result;
+}
+
+// Kernel for transferring `tensor` from host to device.
+// This only supports a single GPU device.
+tfrt::AsyncValueRef<tfrt_stub::FallbackTensor> TransferToDevice(
+    const tfrt_stub::FallbackTensor& tensor,
+    const tfrt::ExecutionContext& exec_ctx) {
+  const absl::StatusOr<Devices>& devices = GetDevices(exec_ctx);
+  if (!devices.ok()) {
+    return tfrt::MakeErrorAsyncValueRef(
+        absl::InternalError(devices.status().message()));
+  }
+  return TransferTensor(exec_ctx, tensor, devices->cpu_device,
+                        devices->gpu_devices.at(0));
+}
+
+// Kernel for transferring `tensor` from device to host.
+// This only supports a single GPU device.
+tfrt::AsyncValueRef<tfrt_stub::FallbackTensor> TransferFromDevice(
+    const tfrt_stub::FallbackTensor& tensor,
+    const tfrt::ExecutionContext& exec_ctx) {
+  const absl::StatusOr<Devices>& devices = GetDevices(exec_ctx);
+  if (!devices.ok()) {
+    return tfrt::MakeErrorAsyncValueRef(
+        absl::InternalError(devices.status().message()));
+  }
+  return TransferTensor(exec_ctx, tensor, devices->gpu_devices.at(0),
+                        devices->cpu_device);
+}
+
+// Kernel for transferring `variable` from host to device. If it has been
+// transferred, the variable will be returned from the variable cache.
+// This only supports a single GPU device.
+tfrt::AsyncValueRef<tfrt_stub::FallbackTensor> MaybeTransferVariable(
+    const tfrt_stub::FallbackTensor& variable,
+    const tfrt::ExecutionContext& exec_ctx) {
+  // For now, we only consider a single GPU device.
+  const int kCopyIndex = 0;
+  auto vars_table = exec_ctx.resource_context()
+                        ->GetOrCreateResource<tfrt::gpu::GpuVariablesTable>(
+                            tfrt::gpu::kGpuVariablesTableResourceName);
+  auto cached_device_variable =
+      vars_table->GetDeviceVariable(variable, kCopyIndex);
+  if (cached_device_variable) {
+    return cached_device_variable.CopyRef();
+  }
+
+  // The variable has not been transferred, so we transfer the variable and save
+  // the device copy in the variable table.
+  const absl::StatusOr<Devices>& devices = GetDevices(exec_ctx);
+  if (!devices.ok()) {
+    return tfrt::MakeErrorAsyncValueRef(
+        absl::InternalError(devices.status().message()));
+  }
+  auto device_variable = TransferTensor(exec_ctx, variable, devices->cpu_device,
+                                        devices->gpu_devices.at(0));
+  if (device_variable.IsError()) return device_variable;
+
+  vars_table->AddOrUpdateDeviceVariable(variable, kCopyIndex,
+                                        std::move(device_variable));
+  return vars_table->GetDeviceVariable(variable, kCopyIndex).CopyRef();
+}
+
+// Kernel for `gpurt.compile_and_execute`.
+// - `args` is a list of AsyncValue* tensors.
+// - `func_name` is the name of the XLA function in the function library.
+// - `resource_indices` are the indices of args that are variables.
+// - `used_output_indices` are the indices of outputs that will be used.
+void CompileAndExecute(tfrt::RemainingArguments args,
+                       tfrt::RemainingResults results,
+                       tfrt::StringAttribute func_name,
+                       tfrt::ArrayAttribute<int64_t> resource_indices,
+                       tfrt::ArrayAttribute<int64_t> used_output_indices,
+                       tfrt::KernelErrorHandler error_handler,
+                       const tfrt::ExecutionContext& exec_ctx) {
+  GpuRunInputs run_inputs;
+  run_inputs.args.reserve(args.size());
+  for (int i = 0; i < args.size(); ++i) {
+    run_inputs.args.push_back(args[i]->get<tfrt_stub::FallbackTensor>());
+  }
+  run_inputs.num_outputs = results.size();
+  run_inputs.resource_indices = {resource_indices.data().begin(),
+                                 resource_indices.data().end()};
+  run_inputs.used_output_indices = {used_output_indices.data().begin(),
+                                    used_output_indices.data().end()};
+  run_inputs.func_name = func_name.str();
+
+  const absl::StatusOr<Devices>& devices = GetDevices(exec_ctx);
+  if (!devices.ok()) {
+    error_handler.ReportError(devices.status().message());
+    return;
+  }
+  run_inputs.cpu_device = devices->cpu_device;
+  run_inputs.gpu_devices = devices->gpu_devices;
+
+  tfrt::RequestContext* req_ctx = exec_ctx.request_ctx();
+  const auto* fallback_request_state =
+      req_ctx->GetDataIfExists<tfd::KernelFallbackCompatRequestState>();
+  if (!fallback_request_state) {
+    error_handler.ReportError("Fallback request state is not found.");
+    return;
+  }
+  run_inputs.fallback_request_state = fallback_request_state;
+  run_inputs.host_ctx = exec_ctx.host();
+
+  // Get GpuRunner from the resource context.
+  std::optional<GpuRunner*> gpu_runner =
+      exec_ctx.resource_context()->GetResource<GpuRunner>(
+          tensorflow::gpu::kGpuRunnerResourceName);
+  if (!gpu_runner) {
+    error_handler.ReportError("Missing GpuRunner in ResourceContext.");
+    return;
+  }
+
+  auto fallback_tensor_results = (*gpu_runner)->Run(std::move(run_inputs));
+  if (!fallback_tensor_results.ok()) {
+    error_handler.ReportError(fallback_tensor_results.status().message());
+    return;
+  }
+  for (auto it : llvm::zip(results.values(), *fallback_tensor_results)) {
+    std::get<0>(it) = std::move(std::get<1>(it));
+  }
+}
+
+}  // namespace
+
+void RegisterGpurtKernels(tfrt::KernelRegistry* registry) {
+  registry->AddKernel("gpurt.transfer_to_device",
+                      TFRT_KERNEL(TransferToDevice));
+  registry->AddKernel("gpurt.transfer_from_device",
+                      TFRT_KERNEL(TransferFromDevice));
+  registry->AddKernel("gpurt.maybe_transfer_variable",
+                      TFRT_KERNEL(MaybeTransferVariable));
+  registry->AddKernel("gpurt.compile_and_execute",
+                      TFRT_KERNEL(CompileAndExecute));
+}
+
+TFRT_STATIC_KERNEL_REGISTRATION(RegisterGpurtKernels);
+
+}  // namespace gpu
+}  // namespace tensorflow

tensorflow/core/tfrt/gpu/kernel/tfrt_gpu_init.cc

@@ -0,0 +1,54 @@
+/* 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/core/tfrt/gpu/kernel/tfrt_gpu_init.h"
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "xla/tsl/framework/serving_device_selector_policies.h"
+#include "tensorflow/core/common_runtime/gpu/gpu_serving_device_selector.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/gpu/kernel/gpu_runner.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace gpu {
+
+absl::Status InitTfrtGpu(const GpuRunnerOptions& options,
+                         tensorflow::tfrt_stub::Runtime& runtime) {
+  auto policy = std::make_unique<tsl::RoundRobinPolicy>();
+  auto serving_device_selector =
+      std::make_unique<tensorflow::gpu::GpuServingDeviceSelector>(
+          options.num_gpu_streams, std::move(policy));
+
+  // We need to move `serving_device_selector` to the heap here, as
+  // `AddCreateRuntimeResourceFn` requires a copyable callback.
+  auto shared_serving_device_selector =
+      std::shared_ptr<tensorflow::gpu::GpuServingDeviceSelector>(
+          serving_device_selector.release());
+  runtime.AddCreateRuntimeResourceFn(
+      [serving_device_selector = std::move(shared_serving_device_selector)](
+          tfrt::ResourceContext* resource_ctx) mutable {
+        resource_ctx->CreateResource<tensorflow::gpu::GpuRunner>(
+            tensorflow::gpu::kGpuRunnerResourceName,
+            serving_device_selector.get());
+      });
+  return absl::OkStatus();
+}
+
+}  // namespace gpu
+}  // namespace tensorflow

tensorflow/core/tfrt/gpu/kernel/tfrt_gpu_init.h

@@ -0,0 +1,37 @@
+/* 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_CORE_TFRT_GPU_KERNEL_TFRT_GPU_INIT_H_
+#define TENSORFLOW_CORE_TFRT_GPU_KERNEL_TFRT_GPU_INIT_H_
+#include "absl/status/status.h"
+#include "xla/tsl/framework/serving_device_selector_policies.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+
+namespace tensorflow {
+namespace gpu {
+
+struct GpuRunnerOptions {
+  int num_gpu_streams = 1;
+  tsl::ServingDeviceSelectorPolicy serving_selector_policy =
+      tsl::ServingDeviceSelectorPolicy::kRoundRobin;
+};
+
+absl::Status InitTfrtGpu(const GpuRunnerOptions& options,
+                         tensorflow::tfrt_stub::Runtime& runtime);
+
+}  // namespace gpu
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_GPU_KERNEL_TFRT_GPU_INIT_H_

tensorflow/core/tfrt/graph_executor/BUILD

@@ -0,0 +1,267 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_python//python:proto.bzl", "py_proto_library")
+load("//tensorflow:tensorflow.bzl", "if_google", "tf_cc_shared_object", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    # TODO(b/476497239): Remove the old JSV codebase once the migration is done.
+    packages = [
+        "//tensorflow/core/tfrt/...",
+        "//tensorflow/core/tfrt/graph_executor/python/...",
+    ],
+)
+
+cc_library(
+    name = "graph_execution_options",
+    srcs = ["graph_execution_options.cc"],
+    hdrs = ["graph_execution_options.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":config",
+        "//tensorflow/compiler/mlir/tfrt:tfrt_compile_options",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "//tensorflow/core/tfrt/runtime:work_queue_interface",
+        "//tensorflow/core/tfrt/utils:bridge_graph_analysis",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
+        "@com_google_absl//absl/types:optional",
+    ],
+)
+
+# For bazel python deps only
+tf_cc_shared_object(
+    name = "graph_execution_options.so",
+    deps = [":graph_execution_options"],
+)
+
+cc_library(
+    name = "executable_context",
+    hdrs = ["executable_context.h"],
+    deps = [
+        "//tensorflow/core/tfrt/mlrt/bytecode",
+        "//tensorflow/core/tfrt/mlrt/interpreter:context",
+        "@tf_runtime//:bef",
+        "@tf_runtime//:befexecutor",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:support",
+    ],
+)
+
+cc_library(
+    name = "graph_executor",
+    srcs = ["graph_executor.cc"],
+    hdrs = ["graph_executor.h"],
+    deps = [
+        ":executable_context",
+        ":export_mlir",
+        ":graph_execution_options",
+        ":sync_resource_state",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:error_util",
+        "//tensorflow/compiler/mlir/tensorflow:import_model",
+        "//tensorflow/compiler/mlir/tensorflow:mlir_roundtrip_flags",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:graph_to_tf_executor",
+        "//tensorflow/compiler/mlir/tfrt:backend_compiler",
+        "//tensorflow/compiler/mlir/tfrt:import_model",
+        "//tensorflow/compiler/mlir/tfrt:tfrt_compile_options",
+        "//tensorflow/compiler/mlir/tfrt:transforms/update_op_cost_in_tfrt_mlir",
+        "//tensorflow/compiler/mlir/tfrt/transforms/mlrt:import_model",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/common_runtime:core_cpu_internal",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/platform:errors",
+        "//tensorflow/core/platform:path",
+        "//tensorflow/core/profiler/lib:connected_traceme",
+        "//tensorflow/core/profiler/lib:traceme_encode",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "//tensorflow/core/public:version",
+        "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_compat_request_state",
+        "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_execute_compat",
+        "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_utils",
+        "//tensorflow/core/tfrt/common:metrics",
+        "//tensorflow/core/tfrt/fallback:cost_recorder",
+        "//tensorflow/core/tfrt/fallback:fallback_state",
+        "//tensorflow/core/tfrt/fallback:op_kernel_runner",
+        "//tensorflow/core/tfrt/mlrt/bytecode",
+        "//tensorflow/core/tfrt/mlrt/bytecode:executable",
+        "//tensorflow/core/tfrt/mlrt/bytecode:function",
+        "//tensorflow/core/tfrt/mlrt/interpreter:context",
+        "//tensorflow/core/tfrt/mlrt/interpreter:execute",
+        "//tensorflow/core/tfrt/mlrt/interpreter:value",
+        "//tensorflow/core/tfrt/mlrt/kernel:context",
+        "//tensorflow/core/tfrt/runtime",
+        "//tensorflow/core/tfrt/runtime:step_id",
+        "//tensorflow/core/tfrt/runtime:stream",
+        "//tensorflow/core/tfrt/runtime:work_queue_interface",
+        "//tensorflow/core/tfrt/stubs:tfrt_native_lowering_stub",
+        "//tensorflow/core/tfrt/utils",
+        "//tensorflow/core/tfrt/utils:fallback_tensor",
+        "//tensorflow/core/tfrt/utils:tfrt_graph_execution_state",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncExtensions",
+        "@llvm-project//mlir:IR",
+        "@tf_runtime//:basic_kernels_alwayslink",
+        "@tf_runtime//:bef",
+        "@tf_runtime//:befexecutor",
+        "@tf_runtime//:core_runtime",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:mlirtobef",
+        "@tf_runtime//:support",
+        "@tsl//tsl/platform:refcount",
+        "@tsl//tsl/platform:status",
+        "@tsl//tsl/profiler/lib:traceme",
+        "@xla//xla/tsl/concurrency:async_value",
+        "@xla//xla/tsl/concurrency:ref_count",
+    ] + if_google(
+        [
+            "//learning/brain/experimental/tfrt/native_lowering/stubs:tfrt_native_lowering_impl",
+        ],
+    ),
+)
+
+tf_cc_test(
+    name = "graph_executor_test",
+    srcs = ["graph_executor_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        ":config",
+        ":graph_execution_options",
+        ":graph_executor",
+        "//tensorflow/cc:array_ops",
+        "//tensorflow/cc:cc_ops",
+        "//tensorflow/cc:const_op",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_lite",
+        "//tensorflow/core:framework_types_hdr",
+        "//tensorflow/core:test",
+        "//tensorflow/core/framework:common_shape_fns",
+        "//tensorflow/core/framework:graph_proto_cc",
+        "//tensorflow/core/framework:op",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/grappler/utils:grappler_test",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:statusor",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "//tensorflow/core/runtime_fallback/kernel:kernel_fallback_compat_request_state",
+        "//tensorflow/core/tfrt/fallback:fallback_state",
+        "//tensorflow/core/tfrt/fallback:op_kernel_runner",
+        "//tensorflow/core/tfrt/mlrt/interpreter:context",
+        "//tensorflow/core/tfrt/mlrt/interpreter:value",
+        "//tensorflow/core/tfrt/mlrt/kernel",
+        "//tensorflow/core/tfrt/saved_model:saved_model_testutil",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/time",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:tensor",
+        "@tf_runtime//cpp_tests:common",
+        "@tsl//tsl/platform:status",
+        "@tsl//tsl/platform:statusor",
+    ] + if_google(
+        [
+            "//learning/brain/experimental/tfrt/native_lowering/kernels:kernels_alwayslink",
+        ],
+    ),
+)
+
+cc_library(
+    name = "config",
+    srcs = ["config.cc"],
+    hdrs = ["config.h"],
+    compatible_with = get_compatible_with_portable(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":config_proto_cc",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_protobuf//:protobuf",
+    ],
+)
+
+tf_proto_library(
+    name = "config_proto",
+    srcs = ["config.proto"],
+    visibility = ["//visibility:public"],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "config_proto_py_pb2",
+#     visibility = ["//visibility:public"],
+#     deps = [":config_proto"],
+# )
+# copybara:uncomment_end
+
+tf_proto_library(
+    name = "test_config_proto",
+    testonly = True,
+    srcs = ["test_config.proto"],
+    visibility = if_google(
+        [":friends"],
+        ["//visibility:public"],
+    ),
+)
+
+tf_cc_test(
+    name = "config_test",
+    srcs = ["config_test.cc"],
+    deps = [
+        ":config",
+        ":config_proto_cc",
+        ":test_config_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:status_matchers",
+        "@tsl//tsl/platform:statusor",
+        "@xla//xla/tsl/lib/core:status_test_util",
+    ],
+)
+
+cc_library(
+    name = "sync_resource_state",
+    hdrs = ["sync_resource_state.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core/tfrt/utils:any_ptr",
+        "@tf_runtime//:tensor",
+    ],
+)
+
+cc_library(
+    name = "export_mlir",
+    hdrs = ["export_mlir.h"],
+    deps = ["@llvm-project//mlir:IR"],
+)

tensorflow/core/tfrt/graph_executor/config.cc

@@ -0,0 +1,40 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/tfrt/graph_executor/config.h"
+
+#include <string>
+#include <utility>
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+absl::StatusOr<RuntimeConfig> RuntimeConfig::CreateFromProto(
+    RuntimeConfigProto proto) {
+  RuntimeConfig model_config;
+  model_config.proto_ = std::move(proto);
+  size_t i = 0;
+  for (const auto& any : model_config.proto_.config()) {
+    std::string full_name;
+    if (!::google::protobuf::Any::ParseAnyTypeUrl(any.type_url(), &full_name)) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Invalid Any proto type url: ", any.type_url()));
+    }
+    model_config.map_[full_name] = i++;
+  }
+  return model_config;
+}
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/graph_executor/config.h

@@ -0,0 +1,84 @@
+/* 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_CORE_TFRT_GRAPH_EXECUTOR_CONFIG_H_
+#define TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_CONFIG_H_
+
+#include <string>
+
+#include "google/protobuf/any.pb.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/core/tfrt/graph_executor/config.pb.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+// The helper class for building RuntimeConfigProto and retrieving configs of
+// certain types from the RuntimeConfigProto.
+class RuntimeConfig {
+ public:
+  RuntimeConfig() = default;
+
+  static absl::StatusOr<RuntimeConfig> CreateFromProto(
+      RuntimeConfigProto proto);
+
+  template <typename ConcreteProto>
+  absl::Status Add(const ConcreteProto& config) {
+    const auto& full_name = config.GetDescriptor()->full_name();
+    if (map_.contains(full_name)) {
+      return absl::AlreadyExistsError(
+          absl::StrCat(full_name, " already exists in ModelConfig."));
+    }
+
+    size_t id = proto_.config_size();
+    if (!proto_.add_config()->PackFrom(config)) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Failed to pack proto to Any: ", full_name));
+    }
+    map_[full_name] = id;
+    return absl::OkStatus();
+  }
+
+  template <typename ConcreteProto>
+  absl::StatusOr<ConcreteProto> Get() const {
+    const auto& full_name = ConcreteProto::GetDescriptor()->full_name();
+    auto iter = map_.find(full_name);
+
+    if (iter == map_.end()) {
+      return absl::NotFoundError(
+          absl::StrCat(full_name, " not found in ModelConfig."));
+    }
+
+    ConcreteProto config;
+    if (!proto_.config().at(iter->second).UnpackTo(&config)) {
+      return absl::DataLossError(
+          absl::StrCat("Failed to unpack proto: ", full_name));
+    }
+    return config;
+  }
+
+  const RuntimeConfigProto& ToProto() const { return proto_; }
+
+ private:
+  RuntimeConfigProto proto_;
+  absl::flat_hash_map<std::string, size_t> map_;
+};
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_CONFIG_H_

tensorflow/core/tfrt/graph_executor/config.proto

@@ -0,0 +1,26 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tfrt_stub;
+
+import "google/protobuf/any.proto";
+
+// The serialization format for custom model configs. Though it is using
+// repeated Any protos, only one object of a proto type is allowed in the list.
+message RuntimeConfigProto {
+  repeated google.protobuf.Any config = 1;
+}

tensorflow/core/tfrt/graph_executor/config_test.cc

@@ -0,0 +1,99 @@
+/* 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/core/tfrt/graph_executor/config.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/tfrt/graph_executor/config.pb.h"
+#include "tensorflow/core/tfrt/graph_executor/test_config.pb.h"
+#include "tsl/platform/status_matchers.h"
+#include "tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+namespace {
+
+TEST(ConfigTest, Basic) {
+  TestConfig1 expected_test_config1;
+  expected_test_config1.set_tag("test config1");
+  TestConfig2 expected_test_config2;
+  expected_test_config2.set_tag("test config2");
+
+  RuntimeConfig runtime_config;
+
+  TF_ASSERT_OK(runtime_config.Add(expected_test_config2));
+  TF_ASSERT_OK(runtime_config.Add(expected_test_config1));
+
+  auto test_config1 = runtime_config.Get<TestConfig1>();
+  TF_ASSERT_OK(test_config1.status());
+  auto test_config2 = runtime_config.Get<TestConfig2>();
+  TF_ASSERT_OK(test_config2.status());
+
+  EXPECT_EQ(test_config1->tag(), "test config1");
+  EXPECT_EQ(test_config2->tag(), "test config2");
+}
+
+TEST(ConfigTest, Load) {
+  TestConfig1 expected_test_config1;
+  expected_test_config1.set_tag("test config1");
+  TestConfig2 expected_test_config2;
+  expected_test_config2.set_tag("test config2");
+
+  RuntimeConfigProto runtime_config_proto;
+  runtime_config_proto.add_config()->PackFrom(expected_test_config1);
+  runtime_config_proto.add_config()->PackFrom(expected_test_config2);
+
+  TF_ASSERT_OK_AND_ASSIGN(RuntimeConfig runtime_config,
+                          RuntimeConfig::CreateFromProto(runtime_config_proto));
+
+  auto test_config1 = runtime_config.Get<TestConfig1>();
+  TF_ASSERT_OK(test_config1.status());
+  auto test_config2 = runtime_config.Get<TestConfig2>();
+  TF_ASSERT_OK(test_config2.status());
+
+  EXPECT_EQ(test_config1->tag(), "test config1");
+  EXPECT_EQ(test_config2->tag(), "test config2");
+}
+
+TEST(ConfigTest, NotFound) {
+  TestConfig1 expected_test_config1;
+  expected_test_config1.set_tag("test config1");
+
+  RuntimeConfigProto runtime_config_proto;
+  runtime_config_proto.add_config()->PackFrom(expected_test_config1);
+
+  TF_ASSERT_OK_AND_ASSIGN(RuntimeConfig runtime_config,
+                          RuntimeConfig::CreateFromProto(runtime_config_proto));
+
+  EXPECT_THAT(runtime_config.Get<TestConfig2>(),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(ConfigTest, Duplicate) {
+  TestConfig1 expected_test_config1;
+  expected_test_config1.set_tag("test config1");
+
+  RuntimeConfig runtime_config;
+
+  TF_ASSERT_OK(runtime_config.Add(expected_test_config1));
+  EXPECT_THAT(runtime_config.Add(expected_test_config1),
+              absl_testing::StatusIs(absl::StatusCode::kAlreadyExists));
+}
+
+}  // namespace
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/graph_executor/executable_context.h

@@ -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.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_EXECUTABLE_CONTEXT_H_
+#define TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_EXECUTABLE_CONTEXT_H_
+
+#include <memory>
+#include <utility>
+
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/context.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/bef_executor/bef_file.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 {
+namespace tfrt_stub {
+
+// Stores executable-related data.
+struct ExecutableContext {
+  ExecutableContext(mlrt::bc::Buffer bytecode_buffer,
+                    std::unique_ptr<mlrt::LoadedExecutable> bytecode_executable)
+      : bytecode_buffer(std::move(bytecode_buffer)),
+        bytecode_executable(std::move(bytecode_executable)) {}
+
+  ExecutableContext(tfrt::BefBuffer bef,
+                    tfrt::RCReference<tfrt::BEFFile> bef_file)
+      : bef(std::move(bef)), bef_file(std::move(bef_file)) {}
+
+  bool IsForMlrt() const { return bytecode_executable != nullptr; }
+
+  // Only one set of values will be filled.
+
+  // For the MLRT path.
+  mlrt::bc::Buffer bytecode_buffer;
+  std::unique_ptr<mlrt::LoadedExecutable> bytecode_executable;
+
+  // For the TFRT path.
+  tfrt::BefBuffer bef;
+  tfrt::RCReference<tfrt::BEFFile> bef_file;
+
+  // There are some resources that need re-creating when the executable is
+  // re-created, so a resource context is stored along with the executable.
+  // This resource context is meant to be passed to the op kernels for their
+  // references. See the comment above `GraphExecutor::resource_context_`
+  // about the todo to merge that resource context with this one.
+  tfrt::ResourceContext resource_context;
+};
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_EXECUTABLE_CONTEXT_H_

tensorflow/core/tfrt/graph_executor/export_mlir.h

@@ -0,0 +1,64 @@
+/* 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_CORE_TFRT_GRAPH_EXECUTOR_EXPORT_MLIR_H_
+#define TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_EXPORT_MLIR_H_
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+class XsymbolUploader {
+ public:
+  virtual ~XsymbolUploader() = default;
+
+  virtual std::string MaybeUploadMlirToXsymbol(mlir::ModuleOp module) {
+    return "";
+  }
+};
+
+class XsymbolUploaderRegistry {
+ public:
+  XsymbolUploaderRegistry()
+      : xsymbol_uploader_(std::make_unique<XsymbolUploader>()) {}
+
+  void Register(std::unique_ptr<XsymbolUploader> xsymbol_uploader) {
+    xsymbol_uploader_ = std::move(xsymbol_uploader);
+  }
+
+  XsymbolUploader &Get() const { return *xsymbol_uploader_; }
+
+ private:
+  std::unique_ptr<XsymbolUploader> xsymbol_uploader_;
+};
+
+inline XsymbolUploaderRegistry &GetGlobalXsymbolUploaderRegistry() {
+  static auto *const registry = new XsymbolUploaderRegistry;
+  return *registry;
+}
+
+inline std::string MaybeUploadMlirToXsymbol(mlir::ModuleOp module) {
+  return GetGlobalXsymbolUploaderRegistry().Get().MaybeUploadMlirToXsymbol(
+      module);
+}
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_EXPORT_MLIR_H_

tensorflow/core/tfrt/graph_executor/graph_execution_options.cc

@@ -0,0 +1,136 @@
+/* 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/tfrt/graph_executor/graph_execution_options.h"
+
+#include <ostream>
+
+#include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+#include "tensorflow/core/protobuf/rewriter_config.pb.h"
+// TODO(b/200579737): using FunctionRegistry is simpler than the OSS trick.
+#include "tensorflow/core/tfrt/utils/bridge_graph_analysis.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+tensorflow::SessionOptions CreateDefaultSessionOptions(
+    const GraphExecutionOptions& options) {
+  tensorflow::SessionOptions session_options;
+  auto& config = session_options.config;
+
+  *config.mutable_experimental()->mutable_session_metadata() =
+      options.model_metadata;
+
+  *config.mutable_graph_options() = options.compile_options.graph_options;
+
+  config.mutable_graph_options()
+      ->mutable_rewrite_options()
+      ->set_disable_meta_optimizer(!options.compile_options.enable_grappler);
+
+  // The following configs are constant.
+
+  // Setting use_tfrt to true avoids grappler logic that lowers to v1 control
+  // flow. Note that other function inlining (e.g. on StatefulPartitionedCall)
+  // is still enabled.
+  config.mutable_experimental()->set_use_tfrt(true);
+  if (options.enable_grappler_function_optimizer) {
+    config.mutable_graph_options()
+        ->mutable_rewrite_options()
+        ->set_function_optimization(tensorflow::RewriterConfig::ON);
+  } else {
+    config.mutable_graph_options()
+        ->mutable_rewrite_options()
+        ->set_function_optimization(tensorflow::RewriterConfig::OFF);
+  }
+  // Do not skip grappler optimization even for small graphs.
+  config.mutable_graph_options()
+      ->mutable_rewrite_options()
+      ->set_min_graph_nodes(-1);
+
+  if (options.tfrt_gpu_parallelism > 1) {
+    if (!options.compile_options.use_gpu_compile_and_execute_op) {
+      LOG(WARNING)
+          << "tfrt_gpu_parallelism > 1, but fused GPU kernel is not used. "
+             "Non-fused GPU kernel does not support multiple GPU devices.";
+    }
+    config.mutable_gpu_options()
+        ->mutable_experimental()
+        ->set_num_virtual_devices_per_gpu(options.tfrt_gpu_parallelism);
+  }
+  if (options.gpu_system_memory_size_in_mb > 0) {
+    config.mutable_gpu_options()
+        ->mutable_experimental()
+        ->set_gpu_system_memory_size_in_mb(
+            options.gpu_system_memory_size_in_mb);
+  }
+  return session_options;
+}
+
+void UpdateTpuTargetByBridgeCompatibility(
+    tensorflow::tfrt_stub::GraphExecutionOptions& options,
+    const tensorflow::GraphDef& graph_def) {
+  if (options.compile_options.device_target ==
+      tensorflow::TfrtDeviceInfraTarget::kBridgeFallback) {
+    auto s = tfrt::CheckTpuMlirBridgeCompatibility(graph_def);
+    if (!s.ok()) {
+      LOG(INFO)
+          << "TFRT detected Bridge unsupported feature, using TF fallback";
+      options.compile_options.device_target =
+          tensorflow::TfrtDeviceInfraTarget::kTfFallback;
+    } else {
+      options.compile_options.device_target =
+          tensorflow::TfrtDeviceInfraTarget::kTpurt;
+    }
+  }
+
+  // We don't need to check for SPMD fallback for non TFRT TPU path.
+  //
+  // TODO(b/288096487): Clean up the enums to reflect the device target better.
+  // One option is to use a  custom target enum for the opaque backend.
+  if (options.compile_options.device_target !=
+          tensorflow::TfrtDeviceInfraTarget::kCpu &&
+      options.compile_options.device_target !=
+          tensorflow::TfrtDeviceInfraTarget::kGpu) {
+    // TODO(linchai): Once native support for SPMD models is fully rollout,
+    // remove the fallback logic.
+    if (!(tfrt::CheckSpmdGraph(graph_def).ok() ||
+          options.compile_options.tpu_fuse_ops)) {
+      options.compile_options.device_target =
+          tensorflow::TfrtDeviceInfraTarget::kTfFallback;
+    }
+    LOG(INFO) << "TFRT uses device target "
+              << options.compile_options.device_target;
+  }
+}
+
+std::ostream& operator<<(std::ostream& os,
+                         const GraphExecutionOptions& options) {
+  return os << "{" << "run_placer_grappler_on_functions = "
+            << options.run_placer_grappler_on_functions
+            << ", enable_grappler_function_optimizer = "
+            << options.enable_grappler_function_optimizer
+            << ", enable_tfrt_gpu = " << options.enable_tfrt_gpu
+            << ", use_ifrt = " << options.use_ifrt << ", runtime = "
+            << options.runtime
+            // clang-tidy off
+            << ", model_metadata = "
+            << options.model_metadata.DebugString()
+            // clang-tidy on
+            << ", compile_options = " << options.compile_options << "}";
+}
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/graph_executor/graph_execution_options.h

@@ -0,0 +1,174 @@
+/* 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_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTION_OPTIONS_H_
+#define TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTION_OPTIONS_H_
+
+#include <functional>
+#include <optional>
+#include <ostream>
+#include <string>
+
+#include "absl/time/time.h"
+#include "absl/types/optional.h"
+#include "tensorflow/compiler/mlir/tfrt/translate/tfrt_compile_options.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/tfrt/graph_executor/config.h"
+#include "tensorflow/core/tfrt/runtime/work_queue_interface.h"
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+class Runtime;
+
+// General options for graph execution.
+struct GraphExecutionOptions {
+  explicit GraphExecutionOptions(const tensorflow::tfrt_stub::Runtime* rt)
+      : runtime(rt) {
+    DCHECK(runtime);
+  }
+
+  // If true, when creating an optimized subgraph, Placer and Grappler will
+  // also run on the functions.
+  bool run_placer_grappler_on_functions = false;
+
+  // If true, the function optimizer in the grappler will be enabled, and
+  // optimizations like function inlining will be applied.
+  bool enable_grappler_function_optimizer = false;
+
+  // Whether to enable TFRT GPU.
+  bool enable_tfrt_gpu = false;
+
+  // The number of virtual GPUs to create on a physical GPU.
+  int tfrt_gpu_parallelism = 1;
+
+  // if not zero, override the reserved memory space for gpu system.
+  int gpu_system_memory_size_in_mb = 0;
+
+  // Whether to use gpurt.compile_and_execute for GPU.
+  // TODO(b/294895431): Remove the flag and default to the fused op.
+  bool tfrt_use_fused_gpu_op = false;
+
+  // Runtime configuration. Refer to tensorflow::tfrt_stub::Runtime class for
+  // more details. It must not be nullptr;
+  const tensorflow::tfrt_stub::Runtime* runtime = nullptr;
+
+  // Model metadata used for monitoring and tracing.
+  tensorflow::SessionMetadata model_metadata;
+
+  // The model-specific runtime configurations.
+  tensorflow::tfrt_stub::RuntimeConfig runtime_config;
+
+  // TODO(b/266251216): Maybe flip the default value.
+  [[deprecated(
+      "Use CostAnalysisOptions's `CostAnalysisOptions::ONCE` instead")]] bool
+      enable_online_cost_analysis = false;
+
+  // Determines how often op costs are recorded, and how often these costs
+  // are used to re-compile the executable. Note to users: CostAnalysisOptions
+  // is overwritten when `enable_online_cost_analysis = true`.
+  struct CostAnalysisOptions {
+    enum CostAnalysisVersion {
+      kDisabled,
+      kOnce,  // Cost recording and recompilation occurs on the first run only.
+      kPeriodic,  // This is experimental.
+    };
+    CostAnalysisVersion version = kDisabled;
+
+    // Time between resets in Op cost estimates. Upon reset, the executable
+    // will be recompiled.
+    // However, a reset always occurs after the first execution.
+    absl::Duration reset_interval = absl::ZeroDuration();
+
+    // Number of times to record costs before resetting Op cost estimates.
+    // However, a reset always occurs after the first execution.
+    int updates_per_interval = 1;
+  };
+
+  CostAnalysisOptions cost_analysis_options;
+
+  // If true, the MLRT interpreter will be used instead of the BEF executor.
+  // This option is experimental.
+  bool enable_mlrt = false;
+
+  // If true, the IFRT will be used instead of the TPU Runner.
+  // This option is experimental.
+  bool use_ifrt = false;
+
+  tensorflow::TfrtCompileOptions compile_options;
+};
+
+std::ostream& operator<<(std::ostream& os,
+                         const GraphExecutionOptions& options);
+
+// Per-request options for graph execution.
+struct GraphExecutionRunOptions {
+  std::optional<std::chrono::system_clock::time_point> deadline;
+
+  // Priority of the request. Larger number means higher priority.
+  int priority = 0;
+
+  // If true, the input specs will be checked before running, and an error
+  // will be raised upon mismatch.
+  bool validate_input_specs = false;
+
+  // TODO(b/279197040) Remove after b/279197040 is fixed.
+  // If true, the input specs will be checked before running, and an error
+  // will be logged upon mismatch.
+  bool validate_input_specs_dry_run = false;
+
+  // The thread pool used for this run. If it is nullptr, a default one set
+  // in the tensorflow::tfrt_stub::Runtime will be used.
+  tensorflow::tfrt_stub::WorkQueueInterface* work_queue = nullptr;
+
+  // If true, just-in-time host compilation is disabled, and then if the
+  // specified graph is not compiled, the execution will return an error.
+  bool disable_compilation = false;
+
+  std::function<void(absl::flat_hash_map<std::string, tensorflow::Tensor>)>
+      streamed_output_callback;
+
+  // The optional name for debugging purposes. If empty, the runtime will pick a
+  // name e.g. the joined string of input names and output names.
+  std::string name;
+
+  // RPC deadline for lazy task cancellation in the batch scheduler.
+  // Independent of `deadline` which is gated by `enforce_session_run_timeout`.
+  // This deadline is only used for task-level deadline enforcement in the batch
+  // scheduler.
+  //
+  // TODO(b/510447914): Consider reuse `deadline` for this after the
+  // investigation.
+  std::optional<absl::Time> rpc_deadline_for_batching_task_cancellation;
+
+  // Callback to check if the request RPC has been cancelled by the client.
+  std::function<bool()> is_rpc_cancelled_callback;
+};
+
+// Creates the default `SessionOptions` from a `GraphExecutionOptions`.
+// The created `SessionOptions` contains the Grappler configs.
+tensorflow::SessionOptions CreateDefaultSessionOptions(
+    const GraphExecutionOptions& options);
+
+// Updates TPU target to fallback if bridge uncompatible, otherwise TPU runtime.
+void UpdateTpuTargetByBridgeCompatibility(
+    tensorflow::tfrt_stub::GraphExecutionOptions& options,
+    const tensorflow::GraphDef& graph_def);
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTION_OPTIONS_H_

tensorflow/core/tfrt/graph_executor/graph_executor.h

@@ -0,0 +1,419 @@
+/* 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_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTOR_H_
+#define TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTOR_H_
+
+#include <functional>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/check.h"
+#include "absl/strings/string_view.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "absl/types/span.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/translate/mlir_roundtrip_flags.h"
+#include "tensorflow/compiler/mlir/tfrt/backend_compiler.h"
+#include "xla/tsl/concurrency/ref_count.h"
+#include "xla/tsl/lib/monitoring/sampler.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h"
+#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
+#include "tensorflow/core/tfrt/graph_executor/executable_context.h"
+#include "tensorflow/core/tfrt/graph_executor/graph_execution_options.h"
+#include "tensorflow/core/tfrt/graph_executor/sync_resource_state.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/function.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/context.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/value.h"
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+#include "tensorflow/core/tfrt/runtime/stream.h"
+#include "tensorflow/core/tfrt/runtime/work_queue_interface.h"
+#include "tensorflow/core/tfrt/utils/tfrt_graph_execution_state.h"
+#include "tsl/platform/thread_annotations.h"
+#include "tfrt/bef/bef_buffer.h"  // from @tf_runtime
+#include "tfrt/bef_executor/bef_file.h"  // from @tf_runtime
+#include "tfrt/core_runtime/core_runtime.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/request_deadline_tracker.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 {
+namespace tfrt_stub {
+
+// Contains request related info.
+struct RequestInfo {
+  tfrt::RCReference<tfrt::RequestContext> tfrt_request_context;
+  // If this request needs to create a new queue, it is stored here. Otherwise,
+  // it can be nullptr.
+  std::unique_ptr<WorkQueueInterface> request_queue_owner;
+  // The inter-op thread pool to be used for this request, and it must not be
+  // nullptr. If `request_queue_owner` is not nullptr, then `request_queue` is
+  // the raw pointer inside `request_queue_owner`.
+  WorkQueueInterface* request_queue = nullptr;
+  // The task runner used by tensorflow::OpKernel.
+  std::function<void(std::function<void()>)> runner;
+
+  tensorflow::CancellationManager cancellation_manager;
+};
+
+struct SymbolUids {
+  std::string tf_symbol_uid;
+  std::string tfrt_symbol_uid;
+};
+
+// Creates a `RequestInfo` given relative data.
+// Note: `resource_context` is per-graph-executor and
+// `client_graph_resource_context` is per-loaded-client-graph. See the comment
+// above `GraphExecutor::resource_context_` about the todo to merge these two.
+absl::StatusOr<std::unique_ptr<RequestInfo>> CreateRequestInfo(
+    const GraphExecutionOptions& options,
+    const GraphExecutionRunOptions& run_options,
+    tensorflow::tfrt_stub::WorkQueueInterface* work_queue,
+    tfrt::ResourceContext* resource_context,
+    tfrt::ResourceContext* client_graph_resource_context,
+    OpKernelRunnerTable* runner_table,
+    tfd::FallbackResourceArray* resource_array, FallbackState& fallback_state,
+    const ProcessFunctionLibraryRuntime& process_function_library_runtime,
+    CostRecorder* cost_recorder = nullptr);
+
+// Runs on a function given input/output and other info.
+// Note: `resource_context` is per-graph-executor and
+// `client_graph_resource_context` is per-loaded-client-graph. See the comment
+// above `GraphExecutor::resource_context_` about the todo to merge these two.
+//
+// TODO(chky): Refactor this function to take `LoadedClientGraph` instead of
+// having a long list of parameters.
+absl::Status GraphExecutionRunOnFunction(
+    const GraphExecutionOptions& options,
+    const GraphExecutionRunOptions& run_options,
+    absl::string_view signature_name, const SymbolUids& symbol_uids,
+    const tfrt::Function* func, const mlrt::LoadedExecutable* loaded_executable,
+    absl::Span<const tensorflow::Tensor> inputs,
+    std::vector<tensorflow::Tensor>* outputs,
+    tfrt::ResourceContext* resource_context,
+    tfrt::ResourceContext* client_graph_resource_context,
+    OpKernelRunnerTable* runner_table,
+    tfd::FallbackResourceArray* resource_array, const Runtime& runtime,
+    FallbackState& fallback_state,
+    const tensorflow::ProcessFunctionLibraryRuntime&
+        process_function_library_runtime,
+    tfrt::RequestDeadlineTracker* req_deadline_tracker,
+    std::optional<StreamCallbackId> stream_callback_id,
+    CostRecorder* cost_recorder = nullptr);
+
+// Runs a MLRT function for executing tensorflow graphs.
+absl::Status RunMlrtFunction(
+    mlrt::bc::Function function,
+    const mlrt::LoadedExecutable& loaded_executable,
+    const tsl::RCReference<tfrt::RequestContext>& request_context,
+    tfrt::ConcurrentWorkQueue& work_queue,
+    absl::Span<const tensorflow::Tensor> inputs,
+    std::vector<tensorflow::Tensor>* outputs,
+    SyncResourceState* sync_resource_state);
+
+// Loads (if not yet) and runs a subgraph in a graph as per each request.
+class GraphExecutor {
+ public:
+  using Options = GraphExecutionOptions;
+  using RunOptions = GraphExecutionRunOptions;
+
+  // The loading result of a `ClientGraph`.
+  class LoadedClientGraph {
+   public:
+    LoadedClientGraph(std::string name, SymbolUids symbol_uids,
+                      GraphExecutor* graph_executor,
+                      std::unique_ptr<mlir::MLIRContext> mlir_context,
+                      mlir::OwningOpRef<mlir::ModuleOp> tf_mlir_with_op_keys,
+                      mlir::OwningOpRef<mlir::ModuleOp> tfrt_mlir,
+                      std::shared_ptr<ExecutableContext> executable_context,
+                      std::optional<StreamCallbackId> stream_callback_id,
+                      FunctionLibraryDefinition flib_def,
+                      tsl::monitoring::SamplerCell* latency_sampler);
+
+    // Returns this instance's CostRecorder if it is time to update costs,
+    // else returns nullptr. Only allows one non-null return value at a time
+    // in order to provide thread-safety. If do_recompilation becomes `true`,
+    // then recompiles using updated costs occurs.
+    CostRecorder* MaybeGetCostRecorder(absl::Time now, bool* do_recompilation);
+    // Updates the op cost values in this `LoadedClientGraph` with records from
+    // `cost_recorder`.
+    absl::Status UpdateCost(const CostRecorder& cost_recorder,
+                            const Runtime& runtime);
+    // Updates `cost_analysis_data_` to make it accurate for the next execution.
+    // Assumes a cost update occurred this cycle.
+    void UpdateCostAnalysisData(absl::Time now, bool do_recompilation);
+    // Getters.
+    std::shared_ptr<ExecutableContext> executable_context() const {
+      tensorflow::mutex_lock lock(executable_context_mu_);
+      return executable_context_;
+    }
+    absl::string_view name() const { return name_; }
+    const SymbolUids& symbol_uids() const { return symbol_uids_; }
+
+    OpKernelRunnerTable& runner_table() { return runner_table_; }
+    tfd::FallbackResourceArray& resource_array() { return resource_array_; }
+    SyncResourceState& sync_resource_state() { return sync_resource_state_; }
+
+    std::optional<StreamCallbackId> stream_callback_id() const {
+      return stream_callback_id_;
+    }
+
+    const ProcessFunctionLibraryRuntime& process_function_library_runtime()
+        const {
+      return pflr_;
+    }
+    tsl::monitoring::SamplerCell* latency_sampler() { return latency_sampler_; }
+
+   private:
+    std::string name_;
+    SymbolUids symbol_uids_;
+    GraphExecutor* graph_executor_ = nullptr;
+
+    // `mlir_context_` is declared here because the resources declared later may
+    // hold references to the MLIR objects.
+    std::unique_ptr<mlir::MLIRContext> mlir_context_;
+
+    struct CostAnalysisData {
+      mutable tensorflow::mutex mu;
+      // Ensures only one GraphExecutor thread updates costs at a time.
+      bool is_available TF_GUARDED_BY(mu) = false;
+      // Maintains the book-keeping of op costs.
+      std::unique_ptr<CostRecorder> cost_recorder;
+      // For recompilation in MLRT, TFRT respectively.
+      mlir::OwningOpRef<mlir::ModuleOp> tf_mlir_with_op_keys;
+      mlir::OwningOpRef<mlir::ModuleOp> tfrt_mlir;
+      // Start of current cost measurement cycle.
+      absl::Time start_time TF_GUARDED_BY(mu) = absl::Now();
+      // Cost recordings within the current measurement cycle.
+      int num_cost_updates TF_GUARDED_BY(mu) = 0;
+    };
+    CostAnalysisData cost_analysis_data_;
+
+    OpKernelRunnerTable runner_table_;
+    tfd::FallbackResourceArray resource_array_;
+    mutable tensorflow::mutex executable_context_mu_;
+    // Can be updated if online cost analysis is enabled.
+    std::shared_ptr<ExecutableContext> executable_context_
+        TF_GUARDED_BY(executable_context_mu_);
+    SyncResourceState sync_resource_state_;
+
+    std::optional<StreamCallbackId> stream_callback_id_;
+    FunctionLibraryDefinition flib_def_;
+    ProcessFunctionLibraryRuntime pflr_;
+    tsl::monitoring::SamplerCell* latency_sampler_;
+  };
+
+  // A subgraph constructed by specifying input/output tensors.
+  struct ClientGraph {
+    // The human-readable name for the graph, e.g. the signature_name in the
+    // saved model.
+    std::string name;
+    // The feed nodes for the corresponding inputs, but they might not be in the
+    // original order and if there are more than one original inputs mapped to
+    // the same feed node, only one is picked here.
+    tensorflow::GraphImportConfig::InputArrays input_nodes;
+    // The fetch nodes for the outputs, which should be in the original order.
+    std::vector<std::string> output_nodes;
+    // The target nodes that should be run but not returned as outputs.
+    std::vector<std::string> target_nodes;
+  };
+
+  // Creates a `GraphExecutor` given the args.
+  static absl::StatusOr<std::unique_ptr<GraphExecutor>> Create(
+      Options options, std::unique_ptr<FallbackState> fallback_state,
+      std::unique_ptr<tfrt::ResourceContext> resource_context,
+      tensorflow::GraphDef graph_def,
+      std::unique_ptr<mlrt::KernelRegistry> kernel_registry,
+      tensorflow::tfrt_stub::RuntimeConfig* runtime_config = nullptr);
+
+  // Ctor. Public for `Create()`. Do not use directly.
+  GraphExecutor(Options options, std::unique_ptr<FallbackState> fallback_state,
+                std::unique_ptr<tfrt::ResourceContext> resource_context,
+                std::unique_ptr<tensorflow::tfrt_stub::TfrtGraphExecutionState>
+                    graph_execution_state,
+                std::unique_ptr<mlrt::KernelRegistry> kernel_registry,
+                tensorflow::tfrt_stub::RuntimeConfig* runtime_config = nullptr);
+
+  // Runs on the graph according to given input/output.
+  absl::Status Run(
+      const RunOptions& run_options,
+      absl::Span<const std::pair<std::string, tensorflow::Tensor>> inputs,
+      absl::Span<const std::string> output_tensor_names,
+      absl::Span<const std::string> target_tensor_names,
+      std::vector<tensorflow::Tensor>* outputs);
+
+  // Similar as `Run`, but it requires additional input parameters to specify
+  // the `graph_name`, the sorted input/output names and the original indices of
+  // the inputs/outputs. The caller must guarantee that inputs are in the same
+  // order as of `sorted_input_names`. The sorted input/output names are needed
+  // to consistently build the key for looking up the `LoadedClientGraph` in the
+  // cache. The original indices are needed to map the results to the original
+  // inputs/outputs. The `graph_name` will be used to lookup the compiled graph
+  // in the cache. It is usually the signature name of the graph. If it is
+  // empty, a joined name will be constructed from the sorted input/output names
+  // to lookup the `LoadedClientGraph` in the cache.
+  absl::Status RunWithSortedInputsOutputs(
+      const RunOptions& run_options, absl::string_view graph_name,
+      absl::Span<const std::pair<std::string, tensorflow::Tensor>> inputs,
+      absl::Span<const std::string> sorted_input_names,
+      absl::Span<const tensorflow::DataType> sorted_input_dtypes,
+      absl::Span<const std::string> sorted_output_names,
+      absl::Span<const std::string> sorted_target_node_names,
+      absl::Span<const int> input_original_indices,
+      absl::Span<const int> output_original_indices,
+      std::vector<tensorflow::Tensor>* outputs);
+
+  // Runs the graph identified by `graph_name` using the input `inputs` and
+  // stores the output of the execution in `outputs`. It is the client's
+  // responsibility to ensure `graph_name` corresponds to logically different
+  // graphs, since this name is used to lookup compiled graphs in the cache. The
+  // graph is run synchronously with the TFRT interpreter.
+  absl::Status RunWithSyncInterpreter(
+      const std::string& graph_name, absl::Span<mlrt::Value> input_values,
+      absl::Span<const std::string> input_names,
+      absl::Span<const tensorflow::DataType> input_dtypes,
+      absl::Span<const std::string> output_tensor_names,
+      absl::Span<const std::string> target_tensor_names,
+      absl::Span<mlrt::Value> outputs);
+
+  // Extends the current graph by `graph`.
+  absl::Status Extend(const GraphDef& graph);
+
+  tensorflow::tfrt_stub::TfrtGraphExecutionState& graph_execution_state()
+      const {
+    return *graph_execution_state_;
+  }
+
+  // Returns the underlying runtime.
+  const tensorflow::tfrt_stub::Runtime& runtime() const {
+    DCHECK(options_.runtime);
+    return *options_.runtime;
+  }
+
+  tfrt::ResourceContext& resource_context() { return *resource_context_; }
+
+  const Options& options() const { return options_; }
+  const FallbackState& fallback_state() const { return *fallback_state_; }
+  FallbackState& fallback_state() { return *fallback_state_; }
+
+  // Compiles graph for `graph_name` and runs any initializers.
+  absl::Status CompileGraph(
+      const std::string& graph_name,
+      absl::Span<const std::string> input_tensor_names,
+      absl::Span<const tensorflow::DataType> input_tensor_dtypes,
+      absl::Span<const std::string> output_tensor_names,
+      absl::Span<const std::string> target_tensor_names);
+
+  const mlrt::KernelRegistry& kernel_registry() const {
+    return *kernel_registry_;
+  }
+
+ private:
+  // A set of methods to load a client graph.
+  absl::StatusOr<std::unique_ptr<GraphExecutor::LoadedClientGraph>>
+  LoadClientGraph(
+      const GraphExecutor::ClientGraph& client_graph,
+      tensorflow::tfrt_stub::WorkQueueInterface* work_queue,
+      absl::Span<const std::pair<std::string, tensorflow::Tensor>> inputs);
+  absl::StatusOr<std::unique_ptr<GraphExecutor::LoadedClientGraph>>
+  ImportAndCompileClientGraph(
+      const GraphExecutor::ClientGraph& client_graph,
+      absl::Span<const std::pair<std::string, tensorflow::Tensor>> inputs);
+  absl::StatusOr<
+      std::pair<FunctionLibraryDefinition, mlir::OwningOpRef<mlir::ModuleOp>>>
+  ImportClientGraphToMlirModule(const GraphExecutor::ClientGraph& client_graph,
+                                mlir::MLIRContext* context) const;
+  absl::StatusOr<tfrt::BefBuffer> CompileMlirModuleToBef(
+      mlir::ModuleOp module) const;
+
+  absl::Status InitBef(LoadedClientGraph* loaded_client_graph,
+                       tensorflow::tfrt_stub::WorkQueueInterface* work_queue);
+
+  absl::Status InitBytecode(LoadedClientGraph* loaded_graph);
+
+  // Returns a `LoadedClientGraph` given input/output tensor info. If there is
+  // no existing one yet, creates one first.
+  absl::StatusOr<std::reference_wrapper<GraphExecutor::LoadedClientGraph>>
+  GetOrCreateLoadedClientGraph(
+      const RunOptions& run_options,
+      absl::Span<const std::string> input_tensor_names,
+      absl::Span<const tensorflow::DataType> input_tensor_dtypes,
+      absl::Span<const std::string> output_tensor_names,
+      absl::Span<const std::string> target_tensor_names,
+      tensorflow::tfrt_stub::WorkQueueInterface* work_queue,
+      absl::string_view graph_name = "",
+      absl::Span<const std::pair<std::string, tensorflow::Tensor>> inputs = {})
+      TF_LOCKS_EXCLUDED(loaded_client_graphs_mu_);
+
+  Options options_;
+  std::unique_ptr<FallbackState> fallback_state_;
+
+  std::unique_ptr<tensorflow::tfrt_stub::TfrtGraphExecutionState>
+      graph_execution_state_;
+
+  tfrt::RequestDeadlineTracker req_deadline_tracker_;
+
+  tensorflow::mutex loaded_client_graphs_mu_;
+  // Caches `LoadedClientGraph` by the joined name.
+  // For pointer stability of values in `absl::flat_hash_map<>`, additional
+  // `std::unique_ptr<>` is necessary. (See https://abseil.io/tips/136.)
+  absl::flat_hash_map<std::string /*joined_name*/,
+                      std::unique_ptr<LoadedClientGraph>>
+      loaded_client_graphs_ TF_GUARDED_BY(loaded_client_graphs_mu_);
+
+  std::unique_ptr<mlrt::KernelRegistry> kernel_registry_;
+
+  std::unique_ptr<tfrt::ResourceContext> resource_context_;
+
+ protected:
+  // For testing basic Cost Analysis functionality.
+  absl::Duration simulated_duration_ = absl::ZeroDuration();
+  tensorflow::mutex num_recompilations_mu_;
+  int num_recompilations_ TF_GUARDED_BY(num_recompilations_mu_) = 0;
+};
+
+void RegisterMlirDialect(mlir::DialectRegistry& registry,
+                         tensorflow::BackendCompiler* backend_compiler);
+
+// Sort the strings in `names` and store the results in `sorted_names`. In
+// addition, the original index in `names` for the item `sorted_names[i]` is
+// stored in `original_indices[i]`.
+void CreateSortedNamesAndOriginalIndices(absl::Span<const std::string> names,
+                                         std::vector<std::string>& sorted_names,
+                                         std::vector<int>& original_indices);
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_GRAPH_EXECUTOR_H_

tensorflow/core/tfrt/graph_executor/graph_executor_test.cc

@@ -0,0 +1,645 @@
+/* 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/tfrt/graph_executor/graph_executor.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "learning/brain/experimental/tfrt/native_lowering/kernels/math_kernels.h"
+#include "learning/brain/experimental/tfrt/native_lowering/kernels/sync_fallback_kernels.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "absl/types/span.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/graph.h"
+#include "tensorflow/core/graph/graph_def_builder.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/protobuf/rewriter_config.pb.h"
+#include "tensorflow/core/runtime_fallback/kernel/kernel_fallback_compat_request_state.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
+#include "tensorflow/core/tfrt/graph_executor/config.h"
+#include "tensorflow/core/tfrt/graph_executor/graph_execution_options.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/context.h"
+#include "tensorflow/core/tfrt/mlrt/interpreter/value.h"
+#include "tensorflow/core/tfrt/mlrt/kernel/kernel.h"
+#include "tensorflow/core/tfrt/saved_model/saved_model_testutil.h"
+#include "tsl/platform/statusor.h"
+#include "tfrt/cpp_tests/test_util.h"  // from @tf_runtime
+#include "tfrt/host_context/resource_context.h"  // from @tf_runtime
+#include "tfrt/tensor/dense_host_tensor.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_stub {
+namespace {
+
+using ::testing::status::StatusIs;
+
+class GraphExecutorForTestingCostAnalysis : public GraphExecutor {
+ public:
+  int num_recompilations() {
+    tensorflow::mutex_lock lock(num_recompilations_mu_);
+    return num_recompilations_;
+  }
+  // This method is not thread safe.
+  void AdvanceTime(absl::Duration duration) {
+    simulated_duration_ = simulated_duration_ + duration;
+  }
+};
+
+class GraphExecutorTest : public ::testing::TestWithParam<bool> {};
+
+absl::Status GetSimpleGraphDef(GraphDef& graph_def) {
+  auto scope = tensorflow::Scope::NewRootScope().WithDevice("/device:CPU:0");
+
+  auto input = ops::Placeholder(scope.WithOpName("input"), DT_INT32);
+  auto rank = ops::Rank(scope.WithOpName("rank"), input);
+
+  return scope.ToGraphDef(&graph_def);
+}
+
+std::unique_ptr<mlrt::KernelRegistry> GetKernelRegistry() {
+  auto kernel_registry = std::make_unique<mlrt::KernelRegistry>();
+  tensorflow::tf_mlrt::RegisterTfMlrtKernels(*kernel_registry);
+  tfrt::cpu::RegisterMlrtMathKernels(kernel_registry.get());
+  tfrt::cpu::RegisterMlrtFallbackCompatKernels(kernel_registry.get());
+
+  return kernel_registry;
+}
+
+TEST_P(GraphExecutorTest, Vanilla) {
+  GraphDef graph_def;
+  TF_ASSERT_OK(GetSimpleGraphDef(graph_def));
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  options.enable_mlrt = GetParam();
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()))
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry()));
+
+  // Set input 'x' to [[1, 1, 1]]
+  std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
+  inputs.push_back({"input", CreateTfTensor<int32_t>(
+                                 /*shape=*/{1, 3}, /*data=*/{1, 1, 1})});
+
+  std::vector<tensorflow::Tensor> outputs;
+
+  TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  ASSERT_EQ(outputs.size(), 1);
+
+  EXPECT_THAT(GetTfTensorData<int32_t>(outputs[0]),
+              ::testing::ElementsAreArray({2}));
+}
+
+TEST_P(GraphExecutorTest, OnlineCostAnalysisOptionsOverrideToOnce) {
+  GraphDef graph_def;
+  TF_ASSERT_OK(GetSimpleGraphDef(graph_def));
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  // Make sure `CostAnalysisOptions` is overriden when
+  // `enable_online_cost_analysis` = true.
+  options.enable_online_cost_analysis = true;
+  options.cost_analysis_options.version =
+      GraphExecutionOptions::CostAnalysisOptions::kDisabled;
+  options.enable_mlrt = GetParam();
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()));
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  tensorflow::tfrt_stub::RuntimeConfig runtime_config;
+  tensorflow::tf2xla::v1::MlirBridgeConfig mlir_bridge_config;
+  mlir_bridge_config.set_enable_tf2xla_mlir_bridge(false);
+  TF_ASSERT_OK(runtime_config.Add(mlir_bridge_config));
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor_base,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry(), &runtime_config));
+  auto graph_executor = std::unique_ptr<GraphExecutorForTestingCostAnalysis>(
+      static_cast<GraphExecutorForTestingCostAnalysis*>(
+          graph_executor_base.release()));
+
+  // Set input 'x' to [[1, 1, 1]]
+  std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
+  inputs.push_back({"input", CreateTfTensor<int32_t>(
+                                 /*shape=*/{1, 3}, /*data=*/{1, 1, 1})});
+
+  std::vector<tensorflow::Tensor> outputs;
+
+  // A first run should trigger online cost analysis.
+  EXPECT_EQ(graph_executor->num_recompilations(), 0);
+  TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  ASSERT_EQ(outputs.size(), 1);
+  EXPECT_THAT(GetTfTensorData<int32_t>(outputs[0]),
+              ::testing::ElementsAreArray({2}));
+  EXPECT_EQ(graph_executor->num_recompilations(), 1);
+
+  // A second run should use re-compiled graph with online profiled costs.
+  // A reset does not occur again.
+  TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  ASSERT_EQ(outputs.size(), 1);
+  EXPECT_THAT(GetTfTensorData<int32_t>(outputs[0]),
+              ::testing::ElementsAreArray({2}));
+  EXPECT_EQ(graph_executor->num_recompilations(), 1);
+}
+
+TEST_P(GraphExecutorTest, OnlineCostAnalysisEveryTime) {
+  GraphDef graph_def;
+  TF_ASSERT_OK(GetSimpleGraphDef(graph_def));
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  options.cost_analysis_options.version =
+      GraphExecutionOptions::CostAnalysisOptions::kPeriodic;
+  options.cost_analysis_options.reset_interval = absl::ZeroDuration();
+  options.cost_analysis_options.updates_per_interval = 1;
+  options.enable_mlrt = GetParam();
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()));
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor_base,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry()));
+  auto graph_executor = std::unique_ptr<GraphExecutorForTestingCostAnalysis>(
+      static_cast<GraphExecutorForTestingCostAnalysis*>(
+          graph_executor_base.release()));
+
+  // Set input 'x' to [[1, 1, 1]]
+  std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
+  inputs.push_back({"input", CreateTfTensor<int32_t>(
+                                 /*shape=*/{1, 3}, /*data=*/{1, 1, 1})});
+
+  std::vector<tensorflow::Tensor> outputs;
+
+  for (int i = 0; i < 10; ++i) {
+    TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                     /*output_tensor_names=*/{"rank"},
+                                     /*target_tensor_names=*/{}, &outputs));
+    ASSERT_EQ(outputs.size(), 1);
+    EXPECT_THAT(GetTfTensorData<int32_t>(outputs[0]),
+                ::testing::ElementsAreArray({2}));
+    EXPECT_EQ(graph_executor->num_recompilations(), i + 1);
+  }
+}
+
+TEST_P(GraphExecutorTest, OnlineCostAnalysisDisabled) {
+  GraphDef graph_def;
+  TF_ASSERT_OK(GetSimpleGraphDef(graph_def));
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  options.cost_analysis_options.version =
+      GraphExecutionOptions::CostAnalysisOptions::kDisabled;
+  options.cost_analysis_options.reset_interval = absl::ZeroDuration();
+  options.cost_analysis_options.updates_per_interval = 1;
+  options.enable_mlrt = GetParam();
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()));
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor_base,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry()));
+  auto graph_executor = std::unique_ptr<GraphExecutorForTestingCostAnalysis>(
+      static_cast<GraphExecutorForTestingCostAnalysis*>(
+          graph_executor_base.release()));
+
+  // Set input 'x' to [[1, 1, 1]]
+  std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
+  inputs.push_back({"input", CreateTfTensor<int32_t>(
+                                 /*shape=*/{1, 3}, /*data=*/{1, 1, 1})});
+
+  std::vector<tensorflow::Tensor> outputs;
+
+  TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  EXPECT_EQ(graph_executor->num_recompilations(), 0);
+}
+
+TEST_P(GraphExecutorTest, OnlineCostAnalysisPeriodic) {
+  GraphDef graph_def;
+  TF_ASSERT_OK(GetSimpleGraphDef(graph_def));
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  options.cost_analysis_options.version =
+      GraphExecutionOptions::CostAnalysisOptions::kPeriodic;
+  options.cost_analysis_options.reset_interval = absl::Minutes(10);
+  options.cost_analysis_options.updates_per_interval = 5;
+  options.enable_mlrt = GetParam();
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()));
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor_base,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry()));
+  auto graph_executor = std::unique_ptr<GraphExecutorForTestingCostAnalysis>(
+      static_cast<GraphExecutorForTestingCostAnalysis*>(
+          graph_executor_base.release()));
+
+  // Set input 'x' to [[1, 1, 1]]
+  std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
+  inputs.push_back({"input", CreateTfTensor<int32_t>(
+                                 /*shape=*/{1, 3}, /*data=*/{1, 1, 1})});
+
+  std::vector<tensorflow::Tensor> outputs;
+  // First run always initiates a recompilation.
+  TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  EXPECT_EQ(graph_executor->num_recompilations(), 1);
+
+  // We have specified that the costs should only update every
+  // `reset_interval` / `updates_per_interval` = 2
+  // minutes. So no cost update occurs here.
+  for (int i = 0; i < 10; ++i) {
+    TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                     /*output_tensor_names=*/{"rank"},
+                                     /*target_tensor_names=*/{}, &outputs));
+    EXPECT_EQ(graph_executor->num_recompilations(), 1);
+  }
+  // With 2 minute breaks in-between, 4 runs = 4 cost updates.
+  for (int i = 0; i < 4; ++i) {
+    graph_executor->AdvanceTime(absl::Minutes(2));
+    TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                     /*output_tensor_names=*/{"rank"},
+                                     /*target_tensor_names=*/{}, &outputs));
+    EXPECT_EQ(graph_executor->num_recompilations(), 1);
+  }
+  // A reset occurs on the 5th run.
+  graph_executor->AdvanceTime(absl::Minutes(2));
+  TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  EXPECT_EQ(graph_executor->num_recompilations(), 2);
+
+  // Demonstrate one more reset.
+  for (int i = 0; i < 4; ++i) {
+    graph_executor->AdvanceTime(absl::Minutes(1000));
+    TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                     /*output_tensor_names=*/{"rank"},
+                                     /*target_tensor_names=*/{}, &outputs));
+    EXPECT_EQ(graph_executor->num_recompilations(), 2);
+  }
+  graph_executor->AdvanceTime(absl::Minutes(1000));
+  TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  EXPECT_EQ(graph_executor->num_recompilations(), 3);
+}
+
+REGISTER_OP("TestCancel")
+    .Input("x: T")
+    .Output("z: T")
+    .Attr("T: {int32}")
+    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);
+
+class TestCancelKernel : public OpKernel {
+ public:
+  explicit TestCancelKernel(OpKernelConstruction* context)
+      : OpKernel(context) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    auto status = absl::CancelledError();
+    ctx->cancellation_manager()->StartCancelWithStatus(status);
+    ctx->SetStatus(status);
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("TestCancel").Device(DEVICE_CPU),
+                        TestCancelKernel);
+
+REGISTER_OP("TestIsCancelled").Output("z: T").Attr("T: {bool}").SetIsStateful();
+
+class TestIsCancelledKernel : public OpKernel {
+ public:
+  explicit TestIsCancelledKernel(OpKernelConstruction* context)
+      : OpKernel(context) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    ctx->set_output(
+        0, tensorflow::Tensor(ctx->cancellation_manager()->IsCancelled()));
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("TestIsCancelled").Device(DEVICE_CPU),
+                        TestIsCancelledKernel);
+
+TEST_P(GraphExecutorTest, Cancellation) {
+  GraphDef graph_def;
+
+  tensorflow::GraphDefBuilder builder(
+      tensorflow::GraphDefBuilder::kFailImmediately);
+
+  const tensorflow::TensorShape tensor_shape({10, 9});
+  tensorflow::Node* input = tensorflow::ops::SourceOp(
+      "Placeholder", builder.opts()
+                         .WithName("input")
+                         .WithAttr("dtype", tensorflow::DT_INT32)
+                         .WithAttr("shape", tensor_shape));
+  tensorflow::ops::SourceOp("TestIsCancelled",
+                            builder.opts()
+                                .WithName("is_cancelled")
+                                .WithAttr("T", tensorflow::DT_BOOL));
+  tensorflow::ops::UnaryOp("TestCancel", input,
+                           builder.opts()
+                               .WithName("test_cancel")
+                               .WithAttr("T", tensorflow::DT_INT32));
+
+  TF_ASSERT_OK(builder.ToGraphDef(&graph_def));
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  options.enable_mlrt = GetParam();
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()))
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry()));
+  {
+    std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
+    inputs.push_back({"input", CreateTfTensor<int32_t>(
+                                   /*shape=*/{1, 3}, /*data=*/{1, 1, 1})});
+
+    std::vector<tensorflow::Tensor> outputs;
+    EXPECT_THAT(graph_executor->Run(/*run_options=*/{}, inputs,
+                                    /*output_tensor_names=*/{"test_cancel:0"},
+                                    /*target_tensor_names=*/{}, &outputs),
+                StatusIs(absl::StatusCode::kCancelled));
+  }
+
+  {
+    std::vector<tensorflow::Tensor> outputs;
+    TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, /*inputs=*/{},
+                                     /*output_tensor_names=*/{"is_cancelled:0"},
+                                     /*target_tensor_names=*/{}, &outputs));
+    ASSERT_EQ(outputs.size(), 1);
+
+    EXPECT_THAT(GetTfTensorData<bool>(outputs[0]),
+                ::testing::ElementsAreArray({false}));
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(GraphExecutorTestSuite, GraphExecutorTest,
+                         ::testing::Bool());
+
+TEST_F(GraphExecutorTest, CreateRequestInfoPopulatesRpcCancellationOptions) {
+  GraphDef graph_def;
+  TF_ASSERT_OK(GetSimpleGraphDef(graph_def));
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()));
+
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  auto client_graph_resource_context =
+      std::make_unique<tfrt::ResourceContext>();
+  OpKernelRunnerTable runner_table;
+  tfd::FallbackResourceArray resource_array;
+
+  GraphExecutionRunOptions run_options;
+  absl::Time now = absl::Now();
+  run_options.rpc_deadline_for_batching_task_cancellation = now;
+  bool is_cancelled = false;
+  run_options.is_rpc_cancelled_callback = [&is_cancelled]() {
+    return is_cancelled;
+  };
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto request_info,
+      CreateRequestInfo(
+          options, run_options, /*work_queue=*/nullptr, resource_context.get(),
+          client_graph_resource_context.get(), &runner_table, &resource_array,
+          *fallback_state, fallback_state->process_function_library_runtime()));
+
+  ASSERT_NE(request_info->tfrt_request_context, nullptr);
+  const auto* fallback_request_state =
+      request_info->tfrt_request_context
+          ->GetDataIfExists<tfd::KernelFallbackCompatRequestState>();
+  ASSERT_NE(fallback_request_state, nullptr);
+
+  EXPECT_EQ(
+      fallback_request_state->rpc_deadline_for_batching_task_cancellation(),
+      now);
+  EXPECT_FALSE(fallback_request_state->is_rpc_cancelled_callback()());
+  is_cancelled = true;
+  EXPECT_TRUE(fallback_request_state->is_rpc_cancelled_callback()());
+}
+
+TEST_F(GraphExecutorTest, Extend) {
+  GraphDef graph_def;
+  {
+    auto scope = tensorflow::Scope::NewRootScope().WithDevice("/device:CPU:0");
+
+    Output a = ops::Const(scope.WithOpName("a"), 0.0f, {10, 10});
+
+    Output b = ops::Const(scope.WithControlDependencies(a).WithOpName("b"),
+                          0.0f, {10, 10});
+    Output c = ops::Identity(scope.WithOpName("c"), b);
+
+    TF_ASSERT_OK(scope.ToGraphDef(&graph_def));
+  }
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  auto session_options = CreateDefaultSessionOptions(options);
+  // Disable optimizations for static graph to allow calls to Session::Extend.
+  session_options.config.mutable_experimental()
+      ->set_disable_optimize_for_static_graph(true);
+  TF_ASSERT_OK_AND_ASSIGN(auto fallback_state,
+                          tensorflow::tfrt_stub::FallbackState::Create(
+                              session_options, graph_def.library()));
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry()));
+
+  GraphDef extension;
+  {
+    auto scope = tensorflow::Scope::NewRootScope().WithDevice("/device:CPU:0");
+
+    auto input = ops::Placeholder(scope.WithOpName("input"), DT_INT32);
+    auto rank = ops::Rank(scope.WithOpName("rank"), input);
+
+    TF_ASSERT_OK(scope.ToGraphDef(&extension));
+  }
+
+  TF_ASSERT_OK(graph_executor->Extend(extension));
+
+  // Set input 'x' to [[1, 1, 1]]
+  std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
+  inputs.push_back({"input", CreateTfTensor<int32_t>(
+                                 /*shape=*/{1, 3}, /*data=*/{1, 1, 1})});
+
+  std::vector<tensorflow::Tensor> outputs;
+
+  TF_ASSERT_OK(graph_executor->Run(/*run_options=*/{}, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  ASSERT_EQ(outputs.size(), 1);
+
+  EXPECT_THAT(GetTfTensorData<int32_t>(outputs[0]),
+              ::testing::ElementsAreArray({2}));
+}
+
+TEST_F(GraphExecutorTest, DisableCompilation) {
+  GraphDef graph_def;
+  TF_ASSERT_OK(GetSimpleGraphDef(graph_def));
+
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()));
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry()));
+
+  // Set input 'x' to [[1, 1, 1]]
+  std::vector<std::pair<std::string, tensorflow::Tensor>> inputs;
+  inputs.push_back({"input", CreateTfTensor<int32_t>(
+                                 /*shape=*/{1, 3}, /*data=*/{1, 1, 1})});
+
+  std::vector<tensorflow::Tensor> outputs;
+
+  GraphExecutor::RunOptions run_options;
+  run_options.disable_compilation = true;
+
+  auto status = graph_executor->Run(run_options, inputs,
+                                    /*output_tensor_names=*/{"rank"},
+                                    /*target_tensor_names=*/{}, &outputs);
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(
+      status.ToString(),
+      ::testing::HasSubstr("GraphExecutor: compilation is disabled in "
+                           "execution but the compiled graph is not found"));
+
+  run_options.disable_compilation = false;
+  TF_ASSERT_OK(graph_executor->Run(run_options, inputs,
+                                   /*output_tensor_names=*/{"rank"},
+                                   /*target_tensor_names=*/{}, &outputs));
+  ASSERT_EQ(outputs.size(), 1);
+
+  EXPECT_THAT(GetTfTensorData<int32_t>(outputs[0]),
+              ::testing::ElementsAreArray({2}));
+}
+
+TEST_F(GraphExecutorTest, SyncExecute) {
+  GraphDef graph_def;
+  TF_ASSERT_OK(GetSimpleGraphDef(graph_def));
+  auto runtime = DefaultTfrtRuntime(/*num_threads=*/1);
+  GraphExecutor::Options options(runtime.get());
+  options.compile_options.compile_to_sync_tfrt_dialect = true;
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto fallback_state,
+      tensorflow::tfrt_stub::FallbackState::Create(
+          CreateDefaultSessionOptions(options), graph_def.library()));
+
+  auto resource_context = std::make_unique<tfrt::ResourceContext>();
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto graph_executor,
+      GraphExecutor::Create(std::move(options), std::move(fallback_state),
+                            std::move(resource_context), graph_def,
+                            GetKernelRegistry()));
+
+  std::vector<mlrt::Value> inputs;
+  tfrt::DenseHostTensor dht =
+      tfrt::CreateTensorFromValues<int32_t>({1, 3}, {1, 1, 1});
+  inputs.emplace_back(std::move(dht));
+  std::vector<mlrt::Value> results;
+  results.resize(1);
+
+  TF_ASSERT_OK(graph_executor->RunWithSyncInterpreter(
+      "test_graph", absl::Span<mlrt::Value>(inputs),
+      /*input_names=*/{"input"}, /*input_dtypes=*/{DT_INT32},
+      /*output_tensor_names=*/{"rank"},
+      /*target_tensor_names=*/{}, absl::Span<mlrt::Value>(results)));
+  tfrt::DenseHostTensor expected =
+      tfrt::CreateTensorFromValues<int32_t>({}, {2});
+
+  EXPECT_EQ(expected, results[0].Get<tfrt::DenseHostTensor>());
+}
+
+}  // namespace
+}  // namespace tfrt_stub
+}  // namespace tensorflow

tensorflow/core/tfrt/graph_executor/python/BUILD

@@ -0,0 +1,38 @@
+# This package loads pybind11 dependencies to ensure compatibility with Python API.
+# Files affected: graph_execution_options
+load("//tensorflow:tensorflow.default.bzl", "tf_python_pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    packages = [
+        # Authorized users go here.
+        "//tensorflow/core/tfrt/saved_model/...",
+        "//tensorflow/core/tfrt/graph_executor/...",
+    ],
+)
+
+tf_python_pybind_extension(
+    name = "_pywrap_graph_execution_options",
+    srcs = ["graph_execution_options.cc"],
+    enable_stub_generation = True,
+    module_name = "_pywrap_graph_execution_options",
+    pytype_srcs = [
+        "_pywrap_graph_execution_options.pyi",
+    ],
+    deps = [
+        "//tensorflow/core/tfrt/graph_executor",
+        "//tensorflow/core/tfrt/graph_executor:graph_execution_options",
+        "//tensorflow/core/tfrt/runtime",
+        "//tensorflow/core/tfrt/saved_model",
+        "//tensorflow/core/tfrt/saved_model:saved_model_aot_compile",
+        "@pybind11",
+        "@pybind11_abseil//pybind11_abseil:status_casters",
+        "@xla//third_party/python_runtime:headers",
+    ],
+)

tensorflow/core/tfrt/graph_executor/python/_pywrap_graph_execution_options.pyi

@@ -0,0 +1,17 @@
+# 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.
+# ==============================================================================
+
+class GraphExecutionOptions:
+    def __init__(self, arg0) -> None: ...

tensorflow/core/tfrt/graph_executor/python/graph_execution_options.cc

@@ -0,0 +1,35 @@
+/* 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/core/tfrt/graph_executor/graph_execution_options.h"
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11_abseil/status_casters.h"  // from @pybind11_abseil
+#include "tensorflow/core/tfrt/runtime/runtime.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(GraphExecutionRunOptions, m) {
+  py::class_<tensorflow::tfrt_stub::GraphExecutionRunOptions>(
+      m, "GraphExecutionRunOptions")
+      .def(py::init<>());
+  m.doc() =
+      "pybind11 GraphExecutionRunOptions wrapper";  // optional module docstring
+}
+
+PYBIND11_MODULE(_pywrap_graph_execution_options, m) {
+  py::class_<tensorflow::tfrt_stub::GraphExecutionOptions>(
+      m, "GraphExecutionOptions")
+      .def(py::init<const tensorflow::tfrt_stub::Runtime *>());
+}

tensorflow/core/tfrt/graph_executor/sync_resource_state.h

@@ -0,0 +1,66 @@
+/* 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_CORE_TFRT_GRAPH_EXECUTOR_SYNC_RESOURCE_STATE_H_
+#define TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_SYNC_RESOURCE_STATE_H_
+
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/tfrt/utils/any_ptr.h"
+#include "tfrt/tensor/dense_host_tensor.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace tfrt_stub {
+
+class SyncResourceState {
+ public:
+  // Sets `dht` in the array at `index`. `index` should be dense and
+  // duplicate indices are not allowed.
+  void SetResourceDht(int index, tfrt::DenseHostTensor dht) {
+    if (resource_dht_.size() <= index) {
+      resource_dht_.resize(index + 1);
+    }
+
+    resource_dht_[index] = std::move(dht);
+  }
+
+  tfrt::DenseHostTensor GetResourceDht(int index) const {
+    return resource_dht_.at(index).CopyRef();
+  }
+
+  template <typename T>
+  void Set(int index, T* resource) {
+    if (resources_.size() <= index) {
+      resources_.resize(index + 1);
+    }
+
+    resources_[index] = tfrt::AnyPtr(resource);
+  }
+
+  template <typename T>
+  T* Get(int index) const {
+    return resources_.at(index).get<T>();
+  }
+
+ private:
+  std::vector<tfrt::DenseHostTensor> resource_dht_;
+  // TODO(b/288899457): Consider provide a simpler solution than forking AnyPtr.
+  std::vector<tfrt::AnyPtr> resources_;
+};
+
+}  // namespace tfrt_stub
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_GRAPH_EXECUTOR_SYNC_RESOURCE_STATE_H_

tensorflow/core/tfrt/graph_executor/test_config.proto

@@ -0,0 +1,26 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tfrt_stub;
+
+message TestConfig1 {
+  string tag = 1;
+}
+
+message TestConfig2 {
+  string tag = 1;
+}

tensorflow/core/tfrt/ifrt/BUILD

@@ -0,0 +1,763 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        # copybara:uncomment "//learning/brain/tfrt:__subpackages__",
+        # copybara:uncomment "//third_party/neptune:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/tfrt:__subpackages__",
+        "//tensorflow/core/tfrt/mlrt:__subpackages__",
+    ],
+)
+
+tf_proto_library(
+    name = "ifrt_config_proto",
+    srcs = ["ifrt_config.proto"],
+    protodeps = [
+        "@xla//xla:xla_data_proto",
+        "//tensorflow/core/framework:tensor_proto",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "ifrt_serving_core_selector",
+    srcs = ["ifrt_serving_core_selector.cc"],
+    hdrs = ["ifrt_serving_core_selector.h"],
+    deps = [
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@xla//xla/tsl/framework:serving_device_selector",
+    ],
+)
+
+tf_cc_test(
+    name = "ifrt_serving_core_selector_test",
+    srcs = ["ifrt_serving_core_selector_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_serving_core_selector",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/tsl/framework:serving_device_selector",
+        "@xla//xla/tsl/framework/test_util:mock_serving_device_selector",
+    ],
+)
+
+cc_library(
+    name = "grid",
+    srcs = ["grid.cc"],
+    hdrs = ["grid.h"],
+    deps = [
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ],
+)
+
+cc_library(
+    name = "ifrt_device_utils",
+    srcs = ["ifrt_device_utils.cc"],
+    hdrs = ["ifrt_device_utils.h"],
+    deps = [
+        ":grid",
+        ":ifrt_config_proto_cc",
+        "//tensorflow/compiler/tf2xla:host_compute_metadata_proto_cc",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:attribute_map",
+        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "ifrt_device_utils_test",
+    srcs = [
+        "ifrt_device_utils_test.cc",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_device_utils",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:attribute_map",
+        "@xla//xla/python/ifrt:mock",
+        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/tsl/platform:status_matchers",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "ifrt_serving_executable",
+    srcs = ["ifrt_serving_executable.cc"],
+    hdrs = ["ifrt_serving_executable.h"],
+    deps = [
+        ":ifrt_config_proto_cc",
+        ":ifrt_device_utils",
+        ":ifrt_loaded_variable_registry",
+        ":ifrt_loaded_variable_utils",
+        ":ifrt_persistent_compilation_cache",
+        ":ifrt_restore_tensor_registry",
+        ":ifrt_serving_core_selector",
+        ":ifrt_tensor_utils",
+        ":sharding_utils",
+        ":tf_host_callback",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
+        "//tensorflow/compiler/mlir/tfrt:export",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:extract_callback",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf2hlo",
+        "//tensorflow/compiler/tf2xla:common",
+        "//tensorflow/compiler/tf2xla:host_compute_metadata_proto_cc",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/hash",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:hostcontext",
+        "@tsl//tsl/platform:tstring",
+        "@tsl//tsl/profiler/lib:traceme",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/ir:hlo",
+        "@xla//xla/hlo/translate:stablehlo",
+        "@xla//xla/pjrt:host_callback",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt:pjrt_compiler",
+        "@xla//xla/pjrt:pjrt_executable",
+        "@xla//xla/pjrt:pjrt_layout",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt/hlo:hlo_program",
+        "@xla//xla/python/pjrt_ifrt",
+        "@xla//xla/python/pjrt_ifrt:pjrt_layout",
+        "@xla//xla/service:computation_placer_hdr",
+        "@xla//xla/service:dump",
+        "@xla//xla/tsl/concurrency:future",
+        "@xla//xla/tsl/concurrency:ref_count",
+        "@xla//xla/tsl/framework:serving_device_selector",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "ifrt_executable_registry",
+    srcs = ["ifrt_executable_registry.cc"],
+    hdrs = ["ifrt_executable_registry.h"],
+    deps = [
+        ":ifrt_serving_executable",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
+cc_library(
+    name = "ifrt_restore_tensor_registry",
+    srcs = ["ifrt_restore_tensor_registry.cc"],
+    hdrs = ["ifrt_restore_tensor_registry.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:types_proto_cc",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@xla//xla/tsl/concurrency:future",
+    ],
+)
+
+cc_library(
+    name = "ifrt_persistent_compilation_cache",
+    srcs = ["ifrt_persistent_compilation_cache.cc"],
+    hdrs = ["ifrt_persistent_compilation_cache.h"],
+    deps = [
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf2hlo",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//mlir:IR",
+        "@xla//xla/pjrt:pjrt_executable",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt/hlo:hlo_program",
+        "@xla//xla/python/pjrt_ifrt:xla_ifrt",
+        "@xla//xla/tsl/concurrency:ref_count",
+    ],
+)
+
+cc_library(
+    name = "ifrt_loaded_variable_registry",
+    srcs = ["ifrt_loaded_variable_registry.cc"],
+    hdrs = ["ifrt_loaded_variable_registry.h"],
+    deps = [
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/hash",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla:shape_util",
+        "@xla//xla/hlo/ir:hlo",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/tsl/concurrency:future",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "ifrt_loaded_variable_registry_test",
+    srcs = ["ifrt_loaded_variable_registry_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_loaded_variable_registry",
+        "@com_google_absl//absl/hash",
+        "@com_google_absl//absl/hash:hash_testing",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/ir:hlo",
+    ],
+)
+
+cc_library(
+    name = "ifrt_model_restore_context",
+    hdrs = ["ifrt_model_restore_context.h"],
+    deps = [
+        ":checkpoint_loader",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+cc_library(
+    name = "ifrt_model_context",
+    srcs = ["ifrt_model_context.cc"],
+    hdrs = ["ifrt_model_context.h"],
+    visibility = [
+        "//learning/brain/tfrt:__subpackages__",
+        "//learning/serving/servables/tfrt:__subpackages__",
+        "//tensorflow/compiler/mlir/tfrt:__subpackages__",
+        "//tensorflow/core/tfrt:__subpackages__",
+        "//tensorflow/core/tfrt/mlrt:__subpackages__",
+    ],
+    deps = [
+        ":ifrt_config_proto_cc",
+        ":ifrt_executable_registry",
+        ":ifrt_loaded_variable_registry",
+        ":ifrt_persistent_compilation_cache",
+        ":ifrt_restore_tensor_registry",
+        ":ifrt_serving_core_selector",
+        ":sharding_utils",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf2hlo",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:core_cpu_base",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@tf_runtime//:hostcontext",
+        "@tsl//tsl/platform:protobuf",
+        "@xla//xla/pjrt:pjrt_executable",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "ifrt_tensor_utils",
+    srcs = ["ifrt_tensor_utils.cc"],
+    hdrs = ["ifrt_tensor_utils.h"],
+    deps = [
+        "//tensorflow/compiler/tf2xla:common",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/pjrt_ifrt:pjrt_dtype",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "sharding_utils",
+    srcs = [
+        "sharding_utils.cc",
+    ],
+    hdrs = [
+        "sharding_utils.h",
+    ],
+    deps = [
+        ":ifrt_tensor_utils",
+        "//tensorflow/compiler/tf2xla:common",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/tpu/kernels:sharding_utils",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/container:btree",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@eigen_archive//:eigen3",
+        "@llvm-project//llvm:Support",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/ir:hlo",
+        "@xla//xla/pjrt:pjrt_layout",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/pjrt_ifrt:xla_ifrt",
+        "@xla//xla/tsl/concurrency:future",
+        "@xla//xla/tsl/concurrency:ref_count",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "ifrt_loaded_variable_utils",
+    srcs = ["ifrt_loaded_variable_utils.cc"],
+    hdrs = ["ifrt_loaded_variable_utils.h"],
+    deps = [
+        ":ifrt_config_proto_cc",
+        ":ifrt_loaded_variable_registry",
+        ":ifrt_restore_tensor_registry",
+        ":sharding_utils",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/platform:context",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@tf_runtime//:hostcontext",
+        "@xla//xla:shape_util",
+        "@xla//xla/hlo/ir:hlo",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/tsl/concurrency:future",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "tf_host_callback",
+    srcs = ["tf_host_callback.cc"],
+    hdrs = ["tf_host_callback.h"],
+    deps = [
+        "//tensorflow/c/eager:abstract_tensor_handle",
+        "//tensorflow/c/eager:immediate_execution_context",
+        "//tensorflow/c/eager:immediate_execution_operation",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core/common_runtime/eager:context",
+        "//tensorflow/core/common_runtime/eager:core_no_xla",
+        "//tensorflow/core/common_runtime/eager:tensor_handle",
+        "//tensorflow/core/framework:function_proto_cc",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/cleanup",
+        "@com_google_absl//absl/container:fixed_array",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@tsl//tsl/platform:casts",
+        "@tsl//tsl/platform:refcount",
+        "@tsl//tsl/profiler/lib:traceme",
+        "@xla//xla/tsl/platform:errors",
+    ],
+)
+
+cc_library(
+    name = "ifrt_serving_executable_test_util",
+    testonly = True,
+    srcs = ["ifrt_serving_executable_test_util.cc"],
+    hdrs = ["ifrt_serving_executable_test_util.h"],
+    data = [
+        "//tensorflow/core/tfrt/ifrt/testdata",
+    ],
+    deps = [
+        ":ifrt_loaded_variable_registry",
+        ":ifrt_persistent_compilation_cache",
+        ":ifrt_restore_tensor_registry",
+        ":ifrt_serving_core_selector",
+        ":ifrt_serving_executable",
+        ":sharding_utils",
+        ":tf_host_callback",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:tf2hlo",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/platform:resource_loader",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:RegisterAllDialects",
+        "@tf_runtime//:hostcontext",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:test_util",
+        "@xla//xla/tsl/framework/test_util:mock_serving_device_selector",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:status",
+    ],
+)
+
+tf_cc_test(
+    name = "ifrt_restore_tensor_registry_test",
+    srcs = ["ifrt_restore_tensor_registry_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_restore_tensor_registry",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/framework:types_proto_cc",
+        "@com_google_absl//absl/base:log_severity",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/tsl/concurrency:future",
+        "@xla//xla/tsl/lib/core:status_test_util",
+        "@xla//xla/tsl/platform:status_matchers",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "ifrt_loaded_variable_utils_test",
+    srcs = ["ifrt_loaded_variable_utils_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_config_proto_cc",
+        ":ifrt_loaded_variable_registry",
+        ":ifrt_loaded_variable_utils",
+        ":ifrt_restore_tensor_registry",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/framework:tensor_matcher",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/framework:types_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_googletest//:gtest_main",
+        "@tf_runtime//:hostcontext",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/ir:hlo",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:test_util",
+        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
+        "@xla//xla/tsl/concurrency:future",
+        "@xla//xla/tsl/concurrency:ref_count",
+        "@xla//xla/tsl/lib/core:status_test_util",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:status_matchers",
+        "@xla//xla/tsl/platform:statusor",
+        "@xla//xla/tsl/platform:test",
+    ],
+)
+
+tf_cc_test(
+    name = "tf_host_callback_test",
+    srcs = ["tf_host_callback_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        ":tf_host_callback",
+        "//tensorflow/cc:const_op",
+        "//tensorflow/cc:function_ops",
+        "//tensorflow/cc:functional_ops",
+        "//tensorflow/cc:math_ops",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:resource_variable_ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "//tensorflow/core/framework:tensor_matcher",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/kernels:constant_op",
+        "//tensorflow/core/kernels:math",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "sharding_utils_test",
+    srcs = ["sharding_utils_test.cc"],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_tensor_utils",
+        ":pjrt_cpu_client_test_lib",
+        ":sharding_utils",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:test",
+        "//tensorflow/core/framework:tensor_matcher",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/framework:types_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@tsl//tsl/platform:ml_dtypes",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/hlo/ir:hlo",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:test_util",
+        "@xla//xla/python/pjrt_ifrt:xla_ifrt",
+        "@xla//xla/tsl/concurrency:ref_count",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "ifrt_serving_executable_test",
+    srcs = [
+        "ifrt_serving_executable_test.cc",
+    ],
+    data = [
+        "//tensorflow/core/tfrt/ifrt/testdata",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_restore_tensor_registry",
+        ":ifrt_serving_executable",
+        ":ifrt_serving_executable_test_util",
+        ":sharding_utils",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:portable_tensorflow_test_lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core/common_runtime/eager:core",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:tensor_matcher",
+        "//tensorflow/core/framework:tensor_testutil",
+        "//tensorflow/core/framework:types_proto_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:globals",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Support",
+        "@tf_runtime//:basic_kernels_alwayslink",
+        "@tf_runtime//:core_runtime_alwayslink",
+        "@tf_runtime//:test_kernels_alwayslink",
+        "@tf_runtime//backends/cpu:core_runtime_alwayslink",
+        "@tf_runtime//backends/cpu:tf_ops_alwayslink",
+        "@tsl//tsl/platform:tstring",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
+        "@xla//xla/tsl/concurrency:future",
+        "@xla//xla/tsl/framework:serving_device_selector",
+        "@xla//xla/tsl/framework/test_util:mock_serving_device_selector",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "ifrt_executable_registry_test",
+    srcs = [
+        "ifrt_executable_registry_test.cc",
+    ],
+    data = [
+        "//tensorflow/core/tfrt/ifrt/testdata",
+    ],
+    tags = ["no_oss"],
+    deps = [
+        ":ifrt_executable_registry",
+        ":ifrt_loaded_variable_registry",
+        ":ifrt_restore_tensor_registry",
+        ":ifrt_serving_executable",
+        ":sharding_utils",
+        ":tf_host_callback",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/tf2xla:xla_helpers",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:portable_tensorflow_test_lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core/common_runtime/eager:core",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/platform:resource_loader",
+        "//tensorflow/core/platform:status_matchers",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:RegisterAllDialects",
+        "@tf_runtime//:basic_kernels_alwayslink",
+        "@tf_runtime//:core_runtime_alwayslink",
+        "@tf_runtime//:hostcontext",
+        "@tf_runtime//:test_kernels_alwayslink",
+        "@tf_runtime//backends/cpu:core_runtime_alwayslink",
+        "@tf_runtime//backends/cpu:tf_ops_alwayslink",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:test_util",
+        "@xla//xla/python/pjrt_ifrt:tfrt_cpu_client_test_lib",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "checkpoint_loader",
+    srcs = ["checkpoint_loader.cc"],
+    hdrs = ["checkpoint_loader.h"],
+    deps = [
+        ":ifrt_loaded_variable_utils",
+        ":ifrt_restore_tensor_registry",
+        "//tensorflow/compiler/mlir/tfrt/transforms/ifrt:ifrt_types",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/common_runtime:function",
+        "//tensorflow/core/framework:attr_value_proto_cc",
+        "//tensorflow/core/framework:node_def_util",
+        "//tensorflow/core/framework:tensor",
+        "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/protobuf:for_core_protos_cc",
+        "//tensorflow/core/tfrt/fallback:fallback_state",
+        "//tensorflow/core/tfrt/fallback:op_kernel_runner",
+        "//tensorflow/core/tfrt/mlrt/bytecode",
+        "//tensorflow/core/tfrt/mlrt/kernel:context",
+        "//tensorflow/core/tfrt/mlrt/kernel:kernel_runner_utils",
+        "//tensorflow/core/tfrt/mlrt/kernel:shard_restore_util",
+        "//tensorflow/core/tfrt/utils:fallback_tensor",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//mlir:IR",
+        "@tf_runtime//:hostcontext",
+        "@tsl//tsl/platform:tstring",
+        "@xla//xla/tsl/concurrency:future",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+cc_library(
+    name = "pjrt_cpu_client_test_lib",
+    testonly = True,
+    srcs = ["pjrt_cpu_client_test_lib.cc"],
+    deps = [
+        "@com_google_absl//absl/status:statusor",
+        "@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options",
+        "@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
+        "@xla//xla/python/ifrt",
+        "@xla//xla/python/ifrt:test_util",
+        "@xla//xla/python/pjrt_ifrt",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+    alwayslink = True,
+)

tensorflow/core/tfrt/ifrt/checkpoint_loader.cc

@@ -0,0 +1,437 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/tfrt/ifrt/checkpoint_loader.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/tsl/concurrency/future.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/context.h"
+#include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/mlrt/kernel/context.h"
+#include "tensorflow/core/tfrt/mlrt/kernel/kernel_runner_utils.h"
+#include "tensorflow/core/tfrt/mlrt/kernel/shard_restore_util.h"
+#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
+#include "tsl/platform/tstring.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+namespace {
+
+static constexpr int kNumRestoreClusters = 4;
+
+// A shard of variables to be restored.
+struct RestoreVariableShard {
+  tensorflow::Tensor prefix;
+  tensorflow::Tensor tensor_names;
+  tensorflow::Tensor shape_and_slices;
+  std::vector<tensorflow::tfrt_stub::FallbackTensor> var_handles;
+  tensorflow::AttrValue dtypes_attr_value;
+  std::vector<tensorflow::DataType> restored_dtypes;
+  std::vector<bool> truncate_in_cast;
+};
+
+struct AsyncState {
+  explicit AsyncState(
+      const std::vector<tensorflow::TensorValue>& input_tf_tensor_values,
+      const OpKernelContext::Params& params, int num_outputs,
+      const tensorflow::DeviceMgr& device_manager,
+      const tensorflow::ProcessFunctionLibraryRuntime&
+          process_function_library_runtime)
+      : run_state(input_tf_tensor_values, params),
+        context(&run_state.params, num_outputs),
+        device_manager(device_manager),
+        process_function_library_runtime(process_function_library_runtime) {}
+
+  tfrt_stub::OpKernelRunState run_state;
+  OpKernelContext context;
+  const tensorflow::DeviceMgr& device_manager;
+  const tensorflow::ProcessFunctionLibraryRuntime&
+      process_function_library_runtime;
+
+  std::vector<tsl::Promise<tensorflow::Tensor>> host_tensor_results;
+  // If the dtype and shape is available at restore graph at compile time,
+  // the corresponding promise will be nullopt.
+  std::vector<std::optional<tsl::Promise<DtypeAndShape>>>
+      dtype_and_shape_results;
+};
+
+// Returns a casted tensor if successful.
+absl::StatusOr<tensorflow::Tensor> Cast(
+    tensorflow::Tensor& in_tensor, tensorflow::DataType restored_dtype,
+    tensorflow::DataType cast_dtype, bool truncate_in_cast,
+    const tensorflow::DeviceMgr& device_manager,
+    const tensorflow::ProcessFunctionLibraryRuntime&
+        process_function_library_runtime,
+    OpKernelContext::Params& params) {
+  auto runner =
+      tfrt_stub::OpKernelRunner::Create(
+          /*op_name=*/
+          "Cast", /*node_name=*/"Cast", params.device->name(),
+          /*num_args=*/1,
+          [&](tensorflow::AttrValueMap* attr_value_map) {
+            tensorflow::AttrValue restored_dtype_attr_value;
+            restored_dtype_attr_value.set_type(restored_dtype);
+            attr_value_map->insert({"SrcT", restored_dtype_attr_value});
+
+            tensorflow::AttrValue cast_dtype_attr_value;
+            cast_dtype_attr_value.set_type(cast_dtype);
+            attr_value_map->insert({"DstT", cast_dtype_attr_value});
+
+            tensorflow::AttrValue truncate_attr_value;
+            truncate_attr_value.set_b(truncate_in_cast);
+            attr_value_map->insert({"Truncate", truncate_attr_value});
+            return absl::OkStatus();
+          },
+          device_manager, process_function_library_runtime)
+          .value();
+
+  std::vector<tensorflow::TensorValue> input_tf_tensor_values;
+  input_tf_tensor_values.push_back(tensorflow::TensorValue(&in_tensor));
+
+  tf_mlrt::SetUpParams(runner, input_tf_tensor_values, params);
+  // Use persistent device instead of the per request device.
+
+  OpKernelContext op_kernel_context(&params, /*num_outputs=*/1);
+
+  runner.Run(&op_kernel_context);
+
+  if (!op_kernel_context.status().ok()) {
+    return op_kernel_context.status();
+  }
+  DCHECK_EQ(op_kernel_context.num_outputs(), 1);
+  return *(op_kernel_context.mutable_output(0));
+}
+
+void RunShardHelper(const tfrt_stub::OpKernelRunner& runner,
+                    AsyncState* async_state, RestoreVariableShard shard,
+                    IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry) {
+  // Keep input tensor alive in `shard`.
+  auto* op_kernel_context_ptr = &async_state->context;
+  runner.Run(op_kernel_context_ptr);
+
+  auto& op_kernel_context = async_state->context;
+  if (!op_kernel_context.status().ok()) {
+    LOG(ERROR) << "failed to run restore op: " << op_kernel_context.status();
+    for (auto& result : async_state->host_tensor_results) {
+      std::move(result).Set(op_kernel_context.status());
+    }
+    for (auto& result : async_state->dtype_and_shape_results) {
+      if (result.has_value()) {
+        std::move(*result).Set(op_kernel_context.status());
+      }
+    }
+    return;
+  }
+  DCHECK_EQ(shard.var_handles.size(), op_kernel_context.num_outputs());
+  DCHECK_EQ(shard.truncate_in_cast.size(), op_kernel_context.num_outputs());
+
+  // TODO(b/343964091): consider to run multiple casts in parallel.
+  for (int i = 0; i < op_kernel_context.num_outputs(); ++i) {
+    DCHECK(op_kernel_context.mutable_output(i));
+
+    if (op_kernel_context.mutable_output(i)->dtype() !=
+        shard.restored_dtypes[i]) {
+      LOG(ERROR) << "checkpoint_loader: dtype mismatch: "
+                 << shard.var_handles[i].tensor().DebugString() << ", "
+                 << op_kernel_context.mutable_output(i)->dtype() << " vs. "
+                 << shard.restored_dtypes[i];
+      std::move(async_state->host_tensor_results[i])
+          .Set(absl::InvalidArgumentError(
+              absl::StrCat("The restored tensor has a different dtype than the "
+                           "variable handle: ",
+                           op_kernel_context.mutable_output(i)->dtype(),
+                           " vs. ", shard.restored_dtypes[i])));
+      if (async_state->dtype_and_shape_results[i].has_value()) {
+        std::move(*async_state->dtype_and_shape_results[i])
+            .Set(absl::InvalidArgumentError(absl::StrCat(
+                "The restored tensor has a different dtype than the "
+                "variable handle: ",
+                op_kernel_context.mutable_output(i)->dtype(), " vs. ",
+                shard.restored_dtypes[i])));
+      }
+      return;
+    }
+    const ResourceHandle& var_handle =
+        shard.var_handles[i].tensor().scalar<tensorflow::ResourceHandle>()();
+
+    DtypeAndShape dtype_and_shape{
+        .dtype = op_kernel_context.mutable_output(i)->dtype(),
+        .shape = op_kernel_context.mutable_output(i)->shape(),
+    };
+
+    if (shard.restored_dtypes[i] == var_handle.dtypes_and_shapes()[0].dtype) {
+      std::move(async_state->host_tensor_results[i])
+          .Set(*std::move(op_kernel_context.mutable_output(i)));
+      if (async_state->dtype_and_shape_results[i].has_value()) {
+        std::move(*async_state->dtype_and_shape_results[i])
+            .Set(std::move(dtype_and_shape));
+      }
+    } else {
+      absl::StatusOr<tensorflow::Tensor> cast_output =
+          Cast(*op_kernel_context.mutable_output(i), shard.restored_dtypes[i],
+               var_handle.dtypes_and_shapes()[0].dtype,
+               shard.truncate_in_cast[i], async_state->device_manager,
+               async_state->process_function_library_runtime,
+               async_state->run_state.params);
+      if (!cast_output.ok()) {
+        std::move(async_state->host_tensor_results[i])
+            .Set(cast_output.status());
+        if (async_state->dtype_and_shape_results[i].has_value()) {
+          std::move(*async_state->dtype_and_shape_results[i])
+              .Set(cast_output.status());
+        }
+        return;
+      } else {
+        dtype_and_shape.dtype = cast_output->dtype();
+        std::move(async_state->host_tensor_results[i])
+            .Set(*std::move(cast_output));
+        if (async_state->dtype_and_shape_results[i].has_value()) {
+          std::move(*async_state->dtype_and_shape_results[i])
+              .Set(std::move(dtype_and_shape));
+        }
+      }
+    }
+  }
+}
+
+absl::Status RunShard(RestoreVariableShard shard,
+                      IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry,
+                      tfrt::ConcurrentWorkQueue* checkpoint_loader_work_queue,
+                      tf_mlrt::Context& context, bool use_async_restore) {
+  if (!ifrt_restore_tensor_registry) {
+    return absl::InternalError("ifrt_restore_tensor_registry must not be null");
+  }
+  if (!checkpoint_loader_work_queue) {
+    return absl::InternalError("checkpoint_loader_work_queue must not be null");
+  }
+  const int num_outputs = shard.var_handles.size();
+  DCHECK_EQ(num_outputs, shard.tensor_names.NumElements());
+  auto& fallback_request_state = context.fallback_request_state();
+
+  // Use `tf.RestoreV2` to restore tensor. This will also populate
+  // tensorflow::ResourceManager.
+  // TODO(b/319045348): avoid populating tensorflow::ResourceManager if the
+  // variable is only used by device/IFRT.
+  // TODO(b/319045348): consider directly calling restore function such as that
+  // in /tensorflow/core/kernels/save_restore_v2_ops.cc
+  auto runner =
+      tfrt_stub::OpKernelRunner::Create(
+          /*op_name=*/
+          "RestoreV2", /*node_name=*/"RestoreV2",
+          context.params().device->name(),
+          /*num_args=*/3,
+          [&](tensorflow::AttrValueMap* attr_value_map) {
+            attr_value_map->insert({"dtypes", shard.dtypes_attr_value});
+            return absl::OkStatus();
+          },
+          fallback_request_state.device_manager(),
+          fallback_request_state.process_function_library_runtime())
+          .value();
+
+  // Prepare the input tensors.
+  std::vector<tensorflow::TensorValue> input_tf_tensor_values;
+  static constexpr int kNumInputArgs = 3;
+  input_tf_tensor_values.resize(kNumInputArgs);
+  // We need to keep these tensor alive
+  input_tf_tensor_values[0].tensor = &shard.prefix;
+  input_tf_tensor_values[1].tensor = &shard.tensor_names;
+  input_tf_tensor_values[2].tensor = &shard.shape_and_slices;
+
+  auto& params = context.params();
+  tf_mlrt::SetUpParams(runner, input_tf_tensor_values, params);
+  // Use persistent device instead of the per request device.
+  params.device = context.fallback_request_state().device_manager().HostCPU();
+
+  auto async_state = std::make_unique<AsyncState>(
+      input_tf_tensor_values, params, num_outputs,
+      fallback_request_state.device_manager(),
+      fallback_request_state.process_function_library_runtime());
+
+  for (int i = 0; i < num_outputs; ++i) {
+    auto [host_tensor_promise, host_tensor_future] =
+        tsl::MakePromise<tensorflow::Tensor>();
+    auto [dtype_and_shape_promise, dtype_and_shape_future] =
+        tsl::MakePromise<ifrt_serving::DtypeAndShape>();
+    const ResourceHandle& var_handle =
+        shard.var_handles[i].tensor().scalar<tensorflow::ResourceHandle>()();
+
+    std::string runtime_name =
+        ifrt_serving::GetRuntimeNameFromVarHandle(var_handle);
+
+    ifrt_serving::IfrtRestoreTensorRegistry::RestoredTensorInfo
+        restored_tensor_info = {false, std::move(dtype_and_shape_future),
+                                std::move(host_tensor_future)};
+    if (auto status = ifrt_restore_tensor_registry->TryRegister(
+            runtime_name, restored_tensor_info);
+        !status.ok()) {
+      // Propagate errors so that if already-registered futures are being waited
+      // on, they can be unblocked.
+      for (auto& result : async_state->host_tensor_results) {
+        std::move(result).Set(status);
+      };
+      for (auto& result : async_state->dtype_and_shape_results) {
+        if (result.has_value()) {
+          std::move(*result).Set(status);
+        }
+      };
+      return status;
+    }
+    absl::StatusOr<ifrt_serving::DtypeAndShape> var_handle_dtype_and_shape =
+        ifrt_serving::GetDtypeAndShape(var_handle);
+
+    // If dtype_and_shape is available at restore graph at compile time, set it
+    // in the registry immediately. Otherwise, wait for the restore op to
+    // complete to set it.
+    if (var_handle_dtype_and_shape.ok()) {
+      std::move(dtype_and_shape_promise)
+          .Set(*std::move(var_handle_dtype_and_shape));
+      async_state->dtype_and_shape_results.push_back(std::nullopt);
+    } else {
+      async_state->dtype_and_shape_results.push_back(
+          std::move(dtype_and_shape_promise));
+    }
+    async_state->host_tensor_results.push_back(std::move(host_tensor_promise));
+  }
+  // Run the shard synchronously.
+  if (!use_async_restore) {
+    RunShardHelper(runner, async_state.get(), shard,
+                   ifrt_restore_tensor_registry);
+  } else {
+    tensorflow::Context bg_context(tensorflow::ContextKind::kThread);
+    // Use dedicated work queue for restore operation.
+    checkpoint_loader_work_queue->AddTask(
+        [runner = std::move(runner), async_state = std::move(async_state),
+         shard = std::move(shard), bg_context = std::move(bg_context),
+         ifrt_restore_tensor_registry = ifrt_restore_tensor_registry]() {
+          tensorflow::WithContext wc(bg_context);
+          RunShardHelper(runner, async_state.get(), shard,
+                         ifrt_restore_tensor_registry);
+        });
+  }
+
+  return absl::OkStatus();
+}
+
+int64_t GetSizeFromVarHandle(const ResourceHandle& handle) {
+  int size = 0;
+  for (auto& dtype_and_shape : handle.dtypes_and_shapes()) {
+    size += DataTypeSize(dtype_and_shape.dtype) *
+            dtype_and_shape.shape.num_elements();
+  }
+  return size;
+}
+
+}  // namespace
+
+absl::Status CheckpointLoader::PrepareRestore(const PrepareRestoreArgs& args) {
+  VLOG(1) << "Skip CheckpointLoader::PrepareRestore";
+  return absl::OkStatus();
+}
+
+absl::Status CheckpointLoader::Load(
+    const tensorflow::tfrt_stub::FallbackTensor& prefix,
+    const std::vector<tensorflow::tfrt_stub::FallbackTensor>& var_handles,
+    const tensorflow::tfrt_stub::FallbackTensor& tensor_names,
+    const tensorflow::tfrt_stub::FallbackTensor& shape_and_slices,
+    absl::Span<const tensorflow::DataType> restored_dtypes,
+    const std::vector<bool>& truncate_in_cast, tf_mlrt::Context& context) {
+  std::vector<int64_t> variable_sizes;
+  variable_sizes.reserve(var_handles.size());
+  for (auto& handle : var_handles) {
+    variable_sizes.push_back(GetSizeFromVarHandle(
+        handle.tensor().scalar<tensorflow::ResourceHandle>()()));
+  }
+
+  std::vector<std::vector<int>> sharded_indices = tf_mlrt::ShardVariables(
+      kNumRestoreClusters, absl::MakeSpan(variable_sizes));
+
+  // Converts the names and slices back to the tensor.
+  auto vector_to_tensor = [](const std::vector<tsl::tstring>& vec) {
+    tensorflow::Tensor tensor(tensorflow::DT_STRING,
+                              TensorShape({static_cast<int>(vec.size())}));
+    for (int i = 0; i < vec.size(); ++i) {
+      tensor.flat<tsl::tstring>()(i) = vec[i];
+    }
+    return tensor;
+  };
+
+  const auto& tensor_names_flat = tensor_names.tensor().flat<tsl::tstring>();
+  const auto& shape_and_slices_flat =
+      shape_and_slices.tensor().flat<tsl::tstring>();
+
+  std::vector<RestoreVariableShard> shards;
+  shards.reserve(sharded_indices.size());
+  for (auto& sharded_index : sharded_indices) {
+    RestoreVariableShard shard;
+    shard.var_handles.reserve(sharded_index.size());
+    shard.truncate_in_cast.reserve(sharded_index.size());
+    shard.restored_dtypes.reserve(sharded_index.size());
+    std::vector<tsl::tstring> tensor_names;
+    std::vector<tsl::tstring> shape_and_slices;
+    shape_and_slices.reserve(sharded_index.size());
+    tensor_names.reserve(sharded_index.size());
+    for (int index : sharded_index) {
+      tensor_names.push_back(tensor_names_flat(index));
+      shape_and_slices.push_back(shape_and_slices_flat(index));
+      shard.dtypes_attr_value.mutable_list()->add_type(restored_dtypes[index]);
+      shard.var_handles.push_back(var_handles[index]);
+      shard.restored_dtypes.push_back(restored_dtypes[index]);
+      shard.truncate_in_cast.push_back(truncate_in_cast[index]);
+    }
+    shard.prefix = prefix.tensor();
+    shard.tensor_names = vector_to_tensor(tensor_names);
+    shard.shape_and_slices = vector_to_tensor(shape_and_slices);
+    shards.push_back(std::move(shard));
+  }
+  for (const auto& shard : shards) {
+    TF_RETURN_IF_ERROR(RunShard(shard, ifrt_restore_tensor_registry_,
+                                checkpoint_loader_work_queue_, context,
+                                use_async_restore_));
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/checkpoint_loader.h

@@ -0,0 +1,83 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_CHECKPOINT_LOADER_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_CHECKPOINT_LOADER_H_
+
+#include <string>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/protobuf/meta_graph.pb.h"
+#include "tensorflow/core/tfrt/fallback/fallback_state.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/mlrt/bytecode/bytecode.h"
+#include "tensorflow/core/tfrt/mlrt/kernel/context.h"
+#include "tensorflow/core/tfrt/utils/fallback_tensor.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// TODO(b/352551302) Move the unit test in ifrt_ops_kernel for restore to test
+// this class's APIs.
+// Implement the `CheckpointLoaderInterface` by using RestoreV2.
+class CheckpointLoader {
+ public:
+  struct PrepareRestoreArgs {
+    mlir::MLIRContext* context;
+    tensorflow::MetaGraphDef meta_graph_def;
+    tfrt_stub::FallbackState* fallback_state;
+    std::string saved_model_dir;
+    bool run_placer_grappler_on_functions;
+  };
+
+  explicit CheckpointLoader(
+      IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry,
+      tfrt::ConcurrentWorkQueue* checkpoint_loader_work_queue,
+      bool use_async_restore = true)
+      : ifrt_restore_tensor_registry_(ifrt_restore_tensor_registry),
+        checkpoint_loader_work_queue_(checkpoint_loader_work_queue),
+        use_async_restore_(use_async_restore) {}
+  virtual ~CheckpointLoader() = default;
+
+  // Called before `Load` to do some preparation work.
+  virtual absl::Status PrepareRestore(const PrepareRestoreArgs& args);
+
+  // Load the checkpoint. This API is designed to be compatible with the
+  // `tf_mlrt.ifrt_restore_variable` kernel.
+  virtual absl::Status Load(
+      const tensorflow::tfrt_stub::FallbackTensor& prefix,
+      const std::vector<tensorflow::tfrt_stub::FallbackTensor>& var_handles,
+      const tensorflow::tfrt_stub::FallbackTensor& tensor_names,
+      const tensorflow::tfrt_stub::FallbackTensor& shape_and_slices,
+      absl::Span<const tensorflow::DataType> restored_dtypes,
+      const std::vector<bool>& truncate_in_cast, tf_mlrt::Context& context);
+
+ protected:
+  IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry_;
+  tfrt::ConcurrentWorkQueue* checkpoint_loader_work_queue_;
+  bool use_async_restore_ = true;
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_CHECKPOINT_LOADER_H_

tensorflow/core/tfrt/ifrt/grid.cc

@@ -0,0 +1,30 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/grid.h"
+
+#include <string>
+
+#include "absl/strings/str_cat.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+std::string GridCoords::ToString() const {
+  return absl::StrCat("[", dim[0], ",", dim[1], ",", dim[2], ",", dim[3], "]");
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/grid.h

@@ -0,0 +1,77 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_GRID_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_GRID_H_
+
+#include <ostream>
+#include <string>
+
+#include "absl/log/check.h"
+#include "absl/strings/str_format.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// Coordinates that identify a particular point in a 4-d grid (usually a TPU
+// topology).
+struct GridCoords {
+  int dim[4];
+
+  constexpr GridCoords(int d0, int d1, int d2, int d3) : dim{d0, d1, d2, d3} {}
+  GridCoords() : GridCoords(0, 0, 0, 0) {}
+
+  static GridCoords Zeroes() { return GridCoords(0, 0, 0, 0); }
+  static GridCoords Ones() { return GridCoords(1, 1, 1, 1); }
+
+  int operator[](int i) const {
+    DCHECK_GE(i, 0);
+    DCHECK_LT(i, 4);
+    return dim[i];
+  }
+
+  int& operator[](int i) {
+    DCHECK_GE(i, 0);
+    DCHECK_LT(i, 4);
+    return dim[i];
+  }
+
+  int Product() const { return dim[0] * dim[1] * dim[2] * dim[3]; }
+
+  std::string ToString() const;
+
+  template <typename Sink>
+  friend void AbslStringify(Sink& sink, const GridCoords& value) {
+    absl::Format(&sink, "%s", value.ToString());
+  }
+
+  friend bool operator==(const GridCoords& a, const GridCoords& b) {
+    return a[0] == b[0] && a[1] == b[1] && a[2] == b[2] && a[3] == b[3];
+  }
+
+  friend std::ostream& operator<<(std::ostream& os, const GridCoords& c) {
+    return os << c.ToString();
+  }
+
+  template <typename H>
+  friend H AbslHashValue(H h, const GridCoords& c) {
+    return H::combine(std::move(h), c[0], c[1], c[2], c[3]);
+  }
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_GRID_H_

tensorflow/core/tfrt/ifrt/ifrt_config.proto

@@ -0,0 +1,50 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.ifrt_serving;
+
+import "xla/xla_data.proto";
+import "tensorflow/core/framework/tensor.proto";
+
+enum IfrtPjRtServingPlatformType {
+  IFRT_PJRT_SERVING_PLATFORM_TYPE_UNSPECIFIED = 0;
+  IFRT_PJRT_SERVING_PLATFORM_TYPE_TPU = 1;
+  IFRT_PJRT_SERVING_PLATFORM_TYPE_GPU = 2;
+}
+
+// `tf.IfrtLoadVariableOp`'s configuration message for device and spmd
+// configuration.
+message VariableDeviceShardingConfigProto {
+  xla.OpSharding sharding = 1;
+  repeated int32 device_ids = 2;
+}
+
+enum IfrtServingCoreSelectionPolicy {
+  // Default policy to select the soonest-to-finish core defined at
+  // http://shortn/_XQhI1ASAu0.
+  IFRT_SERVING_CORE_SELECTION_POLICY_DEFAULT = 0;
+  // Policy that round robin with local ordinal http://shortn/_7BtVe4dkp5.
+  IFRT_SERVING_CORE_SELECTION_POLICY_LOCAL_ROUND_ROBIN = 1;
+}
+
+message DefaultSignatureInputConfig {
+  message Signature {
+    map<string, tensorflow.TensorProto> default_inputs = 1;
+  }
+
+  map<string, Signature> signatures = 1;
+}

tensorflow/core/tfrt/ifrt/ifrt_device_utils.cc

@@ -0,0 +1,192 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_device_utils.h"
+
+#include <cstdint>
+#include <optional>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/attribute_map.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/device.h"
+#include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/host_callback.h"
+#include "xla/service/computation_placer.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/example/feature.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tensorflow/core/tfrt/ifrt/grid.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_config.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+static constexpr int kTpuTopologyRank = 4;  // (x, y, z, core).
+
+absl::StatusOr<std::vector<xla::ifrt::Device*>> GetAssignedIfrtDevices(
+    const xla::ifrt::Client& ifrt_client, int num_replicas,
+    int num_cores_per_replica,
+    std::optional<std::vector<int>> device_assignment) {
+  const int num_devices = num_replicas * num_cores_per_replica;
+
+  // No device coordinates from ifrt devices. This disallow the mapping from
+  // device assignment attribute to ifrt devices.
+  bool no_device_coordinates = false;
+  for (auto* device : ifrt_client.devices()) {
+    auto coords_attr = device->Attributes().Get<std::vector<int64_t>>("coords");
+    auto core_on_chip_attr = device->Attributes().Get<int64_t>("core_on_chip");
+    if (!coords_attr.ok() || !core_on_chip_attr.ok()) {
+      no_device_coordinates = true;
+      break;
+    }
+  }
+
+  // If there is no device assignment attribute or no device coordinates,
+  // get the default device assignment from IFRT.
+  if (!device_assignment || device_assignment->empty() ||
+      no_device_coordinates) {
+    TF_ASSIGN_OR_RETURN(xla::DeviceAssignment xla_device_assignment,
+                        ifrt_client.GetDefaultDeviceAssignment(
+                            num_replicas, num_cores_per_replica));
+    VLOG(3) << "Getting default device lists";
+    std::vector<xla::ifrt::Device*> devices;
+    devices.reserve(num_devices);
+    for (int replica_idx = 0; replica_idx < num_replicas; replica_idx++) {
+      for (int core_idx = 0; core_idx < num_cores_per_replica; core_idx++) {
+        // This relies on the IFRT implementation of GetDefaultDeviceAssignment
+        // that keeps device id the same between device assignment  and ifrt
+        // device list.
+        auto device_id = xla_device_assignment(replica_idx, core_idx);
+        TF_ASSIGN_OR_RETURN(
+            xla::ifrt::Device * device,
+            ifrt_client.LookupDevice(xla::ifrt::DeviceId(device_id)));
+        devices.push_back(device);
+      }
+    }
+    return devices;
+  }
+
+  // Devices ordered as in the device assignment attribute.
+  absl::flat_hash_map<GridCoords, int> devices_from_attribute;
+
+  // Each device is encoded by [x,y,z,c] at the attribute.
+  std::vector<int> coord;
+  coord.reserve(kTpuTopologyRank);
+  int device_index = 0;
+
+  for (auto coord_attr : *device_assignment) {
+    coord.push_back(coord_attr);
+    if (coord.size() == kTpuTopologyRank) {
+      devices_from_attribute.insert(
+          {GridCoords(coord[0], coord[1], coord[2], coord[3]), device_index});
+      device_index++;
+      coord.clear();
+    }
+  }
+  if (!coord.empty()) {
+    return absl::FailedPreconditionError(
+        absl::StrCat("Device assignment attribute is expected to be a multiple "
+                     "of 4, but got ",
+                     device_assignment->size()));
+  }
+
+  if (devices_from_attribute.size() != num_devices) {
+    return absl::FailedPreconditionError(
+        absl::StrCat("Device assignment has ", devices_from_attribute.size(),
+                     " devices, but expected ", num_devices));
+  }
+
+  struct IfrtDeviceGrid {
+    xla::ifrt::Device* device;
+    GridCoords grid;
+    int index_at_attribute;
+  };
+  std::vector<IfrtDeviceGrid> ifrt_devices;
+  ifrt_devices.reserve(num_devices);
+
+  for (auto* device : ifrt_client.devices()) {
+    GridCoords grid;
+    auto coords_attr = device->Attributes().Get<std::vector<int64_t>>("coords");
+    auto core_on_chip_attr = device->Attributes().Get<int64_t>("core_on_chip");
+    if (coords_attr.ok() && core_on_chip_attr.ok()) {
+      VLOG(3) << "Adding coords and core_on_chip attributes:"
+              << device->DebugString();
+      auto coords = *coords_attr;
+      auto core_on_chip = *core_on_chip_attr;
+
+      if (coords.size() != 3) {
+        return absl::InternalError(
+            absl::StrCat("Expected coords to be of size 3, but got ",
+                         coords.size(), " for device ", device->DebugString()));
+      }
+      grid = GridCoords(coords[0], coords[1], coords[2], core_on_chip);
+    } else {
+      return absl::InternalError(
+          absl::StrCat("Device ", device->DebugString(),
+                       " does not have coords or core_on_chip attribute."));
+    }
+
+    auto device_it_from_attribute = devices_from_attribute.find(grid);
+    if (device_it_from_attribute == devices_from_attribute.end()) {
+      VLOG(1) << "Device coordinates " << grid.ToString()
+              << " does not match any TPU device assigned "
+              << absl::StrJoin(*device_assignment, " ");
+      continue;
+    }
+    ifrt_devices.push_back(
+        {.device = device,
+         .grid = grid,
+         .index_at_attribute = device_it_from_attribute->second});
+  }
+
+  if (ifrt_devices.size() != num_devices) {
+    return absl::FailedPreconditionError(absl::StrCat(
+        "Match ", ifrt_devices.size(), " devices, but expected ", num_devices));
+  }
+
+  // Sort the devices by the order in the device assignment attribute.
+  absl::c_sort(ifrt_devices, [&](const auto& lhs, const auto& rhs) {
+    return lhs.index_at_attribute < rhs.index_at_attribute;
+  });
+
+  std::vector<xla::ifrt::Device*> result;
+  result.reserve(ifrt_devices.size());
+  for (auto& device_grid : ifrt_devices) {
+    result.push_back(device_grid.device);
+    VLOG(3) << "Device: " << device_grid.device->DebugString()
+            << " is assigned";
+  }
+  return result;
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_device_utils.h

@@ -0,0 +1,63 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_DEVICE_UTILS_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_DEVICE_UTILS_H_
+
+#include <optional>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/compiler/tf2xla/host_compute_metadata.pb.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/device.h"
+#include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/host_callback.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/example/feature.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_config.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// Returns the assigned IFRT devices based on the device assignment attribute.
+//
+// params:
+//   ifrt_client: The ifrt client.
+//   num_replicas: The number of replicas.
+//   num_cores_per_replica: The number of cores per replica.
+//
+//   device_assignment: The device assignment array encoded as
+//   [x0,y0,z0,core0,x1,y1,z1,core1, ...]. Optional. If not provided, the
+//   devices will be assigned based on the default order returned by the IFRT
+//   client.
+//
+// returns:
+//   The assigned devices.
+absl::StatusOr<std::vector<xla::ifrt::Device*>> GetAssignedIfrtDevices(
+    const xla::ifrt::Client& ifrt_client, int num_replicas,
+    int num_cores_per_replica,
+    std::optional<std::vector<int>> device_assignment);
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_DEVICE_UTILS_H_

tensorflow/core/tfrt/ifrt/ifrt_device_utils_test.cc

@@ -0,0 +1,132 @@
+/* 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/core/tfrt/ifrt/ifrt_device_utils.h"
+
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "xla/python/ifrt/attribute_map.h"
+#include "xla/python/ifrt/device.h"
+#include "xla/python/ifrt/mock.h"
+#include "xla/service/computation_placer.h"
+#include "xla/tsl/platform/status_matchers.h"
+#include "xla/tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+using ::testing::ElementsAre;
+using ::testing::Return;
+using ::testing::ReturnRef;
+
+static constexpr int kNumReplicas = 1;
+static constexpr int kNumCoresPerReplica = 2;
+// Intentionally have more devices than kNumReplicas * kNumCoresPerReplica for
+// testing purposes.
+static constexpr int kNumDevices = 4;
+static constexpr int kDeviceIdOffset = 8;
+
+class IfrtDeviceUtilsTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    mocked_devices_.reserve(kNumDevices);
+    devices_.reserve(kNumDevices);
+    for (int i = 0; i < kNumDevices; ++i) {
+      mocked_devices_.push_back(std::make_unique<xla::ifrt::MockDevice>());
+      ON_CALL(*mocked_devices_[i], Attributes())
+          .WillByDefault(ReturnRef(device_attributes_maps_[i]));
+      ON_CALL(*mocked_devices_[i], Id())
+          .WillByDefault(Return(xla::ifrt::DeviceId(kDeviceIdOffset + i)));
+      ON_CALL(client_, LookupDevice(xla::ifrt::DeviceId(kDeviceIdOffset + i)))
+          .WillByDefault(Return(mocked_devices_[i].get()));
+
+      devices_.push_back(mocked_devices_[i].get());
+    };
+
+    ON_CALL(client_, devices()).WillByDefault(Return(devices_));
+
+    // Default use the last two devices.
+    xla::DeviceAssignment assignment(kNumReplicas, kNumCoresPerReplica);
+    assignment(0, 0) = kDeviceIdOffset + 2;
+    assignment(0, 1) = kDeviceIdOffset + 3;
+
+    ON_CALL(client_,
+            GetDefaultDeviceAssignment(kNumReplicas, kNumCoresPerReplica))
+        .WillByDefault(Return(assignment));
+  }
+
+  xla::ifrt::MockClient client_;
+  std::vector<std::unique_ptr<xla::ifrt::MockDevice>> mocked_devices_;
+
+  std::vector<xla::ifrt::Device*> devices_;
+  std::vector<xla::ifrt::AttributeMap> device_attributes_maps_ = {
+      xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{
+          {"coords", xla::ifrt::AttributeMap::Int64ListValue({1, 0, 0})},
+          {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(0)}}),
+      xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{
+          {"coords", xla::ifrt::AttributeMap::Int64ListValue({1, 0, 0})},
+          {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(1)}}),
+      xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{
+          {"coords", xla::ifrt::AttributeMap::Int64ListValue({2, 0, 0})},
+          {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(0)}}),
+      xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{
+          {"coords", xla::ifrt::AttributeMap::Int64ListValue({2, 0, 0})},
+          {"core_on_chip", xla::ifrt::AttributeMap::Int64Value(1)}}),
+  };
+};
+
+TEST_F(IfrtDeviceUtilsTest, Basic) {
+  std::vector<int> device_assignment_attr = {1, 0, 0, 1, 1, 0, 0, 0};
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto devices_from_attribute,
+      GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica,
+                             device_assignment_attr));
+  EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[1], devices_[0]));
+}
+
+TEST_F(IfrtDeviceUtilsTest, SeparateXCoordinates) {
+  std::vector<int> device_assignment_attr = {1, 0, 0, 1, 2, 0, 0, 0};
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto devices_from_attribute,
+      GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica,
+                             device_assignment_attr));
+  EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[1], devices_[2]));
+}
+
+TEST_F(IfrtDeviceUtilsTest, EmptyDeviceAssignmentShallReturnDefault) {
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto devices_from_attribute,
+      GetAssignedIfrtDevices(client_, kNumReplicas, kNumCoresPerReplica,
+                             std::nullopt));
+  EXPECT_THAT(devices_from_attribute, ElementsAre(devices_[2], devices_[3]));
+}
+
+TEST_F(IfrtDeviceUtilsTest, MismatchCoordinatesShallFail) {
+  std::vector<int> device_assignment_attr = {1, 0, 0, 1, 3, 0, 0, 0};
+  auto status = GetAssignedIfrtDevices(client_, 1, 2, device_assignment_attr);
+  EXPECT_THAT(status,
+              absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+}
+
+}  // namespace
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_executable_registry.cc

@@ -0,0 +1,126 @@
+/* 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/core/tfrt/ifrt/ifrt_executable_registry.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <utility>
+
+#include "absl/base/attributes.h"
+#include "absl/base/const_init.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+ServingExecutableRegistry::Handle::Handle(Handle&& other) {
+  *this = std::move(other);
+}
+
+ServingExecutableRegistry::Handle& ServingExecutableRegistry::Handle::operator=(
+    Handle&& other) {
+  if (this != &other) {
+    program_id_ = std::move(other.program_id_);
+    other.program_id_ = std::nullopt;
+  }
+  return *this;
+}
+
+ServingExecutableRegistry::Handle::~Handle() { Release(); }
+
+absl::Status ServingExecutableRegistry::Handle::Freeze() {
+  if (!program_id_.has_value()) {
+    return absl::FailedPreconditionError("Program is not registered");
+  }
+
+  absl::MutexLock l(ServingExecutableRegistry::mu_);
+
+  const auto it = ServingExecutableRegistry::executables_->find(*program_id_);
+  if (it == ServingExecutableRegistry::executables_->end()) {
+    return absl::NotFoundError(
+        absl::StrCat("Program ", *program_id_, " not found in the registry"));
+  }
+
+  VLOG(1) << "Freeze the program " << *program_id_ << " from signature '"
+          << it->second->signature_name() << "' of model '"
+          << it->second->model_name() << "'";
+
+  it->second->Freeze();
+  return absl::OkStatus();
+}
+
+void ServingExecutableRegistry::Handle::Release() {
+  if (!program_id_.has_value()) {
+    return;
+  }
+
+  absl::MutexLock l(ServingExecutableRegistry::mu_);
+
+  const auto it = ServingExecutableRegistry::executables_->find(*program_id_);
+  if (it == ServingExecutableRegistry::executables_->end()) {
+    LOG(ERROR) << "Program " << *program_id_ << " not found in the registry";
+    return;
+  }
+
+  VLOG(1) << "Unregistering program " << *program_id_ << " from signature '"
+          << it->second->signature_name() << "' of model '"
+          << it->second->model_name() << "'";
+  ServingExecutableRegistry::executables_->erase(it);
+
+  program_id_ = std::nullopt;
+}
+
+ServingExecutableRegistry::Handle::Handle(int64_t program_id)
+    : program_id_(program_id) {}
+
+absl::StatusOr<ServingExecutableRegistry::Handle>
+ServingExecutableRegistry::Register(
+    int64_t program_id, std::unique_ptr<IfrtServingExecutable> executable) {
+  absl::MutexLock l(mu_);
+  VLOG(1) << "Registering program " << program_id << " from signature '"
+          << executable->signature_name() << "' of model '"
+          << executable->model_name() << "'"
+          << ", address is " << executable.get();
+  if (!executables_->insert({program_id, std::move(executable)}).second) {
+    return absl::AlreadyExistsError(absl::StrCat(
+        "Program ", program_id, " already exists in the program registry"));
+  }
+  return Handle(program_id);
+}
+
+IfrtServingExecutable* ServingExecutableRegistry::Lookup(int64_t program_id) {
+  absl::ReaderMutexLock l(mu_);
+  VLOG(1) << "Looking up program " << program_id;
+  const auto it = executables_->find(program_id);
+  return it != executables_->end() ? it->second.get() : nullptr;
+}
+
+ABSL_CONST_INIT absl::Mutex ServingExecutableRegistry::mu_(absl::kConstInit);
+
+absl::flat_hash_map<int64_t, std::unique_ptr<IfrtServingExecutable>>* const
+    ServingExecutableRegistry::executables_ =
+        new absl::flat_hash_map<int64_t,
+                                std::unique_ptr<IfrtServingExecutable>>();
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h

@@ -0,0 +1,104 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_EXECUTABLE_REGISTRY_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_EXECUTABLE_REGISTRY_H_
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// Maintains a process-wide map from program ids to executables. Used by the
+// `IfrtCall` TensorFlow op to look up executables and invoke them.
+//
+// Invoking a TPU program inside a `IfrtCall` TF op requires being
+// able to retrieve an executable for the given program. Since there's no easy
+// way to pass non-serializable attributes to TF ops, we encode a program id
+// (that is unique within a process) as an attribute of a `IfrtCall` op and
+// use this registry class to let the `IfrtCall` op look up an executable
+// during TF op execution.
+class ServingExecutableRegistry {
+ public:
+  // RAII handle for registered executables.
+  class Handle {
+   public:
+    Handle();  // Constructs an empty handle.
+
+    // Move only.
+    Handle(Handle&& other);
+    Handle& operator=(Handle&& other);
+    Handle(const Handle&) = delete;
+    Handle& operator=(const Handle&) = delete;
+
+    ~Handle();
+
+    // Returns the program id that the handle represents, or `std::nullopt` if
+    // the handle is empty.
+    std::optional<int64_t> program_id() const { return program_id_; }
+
+    // Unregisters the owned executable, if any, early (before the destructor).
+    // Calling this method multiple times is a no-op.
+    void Release();
+
+    // Freezes the program's compilation. After Freeze() is called, no new model
+    // signature will be compiled. Using a signature or an input shape that
+    // wasn't compiled before the freeze will lead to an error.
+    absl::Status Freeze();
+
+   private:
+    friend class ServingExecutableRegistry;
+
+    // Can only be constructed by `ServingExecutableRegistry::Register()`.
+    explicit Handle(int64_t program_id);
+
+    // Program id. `std::nullopt` if the handle is already released.
+    std::optional<int64_t> program_id_;
+  };
+
+  // Registers an executable under the given program id. Returns an RAII handle
+  // that unregisters the program at its destruction.
+  static absl::StatusOr<Handle> Register(
+      int64_t program_id, std::unique_ptr<IfrtServingExecutable> executable);
+
+  // Looks up an executable registered under the given program id, or returns
+  // nullptr if there's no such program.
+  static IfrtServingExecutable* Lookup(int64_t program_id);
+
+ private:
+  friend class Handle;
+  friend class IfrtBackendCompilerTest;
+
+  static absl::Mutex mu_;
+
+  // Mapping from program ids to executables.
+  static absl::flat_hash_map<int64_t,
+                             std::unique_ptr<IfrtServingExecutable>>* const
+      executables_ ABSL_GUARDED_BY(&mu_);
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_EXECUTABLE_REGISTRY_H_

tensorflow/core/tfrt/ifrt/ifrt_executable_registry_test.cc

@@ -0,0 +1,221 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/test_util.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
+#include "tensorflow/core/tfrt/ifrt/tf_host_callback.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+tsl::thread::ThreadPool& GetThreadPool() {
+  constexpr int kMaxParallelism = 16;
+  static auto* thread_pool =
+      new tsl::thread::ThreadPool(tsl::Env::Default(), tsl::ThreadOptions(),
+                                  "IfrtSharding", kMaxParallelism);
+  return *thread_pool;
+}
+
+absl::StatusOr<std::unique_ptr<IfrtServingExecutable>>
+CreateIfrtServingExecutable(mlir::MLIRContext& context, int64_t program_id) {
+  // Create test input module
+  constexpr absl::string_view kDataDirectory =
+      "tensorflow/core/tfrt/ifrt/testdata";
+  std::string mlir_module_path = tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kDataDirectory, "/executable.mlir"));
+
+  mlir::OwningOpRef<mlir::ModuleOp> mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context);
+
+  if (!mlir_module) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Failed to parse MLIR file: ", mlir_module_path));
+  }
+
+  // Create contexts required for the compiler execution.
+  TF_ASSIGN_OR_RETURN(std::shared_ptr<xla::ifrt::Client> client,
+                      xla::ifrt::test_util::GetClient());
+
+  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry;
+  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry;
+  H2DTransferExecutorFactory h2d_transfer_executor_factory;
+  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue =
+      tfrt::CreateMultiThreadedWorkQueue(
+          /*num_threads=*/4, /*num_blocking_threads=*/4);
+  TF_ASSIGN_OR_RETURN(std::unique_ptr<tensorflow::DynamicDeviceMgr> device_mgr,
+                      CreateTfDynamicDeviceMgr());
+
+  return IfrtServingExecutable::Create(
+      program_id, "test", "main", std::move(mlir_module), client,
+      &GetThreadPool(), &ifrt_loaded_variable_registry,
+      &ifrt_restore_tensor_registry, work_queue.get(), device_mgr.get(),
+      tensorflow::IdentityShapeRepresentationFn(),
+      /*ifrt_serving_core_selector=*/nullptr,
+      /*compilation_environment_proto=*/nullptr,
+      /*tf_to_hlo_compiler=*/nullptr, /*persistent_compilation_cache=*/nullptr,
+      &h2d_transfer_executor_factory);
+}
+
+TEST(IfrtExecutableRegistry, Basic) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+  IfrtServingExecutable* raw_ptr = executable.get();
+
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  IfrtServingExecutable* executable_ptr =
+      ServingExecutableRegistry::Lookup(program_id);
+  ASSERT_EQ(executable_ptr, raw_ptr);
+}
+
+TEST(IfrtExecutableRegistry, DuplicateRegistrationFails) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  EXPECT_THAT(
+      ServingExecutableRegistry::Register(program_id, std::move(executable)),
+      absl_testing::StatusIs(absl::StatusCode::kAlreadyExists));
+}
+
+TEST(IfrtExecutableRegistry, ReleaseOk) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  handle.Release();
+
+  EXPECT_EQ(ServingExecutableRegistry::Lookup(program_id), nullptr);
+}
+
+TEST(IfrtExecutableRegistry, FreezeOk) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+  IfrtServingExecutable* raw_ptr = executable.get();
+
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  ASSERT_OK(handle.Freeze());
+
+  // After the freeze, the lookup should still return the same pointer.
+  IfrtServingExecutable* executable_ptr =
+      ServingExecutableRegistry::Lookup(program_id);
+  ASSERT_EQ(executable_ptr, raw_ptr);
+}
+
+TEST(IfrtExecutableRegistry, FreezeFailedProgramNotRegistered) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  mlir::MLIRContext context(registry);
+
+  int64_t program_id = 1234;
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<IfrtServingExecutable> executable,
+                          CreateIfrtServingExecutable(context, program_id));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto handle, ServingExecutableRegistry::Register(
+                                           program_id, std::move(executable)));
+
+  handle.Release();
+
+  EXPECT_THAT(handle.Freeze(),
+              absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+}
+
+TEST(IfrtExecutableRegistry, InvalidProgramIdShallReturnNull) {
+  int64_t program_id = 1234;
+
+  IfrtServingExecutable* executable_ptr =
+      ServingExecutableRegistry::Lookup(program_id);
+  ASSERT_EQ(executable_ptr, nullptr);
+}
+
+}  // namespace
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.cc

@@ -0,0 +1,55 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/synchronization/mutex.h"
+#include "xla/tsl/platform/statusor.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+absl::Status IfrtLoadedVariableRegistry::TryRegisterLoadedVariable(
+    const Key& key, LoadedVariableConstructor&& loaded_variable_constructor) {
+  absl::MutexLock lock(mutex_);
+  auto& variable = loaded_variable_map_[key];
+  if (variable.array.IsValid()) {
+    return absl::AlreadyExistsError(
+        absl::StrCat("Variable '", key.input_name, "' already registered."));
+  }
+  TF_ASSIGN_OR_RETURN(variable, loaded_variable_constructor());
+  return absl::OkStatus();
+}
+
+absl::StatusOr<IfrtLoadedVariableRegistry::LoadedVariable>
+IfrtLoadedVariableRegistry::GetLoadedVariable(KeyView key_view) const {
+  absl::MutexLock lock(mutex_);
+  auto it = loaded_variable_map_.find(key_view);
+  if (it == loaded_variable_map_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("Variable '", key_view.input_name, "' not found."));
+  }
+  return it->second;
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h

@@ -0,0 +1,165 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_LOADED_VARIABLE_REGISTRY_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_LOADED_VARIABLE_REGISTRY_H_
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/functional/any_invocable.h"
+#include "absl/hash/hash.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
+#include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/shape.h"
+#include "xla/tsl/concurrency/future.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// This class is thread safe.
+class IfrtLoadedVariableRegistry {
+ public:
+  // The key is per variable tensor per device assignment. For single -device
+  // program, variables can be loaded on multiple devices with core selection.
+  // For SPMD program, we currently assume all devices will be used, so we use
+  // vector to make it compatible with SPMD.
+  struct Key {
+    // We use a vector to make it compatible with SPMD because the order of the
+    // devices used for sharding must match the order of the devices used for
+    // xla compilation.
+    std::vector<int> device_ids;
+    std::string input_name;
+    xla::HloSharding hlo_sharding;
+    // Use pointer instead of value to avoid copy when building the key.
+    std::shared_ptr<const xla::Shape> shape_on_device;
+    template <typename H>
+    friend H AbslHashValue(H h, const Key& key) {
+      h = H::combine(std::move(h), key.input_name, key.device_ids,
+                     key.hlo_sharding);
+      if (key.shape_on_device != nullptr) {
+        h = H::combine(std::move(h), *key.shape_on_device);
+      }
+      return h;
+    }
+
+    friend bool operator==(const Key& x, const Key& y) {
+      return KeyView(x) == KeyView(y);
+    }
+
+    std::string ToString() const {
+      return absl::StrCat(
+          input_name, ":", absl::StrJoin(device_ids, ","), ":",
+          hlo_sharding.ToString(), ":",
+          (shape_on_device ? shape_on_device->ToString() : "nullptr"));
+    }
+  };
+
+  // A view of Key that references the data without owning it. This is used for
+  // efficient lookups in the LoadedVariable map without constructing a full Key
+  // object.
+  struct KeyView {
+    KeyView(const Key& key)  // NOLINT
+        : device_ids(key.device_ids),
+          input_name(key.input_name),
+          hlo_sharding(key.hlo_sharding),
+          shape_on_device(key.shape_on_device) {}
+
+    KeyView(absl::Span<const int> device_ids, absl::string_view input_name,
+            const xla::HloSharding& hlo_sharding,
+            std::shared_ptr<const xla::Shape> shape_on_device)
+        : device_ids(device_ids),
+          input_name(input_name),
+          hlo_sharding(hlo_sharding),
+          shape_on_device(std::move(shape_on_device)) {}
+
+    absl::Span<const int> device_ids;
+    absl::string_view input_name;
+    const xla::HloSharding& hlo_sharding;
+    std::shared_ptr<const xla::Shape> shape_on_device;
+
+    template <typename H>
+    friend H AbslHashValue(H h, const KeyView& key) {
+      h = H::combine(std::move(h), key.input_name, key.device_ids,
+                     key.hlo_sharding);
+      if (key.shape_on_device != nullptr) {
+        h = H::combine(std::move(h), *key.shape_on_device);
+      }
+      return h;
+    }
+
+    friend bool operator==(const KeyView& x, const KeyView& y) {
+      bool xla_shape_equal = false;
+      if (x.shape_on_device == nullptr && y.shape_on_device == nullptr) {
+        xla_shape_equal = true;
+      } else if (x.shape_on_device != nullptr && y.shape_on_device != nullptr) {
+        xla_shape_equal = *x.shape_on_device == *y.shape_on_device;
+      }
+      return x.input_name == y.input_name && x.device_ids == y.device_ids &&
+             x.hlo_sharding == y.hlo_sharding && xla_shape_equal;
+    }
+  };
+
+  struct KeyEq {
+    using is_transparent = void;
+    bool operator()(const KeyView& lhs, const KeyView& rhs) const {
+      return rhs == lhs;
+    }
+  };
+
+  struct KeyHash {
+    using is_transparent = void;
+    size_t operator()(const KeyView& key) const { return absl::HashOf(key); }
+  };
+
+  struct LoadedVariable {
+    tsl::Future<xla::ifrt::ArrayRef> array;
+  };
+  using LoadedVariableConstructor =
+      absl::AnyInvocable<absl::StatusOr<LoadedVariable>() const>;
+
+  // Tries to register a loaded variable with the given name.
+  // Returns an error if the named array does not already exists and
+  // loaded_variable_constructor failed to create an array. Note that it returns
+  // OK if the named array already exists.
+  // loaded_variable_constructor is invoked in the caller thread.
+  absl::Status TryRegisterLoadedVariable(
+      const Key& key, LoadedVariableConstructor&& loaded_variable_constructor)
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+  absl::StatusOr<LoadedVariable> GetLoadedVariable(KeyView key_view) const
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+ private:
+  mutable absl::Mutex mutex_;
+  absl::flat_hash_map<Key, LoadedVariable, KeyHash, KeyEq> loaded_variable_map_
+      ABSL_GUARDED_BY(mutex_);
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_LOADED_VARIABLE_REGISTRY_H_

tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry_test.cc

@@ -0,0 +1,195 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+
+#include <memory>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/hash/hash.h"
+#include "absl/hash/hash_testing.h"
+#include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/xla_data.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+TEST(IfrtLoadedVariableRegistryTest, KeyEquality) {
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {2, 2});
+  auto shape_ptr1 = std::make_shared<xla::Shape>(shape);
+  auto shape_ptr2 = std::make_shared<xla::Shape>(shape);
+  // shape_ptr3 points to a different shape value.
+  xla::Shape different_shape = xla::ShapeUtil::MakeShape(xla::F32, {3, 3});
+  auto shape_ptr3 = std::make_shared<xla::Shape>(different_shape);
+
+  IfrtLoadedVariableRegistry::Key key1{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr1,
+  };
+
+  IfrtLoadedVariableRegistry::Key key2{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr2,
+  };
+
+  IfrtLoadedVariableRegistry::Key key3{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr3,
+  };
+
+  EXPECT_EQ(key1, key2);
+  EXPECT_NE(key1, key3);
+}
+
+TEST(IfrtLoadedVariableRegistryTest, KeyHash) {
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {2, 2});
+  auto shape_ptr1 = std::make_shared<xla::Shape>(shape);
+  auto shape_ptr2 = std::make_shared<xla::Shape>(shape);
+  xla::Shape different_shape = xla::ShapeUtil::MakeShape(xla::F32, {3, 3});
+  auto shape_ptr3 = std::make_shared<xla::Shape>(different_shape);
+
+  IfrtLoadedVariableRegistry::Key key1{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr1,
+  };
+
+  IfrtLoadedVariableRegistry::Key key2{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr2,
+  };
+  IfrtLoadedVariableRegistry::Key key3{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr3,
+  };
+
+  EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly({key1, key2, key3}));
+  EXPECT_EQ(absl::Hash<IfrtLoadedVariableRegistry::Key>()(key1),
+            absl::Hash<IfrtLoadedVariableRegistry::Key>()(key2));
+  // Note: hash collision is possible but unlikely for key1 and key3.
+}
+
+TEST(IfrtLoadedVariableRegistryTest, KeyToString) {
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {2, 2});
+  auto shape_ptr = std::make_shared<xla::Shape>(shape);
+
+  IfrtLoadedVariableRegistry::Key key{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr,
+  };
+
+  EXPECT_THAT(key.ToString(),
+              ::testing::HasSubstr("input:0,1:{replicated}:f32[2,2]"));
+}
+
+TEST(IfrtLoadedVariableRegistryTest, KeyViewAndKeyEquality) {
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {2, 2});
+  auto shape_ptr1 = std::make_shared<xla::Shape>(shape);
+  auto shape_ptr2 = std::make_shared<xla::Shape>(shape);
+  // shape_ptr3 points to a different shape value.
+  xla::Shape different_shape = xla::ShapeUtil::MakeShape(xla::F32, {3, 3});
+  auto shape_ptr3 = std::make_shared<xla::Shape>(different_shape);
+
+  IfrtLoadedVariableRegistry::Key key1{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr1,
+  };
+
+  IfrtLoadedVariableRegistry::Key key2{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr2,
+  };
+
+  IfrtLoadedVariableRegistry::Key key3{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr3,
+  };
+
+  IfrtLoadedVariableRegistry::KeyView key_view1(key1);
+  IfrtLoadedVariableRegistry::KeyView key_view2(key2);
+  IfrtLoadedVariableRegistry::KeyView key_view3(key3);
+
+  EXPECT_EQ(key_view1, key_view2);
+  EXPECT_NE(key_view1, key_view3);
+  EXPECT_EQ(key1, key_view1);
+  EXPECT_EQ(key_view1, key1);
+  EXPECT_NE(key1, key_view3);
+  EXPECT_NE(key_view3, key1);
+}
+
+TEST(IfrtLoadedVariableRegistryTest, KeyViewAndKeyHash) {
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::F32, {2, 2});
+  auto shape_ptr1 = std::make_shared<xla::Shape>(shape);
+  auto shape_ptr2 = std::make_shared<xla::Shape>(shape);
+  xla::Shape different_shape = xla::ShapeUtil::MakeShape(xla::F32, {3, 3});
+  auto shape_ptr3 = std::make_shared<xla::Shape>(different_shape);
+
+  IfrtLoadedVariableRegistry::Key key1{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr1,
+  };
+
+  IfrtLoadedVariableRegistry::Key key2{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr2,
+  };
+  IfrtLoadedVariableRegistry::Key key3{
+      .device_ids = {0, 1},
+      .input_name = "input",
+      .hlo_sharding = xla::HloSharding::Replicate(),
+      .shape_on_device = shape_ptr3,
+  };
+  IfrtLoadedVariableRegistry::KeyView key_view1(key1);
+  IfrtLoadedVariableRegistry::KeyView key_view2(key2);
+  IfrtLoadedVariableRegistry::KeyView key_view3(key3);
+
+  EXPECT_TRUE(absl::VerifyTypeImplementsAbslHashCorrectly(
+      {key_view1, key_view2, key_view3}));
+  EXPECT_EQ(absl::Hash<IfrtLoadedVariableRegistry::KeyView>()(key_view1),
+            absl::Hash<IfrtLoadedVariableRegistry::KeyView>()(key_view2));
+  EXPECT_EQ(absl::Hash<IfrtLoadedVariableRegistry::Key>()(key1),
+            absl::Hash<IfrtLoadedVariableRegistry::KeyView>()(key_view1));
+}
+
+}  // namespace
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.cc

@@ -0,0 +1,187 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/device_list.h"
+#include "xla/python/ifrt/layout.h"
+#include "xla/shape.h"
+#include "xla/tsl/concurrency/future.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/context.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+namespace {
+
+absl::StatusOr<xla::ifrt::ArrayRef> LoadIfrtVariable(
+    std::shared_ptr<xla::ifrt::Client> ifrt_client,
+    const tsl::thread::ThreadPool& thread_pool,
+    const tensorflow::Tensor& variable,
+    const VariableDeviceShardingConfig& sharding_config,
+    const xla::ifrt::LayoutRef& xla_input_layout,
+    const xla::ifrt::DeviceListRef& device_list) {
+  TF_ASSIGN_OR_RETURN(
+      auto sharding,
+      tensorflow::ifrt_serving::ToIfrtSharding(
+          *ifrt_client, sharding_config.hlo_sharding, device_list));
+  return tensorflow::ifrt_serving::MakeArrayFromTensor(
+      *ifrt_client, variable, device_list, std::move(sharding), thread_pool,
+      xla_input_layout);
+}
+
+}  // namespace
+
+absl::StatusOr<ifrt_serving::DtypeAndShape> GetDtypeAndShape(
+    const ResourceHandle& resource_handle) {
+  const std::vector<DtypeAndPartialTensorShape>& dtype_and_partial_shapes =
+      resource_handle.dtypes_and_shapes();
+
+  if (dtype_and_partial_shapes.size() != 1) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Expected 1 dtype and shape, got ", dtype_and_partial_shapes.size()));
+  }
+  ifrt_serving::DtypeAndShape dtype_and_shape;
+  if (!dtype_and_partial_shapes.front().shape.AsTensorShape(
+          &dtype_and_shape.shape)) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Failed to convert partial shape to full tensor shape: ",
+                     dtype_and_partial_shapes.front().shape.DebugString()));
+  }
+
+  dtype_and_shape.dtype = dtype_and_partial_shapes.front().dtype;
+  return dtype_and_shape;
+}
+
+std::string GetRuntimeNameFromVarHandle(const ResourceHandle& handle) {
+  return absl::StrCat(handle.container(), "__", handle.name());
+}
+
+absl::Status AsyncLoadRestoredTensorAsIfrtLoadedVariable(
+    absl::string_view tensor_name,
+    std::shared_ptr<xla::ifrt::Client> ifrt_client,
+    const tsl::thread::ThreadPool& thread_pool,
+    const ifrt_serving::IfrtRestoreTensorRegistry& restore_tensor_registry,
+    ifrt_serving::IfrtLoadedVariableRegistry& loaded_variable_registry,
+    tfrt::ConcurrentWorkQueue* checkpoint_loader_queue,
+    const VariableDeviceShardingConfig& sharding_config,
+    const xla::ifrt::LayoutRef& xla_input_layout,
+    std::shared_ptr<const xla::Shape> shape_on_device,
+    const xla::ifrt::DeviceListRef& device_list) {
+  if (loaded_variable_registry
+          .GetLoadedVariable(IfrtLoadedVariableRegistry::KeyView(
+              sharding_config.device_ids, tensor_name,
+              sharding_config.hlo_sharding, shape_on_device))
+          .ok()) {
+    VLOG(1) << "Found alread registered variable for " << tensor_name;
+    return absl::OkStatus();
+  }
+
+  // `GetRestoredTensor` returns a future and wait for the tensor to be
+  // available is not efficient, so we check `GetDtypeAndShape` to find out if
+  // the tensor_name exists in the restore tensor registry.
+  TF_ASSIGN_OR_RETURN(
+      absl::StatusOr<ifrt_serving::DtypeAndShape> dtype_and_shape,
+      restore_tensor_registry.GetDtypeAndShape(tensor_name));
+
+  tsl::Future<tensorflow::Tensor> restored_tensor_future =
+      restore_tensor_registry.GetRestoredTensor(tensor_name);
+
+  if (!restored_tensor_future.IsValid()) {
+    return absl::InternalError(absl::StrCat(
+        "LoadVariableOp: failed to fetch variable tensor: ", tensor_name));
+  }
+  auto [loaded_variable_promise, loaded_variable_future] =
+      tsl::MakePromise<xla::ifrt::ArrayRef>();
+  IfrtLoadedVariableRegistry::Key key{
+      .device_ids = sharding_config.device_ids,
+      .input_name = std::string(tensor_name),
+      .hlo_sharding = sharding_config.hlo_sharding,
+      .shape_on_device = std::move(shape_on_device),
+  };
+  absl::Status status = loaded_variable_registry.TryRegisterLoadedVariable(
+      key,
+      [&]() -> absl::StatusOr<
+                ifrt_serving::IfrtLoadedVariableRegistry::LoadedVariable> {
+        return ifrt_serving::IfrtLoadedVariableRegistry::LoadedVariable(
+            {.array = loaded_variable_future});
+      });
+
+  if (absl::IsAlreadyExists(status)) {
+    // Variable is already registered, so we don't need to load it again. This
+    // can happen in parallel warmup.
+    return absl::OkStatus();
+  }
+  TF_RETURN_IF_ERROR(status);
+
+  tensorflow::Context bg_context(tensorflow::ContextKind::kThread);
+  restored_tensor_future.OnReady(
+      [ifrt_client = std::move(ifrt_client), &thread_pool = thread_pool,
+       checkpoint_loader_queue = checkpoint_loader_queue,
+       sharding_config = sharding_config,
+       loaded_variable_promise = std::move(loaded_variable_promise),
+       bg_context = std::move(bg_context), xla_input_layout = xla_input_layout,
+       device_list = device_list](
+          absl::StatusOr<tensorflow::Tensor> restored_tensor) mutable {
+        if (!restored_tensor.ok()) {
+          loaded_variable_promise.Set(restored_tensor.status());
+          return;
+        }
+
+        // Transfer tensor to array in a separate thread.
+        checkpoint_loader_queue->AddTask(
+            [ifrt_client = ifrt_client, &thread_pool = thread_pool,
+             sharding_config = std::move(sharding_config),
+             restored_tensor = std::move(*restored_tensor),
+             loaded_variable_promise = std::move(loaded_variable_promise),
+             bg_context = std::move(bg_context),
+             xla_input_layout = std::move(xla_input_layout),
+             device_list = std::move(device_list)]() mutable {
+              tensorflow::WithContext wc(bg_context);
+              absl::StatusOr<xla::ifrt::ArrayRef> variable_array =
+                  LoadIfrtVariable(ifrt_client, thread_pool, restored_tensor,
+                                   sharding_config, xla_input_layout,
+                                   device_list);
+              loaded_variable_promise.Set(std::move(variable_array));
+            });
+      });
+  return absl::OkStatus();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h

@@ -0,0 +1,79 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_LOADED_VARIABLE_UTILS_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_LOADED_VARIABLE_UTILS_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/device_list.h"
+#include "xla/python/ifrt/layout.h"
+#include "xla/shape.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_config.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// An index to indicate a non per-core executable bundle cache.
+inline constexpr int kNoCoreSelectedIndex = -1;
+
+// TODO(b/352551302) Delete VariableDeviceShardingConfigProto.
+struct VariableDeviceShardingConfig {
+  std::vector<int> device_ids;
+  xla::HloSharding hlo_sharding;
+};
+
+absl::StatusOr<ifrt_serving::DtypeAndShape> GetDtypeAndShape(
+    const ResourceHandle& resource_handle);
+
+// Returns the runtime name from the resource handle. The name will be concat of
+// handle's container name and handle's name.
+std::string GetRuntimeNameFromVarHandle(const ResourceHandle& handle);
+
+// Loads a restored tensor as an IFRT loaded variable and set the restored
+// tensor in the `restored_tensor_promise` as output. It is an async loading. We
+// look for the restored tensor in `ifrt_restore_tensor_registry` and save a
+// future of IFRT loaded variable in `ifrt_loaded_variable_registry`. The caller
+// can look for the actual loaded variable value in
+// `ifrt_loaded_variable_registry`.
+absl::Status AsyncLoadRestoredTensorAsIfrtLoadedVariable(
+    absl::string_view tensor_name,
+    std::shared_ptr<xla::ifrt::Client> ifrt_client,
+    const tsl::thread::ThreadPool& thread_pool,
+    const ifrt_serving::IfrtRestoreTensorRegistry& ifrt_restore_tensor_registry,
+    ifrt_serving::IfrtLoadedVariableRegistry& ifrt_loaded_variable_registry,
+    tfrt::ConcurrentWorkQueue* checkpoint_loader_queue,
+    const VariableDeviceShardingConfig& sharding_config,
+    const xla::ifrt::LayoutRef& xla_input_layout,
+    std::shared_ptr<const xla::Shape> shape_on_device,
+    const xla::ifrt::DeviceListRef& device_list);
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_LOADED_VARIABLE_UTILS_H_

tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils_test.cc

@@ -0,0 +1,183 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_utils.h"
+
+#include <cstdint>
+#include <memory>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/device.h"
+#include "xla/python/ifrt/test_util.h"
+#include "xla/tsl/concurrency/future.h"
+#include "xla/tsl/concurrency/ref_count.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/status_matchers.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/test.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_matcher.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_config.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+using tensorflow::test::TensorEq;
+
+TEST(ShardingUtilsTest, ShardTensorToIfrtLoadedVariableNotFoundWrongName) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+
+  Tensor variable_handle(DT_RESOURCE, TensorShape({}));
+  ResourceHandle resource_handle;
+  resource_handle.set_name("var_x");
+  resource_handle.set_dtypes_and_shapes({{
+      DT_INT32,
+      TensorShape({2, 2}),
+  }});
+  variable_handle.flat<ResourceHandle>()(0) = std::move(resource_handle);
+
+  IfrtRestoreTensorRegistry restored_tensor_registry;
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+  constexpr int kMaxParallelism = 16;
+  tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), tsl::ThreadOptions(),
+                                      "Resharding", kMaxParallelism);
+  IfrtLoadedVariableRegistry loaded_variable_registry;
+  auto restore_work_queue = tfrt::CreateMultiThreadedWorkQueue(
+      /*num_threads=*/4, /*num_blocking_threads=*/4);
+
+  VariableDeviceShardingConfig sharding_config = {
+      .device_ids = {0},
+      .hlo_sharding = xla::HloSharding::Replicate(),
+  };
+
+  auto [promise, future] = tsl::MakePromise<tensorflow::Tensor>();
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      false,
+      tsl::Future<DtypeAndShape>(
+          GetDtypeAndShape(variable_handle.scalar<ResourceHandle>()()).value()),
+      future};
+  TF_ASSERT_OK(restored_tensor_registry.TryRegister("var_x_wrong",
+                                                    restored_tensor_info));
+  promise.Set(input_tensor);
+  TF_ASSERT_OK_AND_ASSIGN(xla::ifrt::Device * device,
+                          client->LookupDevice(xla::ifrt::DeviceId(0)));
+  TF_ASSERT_OK_AND_ASSIGN(auto device_list, client->MakeDeviceList({device}));
+  EXPECT_THAT(
+      AsyncLoadRestoredTensorAsIfrtLoadedVariable(
+          "var_x", client, thread_pool, restored_tensor_registry,
+          loaded_variable_registry, restore_work_queue.get(), sharding_config,
+          /*xla_input_layout=*/nullptr, /*shape_on_device=*/nullptr,
+          device_list),
+      absl_testing::StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(ShardingUtilsTest, ShardTensorToIfrtLoadedVariableSucceed) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, TensorShape({2, 2}));
+
+  Tensor variable_handle(DT_RESOURCE, TensorShape({}));
+  ResourceHandle resource_handle;
+  resource_handle.set_name("var_x");
+  resource_handle.set_dtypes_and_shapes({{
+      DT_INT32,
+      TensorShape({2, 2}),
+  }});
+  variable_handle.flat<ResourceHandle>()(0) = std::move(resource_handle);
+
+  IfrtRestoreTensorRegistry restored_tensor_registry;
+  TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<xla::ifrt::Client> client,
+                          xla::ifrt::test_util::GetClient());
+  constexpr int kMaxParallelism = 16;
+  tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), tsl::ThreadOptions(),
+                                      "Resharding", kMaxParallelism);
+  IfrtLoadedVariableRegistry loaded_variable_registry;
+  auto restore_work_queue = tfrt::CreateMultiThreadedWorkQueue(
+      /*num_threads=*/4, /*num_blocking_threads=*/4);
+
+  VariableDeviceShardingConfig sharding_config{
+      .device_ids = {0},
+      .hlo_sharding = xla::HloSharding::Replicate(),
+  };
+
+  auto [promise, future] = tsl::MakePromise<tensorflow::Tensor>();
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      false,
+      tsl::Future<DtypeAndShape>(
+          GetDtypeAndShape(variable_handle.scalar<ResourceHandle>()()).value()),
+      future};
+
+  TF_ASSERT_OK(
+      restored_tensor_registry.TryRegister("var_x", restored_tensor_info));
+  TF_ASSERT_OK_AND_ASSIGN(xla::ifrt::Device * device,
+                          client->LookupDevice(xla::ifrt::DeviceId(0)));
+  TF_ASSERT_OK_AND_ASSIGN(auto device_list, client->MakeDeviceList({device}));
+  TF_ASSERT_OK(AsyncLoadRestoredTensorAsIfrtLoadedVariable(
+      "var_x", client, thread_pool, restored_tensor_registry,
+      loaded_variable_registry, restore_work_queue.get(), sharding_config,
+      /*xla_input_layout=*/nullptr, /*shape_on_device=*/nullptr, device_list));
+  promise.Set(input_tensor);
+  IfrtLoadedVariableRegistry::Key key{
+      .device_ids = {0},
+      .input_name = "var_x",
+      .hlo_sharding = sharding_config.hlo_sharding,
+  };
+  TF_ASSERT_OK_AND_ASSIGN(auto v,
+                          loaded_variable_registry.GetLoadedVariable(key));
+  TF_ASSERT_OK_AND_ASSIGN(auto assembled_array, v.array.Await());
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto disassembled_arrays,
+      assembled_array->DisassembleIntoSingleDeviceArrays(
+          xla::ifrt::ArrayCopySemantics::kAlwaysCopy,
+          xla::ifrt::SingleDeviceShardSemantics::kAddressableShards));
+  ASSERT_EQ(disassembled_arrays.size(), 1);
+  for (int i = 0; i < disassembled_arrays.size(); ++i) {
+    tensorflow::Tensor host_tensor(input_tensor.dtype(), input_tensor.shape());
+    TF_ASSERT_OK(
+        disassembled_arrays[i]
+            ->CopyToHostBuffer(host_tensor.data(), /*byte_strides=*/{},
+                               xla::ifrt::ArrayCopySemantics::kAlwaysCopy)
+            .Await());
+    EXPECT_THAT(host_tensor, TensorEq(input_tensor));
+  }
+}
+}  // namespace
+}  // namespace ifrt_serving
+
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_model_context.cc

@@ -0,0 +1,40 @@
+
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_model_context.h"
+
+#include "absl/status/status.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/threadpool.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+tsl::thread::ThreadPool& IfrtModelContext::GetThreadPool() const {
+  return thread_pool_;
+}
+
+absl::Status IfrtModelContext::Freeze() {
+  restore_tensor_registry_.Freeze();
+  for (auto& program_handle : handles_) {
+    TF_RETURN_IF_ERROR(program_handle.Freeze());
+  }
+  frozen_ = true;
+  return absl::OkStatus();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_model_context.h

@@ -0,0 +1,257 @@
+/* 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_CORE_TFRT_IFRT_IFRT_MODEL_CONTEXT_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_MODEL_CONTEXT_H_
+
+#include <cstddef>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/pjrt/pjrt_executable.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/topology.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_config.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_executable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_persistent_compilation_cache.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
+#include "tsl/platform/protobuf.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+inline constexpr absl::string_view kIfrtModelContextName = "IfrtModelContext";
+
+// Device specific configuration not available through ifrt. This should be
+// rare.
+struct DeviceConfig {
+  tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn =
+      tensorflow::IdentityShapeRepresentationFn();
+};
+
+// The runtime context for ifrt to be used in TFRT serving.
+//
+// This class is thread compatible.
+class IfrtModelContext {
+ public:
+  explicit IfrtModelContext(
+      std::shared_ptr<xla::ifrt::Client> client,
+      IfrtServingCoreSelector* ifrt_serving_core_selector,
+      tsl::thread::ThreadPool* thread_pool,
+      std::variant<std::unique_ptr<tsl::protobuf::Message>,
+                   xla::CompileOptions::EnvironmentOptionOverrides>
+          compilation_env_or_overrides,
+      H2DTransferExecutorFactory* h2d_transfer_executor_factory,
+      bool enable_propagate_static_shapes_pass = true,
+      bool use_output_arena = false)
+      : client_(std::move(client)),
+        ifrt_serving_core_selector_(ifrt_serving_core_selector),
+        thread_pool_(*thread_pool),
+        compilation_env_or_overrides_(std::move(compilation_env_or_overrides)),
+        h2d_transfer_executor_factory_(h2d_transfer_executor_factory),
+        enable_propagate_static_shapes_pass_(
+            enable_propagate_static_shapes_pass),
+        use_output_arena_(use_output_arena) {}
+  IfrtModelContext(
+      std::shared_ptr<xla::ifrt::Client> client,
+      IfrtServingCoreSelector* ifrt_serving_core_selector,
+      tsl::thread::ThreadPool* thread_pool, tensorflow::DeviceMgr* device_mgr,
+      tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn,
+      std::variant<std::unique_ptr<tsl::protobuf::Message>,
+                   xla::CompileOptions::EnvironmentOptionOverrides>
+          compilation_env_or_overrides,
+      std::shared_ptr<const void> topology, TfToHloCompiler* tf_to_hlo_compiler,
+      H2DTransferExecutorFactory* h2d_transfer_executor_factory,
+      IfrtPersistentCompilationCache* persistent_compilation_cache = nullptr,
+      bool enable_propagate_static_shapes_pass = true,
+      bool use_output_arena = false)
+      : client_(std::move(client)),
+        topology_(topology),
+        ifrt_serving_core_selector_(ifrt_serving_core_selector),
+        thread_pool_(*thread_pool),
+        device_mgr_(device_mgr),
+        shape_representation_fn_(shape_representation_fn),
+        compilation_env_or_overrides_(std::move(compilation_env_or_overrides)),
+        tf_to_hlo_compiler_(tf_to_hlo_compiler),
+        h2d_transfer_executor_factory_(h2d_transfer_executor_factory),
+        persistent_compilation_cache_(persistent_compilation_cache),
+        enable_propagate_static_shapes_pass_(
+            enable_propagate_static_shapes_pass),
+        use_output_arena_(use_output_arena) {}
+
+  void RegisterHandle(ServingExecutableRegistry::Handle handle) {
+    handles_.push_back(std::move(handle));
+  }
+
+  std::shared_ptr<xla::ifrt::Client> GetClient() const { return client_; }
+
+  const tensorflow::XlaHelpers::ShapeRepresentationFn&
+  GetShapeRepresentationFn() const {
+    return shape_representation_fn_;
+  }
+
+  tsl::thread::ThreadPool& GetThreadPool() const;
+
+  const IfrtLoadedVariableRegistry& GetLoadedVariableRegistry() const {
+    return loaded_variable_registry_;
+  }
+  IfrtLoadedVariableRegistry& GetLoadedVariableRegistry() {
+    return loaded_variable_registry_;
+  }
+
+  const IfrtRestoreTensorRegistry& GetRestoreTensorRegistry() const {
+    return restore_tensor_registry_;
+  }
+  IfrtRestoreTensorRegistry& GetRestoreTensorRegistry() {
+    return restore_tensor_registry_;
+  }
+
+  IfrtPersistentCompilationCache* GetPersistentCompilationCache() const {
+    return persistent_compilation_cache_;
+  }
+
+  H2DTransferExecutorFactory* GetH2DTransferExecutorFactory() const {
+    return h2d_transfer_executor_factory_;
+  }
+
+  tensorflow::DeviceMgr* GetDeviceMgr() const { return device_mgr_; }
+  IfrtServingCoreSelector* GetIfrtServingCoreSelector() const {
+    return ifrt_serving_core_selector_;
+  }
+
+  tfrt::ConcurrentWorkQueue* checkpoint_loader_queue() const {
+    return checkpoint_loader_queue_;
+  }
+  void set_checkpoint_loader_queue(tfrt::ConcurrentWorkQueue* work_queue) {
+    checkpoint_loader_queue_ = work_queue;
+  }
+
+  void set_default_signature_inputs(
+      const DefaultSignatureInputConfig& default_signature_inputs) {
+    default_signature_inputs_ = default_signature_inputs;
+  }
+
+  void set_h2d_transfer_executor_factory(
+      H2DTransferExecutorFactory* h2d_transfer_executor_factory) {
+    h2d_transfer_executor_factory_ = h2d_transfer_executor_factory;
+  }
+
+  const DefaultSignatureInputConfig& default_signature_inputs() const {
+    return default_signature_inputs_;
+  }
+
+  bool enable_propagate_static_shapes_pass() const {
+    return enable_propagate_static_shapes_pass_;
+  }
+
+  void set_enable_propagate_static_shapes_pass(
+      bool enable_propagate_static_shapes_pass) {
+    enable_propagate_static_shapes_pass_ = enable_propagate_static_shapes_pass;
+  }
+
+  bool use_output_arena() const { return use_output_arena_; }
+
+  void set_use_output_arena(bool use_output_arena) {
+    use_output_arena_ = use_output_arena;
+  }
+
+  tsl::protobuf::Message* GetCompilationEnvironmentProto() const {
+    if (std::holds_alternative<std::unique_ptr<tsl::protobuf::Message>>(
+            compilation_env_or_overrides_)) {
+      return std::get<std::unique_ptr<tsl::protobuf::Message>>(
+                 compilation_env_or_overrides_)
+          .get();
+    }
+    return nullptr;
+  }
+
+  std::variant<tsl::protobuf::Message*,
+               xla::CompileOptions::EnvironmentOptionOverrides>
+  GetCompilationEnvOrOverrides() const {
+    if (std::holds_alternative<std::unique_ptr<tsl::protobuf::Message>>(
+            compilation_env_or_overrides_)) {
+      return std::get<std::unique_ptr<tsl::protobuf::Message>>(
+                 compilation_env_or_overrides_)
+          .get();
+    } else {
+      return std::get<xla::CompileOptions::EnvironmentOptionOverrides>(
+          compilation_env_or_overrides_);
+    }
+  }
+
+  TfToHloCompiler* GetTfToHloCompiler() const { return tf_to_hlo_compiler_; }
+
+  // Freeze the model: release the resources such as host tensors that are used
+  // by the device only. The caller guarantees all resources released in this
+  // function is no longer in use in regular execution path.
+  // After Freeze() is called, no new model signature will be compiled. Using a
+  // signature or an input shape that wasn't compiled before the freeze will
+  // leads to an error.
+  absl::Status Freeze();
+
+  bool IsFrozen() const { return frozen_; }
+
+ private:
+  std::shared_ptr<xla::ifrt::Client> client_;
+  // Keep hardware specific topology info alive. This is currently used for
+  // shape determination.
+  std::shared_ptr<const void> topology_;
+
+  IfrtServingCoreSelector* ifrt_serving_core_selector_;  // May be nullptr
+  tsl::thread::ThreadPool& thread_pool_;
+
+  tensorflow::DeviceMgr* device_mgr_ = nullptr;  // Not owned.
+  tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn_ =
+      tensorflow::IdentityShapeRepresentationFn();
+  std::variant<std::unique_ptr<tsl::protobuf::Message>,
+               xla::CompileOptions::EnvironmentOptionOverrides>
+      compilation_env_or_overrides_;
+
+  // Dedicated work queue for heavy task such as variable tensor restoration.
+  tfrt::ConcurrentWorkQueue* checkpoint_loader_queue_ = nullptr;
+
+  std::vector<ServingExecutableRegistry::Handle> handles_;
+
+  DefaultSignatureInputConfig default_signature_inputs_;
+
+  IfrtLoadedVariableRegistry loaded_variable_registry_;
+  IfrtRestoreTensorRegistry restore_tensor_registry_;
+  TfToHloCompiler* tf_to_hlo_compiler_ = nullptr;
+  H2DTransferExecutorFactory* h2d_transfer_executor_factory_ = nullptr;
+  IfrtPersistentCompilationCache* persistent_compilation_cache_ = nullptr;
+  bool frozen_ = false;
+  bool enable_propagate_static_shapes_pass_ = true;
+  bool use_output_arena_ = false;
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_MODEL_CONTEXT_H_

tensorflow/core/tfrt/ifrt/ifrt_model_restore_context.h

@@ -0,0 +1,50 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_MODEL_RESTORE_CONTEXT_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_MODEL_RESTORE_CONTEXT_H_
+
+#include <memory>
+#include <utility>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/core/tfrt/ifrt/checkpoint_loader.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+inline constexpr absl::string_view kIfrtModelRestoreContextName =
+    "IfrtModelRestoreContext";
+
+// A resource context that holds the `CheckpointLoader` for a model. We need a
+// different context than `IfrtModelContext` because `IfrtModelContext` is too
+// large to be a dependency of other libraries.
+class IfrtModelRestoreContext {
+ public:
+  explicit IfrtModelRestoreContext(
+      std::unique_ptr<CheckpointLoader> checkpoint_loader)
+      : checkpoint_loader_(std::move(checkpoint_loader)) {}
+
+  CheckpointLoader* checkpoint_loader() const {
+    return checkpoint_loader_.get();
+  }
+
+ private:
+  std::unique_ptr<CheckpointLoader> checkpoint_loader_;
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_MODEL_RESTORE_CONTEXT_H_

tensorflow/core/tfrt/ifrt/ifrt_persistent_compilation_cache.cc

@@ -0,0 +1,66 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/tfrt/ifrt/ifrt_persistent_compilation_cache.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/statusor.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.h"
+#include "xla/pjrt/pjrt_executable.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/device_list.h"
+#include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
+#include "xla/python/ifrt/host_callback.h"
+#include "xla/python/ifrt/program.h"
+#include "xla/python/pjrt_ifrt/xla_compiler.h"
+#include "xla/tsl/concurrency/ref_count.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+absl::StatusOr<xla::ifrt::LoadedExecutableRef>
+IfrtPersistentCompilationCache::LookupLoadedExecutableOrCreate(
+    std::unique_ptr<xla::ifrt::HloProgram> hlo_program,
+    xla::ifrt::DeviceListRef device_list,
+    const xla::CompileOptions& xla_compile_options,
+    const std::vector<tsl::RCReference<xla::ifrt::LoadedHostCallback>>&
+        loaded_host_callbacks,
+    xla::ifrt::Client* client,
+    absl::AnyInvocable<absl::StatusOr<xla::ifrt::LoadedExecutableRef>(
+        std::unique_ptr<xla::ifrt::Program> program,
+        std::unique_ptr<xla::ifrt::CompileOptions> options)>
+        value_fn) {
+  // No persistent cache implemented, compile directly.
+  auto ifrt_xla_compile_options =
+      std::make_unique<xla::ifrt::XlaCompileOptions>(
+          xla_compile_options, std::move(device_list), loaded_host_callbacks);
+  return value_fn(std::move(hlo_program), std::move(ifrt_xla_compile_options));
+}
+
+absl::StatusOr<Tf2HloResult>
+IfrtPersistentCompilationCache::LookupTf2HloResultOrCreate(
+    Tf2HloArg tf2hlo_arg, TfToHloCompiler* tf_to_hlo_compiler) {
+  // No tf2xla persistent cache is implemented, compile directly.
+  return tf_to_hlo_compiler->CompileTfToHlo(tf2hlo_arg);
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_persistent_compilation_cache.h

@@ -0,0 +1,74 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_PERSISTENT_COMPILATION_CACHE_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_PERSISTENT_COMPILATION_CACHE_H_
+
+#include <memory>
+#include <vector>
+
+#include "absl/functional/any_invocable.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.h"
+#include "xla/pjrt/pjrt_executable.h"
+#include "xla/python/ifrt/device_list.h"
+#include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
+#include "xla/python/ifrt/host_callback.h"
+#include "xla/python/ifrt/program.h"
+#include "xla/tsl/concurrency/ref_count.h"
+namespace tensorflow {
+namespace ifrt_serving {
+
+class IfrtPersistentCompilationCache {
+ public:
+  IfrtPersistentCompilationCache() = default;
+  virtual ~IfrtPersistentCompilationCache() = default;
+
+  // The implementation of this API should be thread-safe. It generates a key
+  // for looking up the executable in the persistent cache and it will return
+  // the LoadedExecutable if hits cache. Otherwise, it will call the `value_fn`
+  // to generate and return the LoadedExecutable.
+  virtual absl::StatusOr<xla::ifrt::LoadedExecutableRef>
+  LookupLoadedExecutableOrCreate(
+      std::unique_ptr<xla::ifrt::HloProgram> hlo_program,
+      xla::ifrt::DeviceListRef device_list,
+      const xla::CompileOptions& xla_compile_options,
+      const std::vector<tsl::RCReference<xla::ifrt::LoadedHostCallback>>&
+          loaded_host_callbacks,
+      xla::ifrt::Client* client,
+      absl::AnyInvocable<absl::StatusOr<xla::ifrt::LoadedExecutableRef>(
+          std::unique_ptr<xla::ifrt::Program> program,
+          std::unique_ptr<xla::ifrt::CompileOptions> options)>
+          value_fn);
+
+  // The implementation of this API should be thread-safe. It generates a key
+  // for looking up the Tf2HloResult in the persistent cache and it will return
+  // the Tf2HloResult if hits cache. Otherwise, it will call the `value_fn` to
+  // generate and return the Tf2HloResult.
+  virtual absl::StatusOr<Tf2HloResult> LookupTf2HloResultOrCreate(
+      Tf2HloArg tf2hlo_arg, TfToHloCompiler* tf_to_hlo_compiler);
+
+  virtual bool IsXlaCompilationCacheEnabled() const { return false; }
+  virtual bool IsTf2HloCompilationCacheEnabled() const { return false; }
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_PERSISTENT_COMPILATION_CACHE_H_

tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.cc

@@ -0,0 +1,110 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/tsl/concurrency/future.h"
+#include "tensorflow/core/framework/tensor.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+absl::Status IfrtRestoreTensorRegistry::TryRegister(
+    absl::string_view name, RestoredTensorInfo restored_tensor_info) {
+  absl::MutexLock lock(mutex_);
+  auto& info = restored_tensors_[name];
+  if (info.dtype_and_shape.IsValid() || info.tensor_future.IsValid()) {
+    LOG(WARNING)
+        << "Variable named '" << name
+        << "' has been already registered. Ignore request of a new tensor with "
+           "same name, dtype and shape.";
+    return absl::OkStatus();
+  }
+  info = std::move(restored_tensor_info);
+  return absl::OkStatus();
+}
+
+tsl::Future<tensorflow::Tensor> IfrtRestoreTensorRegistry::GetRestoredTensor(
+    absl::string_view name) const {
+  absl::MutexLock lock(mutex_);
+  auto it = restored_tensors_.find(name);
+  if (it == restored_tensors_.end()) {
+    return tsl::Future<tensorflow::Tensor>(
+        absl::NotFoundError(absl::StrCat("Variable '", name, "' not found.")));
+  }
+
+  return it->second.tensor_future;
+}
+
+absl::Status IfrtRestoreTensorRegistry::SetUsedByHost(absl::string_view name) {
+  absl::MutexLock lock(mutex_);
+  auto it = restored_tensors_.find(name);
+  if (it == restored_tensors_.end()) {
+    return absl::NotFoundError(
+        absl::StrCat("Variable '", name, "' not found."));
+  }
+
+  it->second.used_by_host = true;
+  return absl::OkStatus();
+}
+
+void IfrtRestoreTensorRegistry::Freeze() {
+  absl::MutexLock lock(mutex_);
+  tsl::Future<tensorflow::Tensor> release_tensor_future(
+      absl::UnavailableError("Tensor is already release."));
+  for (auto& [name, info] : restored_tensors_) {
+    if (!info.used_by_host) {
+      // Release the tensor by replacing the future containing the tensor with
+      // an future containing a status.
+      info.tensor_future = release_tensor_future;
+    }
+  }
+}
+
+absl::StatusOr<DtypeAndShape> IfrtRestoreTensorRegistry::GetDtypeAndShape(
+    absl::string_view name) const {
+  tsl::Future<DtypeAndShape> dtype_and_shape_future;
+  {
+    absl::MutexLock lock(mutex_);
+    auto it = restored_tensors_.find(name);
+    if (it == restored_tensors_.end()) {
+      return absl::NotFoundError(
+          absl::StrCat("Variable '", name, "' not found."));
+    }
+    dtype_and_shape_future = it->second.dtype_and_shape;
+  }
+
+  if (!dtype_and_shape_future.IsValid()) {
+    // This case should ideally not happen if TryRegister ensures futures are
+    // valid
+    return absl::InternalError(
+        absl::StrCat("Invalid future for variable '", name, "'"));
+  }
+
+  return dtype_and_shape_future.Await();
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h

@@ -0,0 +1,78 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_RESTORE_TENSOR_REGISTRY_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_RESTORE_TENSOR_REGISTRY_H_
+
+#include <string>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/tsl/concurrency/future.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// This class is thread safe.
+class IfrtRestoreTensorRegistry {
+ public:
+  struct RestoredTensorInfo {
+    bool used_by_host = false;
+    // We have to use a separate future for dtype and shape because in
+    // `Freeze()` we will release the host tensor by resetting the
+    // `tensor_future` to an error while we want to keep the `dtype_and_shape`
+    // future unchanged.
+    tsl::Future<DtypeAndShape> dtype_and_shape;
+    tsl::Future<tensorflow::Tensor> tensor_future;
+  };
+  // Tries to register a loaded variable with the given name.
+  // Returns an error if the named tensor already exists.
+  absl::Status TryRegister(absl::string_view name,
+                           RestoredTensorInfo restored_tensor_info)
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+  tsl::Future<tensorflow::Tensor> GetRestoredTensor(
+      absl::string_view name) const ABSL_LOCKS_EXCLUDED(mutex_);
+
+  // Sets the tensor as used by the host. To ensure a tensor's host memory
+  // is released, this function must be called at least once before the Freeze.
+  absl::Status SetUsedByHost(absl::string_view name)
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+  absl::StatusOr<DtypeAndShape> GetDtypeAndShape(absl::string_view name) const
+      ABSL_LOCKS_EXCLUDED(mutex_);
+
+  // Part of freezing the model is to release the host tensors that are used by
+  // the device only. The caller guarantees those tensors are already loaded to
+  // the device.
+  void Freeze() ABSL_LOCKS_EXCLUDED(mutex_);
+
+ private:
+  mutable absl::Mutex mutex_;
+  absl::flat_hash_map<std::string, RestoredTensorInfo> restored_tensors_
+      ABSL_GUARDED_BY(mutex_);
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_RESTORE_TENSOR_REGISTRY_H_

tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry_test.cc

@@ -0,0 +1,195 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+
+#include <cstdint>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "testing/base/public/mock-log.h"
+#include "absl/base/log_severity.h"
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/tsl/concurrency/future.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/status_matchers.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+using ::testing::_;
+using ::testing::AnyNumber;
+using ::testing::HasSubstr;
+using ::testing::kDoNotCaptureLogsYet;
+using ::testing::ScopedMockLog;
+
+TEST(IfrtRestoreTensorRegistryTest, RetrieveNonRegisteredTensorFails) {
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.GetRestoredTensor("input_tensor_1").Await(),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(IfrtRestoreTensorRegistryTest,
+     RetrieveNonRegisteredTensorDTypeAndShapeFails) {
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.GetDtypeAndShape("input_tensor_1"),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(IfrtRestoreTensorRegistryTest, SetNonExistedTensorAsUsedByHostFails) {
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.SetUsedByHost("input_tensor_1"),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound));
+}
+
+TEST(IfrtRestoreTensorRegistryTest,
+     RegisteredExistedTensorSucceedsAndWarningIsLogged) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+  auto [promise, future] = tsl::MakePromise<tensorflow::Tensor>();
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      .used_by_host = false,
+      .dtype_and_shape = tsl::Future<DtypeAndShape>(DtypeAndShape{
+          .dtype = DT_INT32,
+          .shape = tensorflow::TensorShape({2, 2}),
+      }),
+      .tensor_future = future};
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.TryRegister("input_tensor_2", restored_tensor_info),
+              absl_testing::IsOk());
+  promise.Set(input_tensor);
+
+  ScopedMockLog mock_log(kDoNotCaptureLogsYet);
+  EXPECT_CALL(mock_log, Log).Times(AnyNumber());
+  EXPECT_CALL(mock_log,
+              Log(base_logging::WARNING, _,
+                  HasSubstr("Variable named 'input_tensor_2' has been already "
+                            "registered. Ignore request of a new tensor")))
+      .Times(1);
+  mock_log.StartCapturingLogs();
+
+  EXPECT_THAT(registry.TryRegister("input_tensor_2", restored_tensor_info),
+              absl_testing::IsOk());
+}
+
+TEST(IfrtRestoreTensorRegistryTest, SetTensorAsUsedByHost) {
+  auto [promise, future] = tsl::MakePromise<tensorflow::Tensor>();
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      .used_by_host = false,
+      .dtype_and_shape = tsl::Future<DtypeAndShape>(DtypeAndShape{
+          .dtype = DT_INT32,
+          .shape = tensorflow::TensorShape({2, 2}),
+      }),
+      .tensor_future = future};
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.TryRegister("input_tensor_1", restored_tensor_info),
+              absl_testing::IsOk());
+  EXPECT_THAT(registry.SetUsedByHost("input_tensor_1"), absl_testing::IsOk());
+}
+
+TEST(IfrtRestoreTensorRegistryTest, RegisteredTensorCanBeRetrieved) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+  auto [promise, future] = tsl::MakePromise<tensorflow::Tensor>();
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      .used_by_host = false,
+      .dtype_and_shape = tsl::Future<DtypeAndShape>(DtypeAndShape{
+          .dtype = DT_INT32,
+          .shape = tensorflow::TensorShape({2, 2}),
+      }),
+      .tensor_future = future};
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.TryRegister("input_tensor_1", restored_tensor_info),
+              absl_testing::IsOk());
+  promise.Set(input_tensor);
+  TF_ASSERT_OK_AND_ASSIGN(tensorflow::Tensor retrieved,
+                          registry.GetRestoredTensor("input_tensor_1").Await());
+  test::ExpectEqual(retrieved, input_tensor);
+  TF_ASSERT_OK_AND_ASSIGN(DtypeAndShape dtype_and_shape,
+                          registry.GetDtypeAndShape("input_tensor_1"));
+  EXPECT_TRUE(
+      dtype_and_shape.shape.IsSameSize(tensorflow::TensorShape({2, 2})));
+  EXPECT_EQ(dtype_and_shape.dtype, DT_INT32);
+}
+
+TEST(IfrtRestoreTensorRegistryTest,
+     RegisteredTensorDTypeAndShapeCanBeRetrieved) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+  auto [promise, future] = tsl::MakePromise<tensorflow::Tensor>();
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info = {
+      .used_by_host = false,
+      .dtype_and_shape = tsl::Future<DtypeAndShape>(DtypeAndShape{
+          .dtype = DT_INT32,
+          .shape = tensorflow::TensorShape({2, 2}),
+      }),
+      .tensor_future = future};
+  IfrtRestoreTensorRegistry registry;
+  EXPECT_THAT(registry.TryRegister("input_tensor_1", restored_tensor_info),
+              absl_testing::IsOk());
+  TF_ASSERT_OK_AND_ASSIGN(DtypeAndShape dtype_and_shape,
+                          registry.GetDtypeAndShape("input_tensor_1"));
+  EXPECT_TRUE(
+      dtype_and_shape.shape.IsSameSize(tensorflow::TensorShape({2, 2})));
+  EXPECT_EQ(dtype_and_shape.dtype, DT_INT32);
+}
+
+TEST(IfrtRestoreTensorRegistryTest, FeezeTensorRegistry) {
+  auto input_tensor =
+      test::AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({2, 2}));
+  auto [promise1, future1] = tsl::MakePromise<tensorflow::Tensor>();
+  auto [promise2, future2] = tsl::MakePromise<tensorflow::Tensor>();
+
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info1 = {
+      .used_by_host = false,
+      .dtype_and_shape = tsl::Future<DtypeAndShape>(DtypeAndShape{
+          .dtype = DT_INT32,
+          .shape = tensorflow::TensorShape({2, 2}),
+      }),
+      .tensor_future = future1};
+  IfrtRestoreTensorRegistry::RestoredTensorInfo restored_tensor_info2 = {
+      .used_by_host = true,
+      .dtype_and_shape = tsl::Future<DtypeAndShape>(DtypeAndShape{
+          .dtype = DT_INT32,
+          .shape = tensorflow::TensorShape({2, 2}),
+      }),
+      .tensor_future = future2};
+  IfrtRestoreTensorRegistry registry;
+  TF_ASSERT_OK(registry.TryRegister("input_tensor_1", restored_tensor_info1));
+  TF_ASSERT_OK(registry.TryRegister("input_tensor_2", restored_tensor_info2));
+  promise1.Set(input_tensor);
+  promise2.Set(input_tensor);
+  registry.Freeze();
+  // Tensor with `used_by_host` set to false will be freed after freeze.
+  EXPECT_THAT(registry.GetRestoredTensor("input_tensor_1").Await(),
+              absl_testing::StatusIs(absl::StatusCode::kUnavailable));
+  // Tensor with `used_by_host` set to true will be kept after freeze.
+  TF_ASSERT_OK_AND_ASSIGN(tensorflow::Tensor retrieved,
+                          registry.GetRestoredTensor("input_tensor_2").Await());
+  test::ExpectEqual(retrieved, input_tensor);
+}
+}  // namespace
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.cc

@@ -0,0 +1,47 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+
+#include <cstdint>
+
+#include "absl/strings/str_cat.h"
+#include "absl/synchronization/mutex.h"
+#include "xla/tsl/framework/serving_device_selector.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+IfrtServingCoreSelector::IfrtServingCoreSelector(
+    tsl::ServingDeviceSelector* device_selector, int num_cores)
+    : device_selector_(device_selector), num_cores_(num_cores) {}
+
+tsl::DeviceReservation IfrtServingCoreSelector::ReserveDevice(
+    int64_t program_id) {
+  absl::MutexLock lock(mu_);
+  int64_t run_count = run_counter_[program_id]++;
+  if (run_count < num_cores_) {
+    // If run_count is less than the number of TPU cores, we use run_count
+    // as device index to iterate through all the TPU cores for the given
+    // program, so that the program can get warmed up on all the TPU cores
+    // when there are enough warmup requests. Set the selector to nullptr,
+    // since this `DeviceReservation` isn't made through the selector.
+    return tsl::DeviceReservation(run_count, /*selector=*/nullptr);
+  }
+  return device_selector_->ReserveDevice(absl::StrCat(program_id));
+}
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h

@@ -0,0 +1,53 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_CORE_SELECTOR_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_CORE_SELECTOR_H_
+
+#include <cstdint>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/synchronization/mutex.h"
+#include "xla/tsl/framework/serving_device_selector.h"
+namespace tensorflow {
+namespace ifrt_serving {
+
+// A wrapper of a `tsl::ServingDeviceSelector` that will be responsible for the
+// core selection during Ifrt TPU execution.
+class IfrtServingCoreSelector {
+ public:
+  explicit IfrtServingCoreSelector(tsl::ServingDeviceSelector* device_selector,
+                                   int num_cores);
+  // Reserves a device for the given `program_id`. The `program_id` is used to
+  // identify an IFRT executable and should be the key of
+  // `tensorflow::ifrt_serving::ServingExecutableRegistry `.
+  tsl::DeviceReservation ReserveDevice(int64_t program_id);
+
+ private:
+  tsl::ServingDeviceSelector* device_selector_;
+
+  absl::Mutex mu_;
+  // A counter of the number of runs for each program. For a given program, it
+  // is used to determine if the core selector should treat the incoming request
+  // as a warmup request to warm up a core.
+  absl::flat_hash_map<int64_t, int64_t> run_counter_ ABSL_GUARDED_BY(mu_);
+  int num_cores_;
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_CORE_SELECTOR_H_

tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector_test.cc

@@ -0,0 +1,61 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+
+#include <cstdint>
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/str_cat.h"
+#include "xla/tsl/framework/serving_device_selector.h"
+#include "xla/tsl/framework/test_util/mock_serving_device_selector.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+
+class IfrtServingCoreSelectorTest : public ::testing::Test {
+ protected:
+  explicit IfrtServingCoreSelectorTest() {
+    core_selector_ = std::make_unique<IfrtServingCoreSelector>(
+        &serving_device_selector_, num_cores_);
+  }
+
+  tsl::test_util::MockServingDeviceSelector serving_device_selector_;
+  std::unique_ptr<IfrtServingCoreSelector> core_selector_;
+  int num_cores_ = 2;
+};
+
+TEST_F(IfrtServingCoreSelectorTest, ReservedDevicesReturns) {
+  int64_t program_id1 = 111111;
+  EXPECT_CALL(serving_device_selector_,
+              ReserveDevice(absl::StrCat(program_id1)))
+      .WillOnce([this](::testing::Unused) {
+        return tsl::DeviceReservation(0, &serving_device_selector_);
+      });
+  // Warm up each core first.
+  for (int i = 0; i < num_cores_; ++i) {
+    EXPECT_THAT(core_selector_->ReserveDevice(program_id1).device_index(), i);
+  }
+  tsl::DeviceReservation reservation =
+      core_selector_->ReserveDevice(program_id1);
+  EXPECT_THAT(reservation.device_index(), 0);
+}
+
+}  // namespace
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h

@@ -0,0 +1,419 @@
+/* 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_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_H_
+
+#include <stdbool.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/inlined_vector.h"
+#include "absl/hash/hash.h"
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/types/span.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/pjrt_layout.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/device.h"
+#include "xla/python/ifrt/device_list.h"
+#include "xla/python/ifrt/dtype.h"
+#include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/layout.h"
+#include "xla/python/ifrt/shape.h"
+#include "xla/python/ifrt/sharding.h"
+#include "xla/shape.h"
+#include "xla/tsl/concurrency/future.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_persistent_compilation_cache.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
+#include "tensorflow/core/tfrt/ifrt/tf_host_callback.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+// Encodes the layout of `xla_input_shapes` to the module.
+absl::Status EncodeLayout(absl::Span<const xla::Shape> xla_input_shapes,
+                          mlir::ModuleOp module);
+
+class IfrtServingExecutable {
+ public:
+  // TODO(b/514902739): Remove the default value of `use_output_arena`.
+  static absl::StatusOr<std::unique_ptr<IfrtServingExecutable>> Create(
+      int64_t program_id, absl::string_view model_name,
+      absl::string_view signature_name,
+      mlir::OwningOpRef<mlir::ModuleOp> module,
+      std::shared_ptr<xla::ifrt::Client> client,
+      tsl::thread::ThreadPool* thread_pool,
+      IfrtLoadedVariableRegistry* ifrt_loaded_variable_registry,
+      const IfrtRestoreTensorRegistry* ifrt_restore,
+      tfrt::ConcurrentWorkQueue* checkpoint_loader_queue,
+      tensorflow::DeviceMgr* device_mgr,
+      tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn,
+      IfrtServingCoreSelector* ifrt_serving_core_selector,
+      std::variant<tsl::protobuf::Message*,
+                   xla::CompileOptions::EnvironmentOptionOverrides>
+          compilation_env_or_overrides,
+      TfToHloCompiler* tf_to_hlo_compiler,
+      IfrtPersistentCompilationCache* persistent_compilation_cache,
+      H2DTransferExecutorFactory* h2d_transfer_executor_factory,
+      bool use_output_arena = false);
+
+  // Movable but not copyable.
+  IfrtServingExecutable(IfrtServingExecutable&& other) = default;
+  IfrtServingExecutable& operator=(IfrtServingExecutable&& other) = default;
+  IfrtServingExecutable(const IfrtServingExecutable& other) = delete;
+  IfrtServingExecutable& operator=(const IfrtServingExecutable& other) = delete;
+
+  absl::string_view model_name() const { return model_name_; }
+  absl::string_view signature_name() const { return signature_name_; }
+
+  // Executes the computation.
+  // variable_arg_indices are in sorted order.
+  absl::StatusOr<std::vector<tensorflow::Tensor>> Execute(
+      absl::Span<const tensorflow::Tensor> inputs,
+      absl::Span<const int> variable_arg_indices);
+
+  // Executes the computation asynchronously.
+  // variable_arg_indices are in sorted order.
+  absl::StatusOr<tsl::Future<std::vector<tensorflow::Tensor>>> ExecuteAsync(
+      absl::Span<const tensorflow::Tensor> inputs,
+      absl::Span<const int> variable_arg_indices);
+
+  // Freezes the model. After the Freeze(), JIT compile is not supported and
+  // Execute() will return error if inputs contain uncompiled shapes.
+  void Freeze();
+
+  int num_executables() const {
+    absl::MutexLock lock(mutex_);
+    return executable_bundles_.size();
+  }
+
+ private:
+  friend class IfrtBackendCompilerTest;
+  // In memory cache key.
+  struct Key {
+    std::vector<tensorflow::TensorShape> input_shapes;
+    template <typename H>
+    friend H AbslHashValue(H h, const Key& key) {
+      for (const auto& shape : key.input_shapes) {
+        for (auto size : shape.dim_sizes()) {
+          h = H::combine(std::move(h), size);
+        }
+      }
+      return h;
+    }
+
+    friend bool operator==(const Key& x, const Key& y) {
+      return x.input_shapes == y.input_shapes;
+    }
+  };
+
+  // A view of the key. This is used to avoid copying the input shapes and
+  // dtypes and shapes when looking up the cache.
+  struct KeyView {
+    absl::Span<const DtypeAndShape> dtypes_and_shapes;
+
+    template <typename H>
+    friend H AbslHashValue(H h, const KeyView& key) {
+      for (const auto& dtype_and_shape : key.dtypes_and_shapes) {
+        for (auto size : dtype_and_shape.GetShapeForCompilation().dim_sizes()) {
+          h = H::combine(std::move(h), size);
+        }
+      }
+      return h;
+    }
+  };
+
+  // Hash function for the key.
+  struct KeyHash {
+    using is_transparent = void;
+
+    size_t operator()(const Key& key) const { return absl::Hash<Key>()(key); }
+    size_t operator()(const KeyView& key) const {
+      return absl::Hash<KeyView>()(key);
+    }
+  };
+
+  // Equality function for the key.
+  struct KeyEq {
+    using is_transparent = void;
+
+    bool operator()(const Key& lhs, const Key& rhs) const { return lhs == rhs; }
+    bool operator()(const Key& lhs, const KeyView& rhs) const {
+      if (lhs.input_shapes.size() != rhs.dtypes_and_shapes.size()) {
+        return false;
+      }
+      for (int i = 0; i < lhs.input_shapes.size(); ++i) {
+        if (lhs.input_shapes[i] !=
+            rhs.dtypes_and_shapes[i].GetShapeForCompilation()) {
+          return false;
+        }
+      }
+      return true;
+    }
+    bool operator()(const KeyView& lhs, const Key& rhs) const {
+      return this->operator()(rhs, lhs);
+    }
+  };
+
+  struct CachedExecutableBundle {
+    std::vector<xla::ifrt::DType> ifrt_input_dtypes;
+    // If populated, these are the input shapes and layouts that the
+    // executable was compiled with. `xla_input_shapes` and `xla_input_layouts`
+    // are either both populated or both empty and they will have the same size.
+    // The index `i` in these vectors corresponds to the i-th argument in the
+    // executable.
+    std::vector<std::shared_ptr<const xla::Shape>> xla_input_shapes;
+    std::vector<absl::InlinedVector<int64_t, 4>> byte_strides;
+    std::vector<std::shared_ptr<const xla::ifrt::Shape>> ifrt_input_shapes;
+    std::vector<xla::ifrt::LayoutRef> xla_input_layouts;
+    xla::ifrt::LoadedExecutableRef ifrt_executable;
+    tensorflow::tpu::TPUCompileMetadataProto compile_metadata;
+    std::vector<std::unique_ptr<TfHostCallback>> host_callbacks;
+
+    using LoadedVariableList =
+        std::vector<IfrtLoadedVariableRegistry::LoadedVariable>;
+    using DeviceLoadedVariableListMap =
+        absl::flat_hash_map<xla::ifrt::DeviceId, LoadedVariableList>;
+    std::variant<LoadedVariableList, DeviceLoadedVariableListMap>
+        variable_arrays;
+    std::vector<xla::HloSharding> arg_hlo_shardings;
+    std::vector<xla::ifrt::ShardingRef> arg_ifrt_shardings;
+    // Only populated when portable execution is used, currently only single
+    // device sharding is supported.
+    absl::flat_hash_map<xla::ifrt::DeviceId,
+                        std::shared_ptr<xla::ifrt::SingleDeviceSharding>>
+        portable_single_device_shardings;
+    std::vector<xla::HloSharding> retval_hlo_shardings;
+
+    // Input tensor shapes that matches the Tf2Hlo compiled shapes.
+    std::vector<tensorflow::TensorShape> reshaped_input_tensors;
+
+    CachedExecutableBundle() = default;
+    // Move only
+    CachedExecutableBundle(CachedExecutableBundle&& other) = default;
+    CachedExecutableBundle& operator=(CachedExecutableBundle&& other) = default;
+    CachedExecutableBundle(const CachedExecutableBundle& other) = delete;
+    CachedExecutableBundle& operator=(const CachedExecutableBundle& other) =
+        delete;
+  };
+
+  IfrtServingExecutable(
+      int64_t program_id, absl::string_view model_name,
+      absl::string_view signature_name,
+      mlir::OwningOpRef<mlir::ModuleOp> module,
+      std::shared_ptr<xla::ifrt::Client> client,
+      tsl::thread::ThreadPool* thread_pool,
+      IfrtLoadedVariableRegistry* ifrt_loaded_variable_registry,
+      const IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry,
+      tfrt::ConcurrentWorkQueue* checkpoint_loader_queue,
+      tensorflow::DeviceMgr* device_mgr,
+      tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn,
+      IfrtServingCoreSelector* ifrt_serving_core_selector,
+      tensorflow::tpu::TPUCompileMetadataProto original_compile_metadata,
+      xla::ifrt::DeviceListRef assigned_device_list,
+      absl::flat_hash_map<size_t, size_t> static_shape_arg_map,
+      std::variant<tsl::protobuf::Message*,
+                   xla::CompileOptions::EnvironmentOptionOverrides>
+          compilation_env_or_overrides,
+      TfToHloCompiler* tf_to_hlo_compiler,
+      IfrtPersistentCompilationCache* persistent_compilation_cache,
+      H2DTransferExecutorFactory* h2d_transfer_executor_factory,
+      bool use_output_arena = false)
+      : program_id_(program_id),
+        model_name_(std::string(model_name)),
+        signature_name_(std::string(signature_name)),
+        module_(std::move(module)),
+        original_compile_metadata_(std::move(original_compile_metadata)),
+        assigned_device_list_(std::move(assigned_device_list)),
+        assigned_device_ids_([this] {
+          std::vector<int> assigned_device_ids;
+          assigned_device_ids.reserve(assigned_device_list_->size());
+          for (const auto& device : assigned_device_list_->devices()) {
+            assigned_device_ids.push_back(device->Id().value());
+          }
+          return assigned_device_ids;
+        }()),
+        static_shape_arg_map_(std::move(static_shape_arg_map)),
+        ifrt_client_(std::move(client)),
+        thread_pool_(*thread_pool),
+        ifrt_loaded_variable_registry_(*ifrt_loaded_variable_registry),
+        ifrt_restore_tensor_registry_(*ifrt_restore_tensor_registry),
+        checkpoint_loader_queue_(checkpoint_loader_queue),
+        device_mgr_(device_mgr),
+        shape_representation_fn_(std::move(shape_representation_fn)),
+        ifrt_serving_core_selector_(std::move(ifrt_serving_core_selector)),
+        compilation_env_or_overrides_(compilation_env_or_overrides),
+        tf_to_hlo_compiler_(tf_to_hlo_compiler),
+        persistent_compilation_cache_(persistent_compilation_cache),
+        h2d_transfer_executor_factory_(h2d_transfer_executor_factory),
+        use_output_arena_(use_output_arena) {
+    execute_options_.fill_status = true;
+    if (use_output_arena_) {
+      execute_options_.custom_options =
+          xla::ifrt::AttributeMap(xla::ifrt::AttributeMap::Map{
+              {"use_output_arena", xla::ifrt::AttributeMap::BoolValue(true)}});
+    }
+  }
+
+  int64_t program_id_;
+  using SharedCachedExecutableBundle = std::shared_ptr<CachedExecutableBundle>;
+
+  struct ExecutionInfo {
+    xla::ifrt::LoadedExecutable::ExecuteResult execution_result;
+    SharedCachedExecutableBundle executable_bundle;
+    xla::ifrt::DeviceListRef device_list;
+    std::vector<xla::ifrt::ArrayRef> transfer_result;
+    std::shared_ptr<tsl::DeviceReservation> device_reservation;
+  };
+
+  std::string model_name_;
+  std::string signature_name_;
+
+  mlir::OwningOpRef<mlir::ModuleOp> module_ ABSL_GUARDED_BY(mutex_);
+  // The original compile metadata. We need to keep it around to be able to
+  // test portable execution condition even if the Module itself is already
+  // released.
+  tensorflow::tpu::TPUCompileMetadataProto original_compile_metadata_;
+  const xla::ifrt::DeviceListRef assigned_device_list_;
+  // Pre-calculated device IDs to avoid redundant computation on the critical
+  // path within the Execute() call.
+  const std::vector<int> assigned_device_ids_;
+  absl::flat_hash_map<size_t /*original_arg_idx*/,
+                      size_t /*static_shape_arg_idx*/>
+      static_shape_arg_map_;
+
+  std::shared_ptr<xla::ifrt::Client> ifrt_client_;
+  tsl::thread::ThreadPool& thread_pool_;
+
+  IfrtLoadedVariableRegistry& ifrt_loaded_variable_registry_;
+  const IfrtRestoreTensorRegistry& ifrt_restore_tensor_registry_;
+  tfrt::ConcurrentWorkQueue* checkpoint_loader_queue_;
+  tensorflow::DeviceMgr* device_mgr_;  // Not owned. For host callback.
+  tensorflow::XlaHelpers::ShapeRepresentationFn shape_representation_fn_;
+  IfrtServingCoreSelector* ifrt_serving_core_selector_;
+
+  std::variant<tsl::protobuf::Message*,
+               xla::CompileOptions::EnvironmentOptionOverrides>
+      compilation_env_or_overrides_;  // proto is NOT OWNED. can be nullptr.
+
+  mutable absl::Mutex mutex_;
+  absl::flat_hash_map<Key, tsl::Future<SharedCachedExecutableBundle>, KeyHash,
+                      KeyEq>
+      executable_bundles_ ABSL_GUARDED_BY(mutex_);
+
+  bool is_frozen_ ABSL_GUARDED_BY(mutex_) = false;
+
+  // The tf_to_hlo_compiler_ is not owned by this executable. It is expected to
+  // be alive during the lifetime of the executable.
+  TfToHloCompiler* tf_to_hlo_compiler_;
+
+  // The persistent compilation cache is a global cache and is not owned by
+  // this executable. When it is nullptr, the persistent compilation cache is
+  // disabled at ifrt serving level.
+  IfrtPersistentCompilationCache* persistent_compilation_cache_;
+
+  H2DTransferExecutorFactory* h2d_transfer_executor_factory_ = nullptr;
+  bool use_output_arena_ = false;
+
+  xla::ifrt::ExecuteOptions execute_options_;
+
+  // Asynchronously load the restored variable tensors to Ifrt array.
+  absl::Status LoadAndRegisterVariableOnExecutable(
+      absl::Span<const tensorflow::Tensor> inputs,
+      absl::Span<const int> variable_arg_indices,
+      const xla::ifrt::DeviceListRef& device_list,
+      CachedExecutableBundle* executable_bundle);
+  absl::Status AsyncLoadIfrtArray(
+      absl::Span<const tensorflow::Tensor> inputs,
+      absl::Span<const int> variable_arg_indices,
+      const CachedExecutableBundle& executable_bundle,
+      const xla::ifrt::DeviceListRef& devices);
+
+  // Returns the cached executable bundle future if it exists, otherwise creates
+  // a new one by calling xla compiler. When compilation happens, it also calls
+  // `LoadAndRegisterVariableOnExecutable` to load variables on the new
+  // executable.
+  absl::StatusOr<
+      tsl::Future<IfrtServingExecutable::SharedCachedExecutableBundle>>
+  LookUpOrCreateExecutable(absl::Span<const tensorflow::Tensor> inputs,
+                           absl::Span<const DtypeAndShape> dtypes_and_shapes,
+                           absl::Span<const int> variable_arg_indices,
+                           const xla::ifrt::DeviceListRef& device_list);
+
+  // Core implementation for Execute and ExecuteAsync.
+  absl::StatusOr<ExecutionInfo> ExecuteCore(
+      absl::Span<const tensorflow::Tensor> inputs,
+      absl::Span<const int> variable_arg_indices);
+
+  // Creates an executable by calling tf2xla and xla compiler
+  absl::StatusOr<IfrtServingExecutable::SharedCachedExecutableBundle>
+  CreateExecutableSynchronously(
+      mlir::OwningOpRef<mlir::ModuleOp> module_copy,
+      const tensorflow::tpu::TPUCompileMetadataProto& compile_metadata,
+      absl::Span<const DtypeAndShape> dtypes_and_shapes,
+      absl::Span<const int> variable_arg_indices);
+
+  absl::Status PopulateInvariantMetadata(
+      const Tf2HloResult& tf2hlo_result,
+      xla::ifrt::LoadedExecutableRef ifrt_executable,
+      std::vector<std::unique_ptr<TfHostCallback>> host_callbacks,
+      const xla::ifrt::Topology* topology,
+      CachedExecutableBundle& executable_bundle);
+
+  absl::StatusOr<std::unique_ptr<xla::ifrt::Sharding>> CreateSharding(
+      int num_devices, const xla::ifrt::Shape& arg_xla_shape,
+      const xla::ifrt::Shape& sharded_shapes);
+
+  std::vector<xla::ifrt::Shape> GetArgShape(
+      int arg_index, const CachedExecutableBundle& entry);
+
+  bool UsePortableExecution();
+};
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_H_

tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test.cc

@@ -0,0 +1,794 @@
+/* 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/core/tfrt/ifrt/ifrt_serving_executable.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/log/check.h"
+#include "absl/log/globals.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/ifrt_types.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tsl/concurrency/future.h"
+#include "xla/tsl/framework/serving_device_selector.h"
+#include "xla/tsl/framework/test_util/mock_serving_device_selector.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_matcher.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/tensor_util.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h"
+#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
+#include "tsl/platform/tstring.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace {
+using tensorflow::ifrt_serving::test_utils::GetMlirModulePath;
+using ::tensorflow::test::AsTensor;
+using ::tensorflow::test::TensorEq;
+using ::testing::_;
+using ::testing::ElementsAre;
+using ::testing::HasSubstr;
+using ::testing::NiceMock;
+using ::testing::Return;
+
+// Helper to set up a mock expectation for `ReserveDevice`.
+// It returns device reservations in a round-robin fashion, cycling through
+// available cores.
+void SetUpMockDeviceReservation(
+    tsl::test_util::MockServingDeviceSelector& selector, int64_t program_id,
+    int num_cores) {
+  auto device_index = std::make_shared<int>(0);
+  int num_cores_const = std::max(1, num_cores);
+  EXPECT_CALL(selector, ReserveDevice(absl::StrCat(program_id)))
+      .WillRepeatedly([device_index, num_cores_const](::testing::Unused) {
+        return tsl::DeviceReservation((*device_index)++ % num_cores_const,
+                                      /*selector=*/nullptr);
+      });
+}
+
+// Mock class for TpuH2DTransferExecutor.
+// By default, `ScheduledH2DTransfers` pads input tensors to their static shapes
+// based on `ifrt_shape` before calling the base class method.
+class MockH2DTransferExecutor : public H2DTransferExecutor {
+ public:
+  explicit MockH2DTransferExecutor(xla::ifrt::Client& client)
+      : H2DTransferExecutor(client) {
+    ON_CALL(*this, ScheduledH2DTransfers(_, _))
+        .WillByDefault([this](absl::Span<const InputHandle> handles,
+                              tsl::thread::ThreadPool& thread_pool)
+                           -> absl::StatusOr<
+                               tsl::Future<std::vector<xla::ifrt::ArrayRef>>> {
+          std::vector<InputHandle> new_handles;
+          new_handles.reserve(handles.size());
+          for (const auto& handle : handles) {
+            new_handles.push_back(handle);
+            tensorflow::TensorShape static_shape;
+            if (handle.input_xla_shape != nullptr) {
+              static_shape =
+                  tensorflow::TensorShape(handle.input_xla_shape->dimensions());
+            } else if (handle.ifrt_shape != nullptr) {
+              static_shape = tensorflow::TensorShape(handle.ifrt_shape->dims());
+            } else {
+              static_shape = handle.tensor.shape();
+            }
+            if (handle.tensor.shape() != static_shape) {
+              tensorflow::Tensor padded_tensor(handle.tensor.dtype(),
+                                               static_shape);
+              tensorflow::tensor::DeepCopy(handle.tensor, &padded_tensor);
+              new_handles.back().tensor = padded_tensor;
+            }
+          }
+          return H2DTransferExecutor::ScheduledH2DTransfers(new_handles,
+                                                            thread_pool);
+        });
+    ON_CALL(*this, RunH2DTransfers()).WillByDefault(Return(absl::OkStatus()));
+  }
+
+  MOCK_METHOD(absl::StatusOr<tsl::Future<std::vector<xla::ifrt::ArrayRef>>>,
+              ScheduledH2DTransfers,
+              (absl::Span<const InputHandle> handles,
+               tsl::thread::ThreadPool& thread_pool),
+              (override));
+  MOCK_METHOD(absl::Status, RunH2DTransfers, (), (override));
+};
+
+class MockH2DTransferExecutorFactory : public H2DTransferExecutorFactory {
+ public:
+  MockH2DTransferExecutorFactory() {
+    ON_CALL(*this, CreateH2DTransferExecutor(_))
+        .WillByDefault([&](xla::ifrt::Client& client) {
+          return std::make_unique<MockH2DTransferExecutor>(client);
+        });
+  }
+
+  MOCK_METHOD(absl::StatusOr<std::unique_ptr<H2DTransferExecutor>>,
+              CreateH2DTransferExecutor, (xla::ifrt::Client & ifrt_client),
+              (override));
+};
+
+struct VariableInputTestParam {
+  std::vector<tensorflow::Tensor> in_tensors;
+  std::vector<bool>
+      is_variable;  // if is_variable[i] = true, then in_tensor[i] is a variable
+                    // and can be preloaded as an ifrt array.
+  std::vector<tensorflow::Tensor> expected_out_tensors;
+};
+
+using VariableInputTest =
+    ::testing::TestWithParam<std::tuple<VariableInputTestParam, bool>>;
+
+class IfrtServingExecutableTest : public ::testing::TestWithParam<bool> {
+ protected:
+  explicit IfrtServingExecutableTest() {
+    helper_ = std::make_unique<test_utils::IfrtServingExecutableTestHelper>(
+        &selector_);
+  }
+
+  absl::StatusOr<std::vector<tensorflow::Tensor>> Execute(
+      IfrtServingExecutable* executable,
+      absl::Span<const tensorflow::Tensor> inputs,
+      absl::Span<const int> variable_arg_indices = {}) {
+    if (GetParam()) {
+      TF_ASSIGN_OR_RETURN(
+          auto future, executable->ExecuteAsync(inputs, variable_arg_indices));
+      return future.Await();
+    } else {
+      return executable->Execute(inputs, variable_arg_indices);
+    }
+  }
+
+  tsl::test_util::MockServingDeviceSelector selector_;
+  std::unique_ptr<test_utils::IfrtServingExecutableTestHelper> helper_;
+};
+
+INSTANTIATE_TEST_SUITE_P(IfrtServingExecutableTests, IfrtServingExecutableTest,
+                         ::testing::Bool());
+
+TEST_P(IfrtServingExecutableTest, Basic) {
+  int64_t program_id = 123456;
+  SetUpMockDeviceReservation(selector_, program_id, helper_->num_cores());
+  auto executable =
+      helper_->MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
+
+  auto x = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
+  auto y = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
+  std::vector<tensorflow::Tensor> inputs{x, y};
+
+  // Iterate over all cores first for warmup execution.
+  for (int i = 0; i < helper_->num_cores(); i++) {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result, Execute(executable.get(), absl::MakeSpan(inputs), {}));
+  }
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result, Execute(executable.get(), absl::MakeSpan(inputs), {}));
+  const auto expected_out =
+      AsTensor<int32_t>({14}, tensorflow::TensorShape({1, 1}));
+
+  EXPECT_THAT(result, ElementsAre(TensorEq(expected_out)));
+}
+
+TEST_P(IfrtServingExecutableTest, MultipleShapes) {
+  int64_t program_id = 123456;
+  SetUpMockDeviceReservation(selector_, program_id, helper_->num_cores());
+  auto executable =
+      helper_->MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
+
+  auto x1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
+  auto y1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
+  const auto expected_out1 =
+      AsTensor<int32_t>({14}, tensorflow::TensorShape({1, 1}));
+  std::vector<tensorflow::Tensor> inputs1{x1, y1};
+
+  auto x2 = AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({1, 4}));
+  auto y2 = AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({4, 1}));
+  const auto expected_out2 =
+      AsTensor<int32_t>({30}, tensorflow::TensorShape({1, 1}));
+
+  std::vector<tensorflow::Tensor> inputs2{x2, y2};
+
+  std::vector<tensorflow::Tensor> outputs1, outputs2;
+  // Iterate over all cores first for warmup execution.
+  for (int i = 0; i < helper_->num_cores(); i++) {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result, Execute(executable.get(), absl::MakeSpan(inputs1), {}));
+  }
+  for (int i = 0; i < 3; i++) {
+    TF_ASSERT_OK_AND_ASSIGN(
+        outputs1, Execute(executable.get(), absl::MakeSpan(inputs1), {}));
+    TF_ASSERT_OK_AND_ASSIGN(
+        outputs2, Execute(executable.get(), absl::MakeSpan(inputs2), {}));
+  }
+
+  ASSERT_EQ(executable->num_executables(), 2);
+
+  EXPECT_THAT(outputs1, ElementsAre(TensorEq(expected_out1)));
+
+  EXPECT_THAT(outputs2, ElementsAre(TensorEq(expected_out2)));
+}
+
+TEST_P(IfrtServingExecutableTest, ReturnFailOnUncompiledShapeAfterFrozen) {
+  int64_t program_id = 123456;
+  SetUpMockDeviceReservation(selector_, program_id, helper_->num_cores());
+  auto executable =
+      helper_->MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
+
+  auto x1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
+  auto y1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
+  const auto expected_out1 =
+      AsTensor<int32_t>({14}, tensorflow::TensorShape({1, 1}));
+  std::vector<tensorflow::Tensor> inputs1{x1, y1};
+  std::vector<tensorflow::Tensor> outputs1;
+  for (int i = 0; i < helper_->num_cores(); i++) {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result, Execute(executable.get(), absl::MakeSpan(inputs1), {}));
+  }
+  TF_ASSERT_OK_AND_ASSIGN(
+      outputs1, Execute(executable.get(), absl::MakeSpan(inputs1), {}));
+
+  // Freeze the model
+  executable->Freeze();
+
+  // After the freeze(), already compiled shape works ok, but uncompiled shape
+  // shall return failure.
+  outputs1.clear();
+  TF_ASSERT_OK_AND_ASSIGN(
+      outputs1, Execute(executable.get(), absl::MakeSpan(inputs1), {}));
+  EXPECT_THAT(outputs1, ElementsAre(TensorEq(expected_out1)));
+
+  auto x2 = AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({1, 4}));
+  auto y2 = AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({4, 1}));
+  std::vector<tensorflow::Tensor> inputs2{x2, y2};
+
+  std::vector<tensorflow::Tensor> outputs2;
+  auto status = Execute(executable.get(), absl::MakeSpan(inputs2), {});
+
+  EXPECT_THAT(status,
+              absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+}
+
+TEST_P(IfrtServingExecutableTest,
+       FrozenErrorMessageContainsRequestedAndCachedShapes) {
+  int64_t program_id = 123456;
+  SetUpMockDeviceReservation(selector_, program_id, helper_->num_cores());
+  auto executable =
+      helper_->MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
+
+  // Warm up with shape {1, 3} x {3, 1}.
+  auto x1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
+  auto y1 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
+  std::vector<tensorflow::Tensor> inputs1{x1, y1};
+  for (int i = 0; i < helper_->num_cores(); i++) {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result, Execute(executable.get(), absl::MakeSpan(inputs1), {}));
+  }
+
+  // Freeze the model.
+  executable->Freeze();
+
+  // Try to execute with a new, uncompiled shape {1, 4} x {4, 1}.
+  auto x2 = AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({1, 4}));
+  auto y2 = AsTensor<int32_t>({1, 2, 3, 4}, tensorflow::TensorShape({4, 1}));
+  std::vector<tensorflow::Tensor> inputs2{x2, y2};
+
+  auto status = Execute(executable.get(), absl::MakeSpan(inputs2), {});
+
+  ASSERT_THAT(status,
+              absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+
+  // Verify the error message contains the requested (offending) input shapes.
+  std::string error_message(status.status().message());
+  EXPECT_THAT(error_message, HasSubstr("Requested input shapes:"));
+  EXPECT_THAT(error_message, HasSubstr("[1,4]"));
+  EXPECT_THAT(error_message, HasSubstr("[4,1]"));
+
+  // Verify the error message contains the count of cached shape sets.
+  EXPECT_THAT(error_message,
+              HasSubstr("Number of already compiled shape sets: 1"));
+
+  // Verify the error message contains the already compiled shapes.
+  EXPECT_THAT(error_message, HasSubstr("Already compiled:"));
+  EXPECT_THAT(error_message, HasSubstr("[1,3]"));
+  EXPECT_THAT(error_message, HasSubstr("[3,1]"));
+}
+
+TEST_P(IfrtServingExecutableTest, Spmd) {
+  int64_t program_id = 111111;
+  EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id))).Times(0);
+  auto executable = helper_->MakeExecutable(
+      program_id, GetMlirModulePath("spmd_executable.mlir"));
+
+  auto x = AsTensor<int32_t>({1, 2, 3, 4, 5, 6, 7, 8},
+                             tensorflow::TensorShape({4, 2}));
+  auto y = AsTensor<int32_t>({11, 12, 13, 14, 15, 16, 17, 18},
+                             tensorflow::TensorShape({4, 2}));
+
+  auto z = AsTensor<int32_t>({21, 22, 23, 24, 25, 26, 27, 28},
+                             tensorflow::TensorShape({4, 2}));
+
+  const auto expected_out = AsTensor<int32_t>({33, 36, 39, 42, 45, 48, 51, 54},
+                                              tensorflow::TensorShape({4, 2}));
+
+  std::vector<tensorflow::Tensor> inputs{x, y, z};
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result, Execute(executable.get(), absl::MakeSpan(inputs), {}));
+
+  EXPECT_THAT(result, ElementsAre(TensorEq(expected_out)));
+}
+
+TEST_P(IfrtServingExecutableTest, SpmdTwoReturns) {
+  int64_t program_id = 111111;
+  EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id))).Times(0);
+  auto executable = helper_->MakeExecutable(
+      program_id, GetMlirModulePath("spmd_executable_two_returns.mlir"));
+
+  auto x = AsTensor<int32_t>({1, 2, 3, 4, 5, 6, 7, 8},
+                             tensorflow::TensorShape({4, 2}));
+  auto y = AsTensor<int32_t>({11, 12, 13, 14, 15, 16, 17, 18},
+                             tensorflow::TensorShape({4, 2}));
+
+  auto z = AsTensor<int32_t>({21, 22, 23, 24, 25, 26, 27, 28},
+                             tensorflow::TensorShape({4, 2}));
+
+  const auto expected_out0 = AsTensor<int32_t>({33, 36, 39, 42, 45, 48, 51, 54},
+                                               tensorflow::TensorShape({4, 2}));
+  const auto expected_out1 = AsTensor<int32_t>({20, 20, 20, 20, 20, 20, 20, 20},
+                                               tensorflow::TensorShape({4, 2}));
+
+  std::vector<tensorflow::Tensor> inputs{x, y, z};
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result, Execute(executable.get(), absl::MakeSpan(inputs), {}));
+
+  EXPECT_THAT(result,
+              ElementsAre(TensorEq(expected_out0), TensorEq(expected_out1)));
+}
+
+TEST_P(IfrtServingExecutableTest, SpmdXlaCallModuleShardy) {
+  int64_t program_id = 111111;
+  EXPECT_CALL(selector_, ReserveDevice(absl::StrCat(program_id))).Times(0);
+  auto executable = helper_->MakeExecutable(
+      program_id,
+      GetMlirModulePath("spmd_executable_xla_call_module_shardy.mlir"));
+
+  auto x = AsTensor<int32_t>({11, 12, 13, 14, 15, 16, 17, 18},
+                             tensorflow::TensorShape({4, 2}));
+  auto y = AsTensor<int32_t>({8, 7, 6, 5, 4, 3, 2, 1},
+                             tensorflow::TensorShape({4, 2}));
+
+  const auto expected_out0 = AsTensor<int32_t>({3, 5, 7, 9, 11, 13, 15, 17},
+                                               tensorflow::TensorShape({4, 2}));
+  const auto expected_out1 = AsTensor<int32_t>({19, 19, 19, 19, 19, 19, 19, 19},
+                                               tensorflow::TensorShape({4, 2}));
+
+  std::vector<tensorflow::Tensor> inputs{x, y};
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result, Execute(executable.get(), absl::MakeSpan(inputs), {}));
+
+  EXPECT_THAT(result,
+              ElementsAre(TensorEq(expected_out0), TensorEq(expected_out1)));
+}
+
+TEST_F(IfrtServingExecutableTest, EncodeLayout) {
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+  mlir::MLIRContext context(registry);
+  context.loadAllAvailableDialects();
+
+  const char* const kMlirModuleStr = R"(
+    module {
+      func.func @main(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> {
+        %0 = "tf.Add"(%arg0, %arg1) : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
+        return %0 : tensor<2x2xf32>
+      }
+    }
+  )";
+
+  mlir::OwningOpRef<mlir::ModuleOp> module =
+      mlir::parseSourceString<mlir::ModuleOp>(kMlirModuleStr, &context);
+  ASSERT_TRUE(module);
+
+  // Create shapes with layout
+  xla::Shape shape0 = xla::ShapeUtil::MakeShapeWithDenseLayout(
+      xla::PrimitiveType::F32, {2, 2}, {1, 0});
+  xla::Shape shape1 = xla::ShapeUtil::MakeShapeWithDenseLayout(
+      xla::PrimitiveType::F32, {2, 2}, {0, 1});  // Different layout
+  std::vector<xla::Shape> input_shapes = {shape0, shape1};
+
+  TF_ASSERT_OK(EncodeLayout(absl::MakeSpan(input_shapes), *module));
+
+  auto func = module->lookupSymbol<mlir::func::FuncOp>("main");
+  ASSERT_TRUE(func);
+
+  auto attr0 = func.getArgAttr(0, "mhlo.layout_mode");
+  ASSERT_TRUE(attr0);
+  EXPECT_EQ(mlir::cast<mlir::StringAttr>(attr0).getValue(), "{1,0}");
+
+  auto attr1 = func.getArgAttr(1, "mhlo.layout_mode");
+  ASSERT_TRUE(attr1);
+  EXPECT_EQ(mlir::cast<mlir::StringAttr>(attr1).getValue(), "{0,1}");
+}
+
+TEST_P(IfrtServingExecutableTest, NoReturn) {
+  int64_t program_id = 111111;
+  SetUpMockDeviceReservation(selector_, program_id, helper_->num_cores());
+  auto executable = helper_->MakeExecutable(
+      program_id, GetMlirModulePath("executable_no_return.mlir"));
+
+  auto x = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
+  auto y = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({3, 1}));
+  std::vector<tensorflow::Tensor> inputs{x, y};
+  // Iterate over all cores first for warmup execution.
+  for (int i = 0; i < helper_->num_cores(); i++) {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result, Execute(executable.get(), absl::MakeSpan(inputs), {}));
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result, Execute(executable.get(), absl::MakeSpan(inputs), {}));
+
+  ASSERT_EQ(result.size(), 0);
+}
+
+TEST_P(IfrtServingExecutableTest, CompilationFailureFulfillsPromise) {
+  int64_t program_id = 999999;
+  SetUpMockDeviceReservation(selector_, program_id, helper_->num_cores());
+  auto executable =
+      helper_->MakeExecutable(program_id, GetMlirModulePath("executable.mlir"));
+
+  EXPECT_EQ(executable->num_executables(), 0);
+
+  // Pass FLOAT tensors to an MLIR module expecting INT32 (executable.mlir has
+  // %arg0: tensor<*xi32>, %arg1: tensor<*xi32>), causing compilation
+  // (UpdateCompileMetadata inside LookUpOrCreateExecutable) to fail cleanly.
+  auto x = AsTensor<float>({1.0f, 2.0f, 3.0f}, tensorflow::TensorShape({1, 3}));
+  auto y = AsTensor<float>({1.0f, 2.0f, 3.0f}, tensorflow::TensorShape({3, 1}));
+  std::vector<tensorflow::Tensor> inputs{x, y};
+
+  auto result = Execute(executable.get(), absl::MakeSpan(inputs), {});
+
+  // 1. Verify that the compilation error status is returned cleanly (and NOT
+  // "Promise destroyed without being set").
+  EXPECT_THAT(result, absl_testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr("Dtype mismatched!")));
+
+  // 2. Verify that the failed compilation future remains cached.
+  EXPECT_EQ(executable->num_executables(), 1);
+
+  // 3. Verify that subsequent execution attempts immediately return the cleanly
+  // cached error status (and NOT "Promise destroyed without being set").
+  auto second_result = Execute(executable.get(), absl::MakeSpan(inputs), {});
+  EXPECT_THAT(second_result, absl_testing::StatusIs(
+                                 absl::StatusCode::kInvalidArgument,
+                                 ::testing::HasSubstr("Dtype mismatched!")));
+  EXPECT_EQ(executable->num_executables(), 1);
+}
+
+TEST_P(IfrtServingExecutableTest, StaticShape) {
+  absl::SetVLogLevel("tpu_h2d_transfer_executor", 2);
+  int64_t program_id = 789012;
+  SetUpMockDeviceReservation(selector_, program_id, helper_->num_cores());
+
+  auto mock_h2d_factory =
+      std::make_unique<NiceMock<MockH2DTransferExecutorFactory>>();
+  helper_->SetH2DTransferExecutorFactory(std::move(mock_h2d_factory));
+
+  auto executable = helper_->MakeExecutable(
+      program_id, GetMlirModulePath("executable_static_shape.mlir"));
+
+  // The MLIR module defines a MatMul operation where `%arg0` and `%arg1`
+  // use `%arg2` as the static shape argument, as indicated by the
+  // `tf._static_shape_arg_idx = 2` attribute.
+
+  // Test case 1: Dynamic shapes are within the static shape bounds.
+  // `input_x1` has a dynamic shape of {2, 3}.
+  // `input_y1` has a dynamic shape of {2, 3}.
+  auto input_x1 =
+      AsTensor<int32_t>({1, 2, 3, 4, 5, 6}, tensorflow::TensorShape({2, 3}));
+  auto input_y1 =
+      AsTensor<int32_t>({1, 2, 3, 4, 5, 6}, tensorflow::TensorShape({2, 3}));
+  // `shape_tensor` provides the static shape {4, 3}.
+  auto shape_tensor = AsTensor<int64_t>({4, 3}, tensorflow::TensorShape({2}));
+  std::vector<tensorflow::Tensor> inputs1{input_x1, input_y1, shape_tensor};
+
+  // Iterate over all cores first for warmup execution. This ensures the
+  // executable is compiled and cached.
+  for (int i = 0; i < helper_->num_cores(); i++) {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result, Execute(executable.get(), absl::MakeSpan(inputs1), {}));
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result1, Execute(executable.get(), absl::MakeSpan(inputs1), {}));
+
+  // The computation is effectively `input_x1` * `input_y1`^T.
+  // With dynamic shapes {2, 3} and {2, 3}, this results in a {2, 2} matrix
+  // before padding.
+  // Calculation:
+  // input_x1 (2x3) * input_y1^T (3x2) -> 2x2
+  //   1 2 3    *    1 4
+  //   4 5 6         2 5
+  //                 3 6
+  // (1*1 + 2*2 + 3*3) = 1 + 4 + 9 = 14
+  // (1*4 + 2*5 + 3*6) = 4 + 10 + 18 = 32
+  // (4*1 + 5*2 + 6*3) = 4 + 10 + 18 = 32
+  // (4*4 + 5*5 + 6*6) = 16 + 25 + 36 = 77
+  // The resulting 2x2 matrix is:
+  // [[14, 32],
+  //  [32, 77]]
+  // Since the executable was compiled with a static shape, the output is padded
+  // to the static output shape of {4, 4}.
+  ASSERT_EQ(result1.size(), 1);
+  const auto& out_tensor1 = result1[0];
+  EXPECT_EQ(out_tensor1.shape(), tensorflow::TensorShape({4, 4}));
+  auto out_matrix1 = out_tensor1.matrix<int32_t>();
+  EXPECT_EQ(out_matrix1(0, 0), 14);
+  EXPECT_EQ(out_matrix1(0, 1), 32);
+  EXPECT_EQ(out_matrix1(1, 0), 32);
+  EXPECT_EQ(out_matrix1(1, 1), 77);
+
+  executable->Freeze();
+
+  // Test case 2: Different dynamic shape inputs, but the same static bounds
+  // as provided by `shape_tensor`. This should result in a cache hit.
+  // `input_x2`: 1x3. `input_y2`: 1x3.
+  auto input_x2 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
+  auto input_y2 = AsTensor<int32_t>({1, 2, 3}, tensorflow::TensorShape({1, 3}));
+  std::vector<tensorflow::Tensor> inputs2{input_x2, input_y2, shape_tensor};
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto result2, Execute(executable.get(), absl::MakeSpan(inputs2), {}));
+
+  // Calculation: `input_x2` (1x3) * `input_y2`^T (3x1) -> 1x1.
+  //   1 2 3  *  1
+  //             2
+  //             3
+  // Result: (1*1 + 2*2 + 3*3) = 14.
+  // The output is padded to the static output shape of {4, 4}.
+  ASSERT_EQ(result2.size(), 1);
+  const auto& out_tensor2 = result2[0];
+  EXPECT_EQ(out_tensor2.shape(), tensorflow::TensorShape({4, 4}));
+  auto out_matrix2 = out_tensor2.matrix<int32_t>();
+  EXPECT_EQ(out_matrix2(0, 0), 14);
+
+  // Verify that only one executable was created, since both test cases use
+  // the same `shape_tensor` for static shape.
+  EXPECT_EQ(executable->num_executables(), 1);
+}
+
+TEST_P(VariableInputTest, InterleaveVariable) {
+  const auto& param = std::get<0>(GetParam());
+  bool use_async = std::get<1>(GetParam());
+
+  tsl::test_util::MockServingDeviceSelector device_selector;
+  test_utils::IfrtServingExecutableTestHelper helper(&device_selector);
+  int64_t program_id = 111111;
+  SetUpMockDeviceReservation(device_selector, program_id, helper.num_cores());
+  auto executable = helper.MakeExecutable(
+      program_id, GetMlirModulePath("executable_long_inputs.mlir"));
+  IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry =
+      helper.ifrt_restore_tensor_registry();
+
+  std::vector<tensorflow::Tensor> inputs;
+  std::vector<int> loaded_variable_indices;
+  for (int i = 0; i < param.in_tensors.size(); i++) {
+    if (param.is_variable[i]) {
+      auto [input_tensor_promise, input_tensor_future] =
+          tsl::MakePromise<tensorflow::Tensor>();
+      IfrtRestoreTensorRegistry::RestoredTensorInfo restore_tensor_info = {
+          .dtype_and_shape = tsl::Future<DtypeAndShape>(
+              DtypeAndShape{.dtype = param.in_tensors[i].dtype(),
+                            .shape = param.in_tensors[i].shape()}),
+          .tensor_future = input_tensor_future};
+      std::string variable_name = absl::StrCat("variable_", i);
+      ASSERT_OK(ifrt_restore_tensor_registry->TryRegister(variable_name,
+                                                          restore_tensor_info));
+      loaded_variable_indices.push_back(i);
+      input_tensor_promise.Set(param.in_tensors[i]);
+      // Use string tensor containing the key (name) in place of variable
+      // tensor.
+      tensorflow::Tensor key_tensor(tensorflow::DT_STRING, {});
+      key_tensor.scalar<tsl::tstring>()() = variable_name;
+      inputs.push_back(key_tensor);
+    } else {
+      inputs.push_back(param.in_tensors[i]);
+    }
+  }
+
+  ASSERT_EQ(inputs.size(), param.is_variable.size());
+
+  auto execute_fn = [&](absl::Span<const tensorflow::Tensor> inputs,
+                        absl::Span<const int> variable_arg_indices)
+      -> absl::StatusOr<std::vector<tensorflow::Tensor>> {
+    if (use_async) {
+      TF_ASSIGN_OR_RETURN(
+          auto future, executable->ExecuteAsync(inputs, variable_arg_indices));
+      return future.Await();
+    } else {
+      return executable->Execute(inputs, variable_arg_indices);
+    }
+  };
+
+  // Iterate over all cores first for warmup execution.
+  for (int i = 0; i < helper.num_cores(); i++) {
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result, execute_fn(absl::MakeSpan(inputs),
+                                absl::MakeSpan(loaded_variable_indices)));
+  }
+
+  TF_ASSERT_OK_AND_ASSIGN(auto result,
+                          execute_fn(absl::MakeSpan(inputs),
+                                     absl::MakeSpan(loaded_variable_indices)));
+
+  EXPECT_THAT(result, ElementsAre(TensorEq(param.expected_out_tensors[0]),
+                                  TensorEq(param.expected_out_tensors[1]),
+                                  TensorEq(param.expected_out_tensors[2])));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    VariableInputTests, VariableInputTest,
+    ::testing::Combine(
+        ::testing::ValuesIn<VariableInputTestParam>(
+            {
+                // Basic case: all variables or all non-variables.
+                {
+                    .in_tensors =
+                        {
+                            AsTensor<int32_t>({2, 2}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({3, 3}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({4, 4}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({5, 5}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({10, 10}, TensorShape({1, 2})),
+                        },
+                    .is_variable = {true, true, true, true, true},
+                    .expected_out_tensors =
+                        {
+                            AsTensor<int32_t>({12}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({40}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({100}, TensorShape({1, 1})),
+                        },
+                },
+                {
+                    .in_tensors =
+                        {
+                            AsTensor<int32_t>({2, 2}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({3, 3}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({4, 4}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({5, 5}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({10, 10}, TensorShape({1, 2})),
+                        },
+                    .is_variable = {false, false, false, false, false},
+                    .expected_out_tensors =
+                        {
+                            AsTensor<int32_t>({12}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({40}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({100}, TensorShape({1, 1})),
+                        },
+                },
+                // Variable and non-variables are non-interleaved
+                {
+                    .in_tensors =
+                        {
+                            AsTensor<int32_t>({2, 2}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({3, 3}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({4, 4}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({5, 5}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({10, 10}, TensorShape({1, 2})),
+                        },
+                    .is_variable = {false, false, false, true, true},
+                    .expected_out_tensors =
+                        {
+                            AsTensor<int32_t>({12}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({40}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({100}, TensorShape({1, 1})),
+                        },
+                },
+                {
+                    .in_tensors =
+                        {
+                            AsTensor<int32_t>({2, 2}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({3, 3}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({4, 4}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({5, 5}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({10, 10}, TensorShape({1, 2})),
+                        },
+                    .is_variable = {true, true, false, false, false},
+                    .expected_out_tensors =
+                        {
+                            AsTensor<int32_t>({12}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({40}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({100}, TensorShape({1, 1})),
+                        },
+                },
+                // Variable and non-variables are interleaved
+                {
+                    .in_tensors =
+                        {
+                            AsTensor<int32_t>({2, 2}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({3, 3}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({4, 4}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({5, 5}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({10, 10}, TensorShape({1, 2})),
+                        },
+                    .is_variable = {true, false, false, true, false},
+                    .expected_out_tensors =
+                        {
+                            AsTensor<int32_t>({12}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({40}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({100}, TensorShape({1, 1})),
+                        },
+                },
+                {
+                    .in_tensors =
+                        {
+                            AsTensor<int32_t>({2, 2}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({3, 3}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({4, 4}, TensorShape({1, 2})),
+                            AsTensor<int32_t>({5, 5}, TensorShape({2, 1})),
+                            AsTensor<int32_t>({10, 10}, TensorShape({1, 2})),
+                        },
+                    .is_variable = {false, true, true, false, true},
+                    .expected_out_tensors =
+                        {
+                            AsTensor<int32_t>({12}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({40}, TensorShape({1, 1})),
+                            AsTensor<int32_t>({100}, TensorShape({1, 1})),
+                        },
+                },
+            }),
+        ::testing::Bool()));
+
+}  // namespace
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.cc

@@ -0,0 +1,105 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "xla/python/ifrt/test_util.h"
+#include "xla/tsl/framework/test_util/mock_serving_device_selector.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/status.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/resource_loader.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_persistent_compilation_cache.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
+#include "tensorflow/core/tfrt/ifrt/tf_host_callback.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+namespace tensorflow {
+namespace ifrt_serving {
+namespace test_utils {
+
+inline constexpr absl::string_view kMlirModulePath =
+    "tensorflow/core/tfrt/ifrt/testdata/";
+
+std::string GetMlirModulePath(absl::string_view module_name) {
+  return tensorflow::GetDataDependencyFilepath(
+      absl::StrCat(kMlirModulePath, module_name));
+}
+
+IfrtServingExecutableTestHelper::IfrtServingExecutableTestHelper(
+    tsl::test_util::MockServingDeviceSelector* device_selector)
+    : device_selector_(device_selector) {
+  auto client_or = xla::ifrt::test_util::GetClient();
+  TF_CHECK_OK(client_or.status());
+  client_ = std::move(client_or.value());
+  core_selector_ = std::make_unique<IfrtServingCoreSelector>(
+      device_selector_, client_->addressable_device_count());
+
+  thread_pool_ = std::make_unique<tsl::thread::ThreadPool>(
+      tsl::Env::Default(), tsl::ThreadOptions(), "IfrtSharding",
+      kThreadPoolNumThreads);
+  work_queue_ = tfrt::CreateMultiThreadedWorkQueue(
+      /*num_threads=*/4, /*num_blocking_threads=*/4);
+
+  auto device_mgr_or = ifrt_serving::CreateTfDynamicDeviceMgr();
+  TF_CHECK_OK(device_mgr_or.status());
+  device_mgr_ = std::move(device_mgr_or.value());
+
+  mlir::registerAllDialects(registry_);
+  mlir::RegisterAllTensorFlowDialects(registry_);
+  context_ = std::make_unique<mlir::MLIRContext>(registry_);
+  ifrt_persistent_compilation_cache_ =
+      std::make_unique<IfrtPersistentCompilationCache>();
+  h2d_transfer_executor_factory_ =
+      std::make_unique<H2DTransferExecutorFactory>();
+}
+
+std::unique_ptr<IfrtServingExecutable>
+IfrtServingExecutableTestHelper::MakeExecutable(int64_t program_id,
+                                                std::string mlir_module_path) {
+  auto mlir_module =
+      mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, context_.get());
+  auto executable_or = IfrtServingExecutable::Create(
+      program_id, "test", "main", std::move(mlir_module), client_,
+      thread_pool_.get(), &ifrt_loaded_variable_registry_,
+      &ifrt_restore_tensor_registry_, work_queue_.get(), device_mgr_.get(),
+      tensorflow::IdentityShapeRepresentationFn(), core_selector_.get(),
+      /*compilation_environment_proto=*/nullptr, &tf_to_hlo_compiler_,
+      ifrt_persistent_compilation_cache_.get(),
+      h2d_transfer_executor_factory_.get());
+  TF_CHECK_OK(executable_or.status());
+  return std::move(executable_or.value());
+}
+
+}  // namespace test_utils
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_serving_executable_test_util.h

@@ -0,0 +1,96 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_TEST_UTIL_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_TEST_UTIL_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/strings/string_view.h"
+#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tfrt/transforms/ifrt/tf2hlo.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/tsl/framework/test_util/mock_serving_device_selector.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "tensorflow/core/common_runtime/device_mgr.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_persistent_compilation_cache.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_core_selector.h"
+#include "tensorflow/core/tfrt/ifrt/ifrt_serving_executable.h"
+#include "tensorflow/core/tfrt/ifrt/sharding_utils.h"
+#include "tfrt/host_context/concurrent_work_queue.h"  // from @tf_runtime
+
+namespace tensorflow {
+namespace ifrt_serving {
+namespace test_utils {
+
+// A test helper class to create and IfrtServingExecutable.
+class IfrtServingExecutableTestHelper {
+ public:
+  explicit IfrtServingExecutableTestHelper(
+      tsl::test_util::MockServingDeviceSelector* device_selector);
+
+  // Creates an IfrtServingExecutable with the given program id.
+  // Note the instance of this class must outlive the returned
+  // IfrtServingExecutable.
+  std::unique_ptr<IfrtServingExecutable> MakeExecutable(
+      int64_t program_id, std::string mlir_module_path);
+
+  void SetH2DTransferExecutorFactory(
+      std::unique_ptr<H2DTransferExecutorFactory> factory) {
+    h2d_transfer_executor_factory_ = std::move(factory);
+  }
+
+  IfrtRestoreTensorRegistry* ifrt_restore_tensor_registry() {
+    return &ifrt_restore_tensor_registry_;
+  }
+
+  int num_cores() const { return client_->addressable_device_count(); }
+
+ private:
+  static constexpr int kThreadPoolNumThreads = 16;
+
+  tsl::test_util::MockServingDeviceSelector* device_selector_;  // Not owned.
+  std::unique_ptr<IfrtServingCoreSelector> core_selector_;
+  std::shared_ptr<xla::ifrt::Client> client_;
+  std::unique_ptr<tsl::thread::ThreadPool> thread_pool_;
+  IfrtLoadedVariableRegistry ifrt_loaded_variable_registry_;
+  IfrtRestoreTensorRegistry ifrt_restore_tensor_registry_;
+  std::unique_ptr<tfrt::ConcurrentWorkQueue> work_queue_;
+  std::unique_ptr<tensorflow::DynamicDeviceMgr> device_mgr_;
+
+  mlir::DialectRegistry registry_;
+  std::unique_ptr<mlir::MLIRContext> context_;
+  std::unique_ptr<IfrtPersistentCompilationCache>
+      ifrt_persistent_compilation_cache_;
+  TfToHloCompiler tf_to_hlo_compiler_;
+  std::unique_ptr<H2DTransferExecutorFactory> h2d_transfer_executor_factory_;
+};
+
+// Returns the path to the MLIR module for the given module name.
+std::string GetMlirModulePath(absl::string_view module_name);
+
+}  // namespace test_utils
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_SERVING_EXECUTABLE_TEST_UTIL_H_

tensorflow/core/tfrt/ifrt/ifrt_tensor_utils.cc

@@ -0,0 +1,65 @@
+/* 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/core/tfrt/ifrt/ifrt_tensor_utils.h"
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "xla/python/ifrt/dtype.h"
+#include "xla/python/ifrt/shape.h"
+#include "xla/python/pjrt_ifrt/pjrt_dtype.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+absl::StatusOr<tensorflow::DataType> ToTensorDataType(
+    xla::ifrt::DType ifrt_dtype) {
+  if (ifrt_dtype.kind() == xla::ifrt::DType::kString) {
+    return tensorflow::DataType::DT_STRING;
+  }
+  TF_ASSIGN_OR_RETURN(xla::PrimitiveType primitive_type,
+                      xla::ifrt::ToPrimitiveType(ifrt_dtype));
+  return tensorflow::EncodePrimitiveTypeAsDataType(primitive_type);
+}
+
+absl::StatusOr<xla::ifrt::DType> ToIfrtDType(
+    tensorflow::DataType tensor_dtype) {
+  if (tensor_dtype == tensorflow::DataType::DT_STRING) {
+    return xla::ifrt::DType(xla::ifrt::DType::kString);
+  }
+  xla::PrimitiveType primitive_type;
+  TF_RETURN_IF_ERROR(
+      tensorflow::DataTypeToPrimitiveType(tensor_dtype, &primitive_type));
+  return xla::ifrt::ToDType(primitive_type);
+}
+
+xla::ifrt::Shape ToIfrtShape(const tensorflow::TensorShape& shape) {
+  return xla::ifrt::Shape(shape.dim_sizes());
+}
+
+tensorflow::TensorShape ToTensorShape(const xla::ifrt::Shape& shape) {
+  return tensorflow::TensorShape(shape.dims());
+}
+}  // namespace ifrt_serving
+}  // namespace tensorflow

tensorflow/core/tfrt/ifrt/ifrt_tensor_utils.h

@@ -0,0 +1,45 @@
+/* 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_CORE_TFRT_IFRT_IFRT_TENSOR_UTILS_H_
+#define TENSORFLOW_CORE_TFRT_IFRT_IFRT_TENSOR_UTILS_H_
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "xla/python/ifrt/dtype.h"
+#include "xla/python/ifrt/shape.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_shape.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
+
+namespace tensorflow {
+namespace ifrt_serving {
+
+absl::StatusOr<tensorflow::DataType> ToTensorDataType(
+    xla::ifrt::DType ifrt_dtype);
+
+absl::StatusOr<xla::ifrt::DType> ToIfrtDType(tensorflow::DataType tensor_dtype);
+
+xla::ifrt::Shape ToIfrtShape(const tensorflow::TensorShape& shape);
+
+tensorflow::TensorShape ToTensorShape(const xla::ifrt::Shape& shape);
+
+}  // namespace ifrt_serving
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_TFRT_IFRT_IFRT_TENSOR_UTILS_H_

tensorflow/core/tfrt/ifrt/pjrt_cpu_client_test_lib.cc

@@ -0,0 +1,45 @@
+/* Copyright 2022 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/test_util.h"
+#include "xla/python/pjrt_ifrt/pjrt_client.h"
+#include "xla/tsl/platform/statusor.h"
+
+namespace xla {
+namespace ifrt {
+namespace {
+
+const bool kUnused =
+    (test_util::RegisterClientFactory(
+         []() -> absl::StatusOr<std::shared_ptr<Client>> {
+           xla::CpuClientOptions options;
+           options.cpu_device_count = 8;
+           TF_ASSIGN_OR_RETURN(auto pjrt_client,
+                               xla::GetXlaPjrtCpuClient(std::move(options)));
+           return std::shared_ptr<Client>(
+               PjRtClient::Create(std::move(pjrt_client)));
+         }),
+     true);
+
+}  // namespace
+}  // namespace ifrt
+}  // namespace xla