Skip to content

Commit e5c2b05

Browse files
ybaturinatf-text-github-robot
authored andcommitted
Remove hermetic CUDA/NCCL/NVSHMEM rules from XLA repo.
These files are provided by `rules_ml_toolchain` repostory. PiperOrigin-RevId: 855290862
1 parent 7215af0 commit e5c2b05

File tree

1 file changed

+25
-5
lines changed

1 file changed

+25
-5
lines changed

WORKSPACE

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,28 @@ load(
180180

181181
python_wheel_version_suffix_repository(name = "tf_wheel_version_suffix")
182182

183+
# Toolchains for ML projects hermetic builds.
184+
# Details: https://github.com/google-ml-infra/rules_ml_toolchain
185+
http_archive(
186+
name = "rules_ml_toolchain",
187+
sha256 = "9dbee8f24cc1b430bf9c2a6661ab70cbca89979322ddc7742305a05ff637ab6b",
188+
strip_prefix = "rules_ml_toolchain-545c80f1026d526ea9c7aaa410bf0b52c9a82e74",
189+
urls = [
190+
"https://github.com/google-ml-infra/rules_ml_toolchain/archive/545c80f1026d526ea9c7aaa410bf0b52c9a82e74.tar.gz",
191+
],
192+
)
193+
194+
load(
195+
"@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl",
196+
"cc_toolchain_deps",
197+
)
198+
199+
cc_toolchain_deps()
200+
201+
register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")
202+
183203
load(
184-
"@local_xla//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
204+
"@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
185205
"cuda_json_init_repository",
186206
)
187207

@@ -193,7 +213,7 @@ load(
193213
"CUDNN_REDISTRIBUTIONS",
194214
)
195215
load(
196-
"@local_xla//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
216+
"@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
197217
"cuda_redist_init_repositories",
198218
"cudnn_redist_init_repository",
199219
)
@@ -207,21 +227,21 @@ cudnn_redist_init_repository(
207227
)
208228

209229
load(
210-
"@local_xla//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
230+
"@rules_ml_toolchain//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
211231
"cuda_configure",
212232
)
213233

214234
cuda_configure(name = "local_config_cuda")
215235

216236
load(
217-
"@local_xla//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
237+
"@rules_ml_toolchain//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
218238
"nccl_redist_init_repository",
219239
)
220240

221241
nccl_redist_init_repository()
222242

223243
load(
224-
"@local_xla//third_party/nccl/hermetic:nccl_configure.bzl",
244+
"@rules_ml_toolchain//third_party/nccl/hermetic:nccl_configure.bzl",
225245
"nccl_configure",
226246
)
227247

0 commit comments

Comments
 (0)