load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
load("@rules_zig//zig:defs.bzl", "zig_library")

PLATFORMS = {
    "cpu": True,
    "cuda": False,
    "rocm": False,
    "tpu": False,
    "neuron": False,
    "oneapi": False,
    "metal": False,
}

[
    bool_flag(
        name = runtime,
        build_setting_default = default,
    )
    for runtime, default in PLATFORMS.items()
]

config_setting(
    name = "rocm.enabled",
    constraint_values = [
        "@platforms//cpu:x86_64",
        "@platforms//os:linux",
    ],
    flag_values = {":rocm": "True"},
    visibility = ["//visibility:public"],
)

config_setting(
    name = "metal.enabled",
    constraint_values = [
        "@platforms//cpu:aarch64",
        "@platforms//os:macos",
    ],
    flag_values = {":metal": "True"},
    visibility = ["//visibility:public"],
)

[
    config_setting(
        name = "{}.enabled".format(runtime),
        flag_values = {":{}".format(runtime): "True"},
        visibility = ["//visibility:public"],
    )
    for runtime in [
        "cpu",
        "cuda",
        "tpu",
        "neuron",
        "oneapi",
    ]
]

zig_library(
    name = "platforms",
    main = "platforms.zig",
    tags = ["manual"],
    visibility = ["//visibility:public"],
    deps = [
        "//pjrt",
    ] + [
        "//platforms/{}".format(runtime)
        for runtime in PLATFORMS.keys()
    ],
)

platform(
    name = "linux_amd64",
    constraint_values = [
        "@platforms//cpu:x86_64",
        "@platforms//os:linux",
        "@llvm//constraints/libc:gnu.2.28",
    ],
    flags = [
        "--force_pic=true",  # required for protobuf upb on Linux
        "--cpu=x86_64",
    ],
    visibility = ["//visibility:public"],
)

platform(
    name = "linux_arm64",
    constraint_values = [
        "@platforms//cpu:aarch64",
        "@platforms//os:linux",
        "@llvm//constraints/libc:gnu.2.28",
    ],
    flags = [
        "--force_pic=true",  # required for protobuf upb on Linux
        "--cpu=aarch64",
    ],
    visibility = ["//visibility:public"],
)

platform(
    name = "macos_amd64",
    constraint_values = [
        "@platforms//cpu:x86_64",
        "@platforms//os:macos",
    ],
    visibility = ["//visibility:public"],
)

platform(
    name = "macos_arm64",
    constraint_values = [
        "@platforms//cpu:aarch64",
        "@platforms//os:macos",
    ],
    visibility = ["//visibility:public"],
)

constraint_setting(name = "execution_location")

constraint_value(
    name = "local_execution",
    constraint_setting = ":execution_location",
    visibility = ["//visibility:public"],
)

# Special constrained host platform that toolchain can use as exec_compatible_with
# if they want to only be used when executing on the local machine.
#
# This is currently used by Zig toolchains so that they are only resolved for
# the host platform and thus can be used for no-remote-exec actions.
platform(
    name = "local_execution_platform",
    constraint_values = [":local_execution"],
    exec_properties = {
        "no-remote-exec": "1",
    },
    parents = ["@platforms//host"],
    visibility = ["//visibility:public"],
)