metal: choose compile spec based on macos (#9188)

* metal: choose compile spec based on macos

* correction
This commit is contained in:
nimlgen
2025-02-21 00:43:39 +03:00
committed by GitHub
parent 3e22747799
commit f986e12f91

View File

@@ -1,4 +1,4 @@
import os, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal
import os, pathlib, struct, ctypes, tempfile, functools, contextlib, decimal, platform
from typing import Any, Union, cast
from tinygrad.helpers import prod, to_mv, getenv, round_up, cache_dir, T, init_c_struct_t, PROFILE
from tinygrad.device import Compiled, Compiler, CompileError, LRUAllocator, cpu_profile, ProfileDeviceEvent, ProfileRangeEvent
@@ -117,9 +117,14 @@ class MetalCompiler(Compiler):
ret = reply[sum(struct.unpack('<LL', reply[8:16])):]
else:
ret = CompileError(errorMessage.decode())
# no changes for compute in 2.0 - 2.4 specs, use 2.0 as default for old versions.
macos_major = int(platform.mac_ver()[0].split('.')[0])
metal_version = "metal3.1" if macos_major >= 14 else "metal3.0" if macos_major >= 13 else "macos-metal2.0"
# llvm will create modules.timestamp in cache path and cache compilation of metal stdlib (250ms => 8ms compilation time)
# note that llvm won't necessarily create anything else here as apple has prebuilt versions of many standard libraries
params = f'-fno-fast-math -std=metal3.1 --driver-mode=metal -x metal -fmodules-cache-path="{cache_dir}" -fno-caret-diagnostics'
params = f'-fno-fast-math -std={metal_version} --driver-mode=metal -x metal -fmodules-cache-path="{cache_dir}" -fno-caret-diagnostics'
# source blob has to be padded to multiple of 4 but at least one 'b\x00' should be added, params blob just has to be null terminated
src_padded, params_padded = src.encode() + b'\x00'*(round_up(len(src) + 1, 4) - len(src)), params.encode() + b'\x00'
request = struct.pack('<QQ', len(src_padded), len(params_padded)) + src_padded + params_padded