# Copyright 2026 The OpenXLA 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.
# ==============================================================================

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",
    ],
)