# Copyright 2026 The JAX 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
#
#     https://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.

# INTEL ONEAPI plugin extension and GPU kernels
load("@local_config_sycl//sycl:build_defs.bzl", "sycl_library")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@rules_python//python:defs.bzl", "py_library")
load(
    "//jaxlib:jax.bzl",
    "if_oneapi_is_configured",
)
load("//jaxlib/oneapi:oneapi_rpath.bzl", "oneapi_nanobind_extension")

licenses(["notice"])

package(
    default_applicable_licenses = [],
    default_visibility = ["//:__subpackages__"],
)

cc_library(
    name = "oneapi_gpu_runtime_hdrs",
    hdrs = ["oneapi_gpu_runtime.h"],
    deps = [
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@local_config_sycl//sycl:sycl_headers",
    ],
)

cc_library(
    name = "oneapi_gpu_runtime",
    srcs = ["oneapi_gpu_runtime.cc"],
    hdrs = ["oneapi_gpu_runtime.h"],
    copts = ["-fsycl"],
    features = ["-use_header_modules"],
    linkopts = ["-fsycl"],
    deps = [
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@local_config_sycl//sycl:sycl_headers",
    ],
)

cc_library(
    name = "oneapi_vendor",
    defines = ["JAX_GPU_ONEAPI=1"],
    textual_hdrs = [
        "//jaxlib/gpu:vendor.h",
    ],
    deps = [
        ":oneapi_gpu_runtime_hdrs",
    ],
)

cc_library(
    name = "oneapi_gpu_kernel_helpers",
    hdrs = ["//jaxlib/gpu:gpu_kernel_helpers.h"],
    features = ["-use_header_modules"],
    deps = [
        ":oneapi_vendor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:str_format",
    ],
)

sycl_library(
    name = "oneapisolver_kernels_ffi",
    srcs = ["oneapi_solver_kernels_ffi.cc"],
    hdrs = ["//jaxlib/gpu:solver_kernels_ffi.h"],
    features = ["-use_header_modules"],
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_gpu_runtime",
        ":oneapi_vendor",
        "//jaxlib:ffi_helpers",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@xla//xla/ffi/api:ffi",
    ],
)

sycl_library(
    name = "oneapi_prng_kernels_impl",
    srcs = ["//jaxlib/gpu:prng_kernels_oneapi.cc"],
    hdrs = ["//jaxlib/gpu:prng_kernels.h"],
    features = ["-use_header_modules"],
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_vendor",
        "@xla//xla/ffi/api:ffi",
    ],
)

sycl_library(
    name = "oneapi_prng_kernels",
    srcs = ["//jaxlib/gpu:prng_kernels.cc"],
    hdrs = ["//jaxlib/gpu:prng_kernels.h"],
    features = ["-use_header_modules"],
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_gpu_runtime",
        ":oneapi_prng_kernels_impl",
        ":oneapi_vendor",
        "//jaxlib:ffi_helpers",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/status",
        "@xla//xla/ffi/api:c_api",
        "@xla//xla/ffi/api:ffi",
    ],
)

oneapi_nanobind_extension(
    name = "_solver",
    srcs = ["//jaxlib/gpu:solver.cc"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    module_name = "_solver",
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_vendor",
        ":oneapisolver_kernels_ffi",
        "//jaxlib:kernel_nanobind_helpers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@nanobind",
        "@oneapi//:headers",
        "@xla//xla/tsl/python/lib/core:numpy",
    ],
)

oneapi_nanobind_extension(
    name = "_prng",
    srcs = ["//jaxlib/gpu:prng.cc"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    linkopts = ["-link_stage"],
    module_name = "_prng",
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_prng_kernels",
        ":oneapi_vendor",
        "//jaxlib:kernel_nanobind_helpers",
        "@nanobind",
    ],
)

cc_library(
    name = "oneapi_hybrid_kernels",
    srcs = ["//jaxlib/gpu:hybrid_kernels.cc"],
    hdrs = ["//jaxlib/gpu:hybrid_kernels.h"],
    features = ["-use_header_modules"],
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_gpu_runtime",
        ":oneapi_vendor",
        "//jaxlib:ffi_helpers",
        "//jaxlib/cpu:lapack_kernels",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@xla//xla/ffi/api:ffi",
    ],
)

