Refactor the platform and the dedupe logic

Author: finn-ball

cuda/extensions.bzl

@@ -1,6 +1,7 @@
 """Entry point for extensions used by bzlmod."""
 
 load("//cuda:platform_alias_extension.bzl", "platform_alias_repo")
+load("//cuda/private:platforms.bzl", "SUPPORTED_PLATFORMS")
 load("//cuda/private:redist_json_helper.bzl", "redist_json_helper")
 load("//cuda/private:repositories.bzl", "cuda_component", "cuda_toolkit")
 
@@ -96,6 +97,9 @@ def _find_modules(module_ctx):
 def _module_tag_to_dict(t):
     return {attr: getattr(t, attr) for attr in dir(t)}
 
+def _platform_repos_attr(platform):
+    return platform.replace("-", "_") + "_repos"
+
 def _component_attrs_match(existing, current):
     for key, value in current.items():
         if key == "name":
@@ -109,36 +113,43 @@ def _component_attrs_match(existing, current):
             return False
     return True
 
-def _redist_json_impl(module_ctx, attr, generated_components):
+def _component_entry_key(component_name, platform, redist_ver):
+    return "{}|{}|{}".format(component_name, platform, redist_ver)
+
+def _register_redist_components(module_ctx, attr, component_entries):
     url, json_object = redist_json_helper.get(module_ctx, attr)
     redist_ver = redist_json_helper.get_redist_version(module_ctx, attr, json_object)
 
-    platform_mapping = {}
     for platform in attr.platforms:
         component_specs = redist_json_helper.collect_specs(module_ctx, attr, platform, json_object, url)
-        mapping = {}
         for spec in component_specs:
             repo_name = redist_json_helper.get_repo_name(module_ctx, spec)
-            mapping[spec["component_name"]] = repo_name
 
             component_attr = {key: value for key, value in spec.items()}
             component_repo_name = repo_name + "_" + platform.replace("-", "_") + "_" + redist_ver.replace(".", "_")
             component_attr["name"] = component_repo_name
 
-            dedupe_key = "{}|{}|{}".format(spec["component_name"], platform, redist_ver)
-            existing_attr = generated_components.get(dedupe_key)
+            dedupe_key = _component_entry_key(spec["component_name"], platform, redist_ver)
+            existing_entry = component_entries.get(dedupe_key)
+            existing_attr = existing_entry["component_attr"] if existing_entry else None
             if existing_attr == None:
                 cuda_component(**component_attr)
-                generated_components[dedupe_key] = component_attr
+                component_entries[dedupe_key] = {
+                    "component_name": spec["component_name"],
+                    "platform": platform,
+                    "redist_version": redist_ver,
+                    "repo_name": repo_name,
+                    "generated_repo_name": component_repo_name,
+                    "component_attr": component_attr,
+                }
             elif not _component_attrs_match(existing_attr, component_attr):
                 fail(("Conflicting CUDA component definition for {} on {} at version {}. " +
                       "Use distinct component versions when registries are not identical.").format(
                     spec["component_name"],
                     platform,
                     redist_ver,
                 ))
-        platform_mapping[platform] = mapping
-    return redist_ver, platform_mapping
+    return redist_ver
 
 def _impl(module_ctx):
     # Toolchain configuration is only allowed in the root module, or in rules_cuda.
@@ -158,47 +169,47 @@ def _impl(module_ctx):
     for component in components:
         cuda_component(**_module_tag_to_dict(component))
 
-    redist_version = None
     components_mapping = None
     redist_versions = []
-    redist_components_mapping = {}
-
-    # Track all versioned repositories for each component and platform.
-    versioned_repos = {}
-    generated_components = {}
+    component_entries = {}
     for redist_json in redist_jsons:
-        components_mapping = {}
-        redist_version, platform_mapping = _redist_json_impl(module_ctx, redist_json, generated_components)
+        redist_version = _register_redist_components(module_ctx, redist_json, component_entries)
         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
