diff --git a/test/test_linearizer.py b/test/test_linearizer.py index cf033f5a91..3da845c278 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -116,7 +116,8 @@ class TestLinearizer(unittest.TestCase): k.linearize() assert len([uop for uop in k.uops if uop.uop == UOps.WMMA]) == 1, "tensor core not triggered" np_c = np_a @ np_b - np.testing.assert_allclose(np_c, r.numpy(), atol=5e-3, rtol=1e-4) + (tc_atol, tc_rtol) = (1e-2, 1e-3) if tc.dtype_out == dtypes.half else (5e-3, 1e-4) + np.testing.assert_allclose(np_c, r.numpy(), atol=tc_atol, rtol=tc_rtol) def test_limit_dims_to_max_5d_global(self): t = Tensor.rand(3, 4, 5, 6, 7).pad(((1, 1), (1, 1), (1, 1), (1, 1), (1, 1))) + 1 diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index fb7e82a77c..fafe356c27 100644 --- a/tinygrad/codegen/kernel.py +++ b/tinygrad/codegen/kernel.py @@ -44,6 +44,7 @@ tensor_cores: Dict[str, List[TensorCore]] = { ], "HIP": [ TensorCore(dims=[16,16,16], dtype_in=dtypes.half, dtype_out=dtypes.float, wmma_func="__builtin_amdgcn_wmma_f32_16x16x16_f16_w32", upcast_dim=1, threads=[(0,16),(1,2)], thread_local_sizes=[16,16,8], thread_local_aliases=[ [[0],[0],[-1],[1]], [[0],[1],[-1],[0]], [[0],[1],[0],[2,-1]] ]), # noqa: E501 + TensorCore(dims=[16,16,16], dtype_in=dtypes.half, dtype_out=dtypes.half, wmma_func="__hip_wmma_f16_f16", upcast_dim=1, threads=[(0,16),(1,2)], thread_local_sizes=[16,16,8], thread_local_aliases=[ [[0],[0],[-1],[1]], [[0],[1],[-1],[0]], [[0],[1],[0],[2,-1]] ]), # noqa: E501 ] } diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index bd6cb542ec..318ccb6e13 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -273,8 +273,12 @@ class HIPLanguage(CStyleLanguage): __attribute__((device)) __attribute__((const)) _Float16 __ocml_sqrt_f16(_Float16); }\n""" + '\n'.join([_make_hip_dtype(*x) for x in [("signed int", "int", 2), ("_Float16", "half", 2), ("_Float16", "half", 4), ("_Float16", "half", 8), ("_Float16", "half", 16), - ("float", "float", 2), ("float", "float", 4), ("float", "float", 8)]]) + \ - 'extern "C" __attribute__((global))' + ("float", "float", 2), ("float", "float", 4), ("float", "float", 8)]]) + """ + static __attribute__((device)) half8 __hip_wmma_f16_f16(half16 a, half16 b, half8 c) { + half16 c_frag = {}; half8 d; for (int n = 0; n < 8; n++) { c_frag[n*2] = c[n]; } + c_frag = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32(a, b, c_frag, false); + for (int n = 0; n < 8; n++) { d[n] = c_frag[n*2]; } return d; + }\nextern "C" __attribute__((global))""" code_for_workitem = {"g": lambda x: f"__ockl_get_group_id({x})", "l": lambda x: f"__ockl_get_local_id({x})", "i": lambda x: f"(__ockl_get_group_id({x})*__ockl_get_local_size({x})+__ockl_get_local_id({x}))"} code_for_op = {**CStyleLanguage().code_for_op, **code_for_op_hip}