Allow multiple 'local' installs of cuda to be specified

.github/workflows/build-tests.yaml

@@ -75,9 +75,9 @@ jobs:
       - run: bazelisk build @rules_cuda_examples//if_cuda:main
       - run: bazelisk build @rules_cuda_examples//if_cuda:main --enable_cuda=False
 
-      - run: bazelisk build @rules_cuda_examples//basic:main --config=bzlmod
-      - run: bazelisk build @rules_cuda_examples//rdc:main --config=bzlmod
-      - run: bazelisk build @rules_cuda_examples//if_cuda:main --config=bzlmod
-      - run: bazelisk build @rules_cuda_examples//if_cuda:main --enable_cuda=False --config=bzlmod
+      - run: cd examples && bazelisk build //basic:main --config=bzlmod
+      - run: cd examples && bazelisk build //rdc:main --config=bzlmod
+      - run: cd examples && bazelisk build //if_cuda:main --config=bzlmod
+      - run: cd examples && bazelisk build //if_cuda:main --enable_cuda=False --config=bzlmod
 
       - run: bazelisk shutdown

MODULE.bazel

@@ -1,13 +1,13 @@
 module(
     name = "rules_cuda",
     compatibility_level = 1,
-    version = "0.1.1",
+    version = "0.0.0",
 )
 
 bazel_dep(name = "bazel_skylib", version = "1.2.1")
-bazel_dep(name = "platforms", version = "0.0.4")
+bazel_dep(name = "platforms", version = "0.0.6")
 
-toolchain = use_extension("@rules_cuda//cuda:extensions.bzl", "local_toolchain")
+toolchain = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
 use_repo(toolchain, "local_cuda")
 
 register_toolchains(

cuda/extensions.bzl

@@ -2,7 +2,28 @@
 
 load("//cuda/private:repositories.bzl", "local_cuda")
 
+cuda_toolkit = tag_class(attrs = {
+    "name": attr.string(doc = "Name for the toolchain repository", default = "local_cuda"),
+    "toolkit_path": attr.string(doc = "Path to the CUDA SDK, if empty the environment variable CUDA_PATH will be used to deduce this path."),
+})
+
 def _init(module_ctx):
-    local_cuda(name = "local_cuda")
+    registrations = {}
+    for mod in module_ctx.modules:
+        for toolchain in mod.tags.local_toolchain:
+            if not mod.is_root:
+                fail("Only the root module may override the path for the local cuda toolchain")
+            if toolchain.name in registrations.keys():
+                if toolchain.toolkit_path == registrations[toolchain.name]:
+                    # No problem to register a matching toolchain twice
+                    continue
+                fail("Multiple conflicting toolchains declared for name {} ({} and {}".format(toolchain.name, toolchain.toolkit_path, registrations[toolchain.name]))
+            else:
+                registrations[toolchain.name] = toolchain.toolkit_path
+    for name, toolkit_path in registrations.items():
+        local_cuda(name = name, toolkit_path = toolkit_path)
 
-local_toolchain = module_extension(implementation = _init)
+toolchain = module_extension(
+    implementation = _init,
+    tag_classes = {"local_toolchain": cuda_toolkit},
+)

cuda/private/repositories.bzl

@@ -29,18 +29,20 @@ def detect_cuda_toolkit(repository_ctx):
     """Detect CUDA Toolkit.
 
     The path to CUDA Toolkit is determined as:
-
+      - the value of `toolkit_path` passed to local_cuda as an attribute
       - taken from `CUDA_PATH` environment variable or
-      - determined through `which ptxas` or
-      - defaults to `/usr/local/cuda`
+      - determined through 'which ptxas' or
+      - defaults to '/usr/local/cuda'
 
     Args:
         repository_ctx: repository_ctx
 
     Returns:
         A struct contains the information of CUDA Toolkit.
     """
-    cuda_path = repository_ctx.os.environ.get("CUDA_PATH", None)
+    cuda_path = repository_ctx.attr.toolkit_path
+    if cuda_path == "":
+        cuda_path = repository_ctx.os.environ.get("CUDA_PATH", None)
     if cuda_path == None:
         ptxas_path = repository_ctx.which("ptxas")
         if ptxas_path:
@@ -65,13 +67,13 @@ def detect_cuda_toolkit(repository_ctx):
     fatbinary = "@rules_cuda//cuda/dummy:fatbinary"
     if cuda_path != None:
         if repository_ctx.path(cuda_path + "/bin/nvlink" + bin_ext).exists:
-            nvlink = "@local_cuda//:cuda/bin/nvlink" + bin_ext
+            nvlink = ":cuda/bin/nvlink{}".format(bin_ext)
         if repository_ctx.path(cuda_path + "/bin/crt/link.stub").exists:
-            link_stub = "@local_cuda//:cuda/bin/crt/link.stub"
+            link_stub = ":cuda/bin/crt/link.stub"
         if repository_ctx.path(cuda_path + "/bin/bin2c" + bin_ext).exists:
-            bin2c = "@local_cuda//:cuda/bin/bin2c" + bin_ext
+            bin2c = ":cuda/bin/bin2c{}".format(bin_ext)
         if repository_ctx.path(cuda_path + "/bin/fatbinary" + bin_ext).exists:
-            fatbinary = "@local_cuda//:cuda/bin/fatbinary" + bin_ext
+            fatbinary = ":cuda/bin/fatbinary{}".format(bin_ext)
 
     nvcc_version_major = -1
     nvcc_version_minor = -1
@@ -197,6 +199,7 @@ def _local_cuda_impl(repository_ctx):
 
 local_cuda = repository_rule(
     implementation = _local_cuda_impl,
+    attrs = {"toolkit_path": attr.string(mandatory = False)},
     configure = True,
     local = True,
     environ = ["CUDA_PATH", "PATH", "CUDA_CLANG_PATH", "BAZEL_LLVM"],

cuda/repositories.bzl

@@ -1,5 +1,6 @@
-load("//cuda/private:repositories.bzl", _rules_cuda_dependencies = "rules_cuda_dependencies")
+load("//cuda/private:repositories.bzl", _local_cuda = "local_cuda", _rules_cuda_dependencies = "rules_cuda_dependencies")
 load("//cuda/private:toolchain.bzl", _register_detected_cuda_toolchains = "register_detected_cuda_toolchains")
 
 rules_cuda_dependencies = _rules_cuda_dependencies
+local_cuda = _local_cuda
 register_detected_cuda_toolchains = _register_detected_cuda_toolchains

examples/MODULE.bazel

@@ -4,9 +4,11 @@ module(
     version = "0.0.0",
 )
 
-bazel_dep(name = "rules_cuda", version = "0.1.1")
+bazel_dep(name = "rules_cuda", version = "0.1.3")
+local_path_override(module_name = "rules_cuda", path = "..")
 
-cuda_toolchain = use_extension("@rules_cuda//cuda:extensions.bzl", "local_toolchain")
-use_repo(cuda_toolchain, "local_cuda")
+cuda = use_extension("@rules_cuda//cuda:extensions.bzl", "toolchain")
 
-local_path_override(module_name = "rules_cuda", path = "..")
+cuda.local_toolchain(name = "local_cuda", toolkit_path = "")
+
+use_repo(cuda, "local_cuda")