-
-                # Track the versioned repo name for this component/platform/version.
-                if component_name not in versioned_repos:
-                    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(".", "_")
-
-    for component_name in redist_components_mapping.keys():
-        # Build dictionaries mapping versions to repo names for each platform.
-        x86_64_repos = {ver: versioned_repos[component_name]["linux-x86_64"][ver] for ver in redist_versions if "linux-x86_64" in versioned_repos[component_name] and ver in versioned_repos[component_name]["linux-x86_64"]}
-        windows_x86_64_repos = {ver: versioned_repos[component_name]["windows-x86_64"][ver] for ver in redist_versions if "windows-x86_64" in versioned_repos[component_name] and ver in versioned_repos[component_name]["windows-x86_64"]}
-        aarch64_repos = {ver: versioned_repos[component_name]["linux-aarch64"][ver] for ver in redist_versions if "linux-aarch64" in versioned_repos[component_name] and ver in versioned_repos[component_name]["linux-aarch64"]}
-        sbsa_repos = {ver: versioned_repos[component_name]["linux-sbsa"][ver] for ver in redist_versions if "linux-sbsa" in versioned_repos[component_name] and ver in versioned_repos[component_name]["linux-sbsa"]}
-
-        platform_alias_repo(
-            name = redist_components_mapping[component_name],
-            component_name = component_name,
-            linux_x86_64_repos = x86_64_repos,
-            windows_x86_64_repos = windows_x86_64_repos,
-            linux_aarch64_repos = aarch64_repos,
-            linux_sbsa_repos = sbsa_repos,
-            versions = redist_versions,
-        )
-        components_mapping[component_name] = "@" + redist_components_mapping[component_name]
+
+    if len(component_entries) > 0:
+        components_mapping = {}
+        redist_components_mapping = {}
+        versioned_repos = {}
+        for entry in component_entries.values():
+            component_name = entry["component_name"]
+            platform = entry["platform"]
+            redist_version = entry["redist_version"]
+
+            redist_components_mapping[component_name] = entry["repo_name"]
+            if component_name not in versioned_repos:
+                versioned_repos[component_name] = {}
+            if platform not in versioned_repos[component_name]:
+                versioned_repos[component_name][platform] = {}
+            versioned_repos[component_name][platform][redist_version] = entry["generated_repo_name"]
+
+        for component_name in redist_components_mapping.keys():
+            # Build dictionaries mapping versions to repo names for each platform.
+            platform_repo_kwargs = {}
+            for platform in SUPPORTED_PLATFORMS:
+                platform_repo_kwargs[_platform_repos_attr(platform)] = {
+                    ver: versioned_repos[component_name][platform][ver]
+                    for ver in redist_versions
+                    if platform in versioned_repos[component_name] and ver in versioned_repos[component_name][platform]
+                }
+
+            platform_alias_repo(
+                name = redist_components_mapping[component_name],
+                component_name = component_name,
+                versions = redist_versions,
+                **platform_repo_kwargs
+            )
+            components_mapping[component_name] = "@" + redist_components_mapping[component_name]
 
     registrations = {}
     for toolkit in toolkits:

cuda/platform_alias_extension.bzl

@@ -10,6 +10,16 @@ load("//cuda/private:templates/registry.bzl", "REGISTRY")
 # Use REGISTRY as the source of truth for component targets
 TARGET_MAPPING = REGISTRY
 
+def _platform_repos_attr(platform):
+    return platform.replace("-", "_") + "_repos"
+
+_PLATFORM_REPO_ATTRS = {
+    _platform_repos_attr(_platform): attr.string_dict(
+        default = {},
+        doc = "Dictionary mapping versions to repository names for {}".format(_platform),
+    )
+    for _platform in SUPPORTED_PLATFORMS
+}
 def _platform_alias_repo_impl(ctx):
     """Implementation of the platform_alias_repo repository rule.
 
@@ -46,16 +56,13 @@ def _platform_alias_repo_impl(ctx):
     # Build a target for the name of the repo (only if at least one platform is available).
     platform_type = "exec" if ctx.attr.component_name in ["nvcc", "nvvm"] else "target"
 
+    platform_repos_map = {
+        platform: getattr(ctx.attr, _platform_repos_attr(platform))
+        for platform in SUPPORTED_PLATFORMS
+    }
+
     # Check which platforms are available (have at least one version).
-    platforms_available = []
-    if len(ctx.attr.linux_x86_64_repos) > 0:
-        platforms_available.append("linux-x86_64")
-    if len(ctx.attr.windows_x86_64_repos) > 0:
-        platforms_available.append("windows-x86_64")
-    if len(ctx.attr.linux_sbsa_repos) > 0:
-        platforms_available.append("linux-sbsa")
-    if len(ctx.attr.linux_aarch64_repos) > 0:
-        platforms_available.append("linux-aarch64")
+    platforms_available = [platform for platform in SUPPORTED_PLATFORMS if len(platform_repos_map[platform]) > 0]
 
     # Always create unsupported_cuda_platform target - it's used as the default case
     # in select() when no platform condition matches.
@@ -113,13 +120,6 @@ def _platform_alias_repo_impl(ctx):
         # Platforms where it doesn't exist get dummy targets for all versions.
         # This ensures builds on any platform have matching select conditions.
 
-        platform_repos_map = {
-            "linux-x86_64": ctx.attr.linux_x86_64_repos,
-            "windows-x86_64": ctx.attr.windows_x86_64_repos,
-            "linux-sbsa": ctx.attr.linux_sbsa_repos,
-            "linux-aarch64": ctx.attr.linux_aarch64_repos,
-        }
-
         for platform in SUPPORTED_PLATFORMS:
             platform_suffix = platform.replace("-", "_")
             repos_dict = platform_repos_map[platform]
@@ -167,30 +167,14 @@ def _platform_alias_repo_impl(ctx):
 
 platform_alias_repo = repository_rule(
     implementation = _platform_alias_repo_impl,
-    attrs = {
+    attrs = dict({
         "component_name": attr.string(
             mandatory = True,
             doc = "Name of the component",
         ),
-        "linux_x86_64_repos": attr.string_dict(
-            default = {},
-            doc = "Dictionary mapping versions to x86_64 repository names",
-        ),
-        "linux_aarch64_repos": attr.string_dict(
-            default = {},
-            doc = "Dictionary mapping versions to ARM64/Jetpack repository names",
-        ),
-        "windows_x86_64_repos": attr.string_dict(
-            default = {},
-            doc = "Dictionary mapping versions to Windows x86_64 repository names",
-        ),
-        "linux_sbsa_repos": attr.string_dict(
-            default = {},
-            doc = "Dictionary mapping versions to SBSA repository names",
-        ),
         "versions": attr.string_list(
             mandatory = True,
             doc = "List of versions to create aliases for",
         ),
-    },
+    }, **_PLATFORM_REPO_ATTRS),
 )