chore: import upstream snapshot with attribution

third_party/xla/build_tools/pjrt_wheels/BUILD.bazel

@@ -0,0 +1,172 @@
+load("@cuda_cudart//:version.bzl", cuda_major_version = "VERSION")
+load("@nightly_timestamp//:timestamp.bzl", "XLA_NIGHTLY_TIMESTAMP")
+load("@rc_number//:rc_number.bzl", "XLA_RC_NUMBER")
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("@rules_python//python:packaging.bzl", "py_wheel")
+
+# This ensures we can only build plugins for selected CUDA versions.
+cuda_label = "cuda" + cuda_major_version if cuda_major_version else "null"
+
+# If we're building a nightly, append .devYYYYMMDD to the version
+# If we're not, the timestamp is an empty string
+# If we're building a release candidate, append rc# to the version
+# If we're not, the rc number is an empty string
+# Note that PEP 440 supports both of these at the same time, but
+# our CI will never set both.
+wheel_version = "0.0.0" + XLA_RC_NUMBER + XLA_NIGHTLY_TIMESTAMP
+
+wheel_platform = select({
+    "//conditions:default": "manylinux_2_27_x86_64",
+    "//xla/tsl:linux_aarch64": "manylinux_2_27_aarch64",
+})
+
+# Aliases and filegroups don't move files within bazel-out, so we need to use genrules to place them
+# in the correct directory structure for the wheel.
+genrule(
+    name = "init_file_" + cuda_label,
+    srcs = ["__init__.py"],
+    outs = ["xla_plugins/xla_" + cuda_label + "_pjrt/__init__.py"],
+    cmd = "cp $< $@",
+)
+
+genrule(
+    name = "init_file_cpu",
+    srcs = ["__init__.py"],
+    outs = ["xla_plugins/xla_cpu_pjrt/__init__.py"],
+    cmd = "cp $< $@",
+)
+
+# The wheels have the same C API header, but we need to give each their own copy
+genrule(
+    name = "pjrt_c_api_hdr_cuda",
+    srcs = ["//xla/pjrt/c:pjrt_c_api.h"],
+    outs = ["xla_plugins/xla_" + cuda_label + "_pjrt/include/pjrt_c_api.h"],
+    cmd = "cp $< $@",
+)
+
+genrule(
+    name = "pjrt_c_api_hdr_cpu",
+    srcs = ["//xla/pjrt/c:pjrt_c_api.h"],
+    outs = ["xla_plugins/xla_cpu_pjrt/include/pjrt_c_api.h"],
+    cmd = "cp $< $@",
+)
+
+# GPU-specific files
+cc_binary(
+    name = "xla_plugins/xla_" + cuda_label + "_pjrt/xla_gpu_pjrt.so",
+    linkopts = [
+        "-Wl,--version-script,$(location :pjrt_symbols.lds)",
+        "-Wl,--no-undefined",
+    ],
+    linkshared = True,
+    deps = [
+        ":pjrt_symbols.lds",
+        "//xla/pjrt/c:pjrt_c_api_gpu",
+    ],
+)
+
+py_wheel(
+    name = "xla_" + cuda_label + "_pjrt",
+    author = "The OpenXLA Authors",
+    classifiers = ["Development Status :: 1 - Planning"],
+    distribution = "xla_" + cuda_label + "_pjrt",
+    entry_points = {
+        "xla_plugins": [
+            "xla_" + cuda_label + "_pjrt = xla_plugins.xla_" + cuda_label + "_pjrt\n",
+        ],
+    },
+    platform = wheel_platform,
+    python_tag = "py3",
+    strip_path_prefixes = ["build_tools/pjrt_wheels"],
+    summary = "XLA PJRT Plugin",
+    version = wheel_version,
+    deps = [
+        ":init_file_" + cuda_label,
+        ":pjrt_c_api_hdr_cuda",
+        ":xla_plugins/xla_" + cuda_label + "_pjrt/xla_gpu_pjrt.so",
+    ],
+)
+
+alias(
+    name = "xla_cuda_pjrt.dist",
+    actual = ":xla_" + cuda_label + "_pjrt.dist",
+)
+
+# CPU-specific files
+cc_binary(
+    name = "xla_plugins/xla_cpu_pjrt/xla_cpu_pjrt.so",
+    linkopts = [
+        "-Wl,--version-script,$(location :pjrt_symbols.lds)",
+        "-Wl,--no-undefined",
+    ],
+    linkshared = True,
+    deps = [
+        ":pjrt_symbols.lds",
+        "//xla/pjrt/c:pjrt_c_api_cpu",
+    ],
+)
+
+py_wheel(
+    name = "xla_cpu_pjrt",
+    author = "The OpenXLA Authors",
+    classifiers = ["Development Status :: 1 - Planning"],
+    distribution = "xla_cpu_pjrt",
+    entry_points = {
+        "xla_plugins": [
+            "xla_cpu_pjrt = xla_plugins.xla_cpu_pjrt\n",
+        ],
+    },
+    platform = wheel_platform,
+    python_tag = "py3",
+    strip_path_prefixes = ["build_tools/pjrt_wheels"],
+    summary = "XLA PJRT Plugin",
+    version = wheel_version,
+    deps = [
+        ":init_file_cpu",
+        ":pjrt_c_api_hdr_cpu",
+        ":xla_plugins/xla_cpu_pjrt/xla_cpu_pjrt.so",
+    ],
+)
+
+# Tests
+cc_test(
+    name = "cpu_smoke_test",
+    srcs = ["smoke_test.cc"],
+    data = [":xla_plugins/xla_cpu_pjrt/xla_cpu_pjrt.so"],
+    env = {
+        "PJRT_PLUGIN_PATH": "build_tools/pjrt_wheels/xla_plugins/xla_cpu_pjrt/xla_cpu_pjrt.so",
+    },
+    linkopts = ["-ldl"],
+    tags = ["no_windows"],
+    deps = ["//xla/pjrt/c:pjrt_c_api_hdrs"],
+)
+
+cc_test(
+    name = cuda_label + "_smoke_test",
+    srcs = ["smoke_test.cc"],
+    data = [":xla_plugins/xla_" + cuda_label + "_pjrt/xla_gpu_pjrt.so"],
+    env = {
+        "PJRT_PLUGIN_PATH": "build_tools/pjrt_wheels/xla_plugins/xla_" + cuda_label + "_pjrt/xla_gpu_pjrt.so",
+    },
+    linkopts = ["-ldl"],
+    tags = ["no_windows"],
+    deps = ["//xla/pjrt/c:pjrt_c_api_hdrs"],
+)
+
+# The CPU and CUDA test suites are run in CI's build script
+test_suite(
+    name = "cpu_test_suite",
+    tags = ["no_windows"],
+    tests = [
+        ":cpu_smoke_test",
+    ],
+)
+
+test_suite(
+    name = "cuda_test_suite",
+    tags = ["no_windows"],
+    tests = [
+        ":" + cuda_label + "_smoke_test",
+    ],
+)

