From 59800e973659803d942178102d9ab4a3e4ce0a92 Mon Sep 17 00:00:00 2001
From: Brad Larson <bradlarson@modular.com>
Date: Tue, 16 Jun 2026 11:45:40 -0500
Subject: [PATCH] Add separation between architecture and Metal API selection.

---
 mojo/extensions.bzl         | 19 +++++++++++++++----
 mojo/mojo_host_platform.bzl | 33 ++++++++++++++++++++++++---------
 2 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/mojo/extensions.bzl b/mojo/extensions.bzl
index 6ea8fa0..c6be0a9 100644
--- a/mojo/extensions.bzl
+++ b/mojo/extensions.bzl
@@ -201,8 +201,16 @@ _gpu_toolchains_tag = tag_class(
                 "mi325": "amdgpu:gfx942",
                 "mi355": "amdgpu:gfx950",
                 "rtx5090": "nvidia:120a",
-                "metal3": "metal:30",
-                "metal4": "metal:40",
+                "metal1": "metal:1",
+                "metal2": "metal:2",
+                "metal3": "metal:3",
+                "metal4": "metal:4",
+                "metal5": "metal:5",
+                "metal1_metal4": "metal:1-metal4",
+                "metal2_metal4": "metal:2-metal4",
+                "metal3_metal4": "metal:3-metal4",
+                "metal4_metal4": "metal:4-metal4",
+                "metal5_metal4": "metal:5-metal4",
             },
             doc = "The GPUs supported by this toolchain, mapping to Mojo's target accelerators.",
         ),
@@ -227,8 +235,11 @@ _gpu_toolchains_tag = tag_class(
                 "MI355": "mi355",
                 "Navi": "radeon",
                 "AMD Radeon Graphics": "radeon",
-                "Metal 3": "metal3",
-                "Metal 4": "metal4",
+                "Apple M1": "metal1",
+                "Apple M2": "metal2",
+                "Apple M3": "metal3",
+                "Apple M4": "metal4",
+                "Apple M5": "metal5",
             },
             doc = "The output from nvidia-smi or rocm-smi to the corresponding GPU name in SUPPORTED_GPUS.",
         ),
diff --git a/mojo/mojo_host_platform.bzl b/mojo/mojo_host_platform.bzl
index 781004c..603b456 100644
--- a/mojo/mojo_host_platform.bzl
+++ b/mojo/mojo_host_platform.bzl
@@ -87,6 +87,13 @@ def _get_amd_constraints_with_rocm_smi(rctx, rocm_smi, gpu_mapping):
 
     return constraints
 
+def _toolchain_supports_metal4(rctx):
+    result = rctx.execute(["/usr/bin/xcrun", "--sdk", "macosx", "--show-sdk-version"])
+    _log_result(rctx, "/usr/bin/xcrun --sdk macosx --show-sdk-version", result)
+    if result.return_code != 0:
+        return False
+    return int(result.stdout.strip().split(".")[0]) >= 26
+
 def _get_apple_constraint(rctx, gpu_mapping):
     result = rctx.execute(["/usr/bin/sw_vers", "--productVersion"])
     _log_result(rctx, "/usr/sbin/sw_vers --productVersion", result)
@@ -102,21 +109,29 @@ def _get_apple_constraint(rctx, gpu_mapping):
 
     _log_result(rctx, "/usr/sbin/system_profiler SPDisplaysDataType", result)
 
-    metal_version = None
+    chipset_model = None
+    metal_support = None
     for line in result.stdout.splitlines():
-        if "Metal Support:" in line:
-            metal_version = line
-            break
+        if "Chipset Model:" in line:
+            chipset_model = line
+        elif "Metal Support:" in line:
+            metal_support = line
 
-    if not metal_version:  # macOS VMs may not have GPUs attached
+    if not chipset_model:  # macOS VMs may not have GPUs attached
         return None
 
+    metal4 = (
+        metal_support and "Metal 4" in metal_support and
+        _toolchain_supports_metal4(rctx)
+    )
+
     for gpu_name, constraint in gpu_mapping.items():
-        if gpu_name in metal_version:
-            if constraint:
-                return "@mojo_gpu_toolchains//:{}_gpu".format(constraint)
-            else:
+        if gpu_name in chipset_model:
+            if not constraint:
                 return None
+            if metal4:
+                constraint += "_metal4"
+            return "@mojo_gpu_toolchains//:{}_gpu".format(constraint)
 
     _fail(rctx, "Unrecognized system_profiler output, please add it to your gpu_mapping in the MODULE.bazel file: {}".format(result.stdout))
     return None