Add GPU detection logic (#12)

There are a few pieces to GPU detection for supporting Mojo. The final
goal is to have a valid value we can pass as
`--target-accelerator=nvidia:80` to `mojo build`. In order to do this
based on the target platform we have to detect the current GPU with
nvidia-smi or rocm-smi, parse the output, and setup various
config_settings to determine which toolchain should be used. If you use
`--platforms=@mojo_host_toolchain` you get the rest of this logic for
free. The currently known supported GPUs are seeded in the
`mojo.gpu_toolchains` module extension, but new ones can be added in
individual projects as well.

Author: keith

.bazelrc

@@ -1,2 +1,8 @@
 common --incompatible_strict_action_env
 common --test_output=errors
+
+# https://github.com/bazelbuild/bazel/issues/25145
+info --platforms=
+
+common --platforms=@mojo_host_platform
+common --host_platform=@mojo_host_platform

BUILD.bazel

@@ -10,3 +10,8 @@ toolchain_type(
     name = "toolchain_type",
     visibility = ["//visibility:public"],
 )
+
+toolchain_type(
+    name = "gpu_toolchain_type",
+    visibility = ["//visibility:public"],
+)

MODULE.bazel

@@ -11,9 +11,10 @@ bazel_dep(name = "rules_python", version = "1.0.0")
 
 mojo = use_extension("//mojo:extensions.bzl", "mojo")
 mojo.toolchain()
-use_repo(mojo, "mojo_toolchains")
+mojo.gpu_toolchains()
+use_repo(mojo, "mojo_gpu_toolchains", "mojo_host_platform", "mojo_toolchains")
 
-register_toolchains("@mojo_toolchains//...")
+register_toolchains("@mojo_toolchains//...", "@mojo_gpu_toolchains//...")
 
 _DEFAULT_PYTHON_VERSION = "3.12"
 

mojo/extensions.bzl

@@ -1,5 +1,8 @@
 """MODULE.bazel extensions for Mojo toolchains."""
 
+load("//mojo:mojo_host_platform.bzl", "mojo_host_platform")
+load("//mojo/private:mojo_gpu_toolchains_repository.bzl", "mojo_gpu_toolchains_repository")
+
 _PLATFORMS = ["linux_aarch64", "linux_x86_64", "macos_arm64"]
 _DEFAULT_VERSION = "25.4.0.dev2025050902"
 _KNOWN_SHAS = {
@@ -98,20 +101,36 @@ def _mojo_impl(mctx):
         if not module.is_root:
             continue
 
+        toolchains = module.tags.toolchain
         if len(module.tags.toolchain) > 1:
             fail("mojo.toolchain() can only be called once per module.")
+        if toolchains:
+            has_toolchains = True
+            tags = toolchains[0]
+
+            for platform in _PLATFORMS:
+                name = "mojo_toolchain_{}".format(platform)
+                _mojo_toolchain_repository(
+                    name = name,
+                    version = tags.version,
+                    platform = platform,
+                    url_override = tags.url_override,
+                    use_prebuilt_packages = tags.use_prebuilt_packages,
+                )
 
-        has_toolchains = True
-        tags = module.tags.toolchain[0]
-
-        for platform in _PLATFORMS:
-            name = "mojo_toolchain_{}".format(platform)
-            _mojo_toolchain_repository(
-                name = name,
-                version = tags.version,
-                platform = platform,
-                url_override = tags.url_override,
-                use_prebuilt_packages = tags.use_prebuilt_packages,
+        gpu_toolchains = module.tags.gpu_toolchains
+        if len(gpu_toolchains) > 1:
+            fail("mojo.gpu_toolchain() can only be called once per module.")
+        if gpu_toolchains:
+            gpu_toolchain = gpu_toolchains[0]
+            mojo_gpu_toolchains_repository(
+                name = "mojo_gpu_toolchains",
+                supported_gpus = gpu_toolchain.supported_gpus,
+            )
+
+            mojo_host_platform(
+                name = "mojo_host_platform",
+                gpu_mapping = gpu_toolchain.gpu_mapping,
             )
 
     _mojo_toolchain_hub(
@@ -140,10 +159,57 @@ _toolchain_tag = tag_class(
     },
 )
 
+_gpu_toolchains_tag = tag_class(
+    doc = "Tags for configuring Mojo GPU toolchains.",
+    attrs = {
+        "supported_gpus": attr.string_dict(
+            default = {
+                "780M": "amdgpu:gfx1103",
+                "a10": "nvidia:86",
+                "a100": "nvidia:80",
+                "a3000": "nvidia:86",
+                "b100": "nvidia:100a",
+                "b200": "nvidia:100a",
+                "h100": "nvidia:90a",
+                "h200": "nvidia:90a",
+                "l4": "nvidia:89",
+                "mi300x": "amdgpu:gfx942",
+                "mi325": "amdgpu:gfx942",
+                "rtx5090": "nvidia:120a",
+            },
+            doc = "The GPUs supported by this toolchain, mapping to Mojo's target accelerators.",
+        ),
+        "gpu_mapping": attr.string_dict(
+            default = {
+                " A10G": "a10",
+                "A100-": "a100",
+                " H100 ": "h100",
+                " H200 ": "h200",
+                " L4 ": "L4",
+                " Ada ": "L4",
+                " A3000 ": "a3000",
+                "B100": "b100",
+                "B200": "b200",
+                " RTX 5090": "rtx5090",
+                "Laptop GPU": "",
+                "RTX 4070 Ti": "",
+                "RTX 4080 SUPER": "",
+                "NVIDIA GeForce RTX 3090": "",
+                "MI300X": "mi300x",
+                "MI325": "mi325",
+                "Navi": "radeon",
+                "AMD Radeon Graphics": "radeon",
+            },
+            doc = "The output from nvidia-smi or rocm-smi to the corresponding GPU name in SUPPORTED_GPUS.",
+        ),
+    },
+)
+
 mojo = module_extension(
     doc = "Mojo toolchain extension.",
     implementation = _mojo_impl,
     tag_classes = {
         "toolchain": _toolchain_tag,
+        "gpu_toolchains": _gpu_toolchains_tag,
     },
 )

mojo/mojo_host_platform.bzl

@@ -0,0 +1,148 @@
+"""Setup a host platform that takes into account current GPU hardware"""
+
+def _verbose_log(rctx, msg):
+    if rctx.getenv("MOJO_VERBOSE_GPU_DETECT"):
+        # buildifier: disable=print
+        print(msg)
+
+def _log_result(rctx, binary, result):
+    _verbose_log(
+        rctx,
+        "\n------ {}:\nexit status: {}\nstdout: {}\nstderr: {}\n------ end gpu-query info"
+            .format(binary, result.return_code, result.stdout, result.stderr),
+    )
+
+def _get_amdgpu_constraint(series, gpu_mapping):
+    for gpu_name, constraint in gpu_mapping.items():
+        if gpu_name in series:
+            if constraint:
+                return "@mojo_gpu_toolchains//:{}_gpu".format(constraint)
+            else:
+                return None
+
+    fail("Unrecognized amd-smi/rocm-smi output, please add it to your gpu_mapping in the MODULE.bazel file: {}".format(series))
+
+def _get_rocm_constraint(blob, gpu_mapping):
+    for value in blob.values():
+        series = value["Card Series"]
+        return _get_amdgpu_constraint(series, gpu_mapping)
+    fail("Unrecognized rocm-smi output, please report: {}".format(blob))
+
+def _get_amd_constraint(blob, gpu_mapping):
+    for value in blob:
+        series = value["asic"]["market_name"]
+        return _get_amdgpu_constraint(series, gpu_mapping)
+    fail("Unrecognized amd-smi output, please report: {}".format(blob))
+
+def _get_nvidia_constraint(lines, gpu_mapping):
+    line = lines[0]
+    for gpu_name, constraint in gpu_mapping.items():
+        if gpu_name in line:
+            if constraint:
+                return "@mojo_gpu_toolchains//:{}_gpu".format(constraint)
+            else:
+                return None
+
+    fail("Unrecognized nvidia-smi output, please add it to your gpu_mapping in the MODULE.bazel file: {}".format(lines))
+
+def _impl(rctx):
+    constraints = []
+
+    if rctx.os.name == "linux" and rctx.os.arch == "amd64":
+        # A system may have both rocm-smi and nvidia-smi installed, check both.
+        nvidia_smi = rctx.which("nvidia-smi")
+
+        # amd-smi supersedes rocm-smi
+        amd_smi = rctx.which("amd-smi")
+        rocm_smi = rctx.which("rocm-smi")
+
+        _verbose_log(rctx, "nvidia-smi path: {}, rocm-smi path: {}, amd-smi path: {}".format(nvidia_smi, rocm_smi, amd_smi))
+
+        # NVIDIA
+        if nvidia_smi:
+            result = rctx.execute([nvidia_smi, "--query-gpu=gpu_name", "--format=csv,noheader"])
+            _log_result(rctx, nvidia_smi, result)
+            if result.return_code == 0:
+                lines = result.stdout.splitlines()
+                if len(lines) == 0:
+                    fail("nvidia-smi succeeded but had no GPUs, please report this issue")
+
+                constraint = _get_nvidia_constraint(lines, rctx.attr.gpu_mapping)
+                if constraint:
+                    constraints.extend([
+                        "@mojo_gpu_toolchains//:nvidia_gpu",
+                        "@mojo_gpu_toolchains//:has_gpu",
+                        constraint,
+                    ])
+
+                if len(lines) > 1:
+                    constraints.append("@mojo_gpu_toolchains//:has_multi_gpu")
+                if len(lines) >= 4:
+                    constraints.append("@mojo_gpu_toolchains//:has_4_gpus")
+
+        # AMD
+        if amd_smi:
+            result = rctx.execute([amd_smi, "static", "--json"])
+            _log_result(rctx, amd_smi, result)
+
+            if result.return_code == 0:
+                constraints.extend([
+                    "@mojo_gpu_toolchains//:amd_gpu",
+                    "@mojo_gpu_toolchains//:has_gpu",
+                ])
+
+                blob = json.decode(result.stdout)
+                if len(blob) == 0:
+                    fail("amd-smi succeeded but didn't actually have any GPUs, please report this issue")
+
+                constraints.append(_get_amd_constraint(blob, rctx.attr.gpu_mapping))
+                if len(blob) > 1:
+                    constraints.append("@mojo_gpu_toolchains//:has_multi_gpu")
+                if len(blob) >= 4:
+                    constraints.append("@mojo_gpu_toolchains//:has_4_gpus")
+
+        elif rocm_smi:
+            result = rctx.execute([rocm_smi, "--json", "--showproductname"])
+            _log_result(rctx, rocm_smi, result)
+
+            if result.return_code == 0:
+                constraints.extend([
+                    "@mojo_gpu_toolchains//:amd_gpu",
+                    "@mojo_gpu_toolchains//:has_gpu",
+                ])
+
+                blob = json.decode(result.stdout)
+                if len(blob.keys()) == 0:
+                    fail("rocm-smi succeeded but didn't actually have any GPUs, please report this issue")
+
+                constraints.append(_get_rocm_constraint(blob, rctx.attr.gpu_mapping))
+                if len(blob.keys()) > 1:
+                    constraints.append("@mojo_gpu_toolchains//:has_multi_gpu")
+                if len(blob.keys()) >= 4:
+                    constraints.append("@mojo_gpu_toolchains//:has_4_gpus")
+
+    rctx.file("WORKSPACE.bazel", "workspace(name = {})".format(rctx.attr.name))
+    rctx.file("BUILD.bazel", """
+platform(
+    name = "mojo_host_platform",
+    parents = ["@platforms//host"],
+    visibility = ["//visibility:public"],
+    constraint_values = [{constraints}],
+    exec_properties = {{
+        "no-remote-exec": "1",
+    }},
+)
+""".format(constraints = ", ".join(['"{}"'.format(x) for x in constraints])))
+
+mojo_host_platform = repository_rule(
+    implementation = _impl,
+    configure = True,
+    environ = [
+        "MOJO_VERBOSE_GPU_DETECT",
+    ],
+    attrs = {
+        "gpu_mapping": attr.string_dict(
+            doc = "A dictionary of GPU strings from nvidia-smi or amd-smi, mapped to supported GPUs defined by mojo.gpu_toolchains()",
+        ),
+    },
+)

mojo/private/mojo_gpu_toolchain.bzl

@@ -0,0 +1,26 @@
+"""Bazel toolchain representing the currently targeted GPU hardware"""
+
+load("//mojo:providers.bzl", "MojoGPUToolchainInfo")
+
+def _mojo_gpu_toolchain_impl(ctx):
+    brand = ctx.attr.target_accelerator.split(":")[0]
+    return [
+        platform_common.ToolchainInfo(
+            mojo_gpu_toolchain_info = MojoGPUToolchainInfo(
+                brand = brand,
+                has_4_gpus = ctx.attr.has_4_gpus,
+                multi_gpu = ctx.attr.multi_gpu,
+                name = ctx.attr.name,
+                target_accelerator = ctx.attr.target_accelerator,
+            ),
+        ),
+    ]
+
+mojo_gpu_toolchain = rule(
+    implementation = _mojo_gpu_toolchain_impl,
+    attrs = {
+        "target_accelerator": attr.string(mandatory = True),
+        "multi_gpu": attr.bool(mandatory = True),
+        "has_4_gpus": attr.bool(mandatory = True),
+    },
+)