Harden multi-version CUDA toolchain selection on top of rebased #422

- Make redist component repo names unique per redist tag to avoid collisions
  when multiple `cuda.redist_json` entries resolve to the same version
- Deduplicate generated version selectors and namespace version config_settings
  per component to prevent `config_setting` conflicts
- Add platform-aware redist selection in repo rules and align deliverable tool
  labels/aliases (`nvcc`, `nvlink`, `bin2c`, `fatbinary`, `ptxas`, `cicc`)
- Improve clang CUDA path handling for deliverable layouts to avoid invalid
  `--cuda-path` resolution
- Extend dummy/tool alias coverage and templates for missing components/files
- Add multi-version integration test coverage and skip linux-redist tests on
  Windows runners; guard CI redist override to NVIDIA-source cases
- Fix linting errors

Author: finn-ball

.github/workflows/build-tests.yaml

@@ -98,9 +98,9 @@ jobs:
 
       # Override cuda redist_json rule version for CI
       - run: echo "CUDA_REDIST_VERSION_OVERRIDE=${{ matrix.cases.cuda-version }}" >> $GITHUB_ENV
-        if: ${{ !startsWith(matrix.cases.os, 'windows') }}
+        if: ${{ !startsWith(matrix.cases.os, 'windows') && matrix.cases.source == 'nvidia' }}
       - run: echo "CUDA_REDIST_VERSION_OVERRIDE=${{ matrix.cases.cuda-version }}" >> $env:GITHUB_ENV
-        if: ${{ startsWith(matrix.cases.os, 'windows') }}
+        if: ${{ startsWith(matrix.cases.os, 'windows') && matrix.cases.source == 'nvidia' }}
 
       # Use Bazel with version specified in .bazelversion
       - run: echo "USE_BAZEL_VERSION=$(cat .bazelversion)" >> $GITHUB_ENV

README.md

@@ -32,6 +32,46 @@ use_repo(cuda, "cuda")
 
 `rules_cc` provides the C++ toolchain dependency for `rules_cuda`; in Bzlmod, the compatibility repository is handled by `rules_cc` itself.
 
+#### Multi-version hermetic toolchains (Bzlmod)
+
+To select CUDA versions and platforms at build time, define multiple `cuda.redist_json` entries and a single `cuda.toolkit`.
+Then use `@rules_cuda//cuda:version`, `@rules_cuda//cuda:exec_platform`, and `@rules_cuda//cuda:target_platform` flags.
+
+```starlark
+cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
+
+cuda.redist_json(
+    name = "cuda_13_0_2",
+    version = "13.0.2",
+    platforms = [
+        "linux-x86_64",
+        "linux-sbsa",
+    ],
+)
+cuda.redist_json(
+    name = "cuda_13_0_0",
+    version = "13.0.0",
+    platforms = [
+        "linux-x86_64",
+        "linux-sbsa",
+    ],
+)
+
+cuda.toolkit(name = "cuda")
+use_repo(cuda, "cuda")
+```
+
+Example `.bazelrc` entries:
+
+```
+build --@rules_cuda//cuda:exec_platform=linux-x86_64
+build --@rules_cuda//cuda:target_platform=linux-x86_64
+build --@rules_cuda//cuda:version=13.0.0
+```
+
+Note: In Bzlmod, `platforms` is required for `cuda.redist_json` because module extensions don't have access to host OS/arch
+information, so the platforms must be declared explicitly.
+
 <details>
 <summary>Traditional WORKSPACE approach</summary>
 

cuda/BUILD.bazel

@@ -54,7 +54,8 @@ string_flag(
     config_setting(
         name = "target_platform_is_{}".format(platform.replace("-", "_")),
         flag_values = {":target_platform": platform},
-    ) for platform in SUPPORTED_PLATFORMS
+    )
+    for platform in SUPPORTED_PLATFORMS
 ]
 
 string_flag(
@@ -67,7 +68,8 @@ string_flag(
     config_setting(
         name = "exec_platform_is_{}".format(platform.replace("-", "_")),
         flag_values = {":exec_platform": platform},
-    ) for platform in SUPPORTED_PLATFORMS
+    )
+    for platform in SUPPORTED_PLATFORMS
 ]
 
 # Command line flag to specify the list of CUDA architectures to compile for.

