Allow multiple cuda versions in the same toolchain

Author: finn-ball

cuda/extensions.bzl

@@ -72,6 +72,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):
@@ -125,32 +129,51 @@ def _impl(module_ctx):
     for component in components:
         cuda_component(**_module_tag_to_dict(component))
 
-    if len(redist_jsons) > 1:
-        fail("Using multiple cuda.redist_json is not supported yet.")
-
-    redist_version = None
-    components_mapping = None
+    # Process each redist_json, keyed by name.
+    redist_results = {}
     for redist_json in redist_jsons:
+        if redist_json.name in redist_results:
+            fail("Multiple redist_json tags declared with the same name '{}'".format(redist_json.name))
         redist_version, components_mapping = _redist_json_impl(module_ctx, redist_json)
+        redist_results[redist_json.name] = struct(
+            version = redist_version,
+            components_mapping = components_mapping,
+        )
 
+    # Deduplicate toolkit registrations by 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")
-
+    # Instantiate each toolkit, wiring it to the appropriate redist_json.
     for _, toolkit in registrations.items():
-        if components_mapping != None:
-            cuda_toolkit(name = toolkit.name, components_mapping = components_mapping, version = redist_version)
+        rj_name = toolkit.redist_json_name
+        if rj_name:
+            if rj_name not in redist_results:
+                fail("Toolkit '{}' references redist_json_name '{}', but no redist_json with that name was declared".format(
+                    toolkit.name,
+                    rj_name,
+                ))
+            rj = redist_results[rj_name]
+            cuda_toolkit(name = toolkit.name, components_mapping = rj.components_mapping, version = rj.version)
+        elif len(redist_results) == 1:
+            # Backward compatible: single redist_json, use it implicitly.
+            rj = redist_results.values()[0]
+            cuda_toolkit(name = toolkit.name, components_mapping = rj.components_mapping, version = rj.version)
+        elif redist_results:
+            fail("Multiple redist_json tags declared but toolkit '{}' does not specify redist_json_name".format(toolkit.name))
         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,