Try FP16 emulation

Author: davschneller

tensorforge/backend/instructions/compute/primitives/amd.py

@@ -408,18 +408,11 @@ def mfma_emu_bf16_f32(writer: Writer, C, B, A, c, a, b):
     writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k({A[0]}_p1, {B[0]}_p1, {C}, {c}, {a}, {b});')
 
 def mfma_emu_f16_f32(writer: Writer, C, B, A, c, a, b):
-    Ar = writer.varalloc()
-    A1 = writer.varalloc()
-    A2 = writer.varalloc()
-    Br = writer.varalloc()
-    B1 = writer.varalloc()
-    B2 = writer.varalloc()
-    writer(f'const f16x4 {Ar} = f16x4({A});')
-    writer(f'const f16x4 {Br} = f16x4({B});')
-    writer(f'const f16x4 {A1} = f16x4({A});')
-    writer(f'const f16x4 {B1} = f16x4({B});')
-    writer(f'const f16x4 {A2} = f16x4({A} - {A1});')
-    writer(f'const f16x4 {B2} = f16x4({B} - {B1});')
+    writer(f'const auto [{A[0]}_p0, {A[0]}_p1] = tensorforge::splitFloatx4F16({A[0]}, {A[1]}, {A[2]}, {A[3]});')
+    writer(f'const auto [{B[0]}_p0, {B[0]}_p1] = tensorforge::splitFloatx4F16({B[0]}, {B[1]}, {B[2]}, {B[3]});')
+    writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4f16({A[0]}_p0, {B[0]}_p0, {C}, {c}, {a}, {b});')
+    writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4f16({A[0]}_p1, {B[0]}_p0, {C}, {c}, {a}, {b});')
+    writer(f'{C} = __builtin_amdgcn_mfma_f32_4x4x4f16({A[0]}_p0, {B[0]}_p1, {C}, {c}, {a}, {b});')
 
 def matmul32(writer: Writer, C, B, A, M, N, K, kx, threads):
     with writer.AnonymousScope():
@@ -486,10 +479,14 @@ def write_matmul(block, start, cap):
                                             fB[kkk] = B(writer, f'{tmpB}_{kkk}', i, k + kk + kkk)
                                         for kkk in range(dkk, block):
                                             writer(f'float {tmpB}_{kkk} = 0;')
-                                        if True:
+                                        if False:
                                             Ar = [f'{tmpA}_{k // threads}_{kkk}' for kkk in range(4)]
                                             Br = [f'{tmpB}_{kkk}' for kkk in range(4)]
                                             mfma_emu_bf16_f32(writer, tmpacc, Br, Ar, scale, kk // 4, 0)
+                                        elif True:
+                                            Ar = [f'{tmpA}_{k // threads}_{kkk}' for kkk in range(4)]
+                                            Br = [f'{tmpB}_{kkk}' for kkk in range(4)]
+                                            mfma_emu_f16_f32(writer, tmpacc, Br, Ar, scale, kk // 4, 0)
                                         else:
                                             for kkk in range(dkk):
                                                 if fB[kkk]:

tensorforge/include/tensorforge_device/hip.h

@@ -889,4 +889,23 @@ __device__ __forceinline__
           VectorT<short, 4>{i1p2, i2p2, i3p2, i4p2}};
 }
 
+__device__ __forceinline__ std::tuple<_Float16, _Float16>
+splitFloatF16(float input) {
+  const auto i1 = static_cast<_Float16>(input);
+  const auto i1r = input - static_cast<float>(i1);
+  const auto i2 = static_cast<_Float16>(i1r);
+  return {i1, i2};
+}
+
+__device__
+    __forceinline__ std::tuple<VectorT<_Float16, 4>, VectorT<_Float16, 4>>
+    splitFloatx4F16(float i1, float i2, float i3, float i4) {
+  const auto [i1p0, i1p1] = splitFloatF16(i1);
+  const auto [i2p0, i2p1] = splitFloatF16(i2);
+  const auto [i3p0, i3p1] = splitFloatF16(i3);
+  const auto [i4p0, i4p1] = splitFloatF16(i4);
+  return {VectorT<_Float16, 4>{i1p0, i2p0, i3p0, i4p0},
+          VectorT<_Float16, 4>{i1p1, i2p1, i3p1, i4p1}};
+}
+
 } // namespace tensorforge