Add ability to pass custom architecture name to CUDA/CUDNN/NCCL repo …

…rules.

This is useful for cross-compile scenarios. Repositories are created on a host machine, so if we need to download redistributions for other architectures, we need to pass the architecture names to repo rules.

PiperOrigin-RevId: 719020277

third_party/tsl/third_party/gpus/cuda/hermetic/cuda_redist_init_repositories.bzl

@@ -297,7 +297,19 @@ def _download_redistribution(repository_ctx, arch_key, path_prefix):
     repository_ctx.delete(file_name)
 
 def _get_platform_architecture(repository_ctx):
-    host_arch = repository_ctx.os.arch
+    custom_arch = get_env_var(repository_ctx, "CUSTOM_PLATFORM_ARCHITECTURE")
+    if custom_arch:
+        if custom_arch in OS_ARCH_DICT.keys():
+            host_arch = custom_arch
+        else:
+            fail(
+                "Unsupported architecture: {arch}, use one of {supported}".format(
+                    arch = custom_arch,
+                    supported = OS_ARCH_DICT.keys(),
+                ),
+            )
+    else:
+        host_arch = repository_ctx.os.arch
 
     if host_arch == "aarch64":
         uname_result = repository_ctx.execute(["uname", "-a"]).stdout
@@ -379,6 +391,7 @@ cuda_repo = repository_rule(
         "HERMETIC_CUDA_VERSION",
         "TF_CUDA_VERSION",
         "LOCAL_CUDA_PATH",
+        "CUSTOM_PLATFORM_ARCHITECTURE",
     ],
 )
 
@@ -464,6 +477,7 @@ cudnn_repo = repository_rule(
         "HERMETIC_CUDA_VERSION",
         "TF_CUDA_VERSION",
         "LOCAL_CUDNN_PATH",
+        "CUSTOM_PLATFORM_ARCHITECTURE",
     ],
 )
 

third_party/tsl/third_party/nccl/hermetic/nccl_redist_init_repository.bzl

@@ -46,7 +46,19 @@ def _use_downloaded_nccl_wheel(repository_ctx):
         return
 
     # Download archive only when GPU config is used.
-    arch = OS_ARCH_DICT[repository_ctx.os.arch]
+    custom_arch = get_env_var(repository_ctx, "CUSTOM_PLATFORM_ARCHITECTURE")
+    if custom_arch:
+        if custom_arch in OS_ARCH_DICT.keys():
+            arch = OS_ARCH_DICT[custom_arch]
+        else:
+            fail(
+                "Unsupported architecture: {arch}, use one of {supported}".format(
+                    arch = custom_arch,
+                    supported = OS_ARCH_DICT.keys(),
+                ),
+            )
+    else:
+        arch = OS_ARCH_DICT[repository_ctx.os.arch]
     dict_key = "{cuda_version}-{arch}".format(
         cuda_version = cuda_version,
         arch = arch,
@@ -116,7 +128,12 @@ cuda_nccl_repo = repository_rule(
         "build_templates": attr.label_list(mandatory = True),
         "strip_prefix": attr.string(),
     },
-    environ = ["HERMETIC_CUDA_VERSION", "TF_CUDA_VERSION", "LOCAL_NCCL_PATH"],
+    environ = [
+        "HERMETIC_CUDA_VERSION",
+        "TF_CUDA_VERSION",
+        "LOCAL_NCCL_PATH",
+        "CUSTOM_PLATFORM_ARCHITECTURE",
+    ],
 )
 
 def nccl_redist_init_repository(