From 324cf2c6f968296da7222fc99971f6325e3b7a42 Mon Sep 17 00:00:00 2001
From: Han Qi <hanq@google.com>
Date: Fri, 27 Sep 2024 21:08:03 +0000
Subject: [PATCH] hermetic

---
 .bazelrc  |  5 +++++
 WORKSPACE | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/.bazelrc b/.bazelrc
index 694cf3fd125..92fa1a0f760 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -53,6 +53,11 @@ build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
 build:cuda --@local_config_cuda//:enable_cuda
 build:cuda --define=xla_python_enable_gpu=true
 build:cuda --cxxopt=-DXLA_CUDA=1
+build:cuda --repo_env=HERMETIC_CUDA_VERSION="12.3.2"
+build:cuda --repo_env=HERMETIC_CUDNN_VERSION="9.1.1"
+build:cuda --repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_50,sm_60,sm_70,sm_80,compute_90"
+build:cuda --@local_config_cuda//cuda:include_cuda_libs=true
+test:cuda --@cuda_driver//:enable_forward_compatibility=true
 
 # Coverage with cuda/gcc/nvcc requires manually setting coverage flags.
 coverage:cuda --per_file_copt=third_party/.*,torch_xla/.*@--coverage
diff --git a/WORKSPACE b/WORKSPACE
index 423d1f8cd08..50e3ed01fbb 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -137,7 +137,56 @@ load("@xla//:workspace0.bzl", "xla_workspace0")
 
 xla_workspace0()
 
-load("@tsl//third_party/gpus:cuda_configure.bzl", "cuda_configure")
+load(
+    "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
+    "cuda_json_init_repository",
+)
+
+cuda_json_init_repository()
+
+load(
+   "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
+   "cuda_json_init_repository",
+)
+
+cuda_json_init_repository()
+
+load(
+   "@cuda_redist_json//:distributions.bzl",
+   "CUDA_REDISTRIBUTIONS",
+   "CUDNN_REDISTRIBUTIONS",
+)
+load(
+   "@tsl//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl",
+   "cuda_redist_init_repositories",
+   "cudnn_redist_init_repository",
+)
+
+cuda_redist_init_repositories(
+   cuda_redistributions = CUDA_REDISTRIBUTIONS,
+)
+
+cudnn_redist_init_repository(
+   cudnn_redistributions = CUDNN_REDISTRIBUTIONS,
+)
+
+load(
+   "@tsl//third_party/gpus/cuda/hermetic:cuda_configure.bzl",
+   "cuda_configure",
+)
+
 cuda_configure(name = "local_config_cuda")
-load("@tsl//third_party/nccl:nccl_configure.bzl", "nccl_configure")
+
+load(
+   "@tsl//third_party/nccl/hermetic:nccl_redist_init_repository.bzl",
+   "nccl_redist_init_repository",
+)
+
+nccl_redist_init_repository()
+
+load(
+   "@tsl//third_party/nccl/hermetic:nccl_configure.bzl",
+   "nccl_configure",
+)
+
 nccl_configure(name = "local_config_nccl")