Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 72 additions & 6 deletions cuda/platform_alias_extension.bzl
Original file line number Diff line number Diff line change
Expand Up @@ -111,13 +111,60 @@ def _emit_outer_constraint_alias(build_content, alias_name, target_name, platfor
build_content.append(")")
build_content.append("")

def _detect_host_platform(ctx):
"""Detect the host platform where Bazel is running.

This is used to set a sensible default for exec platform when auto-detection
fallback is used.

Args:
ctx: Repository rule context

Returns:
A string representing the detected platform (e.g., "linux-x86_64")
"""
os_name = ctx.os.name.lower()
arch = ctx.os.arch.lower()

# Normalize OS name
if os_name.startswith("linux"):
os_name = "linux"
elif os_name.startswith("windows"):
os_name = "windows"
elif os_name.startswith("mac"):
os_name = "macos"
else:
# Unknown OS, default to linux
os_name = "linux"

# Normalize architecture
if arch in ["x86_64", "amd64", "x64"]:
arch = "x86_64"
elif arch in ["aarch64", "arm64"]:
arch = "aarch64"
else:
# Unknown arch, default to x86_64
arch = "x86_64"

platform = "{}-{}".format(os_name, arch)

# Map to a supported CUDA platform
# For aarch64 on Linux, default to sbsa (server-class)
if platform == "linux-aarch64":
platform = "linux-sbsa"

return platform

def _platform_alias_repo_impl(ctx):
"""Implementation of the platform_alias_repo repository rule.

Args:
ctx: Repository context with attributes x86_repo, arm64_repo, and targets.
"""

# Detect the host platform for exec platform default
host_platform = _detect_host_platform(ctx)

# Generate BUILD.bazel content with platform-specific aliases
build_content = ["# Generated by platform_alias_repo rule", ""]

Expand Down Expand Up @@ -197,12 +244,22 @@ def _platform_alias_repo_impl(ctx):
)

if platform_type == "exec":
# Outer alias: explicit --exec_platform flag wins; the default
# branch falls through to constraint-based auto-detection in
# the private :auto_<name> alias below.
# Outer alias for exec components (nvcc, cicc, etc.)
# Priority order:
# 1. Explicit --@rules_cuda//cuda:exec_platform flag
# 2. Detected host platform (where Bazel is running)
# 3. Constraint-based auto-detection (for RBE scenarios)
#
# The host platform detection happens at repository instantiation time
# and provides a sensible default that works for local builds and
# cross-compilation.
host_platform_suffix = host_platform.replace("-", "_")

build_content.append("alias(")
build_content.append(' name = "{}",'.format(target_name))
build_content.append(" actual = select({")

# Explicit flag overrides
for platform in SUPPORTED_PLATFORMS:
platform_suffix = platform.replace("-", "_")
build_content.append(
Expand All @@ -212,14 +269,23 @@ def _platform_alias_repo_impl(ctx):
build_content.append(' ":{}_{}",'.format(platform_suffix, target_name))
else:
build_content.append(' "{}",'.format(dummy_target))
build_content.append(' "//conditions:default": ":auto_{}",'.format(target_name))

# Default: use detected host platform
if host_platform in platforms_available:
build_content.append(' "//conditions:default": ":{}_{}",'.format(host_platform_suffix, target_name))
else:
# Host platform not available, fall back to auto-detection alias
build_content.append(' "//conditions:default": ":auto_{}",'.format(target_name))

build_content.append(" }),")
build_content.append(' visibility = ["//visibility:public"],')
build_content.append(")")
build_content.append("")

# Auto-detection fallback. Private — only the outer flag-select
# in this same package needs to reach it.
# Auto-detection fallback for when host platform is not available.
# This is a private target only used when the detected host platform
# doesn't have this component available.
# Uses constraint-based detection which works for RBE scenarios.
_emit_outer_constraint_alias(
build_content,
"auto_" + target_name,
Expand Down
Loading