Initial remote hermetic cuda toolchain

.github/workflows/build-tests.yaml

@@ -67,9 +67,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"),
+})
+
 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,6 +29,7 @@ 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'
@@ -39,7 +40,9 @@ def detect_cuda_toolkit(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:
@@ -58,13 +61,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
@@ -184,10 +187,10 @@ 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"],
-    # remotable = True,
 )
 
 def rules_cuda_dependencies():
@@ -212,4 +215,31 @@ def rules_cuda_dependencies():
         ],
     )
 
+    maybe(
+        name = "libcudacxx",
+        repo_rule = http_archive,
+        build_file = "@rules_cuda//third_party:libcu++.BUILD",
+        sha256 = "05b1435ad65f5bdef1bb8d1eb29dc8f0f7df34c01d77bf8686687a4603588bce",
+        strip_prefix = "libcudacxx-1.9.0",
+        urls = ["https://github.com/NVIDIA/libcudacxx/archive/refs/tags/1.9.0.tar.gz"],
+    )
+
+    maybe(
+        name = "thrust",
+        repo_rule = http_archive,
+        build_file = "@rules_cuda//third_party:thrust.BUILD",
+        sha256 = "d021e37f5aac30fd1b9737865399feb57db8e601ae2fc0af3cd41784435e9523",
+        strip_prefix = "thrust-1.17.2",
+        urls = ["https://github.com/NVIDIA/thrust/archive/refs/tags/1.17.2.tar.gz"],
+    )
+
+    maybe(
+        name = "cub",
+        repo_rule = http_archive,
+        build_file = "@rules_cuda//third_party:cub.BUILD",
+        sha256 = "1013a595794548c359f22c07e1f8c620b97e3a577f7e8496d9407f74566a3e2a",
+        strip_prefix = "cub-1.17.2",
+        urls = ["https://github.com/NVIDIA/cub/archive/refs/tags/1.17.2.tar.gz"],
+    )
+
     local_cuda(name = "local_cuda")

cuda/private/toolchain.bzl

@@ -38,10 +38,43 @@ cuda_toolchain = rule(
         "_cc_toolchain": attr.label(default = "@bazel_tools//tools/cpp:current_cc_toolchain"),
     },
 )
-
 CPP_TOOLCHAIN_TYPE = "@bazel_tools//tools/cpp:toolchain_type"
 CUDA_TOOLCHAIN_TYPE = "//cuda:toolchain_type"
 
+def _remote_cuda_impl(repository_ctx):
+    redist = repository_ctx.read(Label(repository_ctx.attr.json_path))
+    repos = json.decode(redist)
+    repos_to_define = dict()
+    base_url = repository_ctx.attr.base_url
+
+    for key in repos:
+        if key == "release_date":
+            continue
+        for arch in repos[key]:
+            if arch == "name" or arch == "license" or arch == "version":
+                continue
+            repos_to_define[key + "-%s" % arch] = {
+                "repo_rule": "http_archive",
+                "name": key + "-%s" % arch,
+                "sha256": repos[key][arch]["sha256"],
+                "url": base_url + repos[key][arch]["relative_path"],
+            }
+
+    repo_defs = "\n".join(["    maybe(name = \"{}\", build_file = \"@rules_cuda//cuda:templates/BUILD.remote_nvcc\", sha256 = \"{}\", repo_rule = {}, urls = [\"{}\"], strip_prefix=\"{}\")\n".format(repos_to_define[repo_name]["name"], repos_to_define[repo_name]["sha256"], repos_to_define[repo_name]["repo_rule"], repos_to_define[repo_name]["url"], repos_to_define[repo_name]["url"].split("/")[-1][:-7]) for repo_name in repos_to_define])
+
+    repository_ctx.template("repositories.bzl", Label("//cuda:templates/BUILD.repo_template"), substitutions = {"%{repos}": repo_defs}, executable = False)
+
+    repository_ctx.symlink(Label("//cuda:runtime/BUILD.remote_cuda"), "BUILD.bazel")
+    repository_ctx.symlink(Label("//cuda:templates/BUILD.remote_toolchain_nvcc"), "toolchain/BUILD.bazel")
+
+_remote_cuda = repository_rule(
+    implementation = _remote_cuda_impl,
+    attrs = {
+        "base_url": attr.string(default = "https://developer.download.nvidia.com/compute/cuda/redist/"),
+        "json_path": attr.string(mandatory = True),
+    },
+)
+
 # buildifier: disable=unused-variable
 def use_cpp_toolchain(mandatory = True):
     """Helper to depend on the C++ toolchain.
@@ -77,6 +110,24 @@ def find_cuda_toolkit(ctx):
     """
     return ctx.toolchains[CUDA_TOOLCHAIN_TYPE].cuda_toolkit[CudaToolkitInfo]
 
+CUDA_VERSIONS_JSON = {
+    "11.4.4": "//cuda:redistrib/redistrib_11.4.4.json",
+    "11.5.2": "//cuda:redistrib/redistrib_11.5.2.json",
+    "11.6.2": "//cuda:redistrib/redistrib_11.6.2.json",
+    "11.7.1": "//cuda:redistrib/redistrib_11.7.1.json",
+    "11.8.0": "//cuda:redistrib/redistrib_11.8.0.json",
+    "12.0.0": "//cuda:redistrib/redistrib_12.0.0.json",
+    "12.0.1": "//cuda:redistrib/redistrib_12.0.1.json",
+    "12.1.0": "//cuda:redistrib/redistrib_12.1.0.json",
+}
+
+def register_cuda_toolchains(name = "remote_cuda_toolchain", version = "12.0.0", cuda_versions = CUDA_VERSIONS_JSON):
+    _remote_cuda(name = name, json_path = cuda_versions[version])
+
+    native.register_toolchains(
+        "@%s//toolchain:nvcc-local-toolchain" % name,
+    )
+
 # buildifier: disable=unnamed-macro
 def register_detected_cuda_toolchains():
     """Helper to register the automatically detected CUDA toolchain(s).