Fix CUDA build

MODULE.bazel

@@ -61,6 +61,9 @@ use_repo(
     "cuda_profiler_api",
 )
 
+nvshmem_redist = use_extension("//extensions:nvshmem_redist.bzl", "nvshmem_redist_ext")
+use_repo(nvshmem_redist, "nvidia_nvshmem")
+
 ##############################################################
 # NCCL configuration
 

extensions/nvshmem_redist.bzl

@@ -0,0 +1,19 @@
+"""Module extension for nvshmem_redist_init_repository."""
+
+load(
+    "//gpu/nvshmem:nvshmem_redist_init_repository.bzl",
+    "nvshmem_redist_init_repository",
+)
+load(
+    "@cuda_redist_json//:distributions.bzl",
+    "CUDA_REDISTRIBUTIONS",
+)
+
+def _nvshmem_redist_ext_impl(mctx):
+    nvshmem_redist_init_repository(
+        nvshmem_redistributions = CUDA_REDISTRIBUTIONS,
+    )
+
+nvshmem_redist_ext = module_extension(
+    implementation = _nvshmem_redist_ext_impl,
+)

gpu/cuda/build_templates/cuda_cccl.BUILD.tpl

@@ -12,7 +12,7 @@ filegroup(
         %{comment}"include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", "") + "/cuda/**",
         %{comment}"include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", "") + "/nv/**",
         %{comment}"include" + cuda_lib_header_prefix(_cudart_version, 13, "/cccl", "") + "/thrust/**",
-    ]),
+    ], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/cuda/build_templates/cuda_cccl_github.BUILD.tpl

@@ -16,7 +16,7 @@ filegroup(
     name = "thrust_header_list",
     srcs = glob([
         %{comment}"thrust/thrust/**",
-    ]),
+    ], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 
@@ -33,7 +33,7 @@ filegroup(
     name = "cuda_header_list",
     srcs = glob([
         %{comment}"libcudacxx/include/cuda/**",
-    ]),
+    ], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 
@@ -50,7 +50,7 @@ filegroup(
     name = "nv_header_list",
     srcs = glob([
         %{comment}"libcudacxx/include/nv/**",
-    ]),
+    ], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 
@@ -67,7 +67,7 @@ filegroup(
     name = "cub_header_list",
     srcs = glob([
         %{comment}"cub/cub/**",
-    ]),
+    ], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/cuda/build_templates/cuda_crt.BUILD.tpl

@@ -7,7 +7,7 @@ filegroup(
     name = "header_list",
     %{comment}srcs = if_cuda_newer_than(
         %{comment}"13_0",
-        %{comment}if_true = glob(["include/crt/**"]),
+        %{comment}if_true = glob(["include/crt/**"], allow_empty = True),
         %{comment}if_false = [],
     %{comment}),
     visibility = ["@local_config_cuda//cuda:__pkg__"],

gpu/cuda/build_templates/cuda_cudart.BUILD.tpl

@@ -139,7 +139,7 @@ filegroup(
         %{comment}"include/vector_functions.h",
         %{comment}"include/vector_functions.hpp",
         %{comment}"include/vector_types.h",
-    %{comment}]),
+    %{comment}], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/cuda/build_templates/cuda_cudnn.BUILD.tpl

@@ -75,7 +75,7 @@ filegroup(
     name = "header_list",
     %{comment}srcs = glob([
         %{comment}"include/cudnn*.h",
-    %{comment}]),
+    %{comment}], allow_empty = True),
 )
 
 

gpu/cuda/build_templates/cuda_cudnn8.BUILD.tpl

@@ -67,7 +67,7 @@ filegroup(
     name = "header_list",
     %{comment}srcs = glob([
         %{comment}"include/cudnn*.h",
-    %{comment}]),
+    %{comment}], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/cuda/build_templates/cuda_cufft.BUILD.tpl

@@ -31,7 +31,7 @@ filegroup(
     %{comment}srcs = glob([
         %{comment}"include/cudalibxt.h",
         %{comment}"include/cufft*.h"
-    %{comment}]),
+    %{comment}], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/cuda/build_templates/cuda_curand.BUILD.tpl

@@ -28,7 +28,7 @@ cc_library(
 
 filegroup(
     name = "header_list",
-    %{comment}srcs = glob(["include/curand*.h"]),
+    %{comment}srcs = glob(["include/curand*.h"], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/cuda/build_templates/cuda_cusolver.BUILD.tpl

@@ -37,7 +37,7 @@ filegroup(
     name = "header_list",
     %{comment}srcs = glob([
         %{comment}"include/cusolver*.h",
-    %{comment}]),
+    %{comment}], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/cuda/build_templates/cuda_nvcc.BUILD.tpl

@@ -61,7 +61,7 @@ filegroup(
     srcs = glob([
         "bin/**",
         "nvvm/bin/**",
-    ]),
+    ], allow_empty = True),
     visibility = ["//visibility:public"],
 )
 
@@ -89,10 +89,10 @@ filegroup(
     %{comment}srcs = glob([
         %{comment}"include/fatbinary_section.h",
         %{comment}"include/nvPTXCompiler.h",
-    %{comment}]) + if_cuda_newer_than(
+    %{comment}], allow_empty = True) + if_cuda_newer_than(
         %{comment}"13_0",
         %{comment}if_true = [],
-        %{comment}if_false = glob(["include/crt/**"]),
+        %{comment}if_false = glob(["include/crt/**"], allow_empty = True),
     %{comment}),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )

gpu/cuda/build_templates/cuda_nvtx.BUILD.tpl

@@ -5,7 +5,7 @@ filegroup(
     %{comment}srcs = glob([
         %{comment}"include/nvToolsExt*.h",
         %{comment}"include/nvtx3/**",
-    %{comment}]),
+    %{comment}], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/cuda/build_templates/cuda_profiler.BUILD.tpl

@@ -2,7 +2,7 @@ licenses(["restricted"])  # NVIDIA proprietary license
 
 filegroup(
     name = "header_list",
-    %{comment}srcs = glob(["include/**"]),
+    %{comment}srcs = glob(["include/**"], allow_empty = True),
     visibility = ["@local_config_cuda//cuda:__pkg__"],
 )
 

gpu/nvidia_common_rules.bzl

@@ -66,8 +66,9 @@ def _get_orig_repo_name(repository_ctx):
         return repository_ctx.original_name
 
     # With Bzlmod, the repo name will be something like `_main~cuda_redist_init_ext~cuda_nvml`,
+    # or `rules_ml_toolchain++cuda_redist_init_ext+cuda_nvml` in Bazel 8.
     # we need to extract the original repo name.
-    return repository_ctx.name.split("~")[-1]
+    return repository_ctx.name.split("~")[-1].split("+")[-1]
 
 def get_archive_name(url):
     # buildifier: disable=function-docstring-return