Provide Target machine architecture for cross-compile scenarios.

CUDA/CUDNN/NCCL repositories are created on a host machine, so if we need to download redistributions for other architectures in cross-compile scenario, we need to pass the target architecture name to repo rules, e.g. `--repo_env=TARGET_PLATFORM_ARCHITECTURE=aarch64`

PiperOrigin-RevId: 719020277

third_party/gpus/cuda/hermetic/cuda_redist_init_repositories.bzl

@@ -297,7 +297,21 @@ 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
+    target_arch = get_env_var(repository_ctx, "TARGET_PLATFORM_ARCHITECTURE")
+
+    # We use NVCC compiler as the host compiler.
+    if target_arch and repository_ctx.name != "cuda_nvcc":
+        if target_arch in OS_ARCH_DICT.keys():
+            host_arch = target_arch
+        else:
+            fail(
+                "Unsupported architecture: {arch}, use one of {supported}".format(
+                    arch = target_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 +393,7 @@ cuda_repo = repository_rule(
         "HERMETIC_CUDA_VERSION",
         "TF_CUDA_VERSION",
         "LOCAL_CUDA_PATH",
+        "TARGET_PLATFORM_ARCHITECTURE",
     ],
 )
 
@@ -464,6 +479,7 @@ cudnn_repo = repository_rule(
         "HERMETIC_CUDA_VERSION",
         "TF_CUDA_VERSION",
         "LOCAL_CUDNN_PATH",
+        "TARGET_PLATFORM_ARCHITECTURE",
     ],
 )
 

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]
+    target_arch = get_env_var(repository_ctx, "TARGET_PLATFORM_ARCHITECTURE")
+    if target_arch:
+        if target_arch in OS_ARCH_DICT.keys():
+            arch = OS_ARCH_DICT[target_arch]
+        else:
+            fail(
+                "Unsupported architecture: {arch}, use one of {supported}".format(
+                    arch = target_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",
+        "TARGET_PLATFORM_ARCHITECTURE",
+    ],
 )
 
 def nccl_redist_init_repository(