chore: import upstream snapshot with attribution

tensorflow/c/experimental/next_pluggable_device/BUILD

@@ -0,0 +1,111 @@
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "c_api",
+    srcs = ["c_api.cc"],
+    hdrs = ["c_api.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":tensor_pjrt_buffer_util",
+        "//tensorflow/c:c_api_macros_hdrs",
+        "//tensorflow/c:kernels_experimental_hdrs",
+        "//tensorflow/c:kernels_hdrs",
+        "//tensorflow/c:tf_buffer",
+        "//tensorflow/c:tf_status_internal",
+        "//tensorflow/c:tf_tensor_internal",
+        "//tensorflow/compiler/jit:variable_info",
+        "//tensorflow/compiler/jit:variable_info_util",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/common_runtime/next_pluggable_device:plugin_resource",
+        "//tensorflow/core/platform:refcount",
+        "//tensorflow/core/platform:status",
+        "//tensorflow/core/tfrt/common:pjrt_util",
+        "@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/time",
+        "@com_google_absl//absl/types:span",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt/c:pjrt_c_api_hdrs",
+        "@xla//xla/pjrt/c:pjrt_c_api_helpers",
+        "@xla//xla/pjrt/c_api_client:pjrt_c_api_client",
+        "@xla//xla/tsl/distributed_runtime/coordination:coordination_service_agent",
+    ],
+)
+
+# Plugin should include this target to avoid linking the C API implementation.
+cc_library(
+    name = "c_api_hdrs",
+    hdrs = ["c_api.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/c:c_api_macros_hdrs",
+        "//tensorflow/c:kernels_hdrs",
+        "//tensorflow/c:tf_buffer_internal",
+        "//tensorflow/c:tf_status_headers",
+        "@xla//xla/pjrt/c:pjrt_c_api_hdrs",
+    ],
+)
+
+cc_library(
+    name = "tensor_pjrt_buffer_util",
+    srcs = ["tensor_pjrt_buffer_util.cc"],
+    hdrs = ["tensor_pjrt_buffer_util.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/compiler/jit:pjrt_tensor_buffer_util",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/tfrt/common:async_value_tensor",
+        "//tensorflow/core/tfrt/common:global_state",
+        "//tensorflow/core/tfrt/common:pjrt_state",
+        "//tensorflow/core/tfrt/common:pjrt_util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@xla//xla/pjrt:pjrt_client",
+        "@xla//xla/pjrt/c:pjrt_c_api_hdrs",
+        "@xla//xla/pjrt/c_api_client:pjrt_c_api_client",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:statusor",
+    ],
+)
+
+tf_cc_test(
+    name = "tensor_pjrt_buffer_util_test",
+    srcs = ["tensor_pjrt_buffer_util_test.cc"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":tensor_pjrt_buffer_util",
+        "//tensorflow/core:framework_types_hdr",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/tfrt/common:async_value_tensor",
+        "//tensorflow/core/tfrt/common:pjrt_util",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:casts",
+        "@xla//xla:shape_util",
+        "@xla//xla:xla_data_proto_cc",
+        "@xla//xla/pjrt:pjrt_api",
+        "@xla//xla/pjrt/c:pjrt_c_api_cpu",
+        "@xla//xla/pjrt/c:pjrt_c_api_hdrs",
+        "@xla//xla/pjrt/c:pjrt_c_api_wrapper_impl",
+        "@xla//xla/pjrt/c_api_client:pjrt_c_api_client",
+        "@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/platform:status_matchers",
+        "@xla//xla/tsl/platform:statusor",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)

tensorflow/c/experimental/next_pluggable_device/c_api.cc

@@ -0,0 +1,351 @@
+/* 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/c/experimental/next_pluggable_device/c_api.h"
+
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/time/time.h"
+#include "absl/types/span.h"
+#include "tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h"
+#include "tensorflow/c/kernels.h"
+#include "tensorflow/c/kernels_experimental.h"
+#include "tensorflow/c/tf_buffer.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_status_internal.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/c/tf_tensor_internal.h"
+#include "tensorflow/compiler/jit/variable_info.h"
+#include "tensorflow/compiler/jit/variable_info_util.h"
+#include "xla/pjrt/c/pjrt_c_api.h"
+#include "xla/pjrt/c/pjrt_c_api_helpers.h"
+#include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/tsl/distributed_runtime/coordination/coordination_service_agent.h"
+#include "tensorflow/core/common_runtime/next_pluggable_device/plugin_resource.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/tfrt/common/pjrt_util.h"
+
+TF_Device* TF_GetDevice(TF_OpKernelContext* ctx) {
+  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
+  return reinterpret_cast<TF_Device*>(cc_ctx->device());
+}
+
+// --------------------------  Resource  ---------------------------------------
+void TF_CreatePluginResource(TF_OpKernelContext* ctx,
+                             const char* container_name,
+                             const char* plugin_resource_name,
+                             void* plugin_resource, void (*delete_func)(void*),
+                             TF_Status* status) {
+  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
+  tensorflow::PluginResource* cc_resource_ptr = new tensorflow::PluginResource(
+      plugin_resource, plugin_resource_name, delete_func);
+  auto cc_status =
+      cc_ctx->resource_manager()->Create<tensorflow::PluginResource>(
+          container_name, plugin_resource_name, cc_resource_ptr);
+  status->status = cc_status;
+}
+
+void TF_LookupOrCreatePluginResource(
+    TF_OpKernelContext* ctx, const char* container_name,
+    const char* plugin_resource_name, void** result_plugin_resource,
+    void* (*create_func)(void*), void* create_func_args,
+    void (*delete_func)(void*), TF_Status* status) {
+  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
+  auto* resource_mgr = cc_ctx->resource_manager();
+  tensorflow::core::RefCountPtr<tensorflow::PluginResource>
+      tf_plugin_resource_ptr;
+  tensorflow::PluginResource* tf_plugin_resource = nullptr;
+
+  auto cc_status = resource_mgr->LookupOrCreate<tensorflow::PluginResource>(
+      container_name, plugin_resource_name, &tf_plugin_resource,
+      [plugin_resource_name, create_func, create_func_args,
+       delete_func](tensorflow::PluginResource** new_resource) {
+        void* opaque_plugin_resource = create_func(create_func_args);
+        *new_resource = new tensorflow::PluginResource(
+            opaque_plugin_resource, plugin_resource_name, delete_func);
+        return absl::OkStatus();
+      });
+
+  if (cc_status.ok()) {
+    tf_plugin_resource_ptr.reset(tf_plugin_resource);
+    *result_plugin_resource = tf_plugin_resource_ptr->GetOpaquePluginResource();
+  } else {
+    *result_plugin_resource = nullptr;
+  }
+  status->status = cc_status;
+}
+
+// -------------------------  VariableInfo  ------------------------------------
+struct TF_VariableInfo {
+  TF_VariableInfo() = delete;
+  // TF_VariableInfo is constructed here by TensorFlow, and will be passed to
+  // plugin as a opaque pointer. Plugin will need to call C APIs below to
+  // operate on TF_VariableInfo (such as allocate temp tensor for the `var` held
+  // by the underlying tensorflow::VariableInfo.
+  TF_VariableInfo(int index, const std::string& name, tensorflow::Var* var) {
+    var_info = tensorflow::VariableInfo{index, name, var};
+  }
+
+  tensorflow::VariableInfo var_info{0, "", nullptr};
+};
+
+TF_VariableInfo* TF_CreateVariableInfoFromContext(TF_OpKernelContext* ctx,
+                                                  int index,
+                                                  TF_Status* status) {
+  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
+  const tensorflow::Tensor& arg_tensor = cc_ctx->input(index);
+  absl::Status cc_status;
+  if (arg_tensor.dtype() != tensorflow::DT_RESOURCE) {
+    cc_status = absl::InvalidArgumentError(
+        absl::StrCat("Trying to obtain resource handle from Input[", index,
+                     "], which is not type DT_RESOURCE."));
+    status->status = cc_status;
+    return nullptr;
+  }
+  const tensorflow::ResourceHandle& handle =
+      arg_tensor.flat<tensorflow::ResourceHandle>()(0);
+  tensorflow::Var* variable;
+  cc_status = tensorflow::LookupResource(cc_ctx, handle, &variable);
+  return new TF_VariableInfo(index, handle.name(), variable);
+}
+
+void TF_LockVariableInfos(TF_VariableInfo** vars, int num_vars,
+                          TF_Status* status) {
+  std::vector<tensorflow::VariableInfo*> variable_ptrs;
+  variable_ptrs.reserve(num_vars);
+  for (int i = 0; i < num_vars; ++i) {
+    variable_ptrs.push_back(&(vars[i]->var_info));
+  }
+  absl::Status cc_status = LockVariables(absl::MakeSpan(variable_ptrs));
+  status->status = cc_status;
+}
+
+void TF_AllocateTempForVariableInfo(TF_OpKernelContext* ctx,
+                                    TF_VariableInfo* var_info,
+                                    TF_Status* status) {
+  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
+  absl::Status cc_status;
+  if (var_info == nullptr) {
+    cc_status = absl::InvalidArgumentError("TF_VariableInfo is NULL.");
+    status->status = cc_status;
+    return;
+  }
+  if (var_info->var_info.var() == nullptr) {
+    cc_status = absl::InvalidArgumentError(
+        "VariableInfo does not track a resource variable.");
+    status->status = cc_status;
+    return;
+  }
+
+  cc_status = cc_ctx->allocate_temp(var_info->var_info.var()->tensor()->dtype(),
+                                    var_info->var_info.var()->tensor()->shape(),
+                                    var_info->var_info.var()->tensor());
+  status->status = cc_status;
+}
+
+TF_Tensor* TF_GetTensorFromVariableInfo(TF_VariableInfo* var_info,
+                                        TF_Status* status) {
+  absl::Status cc_status;
+  if (var_info == nullptr) {
+    cc_status = absl::InvalidArgumentError("TF_VariableInfo is NULL.");
+    status->status = cc_status;
+    return nullptr;
+  }
+  if (var_info->var_info.var() == nullptr) {
+    cc_status = absl::InvalidArgumentError(
+        "VariableInfo does not track a resource variable.");
+    status->status = cc_status;
+    return nullptr;
+  }
+
+  tensorflow::Tensor* tensor = var_info->var_info.var()->tensor();
+  TF_Tensor* result_tensor =
+      tensorflow::TF_TensorFromTensor(*tensor, &cc_status);
+  status->status = cc_status;
+  return result_tensor;
+}
+
+void TF_DeleteVariableInfo(TF_VariableInfo* var_info) {
+  if (var_info != nullptr) {
+    delete var_info;
+  }
+}
+
+// ---------------------  Coordination service  --------------------------------
+TF_CoordinationServiceAgent* TF_GetCoordinationServiceAgent(
+    TF_OpKernelContext* ctx) {
+  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
+  return reinterpret_cast<TF_CoordinationServiceAgent*>(
+      cc_ctx->coordination_service_agent());
+}
+
+bool TF_CoordinationServiceIsInitialized(TF_CoordinationServiceAgent* agent) {
+  if (agent == nullptr) return false;
+  auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
+  return cc_agent->IsInitialized();
+}
+
+void TF_CoordinationServiceInsertKeyValue(const char* key, int64_t key_size,
+                                          const char* value, int64_t value_size,
+                                          TF_CoordinationServiceAgent* agent,
+                                          TF_Status* status) {
+  auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
+  absl::Status cc_status = cc_agent->InsertKeyValue(
+      absl::string_view(key, key_size), absl::string_view(value, value_size));
+  status->status = cc_status;
+}
+
+TF_Buffer* ProcessGetKeyValueResult(absl::StatusOr<std::string> value,
+                                    TF_Status* status) {
+  status->status = value.status();
+  if (!value.ok()) {
+    return nullptr;
+  }
+  // Caller is responsible to call `TF_DeleteBuffer` to release the buffer.
+  TF_Buffer* result = TF_NewBuffer();
+  const std::string& value_str = *value;
+  void* data = malloc(value_str.length());
+  value_str.copy(static_cast<char*>(data), value_str.length(), 0);
+  result->data = data;
+  result->length = value_str.length();
+  result->data_deallocator = [](void* data, size_t length) { free(data); };
+  return result;
+}
+
+TF_Buffer* TF_CoordinationServiceGetKeyValue(const char* key, int64_t key_size,
+                                             TF_CoordinationServiceAgent* agent,
+                                             TF_Status* status) {
+  auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
+  auto value = cc_agent->GetKeyValue(absl::string_view(key, key_size));
+  return ProcessGetKeyValueResult(value, status);
+}
+
+TF_Buffer* TF_CoordinationServiceGetKeyValueWithTimeout(
+    const char* key, int64_t key_size, int64_t timeout_seconds,
+    TF_CoordinationServiceAgent* agent, TF_Status* status) {
+  if (timeout_seconds <= 0) {
+    status->status = absl::InvalidArgumentError(
+        "TF_CoordinationServiceGetKeyValueWithTimeout invoked with invalid "
+        "timeout_seconds <= 0.");
+    return nullptr;
+  }
+  auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
+  auto value = cc_agent->GetKeyValue(absl::string_view(key, key_size),
+                                     absl::Seconds(timeout_seconds));
+  return ProcessGetKeyValueResult(value, status);
+}
+
+TF_Buffer* TF_CoordinationServiceTryGetKeyValue(
+    const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
+    TF_Status* status) {
+  auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
+  auto value = cc_agent->TryGetKeyValue(absl::string_view(key, key_size));
+  return ProcessGetKeyValueResult(value, status);
+}
+
+void TF_CoordinationServiceDeleteKeyValue(const char* key, int64_t key_size,
+                                          TF_CoordinationServiceAgent* agent,
+                                          TF_Status* status) {
+  auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
+  absl::Status cc_status =
+      cc_agent->DeleteKeyValue(absl::string_view(key, key_size));
+  status->status = cc_status;
+}
+
+// ----------------------------  PJRT  -----------------------------------------
+void TF_CreateAndSetPjRtCApiClient(const char* device_type, TF_Status* status,
+                                   PJRT_NamedValue* create_options,
+                                   int num_options) {
+  absl::StatusOr<std::unique_ptr<xla::PjRtClient>> pjrt_client =
+      xla::GetCApiClient(device_type, pjrt::ConvertFromPjRtNamedValueList(
+                                          create_options, num_options));
+  if (!pjrt_client.ok()) {
+    status->status = pjrt_client.status();
+    return;
+  }
+
+  absl::Status s = tensorflow::SetPjRtClientInTFGlobalResourceManager(
+      tensorflow::DeviceType(device_type), std::move(*pjrt_client));
+  status->status = s;
+}
+
+void TF_ResetPjRtCClient(const char* device_type, TF_Status* status) {
+  status->status =
+      tensorflow::ResetPjRtClient(tensorflow::DeviceType(device_type));
+}
+
+PJRT_Client* TF_GetPjRtCClient(const char* device_type, TF_Status* status) {
+  absl::StatusOr<xla::PjRtCApiClient*> pjrt_c_api_client =
+      tensorflow::GetPjRtCApiClient(tensorflow::DeviceType(device_type));
+  if (!pjrt_c_api_client.ok()) {
+    status->status = pjrt_c_api_client.status();
+    return nullptr;
+  }
+  status->status = absl::OkStatus();
+  return (*pjrt_c_api_client)->pjrt_c_client();
+}
+
+PJRT_Buffer* TF_GetPjRtCBuffer(TF_Tensor* c_tensor, TF_Status* status) {
+  tensorflow::Tensor tensor;
+  auto s = tensorflow::TF_TensorToTensor(c_tensor, &tensor);
+  if (!s.ok()) {
+    status->status = s;
+    return nullptr;
+  }
+  absl::StatusOr<PJRT_Buffer*> c_buffer =
+      tensorflow::GetPjRtCBufferFromTensor(&tensor);
+  if (!c_buffer.ok()) {
+    status->status = c_buffer.status();
+    return nullptr;
+  }
+  status->status = absl::OkStatus();
+  return *c_buffer;
+}
+
+void TF_CreatePjRtBuffer(TF_Tensor* c_tensor, PJRT_Buffer* c_buffer,
+                         const char* device_type, TF_Status* status) {
+  tensorflow::Tensor tensor;
+  auto s = tensorflow::TF_TensorToTensor(c_tensor, &tensor);
+  if (!s.ok()) {
+    status->status = s;
+    return;
+  }
+  absl::StatusOr<xla::PjRtCApiClient*> pjrt_c_api_client =
+      tensorflow::GetPjRtCApiClient(tensorflow::DeviceType(device_type));
+  if (!pjrt_c_api_client.ok()) {
+    status->status = pjrt_c_api_client.status();
+    return;
+  }
+  auto set_buffer_status =
+      SetPjRtCBufferToTensor(c_buffer, *pjrt_c_api_client, &tensor);
+  status->status = set_buffer_status;
+}

tensorflow/c/experimental/next_pluggable_device/c_api.h

@@ -0,0 +1,156 @@
+/* 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_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_C_API_H_
+#define TENSORFLOW_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_C_API_H_
+
+#include <cstdint>
+
+#include "tensorflow/c/c_api_macros.h"
+#include "tensorflow/c/kernels.h"
+#include "tensorflow/c/kernels_experimental.h"
+#include "tensorflow/c/tf_buffer.h"
+#include "tensorflow/c/tf_status.h"
+#include "xla/pjrt/c/pjrt_c_api.h"
+
+// --------------------------------------------------------------------------
+// C API for device. The API is under active development and eventually
+// should allow registering a plugin device with TensorFlow.
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// TF_Device is a C wrapper to the C++ TF Device class. This is to be passed
+// through TF_OpKernelContext, and is opaque to plugin.
+typedef struct TF_Device TF_Device;
+
+typedef struct TF_VariableInfo TF_VariableInfo;
+
+// Returns a `TF_Device` pointer, which actually points to a C++ `Device`.
+// Currently we only allow `NextPluggableDevice` to be casted as `TF_Device`,
+// but in theory every this is a C API for every kind of device.
+TF_CAPI_EXPORT extern TF_Device* TF_GetDevice(TF_OpKernelContext* ctx);
+
+// --------------------------  Resource  ---------------------------------------
+// Create a `tensorflow::PluginResource` to the ResourceMgr provided by the
+// `ctx`. The `tensorflow::PluginResource` wraps a resource by plugin (as a
+// opaque pointer, since TensorFlow cannot parse it). `delete_func` is needed
+// for ResourceMgr to clean up the resource. `status` will be set.
+TF_CAPI_EXPORT extern void TF_CreatePluginResource(
+    TF_OpKernelContext* ctx, const char* container_name,
+    const char* plugin_resource_name, void* plugin_resource,
+    void (*delete_func)(void*), TF_Status* status);
+
+// If the ResourceMgr provided by the `ctx` has a resource
+// `plugin_resource_name`, returns it in `*result_plugin_resource`. Otherwise,
+// invokes create_func to create the resource. `delete_func` is needed for
+// ResourceMgr to clean up the resource. `status` will be set. If `status` is
+// not OK, `*result_plugin_resource` will be set as nullptr.
+//
+// Caller does not take ownership of the `plugin_resource`.
+TF_CAPI_EXPORT extern void TF_LookupOrCreatePluginResource(
+    TF_OpKernelContext* ctx, const char* container_name,
+    const char* plugin_resource_name, void** result_plugin_resource,
+    void* (*create_func)(void*), void* create_func_args,
+    void (*delete_func)(void*), TF_Status* status);
+
+// -------------------------  VariableInfo  ------------------------------------
+TF_CAPI_EXPORT extern TF_VariableInfo* TF_CreateVariableInfoFromContext(
+    TF_OpKernelContext* ctx, int index, TF_Status* status);
+
+TF_CAPI_EXPORT extern void TF_LockVariableInfos(TF_VariableInfo** vars,
+                                                int num_vars,
+                                                TF_Status* status);
+
+TF_CAPI_EXPORT extern void TF_AllocateTempForVariableInfo(
+    TF_OpKernelContext* ctx, TF_VariableInfo* var_info, TF_Status* status);
+
+TF_CAPI_EXPORT extern TF_Tensor* TF_GetTensorFromVariableInfo(
+    TF_VariableInfo* var_info, TF_Status* status);
+
+TF_CAPI_EXPORT extern void TF_DeleteVariableInfo(TF_VariableInfo* var_info);
+
+// ---------------------  Coordination service  --------------------------------
+// Returns a not owning pointer to the coordination service agent, which is
+// opaque to plugin. Plugin OpKernels need to use the accompanying C APIs to
+// access coordination service functionalities.
+TF_CAPI_EXPORT extern TF_CoordinationServiceAgent*
+TF_GetCoordinationServiceAgent(TF_OpKernelContext* ctx);
+
+// Returns true if the coordination service agent has been initialized.
+TF_CAPI_EXPORT extern bool TF_CoordinationServiceIsInitialized(
+    TF_CoordinationServiceAgent* agent);
+
+TF_CAPI_EXPORT extern void TF_CoordinationServiceInsertKeyValue(
+    const char* key, int64_t key_size, const char* value, int64_t value_size,
+    TF_CoordinationServiceAgent* agent, TF_Status* status);
+
+// Obtains key-value from coordination service agent. The returned `TF_Buffer`
+// is a newly allocated buffer to hold the string key-value, and caller is
+// responsible for managing the lifetime. If error, `status` will be set and a
+// nullptr will be returned.
+TF_CAPI_EXPORT extern TF_Buffer* TF_CoordinationServiceGetKeyValue(
+    const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
+    TF_Status* status);
+
+TF_CAPI_EXPORT extern TF_Buffer* TF_CoordinationServiceGetKeyValueWithTimeout(
+    const char* key, int64_t key_size, int64_t timeout_seconds,
+    TF_CoordinationServiceAgent* agent, TF_Status* status);
+
+TF_CAPI_EXPORT extern TF_Buffer* TF_CoordinationServiceTryGetKeyValue(
+    const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
+    TF_Status* status);
+
+TF_CAPI_EXPORT extern void TF_CoordinationServiceDeleteKeyValue(
+    const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
+    TF_Status* status);
+
+// ----------------------------  PJRT  -----------------------------------------
+// Passes the pointer to a vector of PJRT_NamedValue and number of options to
+// set options for creating a PJRT client. Passes nullptr for create_options and
+// 0 for num_options if no options need to be set. You can use
+// ConvertToPjRtNamedValueList in
+// tensorflow/compiler/xla/pjrt/c/pjrt_c_api_helpers.h to generate the options.
+TF_CAPI_EXPORT extern void TF_CreateAndSetPjRtCApiClient(
+    const char* device_type, TF_Status* status, PJRT_NamedValue* create_options,
+    int num_options);
+
+// Resets the PjRt client for a device. After this, `TF_GetPjRtCClient` will
+// returns an error for that device.
+TF_CAPI_EXPORT extern void TF_ResetPjRtCClient(const char* device_type,
+                                               TF_Status* status);
+
+// Gets the `PJRT_Client*` stored in TF global ResourceManager.
+TF_CAPI_EXPORT extern PJRT_Client* TF_GetPjRtCClient(const char* device_type,
+                                                     TF_Status* status);
+
+// Gets the `PJRT_Buffer*` stored in the tensor. The status will contain error
+// if the tensor does not have a `PjRtCApiBuffer`.
+TF_CAPI_EXPORT extern PJRT_Buffer* TF_GetPjRtCBuffer(TF_Tensor* c_tensor,
+                                                     TF_Status* status);
+
+// Creates a `PjRtCApiBuffer` with the `PJRT_Buffer*` passed in and set to the
+// tensor.
+TF_CAPI_EXPORT extern void TF_CreatePjRtBuffer(TF_Tensor* c_tensor,
+                                               PJRT_Buffer* c_buffer,
+                                               const char* device_type,
+                                               TF_Status* status);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_C_API_H_

tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.cc

@@ -0,0 +1,92 @@
+/* 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/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h"
+
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h"
+#include "xla/pjrt/c/pjrt_c_api.h"
+#include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/tfrt/common/async_value_tensor.h"
+#include "tensorflow/core/tfrt/common/global_state.h"
+#include "tensorflow/core/tfrt/common/pjrt_state.h"
+#include "tensorflow/core/tfrt/common/pjrt_util.h"
+
+namespace tensorflow {
+
+absl::StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor) {
+  tensorflow::AsyncValueTensor* av_tensor =
+      tensorflow::AsyncValueTensor::FromTensor(tensor);
+  if (av_tensor == nullptr || av_tensor->GetBuffer() == nullptr) {
+    return absl::InternalError("Input tensor does not have PjRtBuffer.");
+  }
+  auto* c_api_buffer =
+      dynamic_cast<xla::PjRtCApiBuffer*>(av_tensor->GetBuffer().get());
+  if (c_api_buffer == nullptr) {
+    return absl::InternalError(
+        "The PjRtBuffer in the tensor is not type PjRtCApiBuffer.");
+  }
+  return c_api_buffer->c_buffer();
+}
+
+absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
+                                    xla::PjRtCApiClient* c_api_client,
+                                    Tensor* tensor) {
+  auto buffer = std::make_unique<xla::PjRtCApiBuffer>(c_api_client, c_buffer);
+  tensorflow::AsyncValueTensor* av_tensor =
+      tensorflow::AsyncValueTensor::FromTensor(tensor);
+  if (av_tensor == nullptr) {
+    TF_ASSIGN_OR_RETURN(
+        *tensor, MakeTensorFromPjRtBuffer(tensor->dtype(), tensor->shape(),
+                                          std::move(buffer)));
+  } else {
+    av_tensor->SetBuffer(std::move(buffer));
+  }
+  return absl::OkStatus();
+}
+
+absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(
+    const DeviceType& device_type) {
+  TF_ASSIGN_OR_RETURN(absl::StatusOr<xla::PjRtClient*> pjrt_client,
+                      tensorflow::GetPjRtClient(device_type));
+  auto* pjrt_c_api_client = dynamic_cast<xla::PjRtCApiClient*>(*pjrt_client);
+  if (pjrt_c_api_client == nullptr) {
+    return absl::InternalError(absl::StrCat("PjRtClient for ",
+                                            device_type.type_string(),
+                                            " is not type PjRtCApiClient"));
+  }
+  return pjrt_c_api_client;
+}
+
+absl::Status ResetPjRtClient(const DeviceType& device_type) {
+  ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr();
+  PjRtState* pjrt_state;
+  TF_RETURN_IF_ERROR(rmgr->Lookup(rmgr->default_container(),
+                                  kPjRtStateResourceName, &pjrt_state));
+  TF_RETURN_IF_ERROR(pjrt_state->MovePjRtClientToUnused(device_type));
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h

@@ -0,0 +1,39 @@
+/* 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_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_TENSOR_PJRT_BUFFER_UTIL_H_
+#define TENSORFLOW_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_TENSOR_PJRT_BUFFER_UTIL_H_
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "xla/pjrt/c/pjrt_c_api.h"
+#include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
+#include "tensorflow/core/framework/tensor.h"
+
+namespace tensorflow {
+
+absl::StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor);
+
+absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
+                                    xla::PjRtCApiClient* c_api_client,
+                                    Tensor* tensor);
+
+absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(
+    const DeviceType& device_type);
+
+absl::Status ResetPjRtClient(const DeviceType& device_type);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_TENSOR_PJRT_BUFFER_UTIL_H_

tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util_test.cc

@@ -0,0 +1,206 @@
+/* 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/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/base/casts.h"
+#include "absl/log/check.h"
+#include "absl/status/status_matchers.h"
+#include "absl/strings/str_cat.h"
+#include "xla/pjrt/c/pjrt_c_api.h"
+#include "xla/pjrt/c/pjrt_c_api_cpu.h"
+#include "xla/pjrt/c/pjrt_c_api_wrapper_impl.h"
+#include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
+#include "xla/pjrt/pjrt_api.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
+#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/status_matchers.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "xla/xla_data.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/tfrt/common/async_value_tensor.h"
+#include "tensorflow/core/tfrt/common/pjrt_util.h"
+#include "tsl/platform/casts.h"
+
+namespace tensorflow {
+namespace {
+
+using ::testing::HasSubstr;
+using ::testing::NotNull;
+using ::tsl::testing::StatusIs;
+
+PJRT_Buffer* CreateCBuffer() {
+  auto status = pjrt::PjrtApi(DEVICE_CPU);
+  if (!status.ok()) {
+    CHECK_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
+  }
+  auto pjrt_client = xla::GetCApiClient(DEVICE_CPU);
+  CHECK_OK(pjrt_client.status());
+  auto c_api_client = absl::down_cast<xla::PjRtCApiClient*>(pjrt_client->get());
+  std::vector<int32_t> data(1, 0);
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});
+
+  auto buffer = c_api_client->pjrt_c_client()->client->BufferFromHostBuffer(
+      data.data(), shape.element_type(), shape.dimensions(),
+      /*byte_strides=*/std::nullopt,
+      xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, nullptr,
+      c_api_client->pjrt_c_client()->client->memory_spaces()[0],
+      /*device_layout=*/nullptr);
+  CHECK_OK(buffer.status());
+
+  return new PJRT_Buffer{std::move(*buffer), c_api_client->pjrt_c_client()};
+}
+
+TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorNoBuffer) {
+  auto allocator = std::make_unique<AsyncValueAllocator>();
+  tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
+
+  EXPECT_THAT(
+      GetPjRtCBufferFromTensor(&tensor),
+      absl_testing::StatusIs(
+          error::INTERNAL,
+          HasSubstr(absl::StrCat("Input tensor does not have PjRtBuffer"))));
+}
+
+TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorIncoorectType) {
+  auto allocator = std::make_unique<AsyncValueAllocator>();
+  tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
+  xla::CpuClientOptions options;
+  options.asynchronous = true;
+  options.cpu_device_count = 1;
+
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetXlaPjrtCpuClient(options));
+  std::vector<int32_t> data(1, 0);
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto buffer,
+      pjrt_client->BufferFromHostBuffer(
+          data.data(), shape.element_type(), shape.dimensions(),
+          /*byte_strides=*/std::nullopt,
+          xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
+          nullptr, pjrt_client->memory_spaces()[0], /*device_layout=*/nullptr));
+  tensorflow::AsyncValueTensor* av_tensor =
+      tensorflow::AsyncValueTensor::FromTensor(&tensor);
+  av_tensor->SetBuffer(std::move(buffer));
+
+  EXPECT_THAT(
+      GetPjRtCBufferFromTensor(&tensor),
+      absl_testing::StatusIs(
+          error::INTERNAL,
+          HasSubstr(absl::StrCat(
+              "The PjRtBuffer in the tensor is not type PjRtCApiBuffer"))));
+}
+
+TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorSuccess) {
+  auto allocator = std::make_unique<AsyncValueAllocator>();
+  tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
+  auto status = pjrt::PjrtApi(DEVICE_CPU);
+  if (!status.ok()) {
+    TF_ASSERT_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
+  }
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
+  std::vector<int32_t> data(1, 0);
+  xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto buffer,
+      pjrt_client->BufferFromHostBuffer(
+          data.data(), shape.element_type(), shape.dimensions(),
+          /*byte_strides=*/std::nullopt,
+          xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
+          nullptr, pjrt_client->memory_spaces()[0], /*device_layout=*/nullptr));
+  tensorflow::AsyncValueTensor* av_tensor =
+      tensorflow::AsyncValueTensor::FromTensor(&tensor);
+  av_tensor->SetBuffer(std::move(buffer));
+
+  TF_ASSERT_OK_AND_ASSIGN(auto c_buffer, GetPjRtCBufferFromTensor(&tensor));
+
+  EXPECT_THAT(c_buffer, NotNull());
+}
+
+TEST(TensorPjRtBufferUtilTest, SetPjRtCBufferToTensorNotAsyncValueTensor) {
+  tensorflow::Tensor tensor(DT_FLOAT, {1});
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
+  PJRT_Buffer* c_buffer = CreateCBuffer();
+
+  TF_EXPECT_OK(SetPjRtCBufferToTensor(
+      c_buffer, absl::down_cast<xla::PjRtCApiClient*>(pjrt_client.get()),
+      &tensor));
+}
+
+TEST(TensorPjRtBufferUtilTest, SetPjRtCBufferToTensorSuccess) {
+  auto allocator = std::make_unique<AsyncValueAllocator>();
+  tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
+  PJRT_Buffer* c_buffer = CreateCBuffer();
+
+  TF_EXPECT_OK(SetPjRtCBufferToTensor(
+      c_buffer, absl::down_cast<xla::PjRtCApiClient*>(pjrt_client.get()),
+      &tensor));
+}
+
+TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientNotFound) {
+  EXPECT_THAT(GetPjRtCApiClient(tensorflow::DeviceType(DEVICE_CPU)),
+              absl_testing::StatusIs(
+                  error::NOT_FOUND,
+                  HasSubstr(absl::StrCat(
+                      "PjRt client not found for device type ", DEVICE_CPU))));
+}
+
+TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientIncorrectType) {
+  xla::CpuClientOptions options;
+  options.asynchronous = true;
+  options.cpu_device_count = 1;
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetXlaPjrtCpuClient(options));
+
+  TF_ASSERT_OK(SetPjRtClientInTFGlobalResourceManager(DEVICE_CPU,
+                                                      std::move(pjrt_client)));
+
+  EXPECT_THAT(GetPjRtCApiClient(tensorflow::DeviceType(DEVICE_CPU)),
+              absl_testing::StatusIs(
+                  error::INTERNAL,
+                  HasSubstr(absl::StrCat("PjRtClient for ", DEVICE_CPU,
+                                         " is not type PjRtCApiClient"))));
+}
+
+TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientSuccess) {
+  auto status = pjrt::PjrtApi(DEVICE_CPU);
+  if (!status.ok()) {
+    TF_ASSERT_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
+  }
+  TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
+  TF_ASSERT_OK(SetPjRtClientInTFGlobalResourceManager(DEVICE_CPU,
+                                                      std::move(pjrt_client)));
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto pjrt_client_get,
+      GetPjRtCApiClient(tensorflow::DeviceType(DEVICE_CPU)));
+
+  EXPECT_THAT(pjrt_client_get, NotNull());
+}
+
+}  // namespace
+}  // namespace tensorflow