wmma: add HIP FP16 to FP16 tensor core (#3287)

* wmma: add HIP FP16 to FP16 tensor core

* test: fix test_tensor_core to use separate tolerances for half

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

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
   ]
 }
 

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}