cuda/defs.bzl

@@ -3,7 +3,7 @@ Core rules for building CUDA projects.
 """
 
 load("//cuda/private:defs.bzl", _requires_cuda = "requires_cuda")
-load("//cuda/private:errors.bzl", _unsupported_cuda_version = "unsupported_cuda_version", _unsupported_cuda_platform = "unsupported_cuda_platform")
+load("//cuda/private:errors.bzl", _unsupported_cuda_platform = "unsupported_cuda_platform", _unsupported_cuda_version = "unsupported_cuda_version")
 load("//cuda/private:macros/cuda_binary.bzl", _cuda_binary = "cuda_binary")
 load("//cuda/private:macros/cuda_test.bzl", _cuda_test = "cuda_test")
 load("//cuda/private:os_helpers.bzl", _cc_import_versioned_sos = "cc_import_versioned_sos", _if_linux = "if_linux", _if_windows = "if_windows")

cuda/dummy/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@rules_cc//cc:defs.bzl", "cc_binary")
+load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -20,7 +20,10 @@ cc_binary(
     defines = ["TOOLNAME=nvlink"],
 )
 
-exports_files(["link.stub", "libdevice.10.bc"])
+exports_files([
+    "link.stub",
+    "libdevice.10.bc",
+])
 
 cc_binary(
     name = "bin2c",
@@ -34,15 +37,15 @@ cc_binary(
     defines = ["TOOLNAME=fatbinary"],
 )
 
-cc_binary(
-    name = "ptxas",
-    srcs = ["dummy.cpp"],
-    defines = ["TOOLNAME=ptxas"],
-)
-
 # Empty cc_library that provides CcInfo for components not available in this CUDA version.
 cc_library(
     name = "dummy",
     srcs = [],
     hdrs = [],
 )
+
+cc_binary(
+    name = "ptxas",
+    srcs = ["dummy.cpp"],
+    defines = ["TOOLNAME=ptxas"],
+)

cuda/extensions.bzl

@@ -1,8 +1,8 @@
 """Entry point for extensions used by bzlmod."""
 
+load("//cuda:platform_alias_extension.bzl", "platform_alias_repo")
 load("//cuda/private:redist_json_helper.bzl", "redist_json_helper")
 load("//cuda/private:repositories.bzl", "cuda_component", "cuda_toolkit")
-load("//cuda:platform_alias_extension.bzl", "platform_alias_repo")
 
 cuda_component_tag = tag_class(attrs = {
     "name": attr.string(mandatory = True, doc = "Repo name for the deliverable cuda_component"),
@@ -76,6 +76,10 @@ cuda_toolkit_tag = tag_class(attrs = {
     "nvcc_version": attr.string(
         doc = "nvcc version. Required for deliverable toolkit only. Fallback to version if omitted.",
     ),
+    "redist_json_name": attr.string(
+        doc = "Name of the redist_json tag whose components this toolkit should use. " +
+              "If omitted and exactly one redist_json is declared, it is used automatically.",
+    ),
 })
 
 def _find_modules(module_ctx):
@@ -109,7 +113,8 @@ def _redist_json_impl(module_ctx, attr):
             mapping[spec["component_name"]] = repo_name
 
             component_attr = {key: value for key, value in spec.items()}
-            component_attr["name"] = repo_name + "_" + platform.replace("-", "_") + "_" + redist_ver.replace(".", "_")
+            component_repo_name = repo_name + "_" + attr.name + "_" + platform.replace("-", "_") + "_" + redist_ver.replace(".", "_")
+            component_attr["name"] = component_repo_name
             cuda_component(**component_attr)
         platform_mapping[platform] = mapping
     return redist_ver, platform_mapping
@@ -128,6 +133,7 @@ def _impl(module_ctx):
         components = rules_cuda.tags.component
         redist_jsons = rules_cuda.tags.redist_json
         toolkits = rules_cuda.tags.toolkit
+
     for component in components:
         cuda_component(**_module_tag_to_dict(component))
 
@@ -141,7 +147,8 @@ def _impl(module_ctx):
     for redist_json in redist_jsons:
         components_mapping = {}
         redist_version, platform_mapping = _redist_json_impl(module_ctx, redist_json)
-        redist_versions.append(redist_version)
+        if redist_version not in redist_versions:
+            redist_versions.append(redist_version)
         for platform in platform_mapping.keys():
             for component_name, repo_name in platform_mapping[platform].items():
                 redist_components_mapping[component_name] = repo_name
@@ -151,7 +158,7 @@ def _impl(module_ctx):
                     versioned_repos[component_name] = {}
                 if platform not in versioned_repos[component_name]:
                     versioned_repos[component_name][platform] = {}
-                versioned_repos[component_name][platform][redist_version] = repo_name + "_" + platform.replace("-", "_") + "_" + redist_version.replace(".", "_")
+                versioned_repos[component_name][platform][redist_version] = repo_name + "_" + redist_json.name + "_" + platform.replace("-", "_") + "_" + redist_version.replace(".", "_")
 
     for component_name in redist_components_mapping.keys():
         # Build dictionaries mapping versions to repo names for each platform.
@@ -169,19 +176,20 @@ def _impl(module_ctx):
             versions = redist_versions,
         )
         components_mapping[component_name] = "@" + redist_components_mapping[component_name]
+
     registrations = {}
     for toolkit in toolkits:
         if toolkit.name in registrations.keys():
             if toolkit.toolkit_path == registrations[toolkit.name].toolkit_path:
-                # No problem to register a matching toolkit twice
                 continue
-            fail("Multiple conflicting toolkits declared for name {} ({} and {}".format(toolkit.name, toolkit.toolkit_path, registrations[toolkit.name].toolkit_path))
+            fail("Multiple conflicting toolkits declared for name {} ({} and {}".format(
+                toolkit.name,
+                toolkit.toolkit_path,
+                registrations[toolkit.name].toolkit_path,
+            ))
         else:
             registrations[toolkit.name] = toolkit
 
-    if len(registrations) > 1:
-        fail("multiple cuda.toolkit is not supported")
-
     for _, toolkit in registrations.items():
         if components_mapping != None:
             # Always use the maximum version so the toolkit includes all components.
@@ -195,7 +203,8 @@ def _impl(module_ctx):
 
             cuda_toolkit(name = toolkit.name, components_mapping = components_mapping, version = toolkit_version)
         else:
-            cuda_toolkit(**_module_tag_to_dict(toolkit))
+            attrs = {k: v for k, v in _module_tag_to_dict(toolkit).items() if k != "redist_json_name"}
+            cuda_toolkit(**attrs)
 
 toolchain = module_extension(
     implementation = _impl,

cuda/platform_alias_extension.bzl

@@ -26,14 +26,21 @@ def _platform_alias_repo_impl(ctx):
 
     build_content.append("[")
     build_content.append("    config_setting(")
-    build_content.append('        name = "version_is_{}".format(version.replace(".", "_")),')
+    build_content.append('        name = "version_is_{}_" + version.replace(".", "_"),'.format(
+        ctx.attr.component_name,
+    ))
     build_content.append('        flag_values = {"@rules_cuda//cuda:version": "{}".format(version)},')
     build_content.append("    )")
     build_content.append("    for version in {}".format(ctx.attr.versions))
     build_content.append("]")
     build_content.append("")
 
-    build_content.append('unsupported_cuda_version(name = "unsupported_cuda_version", component = "{}", available_versions = {})'.format(ctx.attr.component_name, ctx.attr.versions))
+    build_content.append(
+        'unsupported_cuda_version(name = "unsupported_cuda_version", component = "{}", available_versions = {})'.format(
+            ctx.attr.component_name,
+            ctx.attr.versions,
+        ),
+    )
     build_content.append("")
 
     # Build a target for the name of the repo (only if at least one platform is available).
@@ -50,7 +57,12 @@ def _platform_alias_repo_impl(ctx):
 
     # Always create unsupported_cuda_platform target - it's used as the default case
     # in select() when no platform condition matches.
-    build_content.append('unsupported_cuda_platform(name = "unsupported_cuda_platform", component = "{}", available_platforms = {})'.format(ctx.attr.component_name, platforms_available))
+    build_content.append(
+        'unsupported_cuda_platform(name = "unsupported_cuda_platform", component = "{}", available_platforms = {})'.format(
+            ctx.attr.component_name,
+            platforms_available,
+        ),
+    )
     build_content.append("")
 
     # Only generate target aliases if this component is in TARGET_MAPPING.
@@ -76,10 +88,13 @@ def _platform_alias_repo_impl(ctx):
         build_content.append("alias(")
         build_content.append('    name = "{}",'.format(target_name))
         build_content.append("    actual = select({")
+
         # Add conditions for ALL platforms, using dummy for unavailable ones.
         for platform in SUPPORTED_PLATFORMS:
             platform_suffix = platform.replace("-", "_")
-            build_content.append('        "@rules_cuda//cuda:{}_platform_is_{}":'.format(platform_type, platform_suffix))
+            build_content.append(
+                '        "@rules_cuda//cuda:{}_platform_is_{}":'.format(platform_type, platform_suffix),
+            )
             if platform in platforms_available:
                 build_content.append('            ":{}_{}",'.format(platform_suffix, target_name))
             else:
@@ -106,20 +121,38 @@ def _platform_alias_repo_impl(ctx):
             platform_suffix = platform.replace("-", "_")
             repos_dict = platform_repos_map[platform]
             platform_available = platform in platforms_available
+            default_version = ctx.attr.versions[0] if ctx.attr.versions else None
 
             build_content.append("alias(")
             build_content.append('    name = "{}_{}",'.format(platform_suffix, target_name))
             build_content.append("    actual = select({")
 
             for version in ctx.attr.versions:
-                build_content.append('        ":version_is_{}": '.format(version.replace(".", "_")))
+                build_content.append('        ":version_is_{}_{}": '.format(
+                    ctx.attr.component_name,
+                    version.replace(".", "_"),
+                ))
                 if platform_available and version in repos_dict:
                     repo_name = repos_dict[version]
-                    build_content.append('            "@{}//{}",'.format(repo_name, target if target.find(":") != -1 else ":" + target))
+                    build_content.append(
+                        '            "@{}//{}",'.format(
+                            repo_name,
+                            target if target.find(":") != -1 else ":" + target,
+                        ),
+                    )
                 else:
                     # Platform doesn't have this component for this version, use dummy.
                     build_content.append('            "{}",'.format(dummy_target))
-            build_content.append('        "//conditions:default": ":unsupported_cuda_version",')
+            if platform_available and default_version and default_version in repos_dict:
+                repo_name = repos_dict[default_version]
+                build_content.append(
+                    '        "//conditions:default": "@{}//{}",'.format(
+                        repo_name,
+                        target if target.find(":") != -1 else ":" + target,
+                    ),
+                )
+            else:
+                build_content.append('        "//conditions:default": ":unsupported_cuda_version",')
 
             build_content.append("    }),")
             build_content.append('    visibility = ["//visibility:public"],')

cuda/private/errors.bzl

@@ -1,7 +1,7 @@
 def _unsupported_cuda_version_impl(ctx):
     fail("CUDA component '{}' is not available for the selected CUDA version. Available versions: {}".format(
         ctx.attr.component,
-        ", ".join(ctx.attr.available_versions)
+        ", ".join(ctx.attr.available_versions),
     ))
 
 unsupported_cuda_version = rule(
@@ -15,7 +15,7 @@ unsupported_cuda_version = rule(
 def _unsupported_cuda_platform_impl(ctx):
     fail("CUDA component '{}' is not available for the selected platform. Available platforms: {}".format(
         ctx.attr.component,
-        ", ".join(ctx.attr.available_platforms)
+        ", ".join(ctx.attr.available_platforms),
     ))
 
 unsupported_cuda_platform = rule(

cuda/private/repositories.bzl

@@ -13,6 +13,25 @@ def _is_linux(ctx):
 def _is_windows(ctx):
     return ctx.os.name.lower().startswith("windows")
 
+def _detect_platform(ctx):
+    os = None
+    if _is_linux(ctx):
+        os = "linux"
+    elif _is_windows(ctx):
+        os = "windows"
+    else:
+        fail("Unsupported OS '{}' for CUDA redist.".format(ctx.os.name))
+
+    arch = ctx.os.arch
+    if arch in ["x86_64", "amd64"]:
+        arch = "x86_64"
+    elif arch in ["aarch64", "arm64"]:
+        arch = "aarch64"
+    else:
+        fail("Unsupported arch '{}' for CUDA redist.".format(ctx.os.arch))
+
+    return "{}-{}".format(os, arch)
+
 def _get_nvcc_version(repository_ctx, nvcc_root):
     result = repository_ctx.execute([nvcc_root + "/bin/nvcc", "--version"])
     if result.return_code != 0:
@@ -114,19 +133,18 @@ def _detect_deliverable_cuda_toolkit(repository_ctx):
 
     nvcc_repo = repository_ctx.attr.components_mapping["nvcc"]
 
-    bin_ext = ".exe" if _is_windows(repository_ctx) else ""
-    nvcc = "{}//:nvcc{}".format(nvcc_repo, bin_ext)
-    nvlink = "{}//:nvlink{}".format(nvcc_repo, bin_ext)
+    nvcc = "{}//:nvcc".format(nvcc_repo)
+    nvlink = "{}//:nvlink".format(nvcc_repo)
     link_stub = "{}//:link.stub".format(nvcc_repo)
-    bin2c = "{}//:bin2c{}".format(nvcc_repo, bin_ext)
-    fatbinary = "{}//:fatbinary{}".format(nvcc_repo, bin_ext)
-    ptxas = "{}//:ptxas{}".format(ptxas, bin_ext)
+    bin2c = "{}//:bin2c".format(nvcc_repo)
+    fatbinary = "{}//:fatbinary".format(nvcc_repo)
+    ptxas = "{}//:ptxas".format(nvcc_repo)
 
     cicc = None
     libdevice = None
-    if int(cuda_version_major) >= 13:
+    if "nvvm" in repository_ctx.attr.components_mapping:
         nvvm_repo = repository_ctx.attr.components_mapping["nvvm"]
-        cicc = "{}//:cicc{}".format(nvvm_repo, bin_ext)
+        cicc = "{}//:cicc".format(nvvm_repo)
         libdevice = "{}//:libdevice.10.bc".format(nvvm_repo)
 
     return struct(
@@ -464,7 +482,8 @@ def _cuda_redist_json_impl(repository_ctx):
     attr = repository_ctx.attr
     url, json_object = redist_json_helper.get(repository_ctx, attr)
     redist_ver = redist_json_helper.get_redist_version(repository_ctx, attr, json_object)
-    specs = redist_json_helper.collect_specs(repository_ctx, attr, json_object, url)
+    platform = _detect_platform(repository_ctx)
+    specs = redist_json_helper.collect_specs(repository_ctx, attr, platform, json_object, url)
 
     template_helper.generate_redist_bzl(repository_ctx, specs, redist_ver)
     repository_ctx.symlink(Label("//cuda/private:templates/BUILD.redist_json"), "BUILD")