forked from google-ml-infra/rules_ml_toolchain
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathMODULE.bazel
More file actions
77 lines (63 loc) · 2.47 KB
/
MODULE.bazel
File metadata and controls
77 lines (63 loc) · 2.47 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
module(name = "rules_ml_toolchain")
bazel_dep(name = "bazel_skylib", version = "1.7.1")
bazel_dep(name = "protobuf", version = "32.1", repo_name = "com_google_protobuf")
bazel_dep(name = "rules_proto", version = "6.0.0-rc1")
bazel_dep(name = "pybind11_bazel", version = "2.13.6")
bazel_dep(name = "rules_cc", version = "0.0.9")
bazel_dep(name = "abseil-cpp", version = "20240116.2", repo_name = "com_google_absl")
bazel_dep(name = "googletest", version = "1.14.0.bcr.1", repo_name = "gtest")
bazel_dep(name = "platforms", version = "0.0.11")
pybind11_internal_configure = use_extension(
"@pybind11_bazel//:internal_configure.bzl",
"internal_configure_extension",
)
use_repo(pybind11_internal_configure, "pybind11")
##############################################################
# CUDA
cuda_configure_ext = use_extension("//third_party/extensions:cuda_configure.bzl", "cuda_configure_ext")
use_repo(cuda_configure_ext, "local_config_cuda")
cuda_json_ext = use_extension("//third_party/extensions:cuda_json.bzl", "cuda_json_ext")
use_repo(cuda_json_ext, "cuda_redist_json")
cuda_redist_init_ext = use_extension("//third_party/extensions:cuda_redist_init.bzl", "cuda_redist_init_ext")
use_repo(
cuda_redist_init_ext,
"cuda_cccl",
"cuda_crt",
"cuda_cublas",
"cuda_cudart",
"cuda_cudnn",
"cuda_cufft",
"cuda_cupti",
"cuda_curand",
"cuda_cusolver",
"cuda_cusparse",
"cuda_nvcc",
"cuda_nvdisasm",
"cuda_nvjitlink",
"cuda_nvml",
"cuda_nvtx",
"cuda_nvvm",
)
##############################################################
# NCCL configuration
nccl_redist = use_extension("//third_party/extensions:nccl_redist.bzl", "nccl_redist_ext")
use_repo(nccl_redist, "cuda_nccl")
nccl_configure = use_extension("//third_party/extensions:nccl_configure.bzl", "nccl_configure_ext")
use_repo(nccl_configure, "local_config_nccl")
##############################################################
# SYCL configuration
sycl_configure = use_extension("//third_party/extensions:sycl_configure.bzl", "sycl_configure_ext")
use_repo(sycl_configure, "local_config_sycl")
##############################################################
# Hermetic toolchain configuration
toolchain = use_extension("//third_party/extensions:toolchain.bzl", "toolchain_ext")
use_repo(
toolchain,
"llvm_darwin_aarch64",
"llvm_linux_aarch64",
"llvm_linux_x86_64",
"sysroot_darwin_aarch64",
"sysroot_linux_aarch64",
"sysroot_linux_x86_64",
)
register_toolchains("//cc/...")