diff --git a/mojo/extensions.bzl b/mojo/extensions.bzl index 6d0b6f7..08859a4 100644 --- a/mojo/extensions.bzl +++ b/mojo/extensions.bzl @@ -201,16 +201,8 @@ _gpu_toolchains_tag = tag_class( "mi325": "amdgpu:gfx942", "mi355": "amdgpu:gfx950", "rtx5090": "nvidia:120a", - "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", + "metal3": "metal:30", + "metal4": "metal:40", }, doc = "The GPUs supported by this toolchain, mapping to Mojo's target accelerators.", ), @@ -235,11 +227,8 @@ _gpu_toolchains_tag = tag_class( "MI355": "mi355", "Navi": "radeon", "AMD Radeon Graphics": "radeon", - "Apple M1": "metal1", - "Apple M2": "metal2", - "Apple M3": "metal3", - "Apple M4": "metal4", - "Apple M5": "metal5", + "Metal 3": "metal3", + "Metal 4": "metal4", }, 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 603b456..781004c 100644 --- a/mojo/mojo_host_platform.bzl +++ b/mojo/mojo_host_platform.bzl @@ -87,13 +87,6 @@ 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) @@ -109,29 +102,21 @@ def _get_apple_constraint(rctx, gpu_mapping): _log_result(rctx, "/usr/sbin/system_profiler SPDisplaysDataType", result) - chipset_model = None - metal_support = None + metal_version = None for line in result.stdout.splitlines(): - if "Chipset Model:" in line: - chipset_model = line - elif "Metal Support:" in line: - metal_support = line + if "Metal Support:" in line: + metal_version = line + break - if not chipset_model: # macOS VMs may not have GPUs attached + if not metal_version: # 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 chipset_model: - if not constraint: + if gpu_name in metal_version: + if constraint: + return "@mojo_gpu_toolchains//:{}_gpu".format(constraint) + else: 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