@@ -180,8 +180,28 @@ load(
180180
181181python_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+
183203load (
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)
195215load (
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
209229load (
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
214234cuda_configure (name = "local_config_cuda" )
215235
216236load (
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
221241nccl_redist_init_repository ()
222242
223243load (
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