third_party/xla/build_tools/pjrt_wheels/__init__.py

@@ -0,0 +1 @@
+# This is currently just here to mark the directory as a module.

third_party/xla/build_tools/pjrt_wheels/nightly.bzl

@@ -0,0 +1,35 @@
+"""If we're building a nightly, we use this to pass a timestamp for the wheel version."""
+
+def _nightly_timestamp_impl(rctx):
+    timestamp_val = rctx.getenv("XLA_NIGHTLY_TIMESTAMP", "")  # Default to ""
+
+    # Smuggle the value via a new .bzl file
+    if timestamp_val:
+        rctx.file(
+            "timestamp.bzl",
+            content = 'XLA_NIGHTLY_TIMESTAMP = ".dev{}"'.format(timestamp_val),
+        )
+    else:
+        rctx.file(
+            "timestamp.bzl",
+            content = 'XLA_NIGHTLY_TIMESTAMP = ""',
+        )
+
+    # Create a BUILD file to make timestamp.bzl addressable
+    rctx.file("BUILD.bazel", content = "")
+
+nightly_timestamp_repo = repository_rule(
+    implementation = _nightly_timestamp_impl,
+    environ = ["XLA_NIGHTLY_TIMESTAMP"],
+)
+
+# bzlmod implementation
+def _nightly_timestamp_ext_impl(mctx):  # @unused
+    nightly_timestamp_repo(
+        name = "nightly_timestamp",
+    )
+
+nightly_timestamp_repo_bzlmod = module_extension(
+    implementation = _nightly_timestamp_ext_impl,
+    environ = ["XLA_NIGHTLY_TIMESTAMP"],
+)