oneapi_nanobind_extension(
    name = "_hybrid",
    srcs = ["//jaxlib/gpu:hybrid.cc"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    linkopts = ["-link_stage"],
    module_name = "_hybrid",
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_hybrid_kernels",
        ":oneapi_vendor",
        "//jaxlib:kernel_nanobind_helpers",
        "//jaxlib/cpu:lapack_kernels",
        "@com_google_absl//absl/base",
        "@nanobind",
        "@oneapi//:headers",
        "@xla//xla/ffi/api:ffi",
        "@xla//xla/python:safe_static_init",
    ],
)

sycl_library(
    name = "oneapi_linalg_kernels_impl",
    srcs = ["//jaxlib/gpu:linalg_kernels_oneapi.cc"],
    hdrs = ["//jaxlib/gpu:linalg_kernels.h"],
    copts = [
        "-U_FORTIFY_SOURCE",
        "-D_FORTIFY_SOURCE=0",
    ],
    features = ["-use_header_modules"],
    local_defines = [
        "EIGEN_USE_SYCL=1",
        "EIGEN_DONT_VECTORIZE_SYCL=1",
    ],
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_vendor",
        "//jaxlib:tridiagonal_solve_perturbed",
        "@xla//xla/ffi/api:ffi",
    ],
)

sycl_library(
    name = "oneapi_linalg_kernels",
    srcs = ["//jaxlib/gpu:linalg_kernels.cc"],
    hdrs = ["//jaxlib/gpu:linalg_kernels.h"],
    features = ["-use_header_modules"],
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_gpu_runtime",
        ":oneapi_linalg_kernels_impl",
        ":oneapi_vendor",
        "//jaxlib:ffi_helpers",
        "@com_google_absl//absl/strings:str_format",
        "@xla//xla/ffi/api:ffi",
    ],
)

oneapi_nanobind_extension(
    name = "_linalg",
    srcs = ["//jaxlib/gpu:linalg.cc"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    module_name = "_linalg",
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_linalg_kernels",
        ":oneapi_vendor",
        "//jaxlib:kernel_nanobind_helpers",
        "@nanobind",
    ],
)

py_library(
    name = "oneapi_gpu_support",
    deps = [
        ":_hybrid",
        ":_linalg",
        ":_prng",
        ":_solver",
    ],
)

sycl_library(
    name = "py_client_gpu",
    srcs = ["//jaxlib/gpu:py_client_gpu.cc"],
    hdrs = ["//jaxlib/gpu:py_client_gpu.h"],
    copts = [
        "-fexceptions",
        "-fno-strict-aliasing",
    ],
    features = ["-use_header_modules"],
    deps = [
        ":oneapi_gpu_kernel_helpers",
        ":oneapi_gpu_runtime",
        ":oneapi_vendor",
        "//jaxlib:ffi",
        "//jaxlib:free_threading",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@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:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@dlpack",
        "@nanobind",
        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
        "@xla//xla:comparison_util",
        "@xla//xla:shape_util",
        "@xla//xla:xla_data_proto_cc",
        "@xla//xla/ffi:ffi_api",
        "@xla//xla/ffi/api:ffi",
        "@xla//xla/pjrt:host_callback",
        "@xla//xla/pjrt:transpose",
        "@xla//xla/python:nb_numpy",
        "@xla//xla/python:types",
        "@xla//xla/service:platform_util",
    ],
)

oneapi_nanobind_extension(
    name = "oneapi_plugin_extension",
    srcs = ["oneapi_plugin_extension.cc"],
    module_name = "oneapi_plugin_extension",
    deps = [
        ":py_client_gpu",
        "//jaxlib:kernel_nanobind_helpers",
        "//jaxlib/gpu:gpu_plugin_extension",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/strings",
        "@nanobind",
        "@oneapi//:headers",
    ],
)

py_library(
    name = "gpu_only_test_deps",
    # `if_oneapi_is_configured` will default to `[]`.
    deps = if_oneapi_is_configured([
        ":oneapi_gpu_support",
        ":oneapi_plugin_extension",
    ]),
)