Move XLA:CPU plugin API integration point into the PJRT Plugin direct…

…ory.

PiperOrigin-RevId: 694571064

xla/pjrt/cpu/cpu_client.h

@@ -25,6 +25,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/base/attributes.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/inlined_vector.h"
@@ -723,12 +724,14 @@ struct CpuClientOptions {
   std::function<void(HloModuleConfig&)> customize_hlo_module_config;
 };
 
-absl::StatusOr<std::unique_ptr<PjRtClient>> GetTfrtCpuClient(
-    CpuClientOptions options);
+absl::StatusOr<std::unique_ptr<PjRtClient>> ABSL_DEPRECATED(
+    "Use public XLA:CPU GetXlaPjrtCpuClient instead")
+    GetTfrtCpuClient(CpuClientOptions options);
 
 // Deprecated. Use the overload that takes 'options' instead.
-inline absl::StatusOr<std::unique_ptr<PjRtClient>> GetTfrtCpuClient(
-    bool asynchronous) {
+inline absl::StatusOr<std::unique_ptr<PjRtClient>> ABSL_DEPRECATED(
+    "Use public XLA:CPU GetXlaPjrtCpuClient instead")
+    GetTfrtCpuClient(bool asynchronous) {
   CpuClientOptions options;
   options.asynchronous = asynchronous;
   return GetTfrtCpuClient(std::move(options));

xla/pjrt/plugin/xla_cpu/BUILD

@@ -0,0 +1,34 @@
+load("@tsl//tsl/platform:rules_cc.bzl", "cc_library")
+load("//xla:xla.bzl", "xla_cc_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "xla_cpu_pjrt_client",
+    srcs = [
+        "xla_cpu_pjrt_client.cc",
+    ],
+    hdrs = ["xla_cpu_pjrt_client.h"],
+    deps = [
+        "//xla/pjrt:pjrt_client",
+        "//xla/pjrt/cpu:cpu_client",
+        "@com_google_absl//absl/status:statusor",
+    ],
+)
+
+xla_cc_test(
+    name = "xla_cpu_pjrt_client_test",
+    srcs = ["xla_cpu_pjrt_client_test.cc"],
+    deps = [
+        ":xla_cpu_pjrt_client",
+        "//xla/pjrt:pjrt_client",
+        "//xla/pjrt/cpu:cpu_client",
+        "//xla/tests:xla_internal_test_main",
+        "@tsl//tsl/platform:statusor",
+        "@tsl//tsl/platform:test",
+    ],
+)

xla/pjrt/plugin/xla_cpu/README.md

@@ -0,0 +1,2 @@
+Public PJRT entry point for XLA:CPU. Please use PJRT to access XLA:CPU
+functionality.

xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.cc

@@ -0,0 +1,32 @@
+/* Copyright 2024 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 "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
+
+#include <memory>
+
+#include "absl/status/statusor.h"
+#include "xla/pjrt/cpu/cpu_client.h"
+#include "xla/pjrt/pjrt_client.h"
+
+namespace xla {
+
+absl::StatusOr<std::unique_ptr<PjRtClient>> GetXlaPjrtCpuClient(
+    CpuClientOptions options) {
+  // TODO(masonchang): Wrap the TFRTCPU Client inside the PJRT Sandwich
+  return xla::GetTfrtCpuClient(options);
+}
+
+}  // namespace xla

xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h

@@ -0,0 +1,33 @@
+/* Copyright 2024 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.
+==============================================================================*/
+
+#ifndef XLA_PJRT_PLUGIN_XLA_CPU_XLA_CPU_PJRT_CLIENT_H_
+#define XLA_PJRT_PLUGIN_XLA_CPU_XLA_CPU_PJRT_CLIENT_H_
+
+#include <memory>
+
+#include "absl/status/statusor.h"
+#include "xla/pjrt/cpu/cpu_client.h"
+#include "xla/pjrt/pjrt_client.h"
+
+namespace xla {
+
+// Public entry point to get an XLA:CPU PjRtClient
+absl::StatusOr<std::unique_ptr<PjRtClient>> GetXlaPjrtCpuClient(
+    CpuClientOptions options);
+
+}  // namespace xla
+
+#endif  // XLA_PJRT_PLUGIN_XLA_CPU_XLA_CPU_PJRT_CLIENT_H_

xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client_test.cc

@@ -0,0 +1,30 @@
+/* Copyright 2024 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 "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
+
+#include "xla/pjrt/cpu/cpu_client.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "tsl/platform/statusor.h"
+#include "tsl/platform/test.h"
+
+namespace xla {
+
+TEST(XlaCpuPjrtClientTest, GetXlaPjrtCpuClient) {
+  TF_ASSERT_OK_AND_ASSIGN(auto client, GetXlaPjrtCpuClient(CpuClientOptions()));
+  EXPECT_EQ(client->platform_name(), "cpu");
+}
+
+}  // namespace xla