third_party/xla/build_tools/pjrt_wheels/pjrt_symbols.lds

@@ -0,0 +1,9 @@
+VERS_1.0 {
+  global:
+    extern "C" {
+      GetPjrtApi;
+    };
+
+  local:
+    *;
+};

third_party/xla/build_tools/pjrt_wheels/release_candidate.bzl

@@ -0,0 +1,35 @@
+"""If we're building a release candidate, we use this to pass a rc number for the wheel version."""
+
+def _rc_number_impl(rctx):
+    rc_number = rctx.getenv("XLA_RC_NUMBER", "")  # Default to ""
+
+    # Smuggle the value via a new .bzl file
+    if rc_number:
+        rctx.file(
+            "rc_number.bzl",
+            content = 'XLA_RC_NUMBER = "{}"'.format(rc_number),
+        )
+    else:
+        rctx.file(
+            "rc_number.bzl",
+            content = 'XLA_RC_NUMBER = ""',
+        )
+
+    # Create a BUILD file to make timestamp.bzl addressable
+    rctx.file("BUILD.bazel", content = "")
+
+rc_number_repo = repository_rule(
+    implementation = _rc_number_impl,
+    environ = ["XLA_RC_NUMBER"],
+)
+
+# bzlmod implementation
+def _rc_number_ext_impl(mctx):  # @unused
+    rc_number_repo(
+        name = "rc_number",
+    )
+
+rc_number_repo_bzlmod = module_extension(
+    implementation = _rc_number_ext_impl,
+    environ = ["XLA_RC_NUMBER"],
+)

third_party/xla/build_tools/pjrt_wheels/smoke_test.cc

@@ -0,0 +1,57 @@
+/* Copyright 2025 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 <dlfcn.h>
+
+#include <iostream>
+
+#include "xla/pjrt/c/pjrt_c_api.h"
+
+typedef const PJRT_Api* (*GetPjrtApi_Func)();
+
+int main() {
+  // 1. Open the shared object
+  const char* so_path = std::getenv("PJRT_PLUGIN_PATH");
+  std::cout << "so_path: " << so_path << std::endl;
+  void* handle = dlopen(so_path, RTLD_LAZY);
+  if (!handle) {
+    std::cerr << "Error: Could not open shared object." << std::endl;
+    std::cerr << "Reason: " << dlerror() << std::endl;
+    return 1;
+  }
+
+  // 2. Load the symbol (the function)
+  GetPjrtApi_Func get_pjrt_api = (GetPjrtApi_Func)dlsym(handle, "GetPjrtApi");
+  const char* dlsym_error = dlerror();
+  if (dlsym_error) {
+    std::cerr << "Error: Could not find symbol 'GetPjrtApi'." << std::endl;
+    std::cerr << "Reason: " << dlsym_error << std::endl;
+    dlclose(handle);
+    return 1;
+  }
+
+  // 3. Call the function
+  std::cout << "Successfully loaded symbol. Calling GetPjrtApi()..."
+            << std::endl;
+  const PJRT_Api* api = get_pjrt_api();
+  if (api) {
+    std::cout << "Success! Received PjrtApi struct pointer." << std::endl;
+  } else {
+    std::cerr << "Error: GetPjrtApi() returned a null pointer." << std::endl;
+    return 1;
+  }
+
+  return 0;
+}