From 05aa63e738edeb06ad2697a04f8889f6c0746067 Mon Sep 17 00:00:00 2001 From: Chip Kerchner Date: Tue, 24 Sep 2024 12:54:02 -0500 Subject: [PATCH] More MMA BF16 GEMV code. --- kernel/power/sbgemv_common_power10.c | 179 ++++++++++++++++++- kernel/power/sbgemv_n_power10.c | 246 ++++++++++++++++++++++----- 2 files changed, 372 insertions(+), 53 deletions(-) diff --git a/kernel/power/sbgemv_common_power10.c b/kernel/power/sbgemv_common_power10.c index da088014b..2ee912b9d 100644 --- a/kernel/power/sbgemv_common_power10.c +++ b/kernel/power/sbgemv_common_power10.c @@ -29,6 +29,10 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #define SBGEMV_COMMON_MMA_C #include "sbgemv_common.c" +#if defined(_AIX) || defined(__clang__) +#define USE_MERGE_MMA +#endif + FORCEINLINE void vec_load_mult_mma(__vector_quad *out, vec_bf16 *in, vec_bf16 inp) { vec_bf16 in0 = (vec_bf16)vec_load_vec(in); @@ -69,11 +73,13 @@ FORCEINLINE void vec_mult2_mma(__vector_quad *out, vec_bf16 in0, vec_bf16 inp) __builtin_mma_xvbf16ger2(&out[1], (vec_uc8)inp, (vec_uc8)in01); } +#ifndef USE_MERGE_MMA FORCEINLINE void vec_mult4_mma(__vector_quad *out, vec_bf16 *in0, vec_bf16 inp) { vec_mult2_mma(out + 0, in0[0], inp); vec_mult2_mma(out + 2, in0[1], inp); } +#endif FORCEINLINE void vec_loadN_mult11_mma(__vector_quad *out, vec_bf16 *in, vec_bf16 inp, BLASLONG n) { @@ -96,6 +102,7 @@ FORCEINLINE void vec_load_mult12_mma(__vector_quad *out, vec_bf16 *in, vec_bf16 vec_mult2_mma(out, in0, inp); } +#ifndef USE_MERGE_MMA FORCEINLINE void vec_load_mult18_mma(__vector_quad *out, vec_bf16 *in, vec_bf16 inp) { vec_bf16 in0[4]; @@ -106,6 +113,7 @@ FORCEINLINE void vec_load_mult18_mma(__vector_quad *out, vec_bf16 *in, vec_bf16 vec_mult4_mma(&out[0], in0 + 0, inp); vec_mult4_mma(&out[4], in0 + 2, inp); } +#endif FORCEINLINE void vec_reduce1_mma(__vector_quad *out, vec_f32 *temp, vec_f32 v_alpha, vec_f32 *vy0) { @@ -120,6 +128,7 @@ FORCEINLINE void vec_reduce2_mma(__vector_quad *out, vec_f32 *temp, vec_f32 v_al vec_reduce1_mma(&out[1], &temp[4], v_alpha, &vy0[1]); } +#ifndef USE_MERGE_MMA FORCEINLINE void vec_reduce8_mma(__vector_quad *out, vec_f32 *temp, vec_f32 v_alpha, vec_f32 *vy0) { vec_reduce2_mma(&out[0], &temp[0], v_alpha, vy0 + 0); @@ -127,6 +136,23 @@ FORCEINLINE void vec_reduce8_mma(__vector_quad *out, vec_f32 *temp, vec_f32 v_al vec_reduce2_mma(&out[4], &temp[16], v_alpha, vy0 + 4); vec_reduce2_mma(&out[6], &temp[24], v_alpha, vy0 + 6); } +#else +FORCEINLINE void vec_reduce44_mma(__vector_quad *out, vec_f32 *temp, vec_f32 v_alpha, vec_f32 *vy0) +{ + __builtin_mma_disassemble_acc((void*)temp, &out[0]); + + vy0[0] += (temp[0] * v_alpha); + vy0[2] += (temp[1] * v_alpha); + vy0[4] += (temp[2] * v_alpha); + vy0[6] += (temp[3] * v_alpha); +} + +FORCEINLINE void vec_reduce84_mma(__vector_quad *out, vec_f32 *temp, vec_f32 v_alpha, vec_f32 *vy0) +{ + vec_reduce44_mma(&out[0], &temp[0], v_alpha, vy0 + 0); + vec_reduce44_mma(&out[1], &temp[4], v_alpha, vy0 + 1); +} +#endif FORCEINLINE void vec_mult11a_mma(__vector_quad *out, vec_bf16 in0, vec_bf16 in1, vec_bf16 inp) { @@ -166,18 +192,25 @@ FORCEINLINE void vec_load_mult22a_mma(__vector_quad *out, vec_bf16 *ina, vec_bf1 vec_mult2a_mma(out, in0, in1, inp); } -FORCEINLINE void vec_load_mult28a_mma(__vector_quad *out, vec_bf16 *ina, vec_bf16 *inb, vec_bf16 inp) +FORCEINLINE void vec_load4_mma(vec_bf16 *in0, vec_bf16 *in1, vec_bf16 *ina, vec_bf16 *inb) { - vec_bf16 in0[4], in1[4]; - vec_load_pair((vec_f32 *)(in0 + 0), (vec_f32 *)(ina + 0)); vec_load_pair((vec_f32 *)(in1 + 0), (vec_f32 *)(inb + 0)); vec_load_pair((vec_f32 *)(in0 + 2), (vec_f32 *)(ina + 2)); vec_load_pair((vec_f32 *)(in1 + 2), (vec_f32 *)(inb + 2)); +} + +#ifndef USE_MERGE_MMA +FORCEINLINE void vec_load_mult28a_mma(__vector_quad *out, vec_bf16 *ina, vec_bf16 *inb, vec_bf16 inp) +{ + vec_bf16 in0[4], in1[4]; + + vec_load4_mma(in0, in1, ina, inb); vec_mult4a_mma(&out[0], in0 + 0, in1 + 0, inp); vec_mult4a_mma(&out[4], in0 + 2, in1 + 2, inp); } +#endif FORCEINLINE void vec_loadN_mult22a_mma(__vector_quad *out, vec_bf16 *ina, vec_bf16 *inb, vec_bf16 inp, BLASLONG n) { @@ -209,6 +242,48 @@ FORCEINLINE void vec_mult4b_mma(__vector_quad *out, vec_bf16 *in0, vec_bf16 *in1 vec_mult2b_mma(out + 2, in0[1], in1[1], inp); } +#ifdef USE_MERGE_MMA +FORCEINLINE void vec_mult1c_mma(__vector_quad *out, vec_bf16 in0, vec_bf16 inp) +{ + vec_bf16 in00 = vec_mergeh(in0, in0); + + __builtin_mma_xvbf16ger2pp(out, (vec_uc8)inp, (vec_uc8)in00); +} + +FORCEINLINE void vec_mult2c_mma(__vector_quad *out, vec_bf16 in0, vec_bf16 inp) +{ + vec_bf16 in01 = vec_mergel(in0, in0); + + vec_mult1c_mma(&out[0], in0, inp); + + __builtin_mma_xvbf16ger2pp(&out[1], (vec_uc8)inp, (vec_uc8)in01); +} + +FORCEINLINE void vec_mult44_mma(__vector_quad *out, vec_bf16 *in, vec_bf16 *inp) +{ + vec_mult2_mma(out, in[0], inp[0]); + vec_mult2c_mma(out, in[1], inp[1]); +} + +FORCEINLINE void vec_mult44c_mma(__vector_quad *out, vec_bf16 *in, vec_bf16 *inp) +{ + vec_mult2c_mma(out, in[0], inp[0]); + vec_mult2c_mma(out, in[1], inp[1]); +} + +FORCEINLINE void vec_mult44a_mma(__vector_quad *out, vec_bf16 *in0, vec_bf16 *in1, vec_bf16 *inp) +{ + vec_mult2a_mma(out, in0[0], in1[0], inp[0]); + vec_mult2b_mma(out, in0[1], in1[1], inp[1]); +} + +FORCEINLINE void vec_mult44b_mma(__vector_quad *out, vec_bf16 *in0, vec_bf16 *in1, vec_bf16 *inp) +{ + vec_mult2b_mma(out, in0[0], in1[0], inp[0]); + vec_mult2b_mma(out, in0[1], in1[1], inp[1]); +} +#endif + FORCEINLINE void vec_loadN_mult11b_mma(__vector_quad *out, vec_bf16 *ina, vec_bf16 *inb, vec_bf16 inp, BLASLONG n) { vec_bf16 in0 = vec_loadN(ina, n); @@ -225,18 +300,48 @@ FORCEINLINE void vec_load_mult22b_mma(__vector_quad *out, vec_bf16 *ina, vec_bf1 vec_mult2b_mma(out, in0, in1, inp); } +#ifndef USE_MERGE_MMA FORCEINLINE void vec_load_mult28b_mma(__vector_quad *out, vec_bf16 *ina, vec_bf16 *inb, vec_bf16 inp) { vec_bf16 in0[4], in1[4]; - vec_load_pair((vec_f32 *)(in0 + 0), (vec_f32 *)(ina + 0)); - vec_load_pair((vec_f32 *)(in1 + 0), (vec_f32 *)(inb + 0)); - vec_load_pair((vec_f32 *)(in0 + 2), (vec_f32 *)(ina + 2)); - vec_load_pair((vec_f32 *)(in1 + 2), (vec_f32 *)(inb + 2)); + vec_load4_mma(in0, in1, ina, inb); vec_mult4b_mma(&out[0], in0 + 0, in1 + 0, inp); vec_mult4b_mma(&out[4], in0 + 2, in1 + 2, inp); } +#else +FORCEINLINE void vec_load_mult184_mma(__vector_quad *out, vec_bf16 *in, vec_bf16 *inp) +{ + vec_bf16 in0[4]; + + vec_load_pair((vec_f32 *)(in0 + 0), (vec_f32 *)(in + 0)); + vec_load_pair((vec_f32 *)(in0 + 2), (vec_f32 *)(in + 2)); + + vec_mult44_mma(out, in0 + 0, inp + 0); + vec_mult44c_mma(out, in0 + 2, inp + 2); +} + +FORCEINLINE void vec_load_mult284a_mma(__vector_quad *out, vec_bf16 *ina, vec_bf16 *inb, vec_bf16 *inp) +{ + vec_bf16 in0[4], in1[4]; + + vec_load4_mma(in0, in1, ina, inb); + + vec_mult44a_mma(out, in0 + 0, in1 + 0, inp + 0); + vec_mult44b_mma(out, in0 + 2, in1 + 2, inp + 2); +} + +FORCEINLINE void vec_load_mult284b_mma(__vector_quad *out, vec_bf16 *ina, vec_bf16 *inb, vec_bf16 *inp) +{ + vec_bf16 in0[4], in1[4]; + + vec_load4_mma(in0, in1, ina, inb); + + vec_mult44b_mma(out, in0 + 0, in1 + 0, inp + 0); + vec_mult44b_mma(out, in0 + 2, in1 + 2, inp + 2); +} +#endif FORCEINLINE void vec_loadN_mult22b_mma(__vector_quad *out, vec_bf16 *ina, vec_bf16 *inb, vec_bf16 inp, BLASLONG n) { @@ -262,4 +367,64 @@ FORCEINLINE void vec_store4_pair(vec_f32 *v_y, vec_f32 *vy0) vec_store_pair(v_y + 6, vy0 + 6); } +#ifdef USE_MERGE_MMA +FORCEINLINE void vec_load8_pair(vec_f32 *vy0, vec_f32 *v_y) +{ + vec_load4_pair(vy0 + 0, v_y + 0); + vec_load4_pair(vy0 + 8, v_y + 8); +} + +FORCEINLINE void vec_store8_pair(vec_f32 *v_y, vec_f32 *vy0) +{ + vec_store4_pair(v_y + 0, vy0 + 0); + vec_store4_pair(v_y + 8, vy0 + 8); +} + +#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__ +#define VEC_SHIFT(data, shift) vec_sld(data, data, 16 - shift) +#else +#define VEC_SHIFT(data, shift) vec_sld(data, data, shift) +#endif + +typedef __vector unsigned int vec_ui32; + +static vec_ui32 mask_0 = { 0xffffffff, 0x00000000, 0x00000000, 0x00000000 }; +static vec_ui32 mask_1 = { 0x00000000, 0xffffffff, 0x00000000, 0x00000000 }; +static vec_ui32 mask_2 = { 0x00000000, 0x00000000, 0xffffffff, 0x00000000 }; +static vec_ui32 mask_3 = { 0x00000000, 0x00000000, 0x00000000, 0xffffffff }; + +FORCEINLINE void vec_make_mult1(vec_bf16 *v_x0) +{ + v_x0[ 0] = vec_and(v_x0[0], (vec_bf16)mask_0); + + v_x0[ 1] = VEC_SHIFT(v_x0[ 0], 4); + v_x0[ 2] = VEC_SHIFT(v_x0[ 0], 8); + v_x0[ 3] = VEC_SHIFT(v_x0[ 0], 12); +} + +FORCEINLINE void vec_make_mult2(vec_bf16 *v_x0) +{ + v_x0[ 5] = vec_and(v_x0[0], (vec_bf16)mask_1); + vec_make_mult1(v_x0); + + v_x0[ 4] = VEC_SHIFT(v_x0[ 5], 12); + v_x0[ 6] = VEC_SHIFT(v_x0[ 5], 4); + v_x0[ 7] = VEC_SHIFT(v_x0[ 5], 8); +} + +FORCEINLINE void vec_make_mult4(vec_bf16 *v_x0) +{ + v_x0[10] = vec_and(v_x0[0], (vec_bf16)mask_2); + v_x0[15] = vec_and(v_x0[0], (vec_bf16)mask_3); + vec_make_mult2(v_x0); + + v_x0[ 8] = VEC_SHIFT(v_x0[10], 8); + v_x0[ 9] = VEC_SHIFT(v_x0[10], 12); + v_x0[11] = VEC_SHIFT(v_x0[10], 4); + v_x0[12] = VEC_SHIFT(v_x0[15], 4); + v_x0[13] = VEC_SHIFT(v_x0[15], 8); + v_x0[14] = VEC_SHIFT(v_x0[15], 12); +} +#endif + #endif diff --git a/kernel/power/sbgemv_n_power10.c b/kernel/power/sbgemv_n_power10.c index 7b2beb0c7..f2ed6bf9a 100644 --- a/kernel/power/sbgemv_n_power10.c +++ b/kernel/power/sbgemv_n_power10.c @@ -28,7 +28,9 @@ USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. #ifndef SBGEMV_N_MMA_C #define SBGEMV_N_MMA_C +#if !defined(_AIX) || defined(__clang__) #define USE_BFGEMV_N_MMA +#endif #ifdef USE_BFGEMV_N_MMA #include "sbgemv_common_power10.c" @@ -47,7 +49,7 @@ static void BF16GEMV_N_MMA_1(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA { IFLOAT *a0; __vector_quad temp[2*4]; - vec_f32 temp0[8*4], vy0[2*4]; + vec_f32 temp0[8*4]; vec_f32 v_alpha = { alpha, alpha, alpha, alpha }; a0 = ap[0]; @@ -55,26 +57,61 @@ static void BF16GEMV_N_MMA_1(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA vec_bf16 *va0 = (vec_bf16 *)a0; vec_bf16 *x_bf = (vec_bf16 *)(xo); - vec_bf16 v_x0 = vec_loadN(x_bf, 1); vec_f32 *v_y = (vec_f32 *)y; BLASLONG n8 = n / 8; BLASLONG i = 0; +#ifdef USE_MERGE_MMA + vec_bf16 v_x0[4]; + v_x0[0] = vec_loadN(x_bf, 1); + vec_f32 vy0[2*4*2]; + + vec_make_mult1(v_x0); + + for (; i + 8 <= n8; i += 8) { + vec_load8_pair(vy0, &v_y[(i * 2) + 0]); + + vec_load_mult184_mma(&temp[0], &va0[i + 0], &v_x0[ 0]); + vec_load_mult184_mma(&temp[2], &va0[i + 4], &v_x0[ 0]); + + vec_reduce84_mma(&temp[0], temp0 + 0, v_alpha, vy0 + 0); + vec_reduce84_mma(&temp[2], temp0 + 8, v_alpha, vy0 + 8); + + vec_store8_pair(&v_y[(i * 2) + 0], vy0); + } + + if (n8 & 4) { + vec_load4_pair(vy0, &v_y[(i * 2) + 0]); + + vec_load_mult184_mma(&temp[0], &va0[i + 0], &v_x0[ 0]); + + vec_reduce84_mma(&temp[0], temp0, v_alpha, vy0); + + vec_store4_pair(&v_y[(i * 2) + 0], vy0); + + i += 4; + } +#else + vec_bf16 v_x0[1]; + v_x0[0] = vec_loadN(x_bf, 1); + vec_f32 vy0[2*4]; + for (; i + 4 <= n8; i += 4) { vec_load4_pair(vy0, &v_y[(i * 2) + 0]); - vec_load_mult18_mma(&temp[0], &va0[i + 0], v_x0); + vec_load_mult18_mma(&temp[0], &va0[i + 0], v_x0[ 0]); vec_reduce8_mma(&temp[0], temp0, v_alpha, vy0); vec_store4_pair(&v_y[(i * 2) + 0], vy0); } +#endif for (; i < n8; i++) { vec_load_pair(vy0, &v_y[(i * 2) + 0]); - vec_load_mult12_mma(&temp[0], &va0[i], v_x0); + vec_load_mult12_mma(&temp[0], &va0[i], v_x0[ 0]); vec_reduce2_mma(&temp[0], temp0, v_alpha, vy0); @@ -86,7 +123,7 @@ static void BF16GEMV_N_MMA_1(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA BLASLONG n3 = n & 3; vec_loadN2_f32(vy0, &v_y[(i * 2) + 0], n3); - vec_loadN_mult12_mma(&temp[0], &va0[i], v_x0, n); + vec_loadN_mult12_mma(&temp[0], &va0[i], v_x0[ 0], n); vec_reduce2_mma(&temp[0], temp0, v_alpha, vy0); @@ -94,7 +131,7 @@ static void BF16GEMV_N_MMA_1(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA } else if (n) { vy0[0] = vec_loadN_f32(&v_y[(i * 2) + 0], n); - vec_loadN_mult11_mma(&temp[0], &va0[i], v_x0, n); + vec_loadN_mult11_mma(&temp[0], &va0[i], v_x0[ 0], n); vec_reduce1_mma(&temp[0], temp0, v_alpha, vy0); @@ -106,7 +143,7 @@ static void BF16GEMV_N_MMA_2(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA { IFLOAT *a0, *a1; __vector_quad temp[2*4]; - vec_f32 temp0[8*4], vy0[2*4]; + vec_f32 temp0[8*4]; vec_f32 v_alpha = { alpha, alpha, alpha, alpha }; a0 = ap[0]; @@ -116,26 +153,61 @@ static void BF16GEMV_N_MMA_2(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA vec_bf16 *va1 = (vec_bf16 *)a1; vec_bf16 *x_bf = (vec_bf16 *)(xo); - vec_bf16 v_x0 = vec_loadN(x_bf, 2); vec_f32 *v_y = (vec_f32 *)y; BLASLONG n8 = n / 8; BLASLONG i = 0; +#ifdef USE_MERGE_MMA + vec_bf16 v_x0[4]; + vec_f32 vy0[2*4*2]; + v_x0[0] = vec_loadN(x_bf, 2); + + vec_make_mult1(v_x0); + + for (; i + 8 <= n8; i += 8) { + vec_load8_pair(vy0, &v_y[(i * 2) + 0]); + + vec_load_mult284a_mma(&temp[0], &va0[i + 0], &va1[i + 0], &v_x0[ 0]); + vec_load_mult284a_mma(&temp[2], &va0[i + 4], &va1[i + 4], &v_x0[ 0]); + + vec_reduce84_mma(&temp[0], temp0 + 0, v_alpha, vy0 + 0); + vec_reduce84_mma(&temp[2], temp0 + 8, v_alpha, vy0 + 8); + + vec_store8_pair(&v_y[(i * 2) + 0], vy0); + } + + if (n8 & 4) { + vec_load4_pair(vy0, &v_y[(i * 2) + 0]); + + vec_load_mult284a_mma(&temp[0], &va0[i + 0], &va1[i + 0], &v_x0[ 0]); + + vec_reduce84_mma(&temp[0], temp0, v_alpha, vy0); + + vec_store4_pair(&v_y[(i * 2) + 0], vy0); + + i += 4; + } +#else + vec_bf16 v_x0[1]; + vec_f32 vy0[2*4]; + v_x0[0] = vec_loadN(x_bf, 2); + for (; i + 4 <= n8; i += 4) { vec_load4_pair(vy0, &v_y[(i * 2) + 0]); - vec_load_mult28a_mma(&temp[0], &va0[i + 0], &va1[i + 0], v_x0); + vec_load_mult28a_mma(&temp[0], &va0[i + 0], &va1[i + 0], v_x0[ 0]); vec_reduce8_mma(&temp[0], temp0, v_alpha, vy0); vec_store4_pair(&v_y[(i * 2) + 0], vy0); } +#endif for (; i < n8; i++) { vec_load_pair(vy0, &v_y[(i * 2) + 0]); - vec_load_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x0); + vec_load_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0]); vec_reduce2_mma(&temp[0], temp0, v_alpha, vy0); @@ -147,7 +219,7 @@ static void BF16GEMV_N_MMA_2(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA BLASLONG n3 = n & 3; vec_loadN2_f32(vy0, &v_y[(i * 2) + 0], n3); - vec_loadN_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x0, n); + vec_loadN_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0], n); vec_reduce2_mma(&temp[0], temp0, v_alpha, vy0); @@ -155,7 +227,7 @@ static void BF16GEMV_N_MMA_2(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA } else if (n) { vy0[0] = vec_loadN_f32(&v_y[(i * 2) + 0], n); - vec_loadN_mult11a_mma(&temp[0], &va0[i], &va1[i], v_x0, n); + vec_loadN_mult11a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0], n); vec_reduce1_mma(&temp[0], temp0, v_alpha, vy0); @@ -167,7 +239,7 @@ static void BF16GEMV_N_MMA_4(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA { IFLOAT *a0, *a1, *a2, *a3; __vector_quad temp[2*4]; - vec_f32 temp0[8*4], vy0[2*4]; + vec_f32 temp0[8*4]; vec_f32 v_alpha = { alpha, alpha, alpha, alpha }; a0 = ap[0]; @@ -181,30 +253,68 @@ static void BF16GEMV_N_MMA_4(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA vec_bf16 *va3 = (vec_bf16 *)a3; vec_bf16 *x_bf = (vec_bf16 *)(xo); - vec_bf16 v_x00 = vec_loadN(x_bf, 4); - - vec_bf16 v_x01 = (vec_bf16)vec_splat((vec_f32)v_x00, 1); vec_f32 *v_y = (vec_f32 *)y; BLASLONG n8 = n / 8; BLASLONG i = 0; +#ifdef USE_MERGE_MMA + vec_bf16 v_x0[8]; + vec_f32 vy0[2*4*2]; + v_x0[0] = vec_loadN(x_bf, 4); + + vec_make_mult2(v_x0); + + for (; i + 8 <= n8; i += 8) { + vec_load8_pair(vy0, &v_y[(i * 2) + 0]); + + vec_load_mult284a_mma(&temp[0], &va0[i + 0], &va1[i + 0], &v_x0[ 0]); + vec_load_mult284b_mma(&temp[0], &va2[i + 0], &va3[i + 0], &v_x0[ 4]); + vec_load_mult284a_mma(&temp[2], &va0[i + 4], &va1[i + 4], &v_x0[ 0]); + vec_load_mult284b_mma(&temp[2], &va2[i + 4], &va3[i + 4], &v_x0[ 4]); + + vec_reduce84_mma(&temp[0], temp0 + 0, v_alpha, vy0 + 0); + vec_reduce84_mma(&temp[2], temp0 + 8, v_alpha, vy0 + 8); + + vec_store8_pair(&v_y[(i * 2) + 0], vy0); + } + + if (n8 & 4) { + vec_load4_pair(vy0, &v_y[(i * 2) + 0]); + + vec_load_mult284a_mma(&temp[0], &va0[i + 0], &va1[i + 0], &v_x0[ 0]); + vec_load_mult284b_mma(&temp[0], &va2[i + 0], &va3[i + 0], &v_x0[ 4]); + + vec_reduce84_mma(&temp[0], temp0, v_alpha, vy0); + + vec_store4_pair(&v_y[(i * 2) + 0], vy0); + + i += 4; + } +#else + vec_bf16 v_x0[5]; + vec_f32 vy0[2*4]; + v_x0[0] = vec_loadN(x_bf, 4); + + v_x0[ 4] = (vec_bf16)vec_splat((vec_f32)v_x0[0], 1); + for (; i + 4 <= n8; i += 4) { vec_load4_pair(vy0, &v_y[(i * 2) + 0]); - vec_load_mult28a_mma(&temp[0], &va0[i + 0], &va1[i + 0], v_x00); - vec_load_mult28b_mma(&temp[0], &va2[i + 0], &va3[i + 0], v_x01); + vec_load_mult28a_mma(&temp[0], &va0[i + 0], &va1[i + 0], v_x0[ 0]); + vec_load_mult28b_mma(&temp[0], &va2[i + 0], &va3[i + 0], v_x0[ 4]); vec_reduce8_mma(&temp[0], temp0, v_alpha, vy0); vec_store4_pair(&v_y[(i * 2) + 0], vy0); } +#endif for (; i < n8; i++) { vec_load_pair(vy0, &v_y[(i * 2) + 0]); - vec_load_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x00); - vec_load_mult22b_mma(&temp[0], &va2[i], &va3[i], v_x01); + vec_load_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0]); + vec_load_mult22b_mma(&temp[0], &va2[i], &va3[i], v_x0[ 4]); vec_reduce2_mma(&temp[0], temp0, v_alpha, vy0); @@ -216,8 +326,8 @@ static void BF16GEMV_N_MMA_4(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA BLASLONG n3 = n & 3; vec_loadN2_f32(vy0, &v_y[(i * 2) + 0], n3); - vec_loadN_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x00, n); - vec_loadN_mult22b_mma(&temp[0], &va2[i], &va3[i], v_x01, n); + vec_loadN_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0], n); + vec_loadN_mult22b_mma(&temp[0], &va2[i], &va3[i], v_x0[ 4], n); vec_reduce2_mma(&temp[0], temp0, v_alpha, vy0); @@ -225,8 +335,8 @@ static void BF16GEMV_N_MMA_4(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, FLOA } else if (n) { vy0[0] = vec_loadN_f32(&v_y[(i * 2) + 0], n); - vec_loadN_mult11a_mma(&temp[0], &va0[i], &va1[i], v_x00, n); - vec_loadN_mult11b_mma(&temp[0], &va2[i], &va3[i], v_x01, n); + vec_loadN_mult11a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0], n); + vec_loadN_mult11b_mma(&temp[0], &va2[i], &va3[i], v_x0[ 4], n); vec_reduce1_mma(&temp[0], temp0, v_alpha, vy0); @@ -239,7 +349,7 @@ static void BF16GEMV_N_MMA_8(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, BLAS { IFLOAT *a0, *a1, *a2, *a3, *b0, *b1, *b2, *b3; __vector_quad temp[2*4]; - vec_f32 temp0[8*4], vy0[2*4]; + vec_f32 temp0[8*4]; vec_f32 v_alpha = { alpha, alpha, alpha, alpha }; a0 = ap[0]; @@ -261,36 +371,80 @@ static void BF16GEMV_N_MMA_8(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, BLAS vec_bf16 *vb3 = (vec_bf16 *)b3; vec_bf16 *x_bf = (vec_bf16 *)(xo); - vec_bf16 v_x00 = (vec_bf16)vec_load_vec(x_bf); - - vec_bf16 v_x01 = (vec_bf16)vec_splat((vec_f32)v_x00, 1); - vec_bf16 v_x02 = (vec_bf16)vec_splat((vec_f32)v_x00, 2); - vec_bf16 v_x03 = (vec_bf16)vec_splat((vec_f32)v_x00, 3); vec_f32 *v_y = (vec_f32 *)y; BLASLONG n8 = n / 8; BLASLONG i = 0; +#ifdef USE_MERGE_MMA + vec_bf16 v_x0[16]; + vec_f32 vy0[2*4*2]; + v_x0[0] = (vec_bf16)vec_load_vec(x_bf); + + vec_make_mult4(v_x0); + + for (; i + 8 <= n8; i += 8) { + vec_load8_pair(vy0, &v_y[(i * 2) + 0]); + + vec_load_mult284a_mma(&temp[0], &va0[i + 0], &va1[i + 0], &v_x0[ 0]); + vec_load_mult284b_mma(&temp[0], &va2[i + 0], &va3[i + 0], &v_x0[ 4]); + vec_load_mult284b_mma(&temp[0], &vb0[i + 0], &vb1[i + 0], &v_x0[ 8]); + vec_load_mult284b_mma(&temp[0], &vb2[i + 0], &vb3[i + 0], &v_x0[12]); + vec_load_mult284a_mma(&temp[2], &va0[i + 4], &va1[i + 4], &v_x0[ 0]); + vec_load_mult284b_mma(&temp[2], &va2[i + 4], &va3[i + 4], &v_x0[ 4]); + vec_load_mult284b_mma(&temp[2], &vb0[i + 4], &vb1[i + 4], &v_x0[ 8]); + vec_load_mult284b_mma(&temp[2], &vb2[i + 4], &vb3[i + 4], &v_x0[12]); + + vec_reduce84_mma(&temp[0], temp0 + 0, v_alpha, vy0 + 0); + vec_reduce84_mma(&temp[2], temp0 + 8, v_alpha, vy0 + 8); + + vec_store8_pair(&v_y[(i * 2) + 0], vy0); + } + + if (n8 & 4) { + vec_load4_pair(vy0, &v_y[(i * 2) + 0]); + + vec_load_mult284a_mma(&temp[0], &va0[i + 0], &va1[i + 0], &v_x0[ 0]); + vec_load_mult284b_mma(&temp[0], &va2[i + 0], &va3[i + 0], &v_x0[ 4]); + vec_load_mult284b_mma(&temp[0], &vb0[i + 0], &vb1[i + 0], &v_x0[ 8]); + vec_load_mult284b_mma(&temp[0], &vb2[i + 0], &vb3[i + 0], &v_x0[12]); + + vec_reduce84_mma(&temp[0], temp0, v_alpha, vy0); + + vec_store4_pair(&v_y[(i * 2) + 0], vy0); + + i += 4; + } +#else + vec_bf16 v_x0[13]; + vec_f32 vy0[2*4]; + v_x0[0] = (vec_bf16)vec_load_vec(x_bf); + + v_x0[ 4] = (vec_bf16)vec_splat((vec_f32)v_x0[0], 1); + v_x0[ 8] = (vec_bf16)vec_splat((vec_f32)v_x0[0], 2); + v_x0[12] = (vec_bf16)vec_splat((vec_f32)v_x0[0], 3); + for (; i + 4 <= n8; i += 4) { vec_load4_pair(vy0, &v_y[(i * 2) + 0]); - vec_load_mult28a_mma(&temp[0], &va0[i + 0], &va1[i + 0], v_x00); - vec_load_mult28b_mma(&temp[0], &va2[i + 0], &va3[i + 0], v_x01); - vec_load_mult28b_mma(&temp[0], &vb0[i + 0], &vb1[i + 0], v_x02); - vec_load_mult28b_mma(&temp[0], &vb2[i + 0], &vb3[i + 0], v_x03); + vec_load_mult28a_mma(&temp[0], &va0[i + 0], &va1[i + 0], v_x0[ 0]); + vec_load_mult28b_mma(&temp[0], &va2[i + 0], &va3[i + 0], v_x0[ 4]); + vec_load_mult28b_mma(&temp[0], &vb0[i + 0], &vb1[i + 0], v_x0[ 8]); + vec_load_mult28b_mma(&temp[0], &vb2[i + 0], &vb3[i + 0], v_x0[12]); vec_reduce8_mma(&temp[0], temp0, v_alpha, vy0); vec_store4_pair(&v_y[(i * 2) + 0], vy0); } +#endif for (; i < n8; i++) { vec_load_pair(vy0, &v_y[(i * 2) + 0]); - vec_load_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x00); - vec_load_mult22b_mma(&temp[0], &va2[i], &va3[i], v_x01); - vec_load_mult22b_mma(&temp[0], &vb0[i], &vb1[i], v_x02); - vec_load_mult22b_mma(&temp[0], &vb2[i], &vb3[i], v_x03); + vec_load_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0]); + vec_load_mult22b_mma(&temp[0], &va2[i], &va3[i], v_x0[ 4]); + vec_load_mult22b_mma(&temp[0], &vb0[i], &vb1[i], v_x0[ 8]); + vec_load_mult22b_mma(&temp[0], &vb2[i], &vb3[i], v_x0[12]); vec_reduce2_mma(&temp[0], temp0, v_alpha, vy0); @@ -302,10 +456,10 @@ static void BF16GEMV_N_MMA_8(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, BLAS BLASLONG n3 = n & 3; vec_loadN2_f32(vy0, &v_y[(i * 2) + 0], n3); - vec_loadN_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x00, n); - vec_loadN_mult22b_mma(&temp[0], &va2[i], &va3[i], v_x01, n); - vec_loadN_mult22b_mma(&temp[0], &vb0[i], &vb1[i], v_x02, n); - vec_loadN_mult22b_mma(&temp[0], &vb2[i], &vb3[i], v_x03, n); + vec_loadN_mult22a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0], n); + vec_loadN_mult22b_mma(&temp[0], &va2[i], &va3[i], v_x0[ 4], n); + vec_loadN_mult22b_mma(&temp[0], &vb0[i], &vb1[i], v_x0[ 8], n); + vec_loadN_mult22b_mma(&temp[0], &vb2[i], &vb3[i], v_x0[12], n); vec_reduce2_mma(&temp[0], temp0, v_alpha, vy0); @@ -313,10 +467,10 @@ static void BF16GEMV_N_MMA_8(BLASLONG n, IFLOAT **ap, IFLOAT *xo, FLOAT *y, BLAS } else if (n) { vy0[0] = vec_loadN_f32(&v_y[(i * 2) + 0], n); - vec_loadN_mult11a_mma(&temp[0], &va0[i], &va1[i], v_x00, n); - vec_loadN_mult11b_mma(&temp[0], &va2[i], &va3[i], v_x01, n); - vec_loadN_mult11b_mma(&temp[0], &vb0[i], &vb1[i], v_x02, n); - vec_loadN_mult11b_mma(&temp[0], &vb2[i], &vb3[i], v_x03, n); + vec_loadN_mult11a_mma(&temp[0], &va0[i], &va1[i], v_x0[ 0], n); + vec_loadN_mult11b_mma(&temp[0], &va2[i], &va3[i], v_x0[ 4], n); + vec_loadN_mult11b_mma(&temp[0], &vb0[i], &vb1[i], v_x0[ 8], n); + vec_loadN_mult11b_mma(&temp[0], &vb2[i], &vb3[i], v_x0[12], n); vec_reduce1_mma(&temp[0], temp0, v_alpha, vy0);