github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 23:48:01 -05:00
Code Issues Packages Projects Releases Wiki Activity

hotfix amd bf16 is supported case (#10039)

Browse Source 
* hotfix amd and amd_llvm

* bf16 not supported in ci

* hotfix amd_llvm is not a device

* remove default

* dont gate on ci and amd_llvm

* minor cleanup

* skip bf16 tc test for amd_llvm
This commit is contained in:
Ignacio Sica
2025-04-24 21:29:27 -03:00
committed by GitHub
parent b4f823acbe
commit 76a86735c0
2 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
test/test_linearizer.py
View File

@@ -1061,6 +1061,7 @@ class TestLinearizer(unittest.TestCase):
@unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
def test_tensor_cores(self):
for tc in Device[Device.DEFAULT].renderer.tensor_cores:
if CI and getenv("AMD_LLVM") and tc.dtype_in is dtypes.bfloat16: continue # TODO: compilation error in CI
if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue
# for AMX, tc.dims[2] == 1 so reduceop is None thus tensor_cores are not triggered
helper_tc_allclose(tc.dims[0], tc.dims[1], 2 if AMX else tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
@@ -1099,6 +1100,7 @@ class TestLinearizer(unittest.TestCase):
@unittest.expectedFailure
def test_tensor_cores_padded_amd(self):
for tc in Device[Device.DEFAULT].renderer.tensor_cores:
if CI and getenv("AMD_LLVM") and tc.dtype_in is dtypes.bfloat16: continue # TODO: compilation error in CI
if not is_dtype_supported(tc.dtype_in) or not is_dtype_supported(tc.dtype_out): continue
helper_tc_allclose(tc.dims[0]+(pad:=1), tc.dims[1]+pad, tc.dims[2]+pad, tc.dtype_in, tc.dtype_out, tc_opt=2)

4
tinygrad/device.py
View File

@@ -330,10 +330,10 @@ class Compiled:
def is_dtype_supported(dtype:DType, device:Optional[str]=None) -> bool:
if device is None: device = Device.DEFAULT
if dtype == dtypes.bfloat16:
if device in {"METAL", "AMD", "AMD_LLVM"}: return not CI
if device == "METAL": return not CI
if device in {"CUDA", "NV"}: return not CI and not getenv("PTX")
if device in {"CPU", "LLVM"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"}
return False
return device == "AMD"
if dtype in dtypes.fp8s:
# not supported yet - in progress
return False