From aed08036062e0898b415b4e194086ad7952df428 Mon Sep 17 00:00:00 2001 From: Meghana Vankadari Date: Tue, 1 Oct 2024 07:07:46 +0530 Subject: [PATCH] Introduced un-reorder API for bf16bf16f32of32 Details: - Added a new API called unreorder that converts a matrix from reordered format to it's original format( row-major or col-major ). - Currently this API only supports bf16 datatype. - Added corresponding bench and input file to test accuracy of the API. - The new API is only supported for 'B' matrix. - Modified input validation checks in reorder API to account for row Vs col storage of matrix and transposes for bf16 datatype. Change-Id: Ifb9c53b7e6da6f607939c164eb016e82514581b7 --- addon/aocl_gemm/aocl_gemm_bf16_utils.c | 123 ++- addon/aocl_gemm/aocl_gemm_interface_apis.h | 16 +- addon/aocl_gemm/config/lpgemm_config.c | 2 + addon/aocl_gemm/config/lpgemm_func_map.h | 3 + .../frame/bf16bf16f32/lpgemm_reorder_bf16.c | 87 +- .../frame/bf16bf16f32/lpgemm_reorder_bf16.h | 10 +- addon/aocl_gemm/frame/lpgemm_types.h | 1 + .../kernels/bf16bf16f32/lpgemm_pack_bf16.h | 21 + bench/bench_aocl_gemm/Makefile | 3 +- bench/bench_aocl_gemm/bench_lpgemm.c | 35 - bench/bench_aocl_gemm/bench_lpgemm_helpers.h | 36 + bench/bench_aocl_gemm/bench_lpgemm_unpack.c | 276 ++++++ bench/bench_aocl_gemm/bench_unpack_input.txt | 6 + .../lpgemm_unpackb_bf16_amd512vnni.c | 921 ++++++++++++++++++ 14 files changed, 1495 insertions(+), 45 deletions(-) create mode 100644 bench/bench_aocl_gemm/bench_lpgemm_unpack.c create mode 100644 bench/bench_aocl_gemm/bench_unpack_input.txt create mode 100644 kernels/zen4/lpgemm/bf16bf16f32/lpgemm_unpackb_bf16_amd512vnni.c diff --git a/addon/aocl_gemm/aocl_gemm_bf16_utils.c b/addon/aocl_gemm/aocl_gemm_bf16_utils.c index 146fd97d8..25b581b6e 100644 --- a/addon/aocl_gemm/aocl_gemm_bf16_utils.c +++ b/addon/aocl_gemm/aocl_gemm_bf16_utils.c @@ -112,8 +112,7 @@ AOCL_GEMM_REORDER(bfloat16, bf16bf16f32of32) bli_param_map_netlib_to_blis_trans( trans, &blis_trans ); if ( ( input_buf_addr == NULL ) || ( reorder_buf_addr == NULL ) || - ( k <= 0 ) || ( n <= 0 ) || ( bli_is_notrans( blis_trans ) && ( ldb < n ) ) || - ( bli_is_trans( blis_trans ) && ( ldb < k ) ) ) + ( k <= 0 ) || ( n <= 0 ) ) { return; // Error. } @@ -121,13 +120,29 @@ AOCL_GEMM_REORDER(bfloat16, bf16bf16f32of32) inc_t rs_b, cs_b; if( ( order == 'r') || ( order == 'R' ) ) { - rs_b = bli_is_notrans( blis_trans ) ? ldb : 1; - cs_b = bli_is_notrans( blis_trans ) ? 1 : ldb; + if( ( bli_is_notrans( blis_trans ) && ( ldb < n ) ) || + ( bli_is_trans( blis_trans ) && ( ldb < k ) ) ) + { + return; // Error. + } + else + { + rs_b = bli_is_notrans( blis_trans ) ? ldb : 1; + cs_b = bli_is_notrans( blis_trans ) ? 1 : ldb; + } } else if ( ( order == 'c' ) || ( order == 'C' ) ) { - rs_b = bli_is_notrans( blis_trans ) ? 1 : ldb; - cs_b = bli_is_notrans( blis_trans ) ? ldb : 1; + if( ( bli_is_notrans( blis_trans ) && ( ldb < k ) ) || + ( bli_is_trans( blis_trans ) && ( ldb < n ) ) ) + { + return; // Error. + } + else + { + rs_b = bli_is_notrans( blis_trans ) ? 1 : ldb; + cs_b = bli_is_notrans( blis_trans ) ? ldb : 1; + } } else { @@ -195,6 +210,102 @@ AOCL_GEMM_REORDER(bfloat16, bf16bf16f32of32) reorderb_nr64_bf16bf16f32of32( &b, &b_reorder, &rntm_g, lcntx_g ); } + +AOCL_GEMM_UNREORDER(bfloat16, bf16bf16f32of32) +{ + if ( ( output_buf_addr == NULL ) || ( reorder_buf_addr == NULL ) || + ( k <= 0 ) || ( n <= 0 ) ) + { + return; // Error. + } + + inc_t rs_b, cs_b; + + // Check for the validity of strides. + if( ( order == 'r' ) || ( order == 'R' ) ) + { + if( ldb < n ) return; // Error + else + { + rs_b = ldb; + cs_b = 1; + } + } + else if( ( order == 'c' ) || ( order == 'C' ) ) + { + if( ldb < k ) return; // Error. + else + { + rs_b = 1; + cs_b = ldb; + } + } + else + { + return; // Error. + } + + // Check if avx512_bf16 ISA is supported, lpgemm matmul only works with it. + if ( bli_cpuid_is_avx512bf16_supported() == FALSE ) + { + bli_print_msg(" AVX512_BF16 ISA not supported by processor, " + "cannot perform bf16bf16f32 gemm.", __FILE__, __LINE__ ); + return; // Error. + } + + /* Initialize BLIS. */ + bli_init_auto(); + + // Set MC, NC, KC, NR, MR. + aocl_lpgemm_init_global_cntx(); + + AOCL_MATRIX_TYPE input_mat_type; + bli_param_map_char_to_lpmat_type( mat_type, &input_mat_type ); + + if ( input_mat_type == A_MATRIX ) + { + return; // A reorder not supported. + } +#if (defined(BLIS_KERNELS_ZEN4) && (!defined(LPGEMM_BF16_JIT))) + if( n == 1 ) + { + if( rs_b == 1 ) + { + memcpy( output_buf_addr, reorder_buf_addr, ( k * sizeof( bfloat16 ) ) ); + } + else + { + for( dim_t k0 = 0; k0 < k; k0++ ) + { + output_buf_addr[k0*rs_b] = reorder_buf_addr[k0]; + } + } + return; + } +#endif + // Initialize a local runtime with global settings if necessary. Note + // that in the case that a runtime is passed in, we make a local copy. + rntm_t rntm_g; + bli_rntm_init_from_global( &rntm_g ); + bli_pba_rntm_set_pba( &rntm_g ); + + lpgemm_cntx_t* lcntx_g = lpgemm_get_global_cntx_obj( BF16BF16F32OF32 ); + + // create dummy b_reorder obj. + lpgemm_obj_t b_reorder; + b_reorder.storage.aligned_buffer = ( void* )reorder_buf_addr; + + // create dummy b obj. + lpgemm_obj_t b; + b.storage.aligned_buffer = ( void* )output_buf_addr; + b.rs = rs_b; + b.cs = cs_b; + b.width = n; + b.length = k; + + unreorderb_nr64_bf16bf16f32of32( &b, &b_reorder, &rntm_g, lcntx_g ); +} + AOCL_GEMM_GET_REORDER_BUF_SIZE(bf16s4f32of32) { if ((k <= 0) || (n <= 0)) diff --git a/addon/aocl_gemm/aocl_gemm_interface_apis.h b/addon/aocl_gemm/aocl_gemm_interface_apis.h index c1c870936..0ab6daa45 100644 --- a/addon/aocl_gemm/aocl_gemm_interface_apis.h +++ b/addon/aocl_gemm/aocl_gemm_interface_apis.h @@ -4,7 +4,7 @@ An object-based framework for developing high-performance BLAS-like libraries. - Copyright (C) 2022 - 2023, Advanced Micro Devices, Inc. All rights reserved. + Copyright (C) 2022 - 2024, Advanced Micro Devices, Inc. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are @@ -83,6 +83,20 @@ AOCL_GEMM_REORDER(int8_t,s8s8s16os16); AOCL_GEMM_REORDER(int8_t,u8s4s32os32); AOCL_GEMM_REORDER(int8_t, bf16s4f32of32); +#define AOCL_GEMM_UNREORDER(B_type, LP_SFX) \ +BLIS_EXPORT_ADDON void aocl_unreorder_ ## LP_SFX \ + ( \ + const char order, \ + const char mat_type, \ + const B_type* reorder_buf_addr, \ + B_type* output_buf_addr, \ + const dim_t k, \ + const dim_t n, \ + const dim_t ldb \ + ) \ + +AOCL_GEMM_UNREORDER(bfloat16, bf16bf16f32of32); + #define AOCL_GEMM_MATMUL(A_type,B_type,C_type,Sum_type,LP_SFX) \ BLIS_EXPORT_ADDON void aocl_gemm_ ## LP_SFX \ ( \ diff --git a/addon/aocl_gemm/config/lpgemm_config.c b/addon/aocl_gemm/config/lpgemm_config.c index 3fee7d2cb..b4c14c4ad 100644 --- a/addon/aocl_gemm/config/lpgemm_config.c +++ b/addon/aocl_gemm/config/lpgemm_config.c @@ -121,6 +121,7 @@ static void _lpgemm_cntx_init_func_map() #define KMACRO(ID,FUNC_PTR) global_cntx_t_list[ID].kern_fun_ptr = FUNC_PTR; #define PAMACRO(ID,FUNC_PTR) global_cntx_t_list[ID].packa_fun_ptr = FUNC_PTR; #define PBMACRO(ID,FUNC_PTR) global_cntx_t_list[ID].packb_fun_ptr = FUNC_PTR; +#define UBMACRO(ID, FUNC_PTR) global_cntx_t_list[ID].unpackb_fun_ptr = FUNC_PTR; #define PBSMACRO(ID, FUNC_PTR) global_cntx_t_list[ID].packsclb_fun_ptr = FUNC_PTR; #define JITMACRO(ID, FUNC_PTR) global_cntx_t_list[ID].jit_kernel = FUNC_PTR; //TODO: Default initialize with reference kernels so that kernel pointer @@ -141,6 +142,7 @@ static void _lpgemm_cntx_init_func_map() LPGEMM_KERN_FUNC_MAP_AVX512_VNNI_BF16 LPGEMM_PACKA_FUNC_MAP_AVX512_VNNI_BF16 LPGEMM_PACKB_FUNC_MAP_AVX512_VNNI_BF16 + LPGEMM_UNPACKB_FUNC_MAP_AVX512_VNNI_BF16 LPGEMM_PACKSCLB_FUNC_MAP_AVX512_VNNI_BF16 #ifdef LPGEMM_BF16_JIT diff --git a/addon/aocl_gemm/config/lpgemm_func_map.h b/addon/aocl_gemm/config/lpgemm_func_map.h index f4d2f2b83..e93246f3b 100644 --- a/addon/aocl_gemm/config/lpgemm_func_map.h +++ b/addon/aocl_gemm/config/lpgemm_func_map.h @@ -71,6 +71,9 @@ PBMACRO(U8S4S32OS32, packb_nr64_u8s4s32o32) \ PBMACRO(BF16S4F32OF32, packb_nr64_bf16s4f32of32) +#define LPGEMM_UNPACKB_FUNC_MAP_AVX512_VNNI_BF16 \ + UBMACRO(BF16BF16F32OF32, unpackb_nr64_bf16bf16f32of32) + #define LPGEMM_PACKSCLB_FUNC_MAP_AVX512_VNNI_BF16 \ PBSMACRO(U8S8S16OS16, NULL) \ PBSMACRO(U8S8S32OS32, NULL) \ diff --git a/addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.c b/addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.c index 9305b142c..68c1bb350 100644 --- a/addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.c +++ b/addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.c @@ -4,7 +4,7 @@ An object-based framework for developing high-performance BLAS-like libraries. - Copyright (C) 2022 - 2023, Advanced Micro Devices, Inc. All rights reserved. + Copyright (C) 2022 - 2024, Advanced Micro Devices, Inc. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are @@ -169,6 +169,91 @@ void reorderb_nr64_bf16bf16f32of32 b_reorder->mtag = REORDERED; } +void unreorderb_nr64_bf16bf16f32of32 + ( + lpgemm_obj_t * b, + lpgemm_obj_t * b_reorder, + rntm_t* rntm, + lpgemm_cntx_t* lcntx + ) +{ + dim_t NC = lcntx->blksz.NC; + dim_t KC = lcntx->blksz.KC; + dim_t NR = lcntx->blksz.NR; + + // Extracting the matrix properties from the lpgemm object + dim_t rs_b = b->rs; + dim_t cs_b = b->cs; + dim_t n = b->width; + dim_t k = b->length; + + dim_t k_updated = k; + k_updated += (k_updated & 0x1); + + dim_t n_threads = bli_rntm_num_threads( rntm ); + n_threads = ( n_threads > 0 ) ? n_threads : 1; + +#ifdef BLIS_ENABLE_OPENMP + _Pragma( "omp parallel num_threads(n_threads)" ) + { + // Initialise a local thrinfo obj for work split across threads. + thrinfo_t thread_jc; + bli_thrinfo_set_n_way( n_threads, &thread_jc ); + bli_thrinfo_set_work_id( omp_get_thread_num(), &thread_jc ); +#else + { + // Initialise a local thrinfo obj for work split across threads. + thrinfo_t thread_jc; + bli_thrinfo_set_n_way( 1, &thread_jc ); + bli_thrinfo_set_work_id( 0, &thread_jc ); +#endif + + // Compute the JC loop thread range for the current thread. + dim_t jc_start, jc_end; + bli_thread_range_sub( &thread_jc, n, NR, FALSE, &jc_start, &jc_end ); + + for ( dim_t jc = jc_start; jc < jc_end; jc += NC ) + { + dim_t nc0 = bli_min( ( jc_end - jc ), NC ); + + dim_t jc_cur_loop = jc; + dim_t jc_cur_loop_rem = 0; + dim_t n_sub_updated; + + get_B_panel_reordered_start_offset_width + ( + jc, n, NC, 16, + &jc_cur_loop, &jc_cur_loop_rem, + &nc0, &n_sub_updated + ); + + for ( dim_t pc = 0; pc < k; pc += KC ) + { + dim_t kc0 = bli_min( ( k - pc ), KC ); + + // k needs to be a multiple of 2 so that it can be used with dpbf + // instruction. Padding is added in cases this condition is not + // satisfied, and therefore the k offset used for packed/reordered + // buffer needs to be updated. + dim_t kc0_updated = kc0; + kc0_updated += (kc0_updated & 0x1); + + ( ( unpack_bf16 )lcntx->unpackb_fun_ptr ) + ( + ( ( const bfloat16* )b_reorder->storage.aligned_buffer ) + + ( jc_cur_loop * k_updated ) + ( n_sub_updated * pc ) + + ( jc_cur_loop_rem * kc0_updated ), + ( ( ( bfloat16* )b->storage.aligned_buffer ) + + ( rs_b * pc ) + (jc * cs_b)), + nc0, kc0, rs_b, cs_b + ); + } + + adjust_B_panel_reordered_jc( &jc, jc_cur_loop ); + } + } +} + void reorderb_nr64_bf16s4f32of32( lpgemm_obj_t *b, lpgemm_obj_t *b_reorder, diff --git a/addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.h b/addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.h index 6595753dc..912052a66 100644 --- a/addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.h +++ b/addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.h @@ -4,7 +4,7 @@ An object-based framework for developing high-performance BLAS-like libraries. - Copyright (C) 2022 - 2023, Advanced Micro Devices, Inc. All rights reserved. + Copyright (C) 2022 - 2024, Advanced Micro Devices, Inc. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are @@ -53,4 +53,12 @@ void reorderb_nr64_bf16s4f32of32 lpgemm_cntx_t* lcntx ); +void unreorderb_nr64_bf16bf16f32of32 + ( + lpgemm_obj_t * b, + lpgemm_obj_t * b_reorder, + rntm_t* rntm, + lpgemm_cntx_t* lcntx + ); + #endif // LPGEMM_REORDER_H diff --git a/addon/aocl_gemm/frame/lpgemm_types.h b/addon/aocl_gemm/frame/lpgemm_types.h index b0c69079b..99065255c 100644 --- a/addon/aocl_gemm/frame/lpgemm_types.h +++ b/addon/aocl_gemm/frame/lpgemm_types.h @@ -154,6 +154,7 @@ typedef struct void_fp kern_fun_ptr; void_fp packa_fun_ptr; void_fp packb_fun_ptr; + void_fp unpackb_fun_ptr; void_fp packsclb_fun_ptr; lpgemm_pack_strides_t pack_s; } lpgemm_cntx_t; diff --git a/addon/aocl_gemm/kernels/bf16bf16f32/lpgemm_pack_bf16.h b/addon/aocl_gemm/kernels/bf16bf16f32/lpgemm_pack_bf16.h index 5073dbbab..12319fee1 100644 --- a/addon/aocl_gemm/kernels/bf16bf16f32/lpgemm_pack_bf16.h +++ b/addon/aocl_gemm/kernels/bf16bf16f32/lpgemm_pack_bf16.h @@ -70,6 +70,16 @@ typedef void (*pack_bf16) dim_t* ); +typedef void (*unpack_bf16) + ( + const bfloat16*, + bfloat16*, + const dim_t, + const dim_t, + const dim_t, + const dim_t + ); + typedef void (*pack_s4) ( int8_t*, @@ -131,4 +141,15 @@ void packa_mr16_bf16bf16f32of32 dim_t* rs_p, dim_t* cs_p ); + +void unpackb_nr64_bf16bf16f32of32 + ( + const bfloat16* unpack_b_buffer_bf16bf16f32of32, + bfloat16* b, + const dim_t NC, + const dim_t KC, + dim_t rs_b, + dim_t cs_b + ); + #endif //BLIS_GEMM_BF16_PACKB diff --git a/bench/bench_aocl_gemm/Makefile b/bench/bench_aocl_gemm/Makefile index c8c2b732a..8a07d9fe4 100644 --- a/bench/bench_aocl_gemm/Makefile +++ b/bench/bench_aocl_gemm/Makefile @@ -88,7 +88,7 @@ TEST_OBJS := $(patsubst $(TEST_SRC_PATH)/%.c, \ CINCFLAGS := -I$(INC_PATH) -I$(CBLAS_HEADER_PATH) # Use the CFLAGS for the configuration family. -CFLAGS := $(call get-user-cflags-for,$(CONFIG_NAME)) +CFLAGS := $(call get-user-cflags-for,$(CONFIG_NAME)) # Add local header paths to CFLAGS CFLAGS += -I$(TEST_SRC_PATH) @@ -109,6 +109,7 @@ all: blis blis: \ bench_lpgemm_blis.x \ bench_lpgemm_utils_blis.x \ + bench_lpgemm_unpack_blis.x \ bench_lpgemm_eltwise_ops_blis.x diff --git a/bench/bench_aocl_gemm/bench_lpgemm.c b/bench/bench_aocl_gemm/bench_lpgemm.c index 1f757dbd1..ff5f5212f 100644 --- a/bench/bench_aocl_gemm/bench_lpgemm.c +++ b/bench/bench_aocl_gemm/bench_lpgemm.c @@ -42,41 +42,6 @@ CONVERT_TO_FLOAT(int16_t) CONVERT_TO_FLOAT(float) CONVERT_TO_FLOAT(int32_t) -/* Helper functions to print matrices when debugging */ -void print_matrix_bfloat16 - ( - bfloat16* a, - dim_t m, - dim_t n, - dim_t rs_a, - dim_t cs_a - ) -{ - for(dim_t i = 0; i < m; i++) - { - for(dim_t j = 0; j < n; j++) - { - float temp; - bfloat16_to_float(*(a + i*(rs_a) + j *cs_a), &temp); - printf("%f ", temp); - } - printf("\n"); - } -} - -#define PRINT_MATRIX(ctype) \ -void print_matrix_## ctype ( ctype* a, dim_t m, dim_t n, dim_t rs, dim_t cs) \ -{ \ - for(dim_t i = 0; i < m; i++) \ - { \ - for(dim_t j = 0; j < n; j++) \ - { \ - printf("%f ", (float) (*(a + i * ( rs ) + j * cs ) ) ); \ - } \ - printf("\n"); \ - } \ -} \ - PRINT_MATRIX(uint8_t) PRINT_MATRIX(int8_t) PRINT_MATRIX(int16_t) diff --git a/bench/bench_aocl_gemm/bench_lpgemm_helpers.h b/bench/bench_aocl_gemm/bench_lpgemm_helpers.h index d3464c8c3..6c0a53df8 100644 --- a/bench/bench_aocl_gemm/bench_lpgemm_helpers.h +++ b/bench/bench_aocl_gemm/bench_lpgemm_helpers.h @@ -437,4 +437,40 @@ static inline void lpgemm_destroy_post_ops_struct( aocl_post_op* post_ops ) free( post_ops ); } + +#define PRINT_MATRIX(ctype) \ +void print_matrix_## ctype ( ctype* a, dim_t m, dim_t n, dim_t rs, dim_t cs) \ +{ \ + for(dim_t i = 0; i < m; i++) \ + { \ + for(dim_t j = 0; j < n; j++) \ + { \ + printf("%f ", (float) (*(a + i * ( rs ) + j * cs ) ) ); \ + } \ + printf("\n"); \ + } \ +} \ + +/* Helper functions to print matrices when debugging */ +void print_matrix_bfloat16 + ( + bfloat16* a, + dim_t m, + dim_t n, + dim_t rs_a, + dim_t cs_a + ) +{ + for(dim_t i = 0; i < m; i++) + { + for(dim_t j = 0; j < n; j++) + { + float temp; + bfloat16_to_float(*(a + i*(rs_a) + j *cs_a), &temp); + printf("%3d ", (int)temp); + } + printf("\n"); + } +} + #endif //LPGEMM_BENCH_UTILS_H diff --git a/bench/bench_aocl_gemm/bench_lpgemm_unpack.c b/bench/bench_aocl_gemm/bench_lpgemm_unpack.c new file mode 100644 index 000000000..4cb705a74 --- /dev/null +++ b/bench/bench_aocl_gemm/bench_lpgemm_unpack.c @@ -0,0 +1,276 @@ +/* + + BLIS + An object-based framework for developing high-performance BLAS-like + libraries. + + Copyright (C) 2024, Advanced Micro Devices, Inc. All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions are + met: + - Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + - Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + - Neither the name(s) of the copyright holder(s) nor the names of its + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +*/ + +#include "bench_lpgemm_helpers.h" + +#define GEN_UNPACK_ACC_CHK_DRV_FUNC( B_type, LP_SFX ) \ +void unpack_accuracy_check_driver_ ## LP_SFX \ + ( \ + FILE* fout, \ + const char stor_order, \ + const char mat_type, \ + dim_t m, \ + dim_t n, \ + dim_t rs, \ + dim_t cs, \ + B_type* b, \ + B_type* b_unpacked \ + ) \ +{ \ + for( dim_t i = 0; i < m; i++ ) \ + { \ + for( dim_t j = 0; j < n; j++ ) \ + { \ + if ( b[ i * ( rs ) + j * cs ] != b_unpacked[ i * ( rs ) + j * cs ] ) \ + { \ + if( fout ) \ + { \ + fprintf( fout, "failure, m: %ld, n: %ld, computed:%f, ref:%f, diff:%f\n", \ + i, j, ( float )b[ i * ( rs ) + j * cs ], ( float )b_unpacked[ i * ( rs ) + j * cs ], \ + ( float )( b[ i * ( rs ) + j * cs ] - b_unpacked[ i * ( rs ) + j * cs ] ) ); \ + } \ + printf( "failure, m: %ld, n: %ld, computed:%f, ref:%f, diff:%f\n", \ + i, j, ( float )b[ i * ( rs ) + j * cs ], ( float )b_unpacked[ i * ( rs ) + j * cs ], \ + ( float )( b[ i * ( rs ) + j * cs ] - b_unpacked[ i * ( rs ) + j * cs ] ) ); \ + goto cleanup_acc; \ + } \ + } \ + } \ + \ +cleanup_acc: \ + fflush(stdout); \ + return; \ +} \ + +GEN_UNPACK_ACC_CHK_DRV_FUNC( bfloat16, bf16bf16f32of32 ) + +#define GEN_UNPACK_BENCH_MAIN_FUNC( B_type, LP_SFX ) \ +void unpack_bench_main_ ## LP_SFX \ + ( \ + FILE* fout, \ + char stor_order, \ + char mat_type, \ + dim_t m, \ + dim_t n, \ + dim_t stride \ + ) \ +{ \ + dim_t size_B = 0; \ + \ + dim_t rs, cs; \ + if ( ( stor_order == 'r' ) || ( stor_order == 'R' ) ) \ + { \ + size_B = m * stride; \ + rs = stride; \ + cs = 1; \ + } \ + else \ + { \ + size_B = stride * n; \ + rs = 1; \ + cs = stride; \ + } \ + \ + /* Original B matrix */ \ + B_type* b = ( B_type* ) lpgemm_malloc( sizeof( B_type ) * size_B ); \ + GEN_FUNC_NAME(fill_array_,B_type)(b, size_B ); \ + \ + /* Matrix to be unpacked into */ \ + B_type* b_unpacked = ( B_type* ) lpgemm_malloc( sizeof( B_type ) * size_B ); \ + memset( ( void* ) b_unpacked, 0, sizeof( B_type ) * size_B ); \ + \ + /* Reorder matrix */ \ + B_type* b_reorder = NULL; \ + \ + siz_t b_reorder_buf_siz_req = \ + GEN_FUNC_NAME(aocl_get_reorder_buf_size_,LP_SFX)( stor_order, 'n', mat_type, m, n ); \ + \ + b_reorder = ( B_type* ) lpgemm_malloc( b_reorder_buf_siz_req ); \ + /* reorder B. */ \ + GEN_FUNC_NAME(aocl_reorder_,LP_SFX)( stor_order, 'n', mat_type, \ + b, b_reorder, m, n, stride ); \ + \ + /* Unpack B. */ \ + GEN_FUNC_NAME(aocl_unreorder_,LP_SFX)( stor_order, mat_type, \ + b_reorder, b_unpacked, m, n, stride ); \ + \ + /* Accuracy check */ \ + printf("Running accuracy check\n"); \ + printf("%s m: %ld, n: %ld, stride: %ld, stor_order: %c\n", \ + XSTR(LP_SFX), m, n, stride, stor_order ); \ + GEN_FUNC_NAME( unpack_accuracy_check_driver_, LP_SFX ) \ + ( \ + fout, stor_order, mat_type, m, n, rs, cs,\ + b, b_unpacked \ + ); \ + \ + lpgemm_free( b ); \ + lpgemm_free( b_unpacked ); \ + lpgemm_free( b_reorder ); \ +} \ + +GEN_UNPACK_BENCH_MAIN_FUNC( bfloat16, bf16bf16f32of32 ) + +int main( int argc, char** argv ) +{ + // By default bench mode is set to accuracy. + bench_mode = 'a'; + + FILE* fin = NULL; + if ( argc < 3 ) + { + printf + ( + "Usage: ./bench_lpgemm_unpack -i input.txt \n" \ + ); + exit( 1 ); + } + + char* file_name = NULL; + + // Parse CLI arguments. + getopt_t state; + // Initialize the state for running bli_getopt(). Here, 0 is the + // initial value for opterr, which suppresses error messages. + bli_getopt_init_state( 0, &state ); + +#define UNPACK_OPS_TYPE_STR_LEN 24 + char unpack_ops_type_str[UNPACK_OPS_TYPE_STR_LEN]; + + int opt; + // Process all option arguments until we get a -1, which means we're done. + while( (opt = bli_getopt( argc, argv, "i:", &state )) != -1 ) + { + char opt_ch = ( char )opt; + switch( opt_ch ) + { + case 'i': + file_name = state.optarg; + break; + default: + break; + } + } + + if ( file_name == NULL ) + { + printf( " File name provided is invalid.\n" ); + exit( 1 ); + } + + fin = fopen( file_name, "r" ); + if (fin == NULL) + { + printf( "Error opening the file %s\n", argv[1] ); + exit( 1 ); + } + + FILE* fout = NULL; + + fout = fopen( "lpgemm_unpack_accuracy_test_failures.txt", "w" ); + + char stor_order; + char mat_type; + dim_t m, n; + dim_t stride; + + const dim_t len_list_omp_cores_for_testing = 2; + const dim_t list_omp_cores_for_testing[2] = { 1, 64 }; + + dim_t core_index = 0; + bool can_run = TRUE; + while ( ( can_run == TRUE ) && ( fseek( fin, 0L, SEEK_SET ) == 0 ) ) + { + if ( bench_mode == 'p' ) + { + can_run = FALSE; + } + else if ( bench_mode == 'a' ) + { + // For accuracy testing, we test accuracy using multiple different + // number of cores. This helps uncover any bugs related to over + // subscription or varying thread factorizations. + // Set current number of cores. +#ifdef BLIS_ENABLE_OPENMP + omp_set_num_threads( list_omp_cores_for_testing[core_index] ); +#endif + printf( "Accuracy test using %ld threads.\n", + list_omp_cores_for_testing[core_index] ); + + core_index++; + if ( core_index < len_list_omp_cores_for_testing ) + { + can_run = TRUE; + } + else + { + can_run = FALSE; + } + } + + // Input format: stor_type, mat_type, m, n, stride + while ( fscanf( fin, "%c %c %ld %ld %ld %s\n", + &stor_order, &mat_type, &m, &n, + &stride, unpack_ops_type_str ) == 6 ) + { + str_tolower( unpack_ops_type_str ); + + stor_order = ( ( stor_order == 'r' ) || ( stor_order == 'R' ) || + ( stor_order == 'c' ) || ( stor_order == 'C' ) ) ? + stor_order : 'r'; + + if ( ( strcmp( unpack_ops_type_str, "bf16bf16f32of32" ) == 0 ) || + ( strcmp( unpack_ops_type_str, "bf16bf16f32ofbf16" ) == 0 ) || + ( strcmp( unpack_ops_type_str, "*" ) == 0 ) ) + { + GEN_FUNC_NAME(unpack_bench_main_, bf16bf16f32of32) + ( + fout, stor_order, mat_type, + m, n, stride + ); + } + } + } + + if( fin ) + { + fclose( fin ); + } + if( fout ) + { + fclose( fout ); + } + + return 0; +} \ No newline at end of file diff --git a/bench/bench_aocl_gemm/bench_unpack_input.txt b/bench/bench_aocl_gemm/bench_unpack_input.txt new file mode 100644 index 000000000..cfd3b4862 --- /dev/null +++ b/bench/bench_aocl_gemm/bench_unpack_input.txt @@ -0,0 +1,6 @@ +r b 23 45 45 bf16bf16f32of32 +r b 1 118 118 bf16bf16f32of32 +c b 51 75 51 bf16bf16f32of32 +c b 1 65 1 bf16bf16f32of32 +c b 39 45 39 bf16bf16f32of32 +r b 12 63 63 bf16bf16f32of32 \ No newline at end of file diff --git a/kernels/zen4/lpgemm/bf16bf16f32/lpgemm_unpackb_bf16_amd512vnni.c b/kernels/zen4/lpgemm/bf16bf16f32/lpgemm_unpackb_bf16_amd512vnni.c new file mode 100644 index 000000000..57e9d4bc3 --- /dev/null +++ b/kernels/zen4/lpgemm/bf16bf16f32/lpgemm_unpackb_bf16_amd512vnni.c @@ -0,0 +1,921 @@ +/* + + BLIS + An object-based framework for developing high-performance BLAS-like + libraries. + + Copyright (C) 2024, Advanced Micro Devices, Inc. All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions are + met: + - Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + - Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + - Neither the name(s) of the copyright holder(s) nor the names of its + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +*/ + +#include +#include +#include "blis.h" + +#ifdef BLIS_ADDON_LPGEMM + +void unpackb_nr48_bf16bf16f32of32_row_major + ( + const bfloat16* b, + bfloat16* unpack_b_buffer_bf16bf16f32of32, + const dim_t KC, + dim_t ldb + ) +{ + dim_t NR1 = 32; + + // Used for permuting the mm512i elements. + __m512i selector_even = _mm512_set_epi16( 0x1E, 0x1C, 0x1A, 0x18, 0x16, 0x14, 0x12, 0x10, 0xE, 0xC, 0xA, 0x8, 0x6, 0x4, 0x2, 0x0, + 0x3E, 0x3C, 0x3A, 0x38, 0x36, 0x34, 0x32, 0x30, 0x2E, 0x2C, 0x2A, 0x28, 0x26, 0x24, 0x22, 0x20); + + __m512i selector_odd = _mm512_set_epi16( 0x1F, 0x1D, 0x1B, 0x19, 0x17, 0x15, 0x13, 0x11, 0xF, 0xD, 0xB, 0x9, 0x7, 0x5, 0x3, 0x1, + 0x3F, 0x3D, 0x3B, 0x39, 0x37, 0x35, 0x33, 0x31, 0x2F, 0x2D, 0x2B, 0x29, 0x27, 0x25, 0x23, 0x21); + + __m512i a0, a01, b0; + __m512i c0, d0, c01; + + dim_t k_full_pieces_blks = KC / 2; + dim_t k_full_pieces = k_full_pieces_blks * 2; + dim_t k_partial_pieces = KC % 2; + + dim_t kr_new = 0; + + for ( dim_t kr = 0; kr < k_full_pieces; kr += 2 ) + { + // First 2x32 elements + a0 = _mm512_loadu_si512( b + ( ( kr_new + 0 ) * NR1 ) ); + b0 = _mm512_loadu_si512( b + ( ( kr_new + 1 ) * NR1 ) ); + + a01 = _mm512_permutex2var_epi16( b0, selector_even, a0 ); + b0 = _mm512_permutex2var_epi16( b0, selector_odd, a0 ); + + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 0 ) ), a01 ); + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 1 ) ), b0 ); + + c0 = _mm512_loadu_epi16( b + ( ( kr_new + 2 ) * NR1 ) ); + d0 = _mm512_setzero_si512(); + + c01 = _mm512_permutex2var_epi16( d0, selector_even, c0 ); + d0 = _mm512_permutex2var_epi16( d0, selector_odd, c0 ); + + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 0 ) ) + NR1, 0xFFFF, c01 ); + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 1 ) ) + NR1, 0xFFFF, d0 ); + + kr_new += 3; + } + + if( k_partial_pieces > 0 ) + { + // First 2x32 elements + a0 = _mm512_loadu_si512( b + ( ( kr_new + 0 ) * NR1 ) ); + b0 = _mm512_loadu_si512( b + ( ( kr_new + 1 ) * NR1 ) ); + + a01 = _mm512_permutex2var_epi16( b0, selector_even, a0 ); + + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( k_full_pieces + 0 ) ), a01 ); + + c0 = _mm512_loadu_epi16( b + ( ( kr_new + 2 ) * NR1 ) ); + c01 = _mm512_permutex2var_epi16( c0, selector_even, c0 ); + + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( k_full_pieces + 0 ) ) + NR1, 0xFFFF, c01 ); + } +} +void unpackb_nr32_bf16bf16f32of32_row_major + ( + const bfloat16* b, + bfloat16* unpack_b_buffer_bf16bf16f32of32, + const dim_t KC, + dim_t ldb + ) +{ + dim_t NR = 32; + + // Used for permuting the mm512i elements. + __m512i selector_even = _mm512_set_epi16( 0x1E, 0x1C, 0x1A, 0x18, 0x16, 0x14, 0x12, 0x10, 0xE, 0xC, 0xA, 0x8, 0x6, 0x4, 0x2, 0x0, + 0x3E, 0x3C, 0x3A, 0x38, 0x36, 0x34, 0x32, 0x30, 0x2E, 0x2C, 0x2A, 0x28, 0x26, 0x24, 0x22, 0x20); + + __m512i selector_odd = _mm512_set_epi16( 0x1F, 0x1D, 0x1B, 0x19, 0x17, 0x15, 0x13, 0x11, 0xF, 0xD, 0xB, 0x9, 0x7, 0x5, 0x3, 0x1, + 0x3F, 0x3D, 0x3B, 0x39, 0x37, 0x35, 0x33, 0x31, 0x2F, 0x2D, 0x2B, 0x29, 0x27, 0x25, 0x23, 0x21); + + dim_t k_full_pieces_blks = KC / 2; + dim_t k_full_pieces = k_full_pieces_blks * 2; + dim_t k_partial_pieces = KC % 2; + + // KC when not multiple of 2 will have padding to make it multiple of 2 in packed buffer. + dim_t KC_updated = KC; + if ( k_partial_pieces > 0 ) + { + KC_updated += ( 2 - k_partial_pieces ); + } + + __m512i a0, c0; + __m512i a01; + + for ( dim_t kr = 0; kr < k_full_pieces; kr += 2 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 64 elements in each row. + a0 = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + c0 = _mm512_loadu_si512( b + ( ( kr + 1 ) * NR ) ); + + a01 = _mm512_permutex2var_epi16( c0, selector_even, a0 ); + c0 = _mm512_permutex2var_epi16( c0, selector_odd, a0 ); + + // Store to unpack buffer + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 0 ) ), a01 ); + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 1 ) ), c0 ); + + } + if( k_partial_pieces > 0 ) + { + a0 = _mm512_loadu_si512( b + ( ( k_full_pieces + 0 ) * NR ) ); + c0 = _mm512_loadu_si512( b + ( ( k_full_pieces + 1 ) * NR ) ); + + a0 = _mm512_permutex2var_epi16( c0, selector_even, a0 ); + + // Store to unpack buffer + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( k_full_pieces + 0 ) ), a0 ); + } +} +void unpackb_nr16_bf16bf16f32of32_row_major + ( + const bfloat16* b, + bfloat16* unpack_b_buffer_bf16bf16f32of32, + const dim_t KC, + dim_t ldb + ) +{ + dim_t NR = 16; + + // Used for permuting the mm512i elements. + __m512i selector_even = _mm512_set_epi16( 0x1E, 0x1C, 0x1A, 0x18, 0x16, 0x14, 0x12, 0x10, 0xE, 0xC, 0xA, 0x8, 0x6, 0x4, 0x2, 0x0, + 0x3E, 0x3C, 0x3A, 0x38, 0x36, 0x34, 0x32, 0x30, 0x2E, 0x2C, 0x2A, 0x28, 0x26, 0x24, 0x22, 0x20); + + __m512i selector_odd = _mm512_set_epi16( 0x1F, 0x1D, 0x1B, 0x19, 0x17, 0x15, 0x13, 0x11, 0xF, 0xD, 0xB, 0x9, 0x7, 0x5, 0x3, 0x1, + 0x3F, 0x3D, 0x3B, 0x39, 0x37, 0x35, 0x33, 0x31, 0x2F, 0x2D, 0x2B, 0x29, 0x27, 0x25, 0x23, 0x21); + + __m512i a0; + __m512i c0; + __m512i a01; + + dim_t k_full_pieces_blks = KC / 2; + dim_t k_full_pieces = k_full_pieces_blks * 2; + dim_t k_partial_pieces = KC % 2; + + for ( dim_t kr = 0; kr < k_full_pieces; kr += 2 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 16 elements in each row. + a0 = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + + a01 = _mm512_permutex2var_epi16( a0, selector_even, a0 ); + c0 = _mm512_permutex2var_epi16( a0, selector_odd, a0 ); + + // Store to unpack buffer + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 0 ) ), 0xFFFF, a01 ); + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 1 ) ), 0xFFFF, c0 ); + } + if( k_partial_pieces > 0 ) + { + a0 = _mm512_loadu_si512( b + ( ( k_full_pieces + 0 ) * NR ) ); + + a0 = _mm512_permutex2var_epi16( a0, selector_even, a0 ); + + // Store to unpack buffer + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( k_full_pieces + 0 ) ), 0xFFFF, a0 ); + } +} +void unpackb_nrlt16_bf16bf16f32of32_row_major + ( + const bfloat16* b, + bfloat16* unpack_b_buffer_bf16bf16f32of32, + const dim_t KC, + dim_t ldb, + dim_t n0_partial_rem + ) +{ + dim_t NR = 16; + + // Used for permuting the mm512i elements. + __m512i selector_even = _mm512_set_epi16( 0x1E, 0x1C, 0x1A, 0x18, 0x16, 0x14, 0x12, 0x10, 0xE, 0xC, 0xA, 0x8, 0x6, 0x4, 0x2, 0x0, + 0x3E, 0x3C, 0x3A, 0x38, 0x36, 0x34, 0x32, 0x30, 0x2E, 0x2C, 0x2A, 0x28, 0x26, 0x24, 0x22, 0x20); + + __m512i selector_odd = _mm512_set_epi16( 0x1F, 0x1D, 0x1B, 0x19, 0x17, 0x15, 0x13, 0x11, 0xF, 0xD, 0xB, 0x9, 0x7, 0x5, 0x3, 0x1, + 0x3F, 0x3D, 0x3B, 0x39, 0x37, 0x35, 0x33, 0x31, 0x2F, 0x2D, 0x2B, 0x29, 0x27, 0x25, 0x23, 0x21); + + __m512i a0; + __m512i c0; + __m512i a01; + + __mmask32 store_mask = _cvtu32_mask32( 0xFFFF >> ( 16 - n0_partial_rem ) ); + + dim_t k_full_pieces_blks = KC / 2; + dim_t k_full_pieces = k_full_pieces_blks * 2; + dim_t k_partial_pieces = KC % 2; + + for ( dim_t kr = 0; kr < k_full_pieces; kr += 2 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 16 elements in each row. + a0 = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + + a01 = _mm512_permutex2var_epi16( a0, selector_even, a0 ); + c0 = _mm512_permutex2var_epi16( a0, selector_odd, a0 ); + + // Store to unpack buffer + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 0 ) ), store_mask, a01 ); + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 1 ) ), store_mask, c0 ); + } + if( k_partial_pieces > 0 ) + { + a0 = _mm512_loadu_si512( b + ( ( k_full_pieces + 0 ) * NR ) ); + + a0 = _mm512_permutex2var_epi16( a0, selector_even, a0 ); + + // Store to unpack buffer + _mm512_mask_storeu_epi16( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( k_full_pieces + 0 ) ), store_mask, a0 ); + } +} + +void unpackb_nr64_bf16bf16f32of32_row_major + ( + const bfloat16* b, + bfloat16* unpack_b_buffer_bf16bf16f32of32, + const dim_t NC, + const dim_t KC, + dim_t ldb + ) +{ + dim_t NR = 64; + + // Used for permuting the mm512i elements. + __m512i selector_even = _mm512_set_epi16( 0x1E, 0x1C, 0x1A, 0x18, 0x16, 0x14, 0x12, 0x10, 0xE, 0xC, 0xA, 0x8, 0x6, 0x4, 0x2, 0x0, + 0x3E, 0x3C, 0x3A, 0x38, 0x36, 0x34, 0x32, 0x30, 0x2E, 0x2C, 0x2A, 0x28, 0x26, 0x24, 0x22, 0x20); + + __m512i selector_odd = _mm512_set_epi16( 0x1F, 0x1D, 0x1B, 0x19, 0x17, 0x15, 0x13, 0x11, 0xF, 0xD, 0xB, 0x9, 0x7, 0x5, 0x3, 0x1, + 0x3F, 0x3D, 0x3B, 0x39, 0x37, 0x35, 0x33, 0x31, 0x2F, 0x2D, 0x2B, 0x29, 0x27, 0x25, 0x23, 0x21); + + dim_t n_full_pieces = NC / NR; + dim_t n_full_pieces_loop_limit = n_full_pieces * NR; + dim_t n_partial_pieces = NC % NR; + + dim_t k_full_pieces_blks = KC / 2; + dim_t k_full_pieces = k_full_pieces_blks * 2; + dim_t k_partial_pieces = KC % 2; + + // KC when not multiple of 2 will have padding to make it multiple of 2 in packed buffer. + dim_t KC_updated = KC; + if ( k_partial_pieces > 0 ) + { + KC_updated += ( 2 - k_partial_pieces ); + } + + __m512i a0, b0, c0, d0; + __m512i a01, c01; + + for ( dim_t jc = 0; jc < n_full_pieces_loop_limit; jc += NR ) + { + for ( dim_t kr = 0; kr < k_full_pieces; kr += 2 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 64 elements in each row. + a0 = _mm512_loadu_si512( b + ( jc * KC_updated ) + ( ( kr + 0 ) * NR ) ); + b0 = _mm512_loadu_si512( b + ( jc * KC_updated ) + ( ( kr + 0 ) * NR ) + 32 ); + c0 = _mm512_loadu_si512( b + ( jc * KC_updated ) + ( ( kr + 1 ) * NR ) ); + d0 = _mm512_loadu_si512( b + ( jc * KC_updated ) + ( ( kr + 1 ) * NR ) + 32 ); + + a01 = _mm512_permutex2var_epi16( b0, selector_even, a0 ); + b0 = _mm512_permutex2var_epi16( b0, selector_odd, a0 ); + + c01 = _mm512_permutex2var_epi16( d0, selector_even, c0 ); + d0 = _mm512_permutex2var_epi16( d0, selector_odd, c0 ); + + // Store to unpack buffer + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 0 ) ) + jc, a01 ); + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 0 ) ) + jc + 32, c01 ); + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 1 ) ) + jc, b0 ); + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( kr + 1 ) ) + jc + 32, d0 ); + + } + if( k_partial_pieces > 0 ) + { + a0 = _mm512_loadu_si512( b + ( jc * KC_updated ) + ( ( k_full_pieces + 0 ) * NR ) ); + b0 = _mm512_loadu_si512( b + ( jc * KC_updated ) + ( ( k_full_pieces + 0 ) * NR ) + 32 ); + c0 = _mm512_loadu_si512( b + ( jc * KC_updated ) + ( ( k_full_pieces + 1 ) * NR ) ); + d0 = _mm512_loadu_si512( b + ( jc * KC_updated ) + ( ( k_full_pieces + 1 ) * NR ) + 32 ); + + a01 = _mm512_permutex2var_epi16( b0, selector_even, a0 ); + + c01 = _mm512_permutex2var_epi16( d0, selector_even, c0 ); + + // Store to unpack buffer + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( k_full_pieces + 0 ) ) + jc, a01 ); + _mm512_storeu_si512( unpack_b_buffer_bf16bf16f32of32 + ( ldb * ( k_full_pieces + 0 ) ) + jc + 32, c01 ); + } + } + + if( n_partial_pieces > 0 ) + { + dim_t n0_partial_rem = n_partial_pieces % 16; + dim_t n0_partial_unpack = 0; + + // Split into multiple smaller fringe kernels, so as to maximize + // vectorization after packing. Any n0 < NR(64) can be expressed + // as n0 = 48 + n` / n0 = 32 + n` / n0 = 16 + n`, where n` < 16. + dim_t n0_48 = n_partial_pieces / 48; + dim_t n0_32 = n_partial_pieces / 32; + dim_t n0_16 = n_partial_pieces / 16; + + if ( n0_48 == 1 ) + { + unpackb_nr48_bf16bf16f32of32_row_major + ( + ( b + ( n_full_pieces_loop_limit * KC_updated ) ), + ( unpack_b_buffer_bf16bf16f32of32 + n_full_pieces_loop_limit ), KC, ldb + ); + + n0_partial_unpack = 48; + } + else if ( n0_32 == 1 ) + { + unpackb_nr32_bf16bf16f32of32_row_major + ( + ( b + ( n_full_pieces_loop_limit * KC_updated ) ), + ( unpack_b_buffer_bf16bf16f32of32 + n_full_pieces_loop_limit ), KC, ldb + ); + + n0_partial_unpack = 32; + } + else if ( n0_16 == 1 ) + { + unpackb_nr16_bf16bf16f32of32_row_major + ( + ( b + ( n_full_pieces_loop_limit * KC_updated ) ), + ( unpack_b_buffer_bf16bf16f32of32 + n_full_pieces_loop_limit ), KC, ldb + ); + + n0_partial_unpack = 16; + } + + if ( n0_partial_rem > 0 ) + { + unpackb_nrlt16_bf16bf16f32of32_row_major + ( + ( b + ( n_full_pieces_loop_limit * KC_updated ) + + ( n0_partial_unpack * KC_updated ) ), + ( unpack_b_buffer_bf16bf16f32of32 + n_full_pieces_loop_limit + n0_partial_unpack ), KC, ldb, + n0_partial_rem + ); + } + } +} + +#define STORE_16_COLS_AVX512 \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 0 ) ) + kr, a_reg[0] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 1 ) ) + kr, a_reg[1] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 2 ) ) + kr, a_reg[2] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 3 ) ) + kr, a_reg[3] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 4 ) ) + kr, a_reg[4] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 5 ) ) + kr, a_reg[5] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 6 ) ) + kr, a_reg[6] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 7 ) ) + kr, a_reg[7] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 8 ) ) + kr, a_reg[8] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 9 ) ) + kr, a_reg[9] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 10 ) ) + kr, a_reg[10] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 11 ) ) + kr, a_reg[11] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 12 ) ) + kr, a_reg[12] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 13 ) ) + kr, a_reg[13] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 14 ) ) + kr, a_reg[14] ); \ + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 15 ) ) + kr, a_reg[15] ); + +#define MASK_STORE_16_COLS_AVX512( mask ) \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 0 ) ) + kr, mask, a_reg[0] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 1 ) ) + kr, mask, a_reg[1] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 2 ) ) + kr, mask, a_reg[2] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 3 ) ) + kr, mask, a_reg[3] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 4 ) ) + kr, mask, a_reg[4] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 5 ) ) + kr, mask, a_reg[5] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 6 ) ) + kr, mask, a_reg[6] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 7 ) ) + kr, mask, a_reg[7] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 8 ) ) + kr, mask, a_reg[8] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 9 ) ) + kr, mask, a_reg[9] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 10 ) ) + kr, mask, a_reg[10] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 11 ) ) + kr, mask, a_reg[11] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 12 ) ) + kr, mask, a_reg[12] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 13 ) ) + kr, mask, a_reg[13] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 14 ) ) + kr, mask, a_reg[14] ); \ + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 15 ) ) + kr, mask, a_reg[15] ); \ + + + +#define UNPACKHILO32_AVX512 \ + b_reg[0] = _mm512_unpacklo_epi32(a_reg[0], a_reg[1]); \ + b_reg[2] = _mm512_unpacklo_epi32(a_reg[2], a_reg[3]); \ + b_reg[4] = _mm512_unpacklo_epi32(a_reg[4], a_reg[5]); \ + b_reg[6] = _mm512_unpacklo_epi32(a_reg[6], a_reg[7]); \ + b_reg[8] = _mm512_unpacklo_epi32(a_reg[8], a_reg[9]); \ + b_reg[10] = _mm512_unpacklo_epi32(a_reg[10], a_reg[11]); \ + b_reg[12] = _mm512_unpacklo_epi32(a_reg[12], a_reg[13]); \ + b_reg[14] = _mm512_unpacklo_epi32(a_reg[14], a_reg[15]); \ +\ + b_reg[1] = _mm512_unpackhi_epi32(a_reg[0], a_reg[1]); \ + b_reg[3] = _mm512_unpackhi_epi32(a_reg[2], a_reg[3]); \ + b_reg[5] = _mm512_unpackhi_epi32(a_reg[4], a_reg[5]); \ + b_reg[7] = _mm512_unpackhi_epi32(a_reg[6], a_reg[7]); \ + b_reg[9] = _mm512_unpackhi_epi32(a_reg[8], a_reg[9]); \ + b_reg[11] = _mm512_unpackhi_epi32(a_reg[10], a_reg[11]); \ + b_reg[13] = _mm512_unpackhi_epi32(a_reg[12], a_reg[13]); \ + b_reg[15] = _mm512_unpackhi_epi32(a_reg[14], a_reg[15]); + +#define UNPACKHILO64_AVX512 \ + a_reg[0] = _mm512_unpacklo_epi64(b_reg[0], b_reg[2]); \ + a_reg[1] = _mm512_unpacklo_epi64(b_reg[4], b_reg[6]); \ + a_reg[2] = _mm512_unpacklo_epi64(b_reg[8], b_reg[10]); \ + a_reg[3] = _mm512_unpacklo_epi64(b_reg[12], b_reg[14]); \ + a_reg[4] = _mm512_unpacklo_epi64(b_reg[1], b_reg[3]); \ + a_reg[5] = _mm512_unpacklo_epi64(b_reg[5], b_reg[7]); \ + a_reg[6] = _mm512_unpacklo_epi64(b_reg[9], b_reg[11]); \ + a_reg[7] = _mm512_unpacklo_epi64(b_reg[13], b_reg[15]); \ +\ + a_reg[8] = _mm512_unpackhi_epi64(b_reg[0], b_reg[2]); \ + a_reg[9] = _mm512_unpackhi_epi64(b_reg[4], b_reg[6]); \ + a_reg[10] = _mm512_unpackhi_epi64(b_reg[8], b_reg[10]); \ + a_reg[11] = _mm512_unpackhi_epi64(b_reg[12], b_reg[14]); \ + a_reg[12] = _mm512_unpackhi_epi64(b_reg[1], b_reg[3]); \ + a_reg[13] = _mm512_unpackhi_epi64(b_reg[5], b_reg[7]); \ + a_reg[14] = _mm512_unpackhi_epi64(b_reg[9], b_reg[11]); \ + a_reg[15] = _mm512_unpackhi_epi64(b_reg[13], b_reg[15]); + +#define PERMUTEX2_VAR64_AVX512 \ + b_reg[0] = _mm512_permutex2var_epi64(a_reg[0], selector1, a_reg[1]); \ + b_reg[1] = _mm512_permutex2var_epi64(a_reg[2], selector1, a_reg[3]); \ + b_reg[2] = _mm512_permutex2var_epi64(a_reg[8], selector1, a_reg[9]); \ + b_reg[3] = _mm512_permutex2var_epi64(a_reg[10], selector1, a_reg[11]); \ + b_reg[4] = _mm512_permutex2var_epi64(a_reg[4], selector1, a_reg[5]); \ + b_reg[5] = _mm512_permutex2var_epi64(a_reg[6], selector1, a_reg[7]); \ + b_reg[6] = _mm512_permutex2var_epi64(a_reg[12], selector1, a_reg[13]); \ + b_reg[7] = _mm512_permutex2var_epi64(a_reg[14], selector1, a_reg[15]); \ + b_reg[8] = _mm512_permutex2var_epi64(a_reg[0], selector2, a_reg[1]); \ + b_reg[9] = _mm512_permutex2var_epi64(a_reg[2], selector2, a_reg[3]); \ + b_reg[10] = _mm512_permutex2var_epi64(a_reg[8], selector2, a_reg[9]); \ + b_reg[11] = _mm512_permutex2var_epi64(a_reg[10], selector2, a_reg[11]); \ + b_reg[12] = _mm512_permutex2var_epi64(a_reg[4], selector2, a_reg[5]); \ + b_reg[13] = _mm512_permutex2var_epi64(a_reg[6], selector2, a_reg[7]); \ + b_reg[14] = _mm512_permutex2var_epi64(a_reg[12], selector2, a_reg[13]); \ + b_reg[15] = _mm512_permutex2var_epi64(a_reg[14], selector2, a_reg[15]); + +#define SHUFFLE64x2_AVX512 \ + a_reg[0] = _mm512_shuffle_i64x2(b_reg[0], b_reg[1], 0x44); \ + a_reg[1] = _mm512_shuffle_i64x2(b_reg[2], b_reg[3], 0x44); \ + a_reg[2] = _mm512_shuffle_i64x2(b_reg[4], b_reg[5], 0x44); \ + a_reg[3] = _mm512_shuffle_i64x2(b_reg[6], b_reg[7], 0x44); \ + a_reg[4] = _mm512_shuffle_i64x2(b_reg[8], b_reg[9], 0x44); \ + a_reg[5] = _mm512_shuffle_i64x2(b_reg[10], b_reg[11], 0x44); \ + a_reg[6] = _mm512_shuffle_i64x2(b_reg[12], b_reg[13], 0x44); \ + a_reg[7] = _mm512_shuffle_i64x2(b_reg[14], b_reg[15], 0x44); \ + a_reg[8] = _mm512_shuffle_i64x2(b_reg[0], b_reg[1], 0xEE); \ + a_reg[9] = _mm512_shuffle_i64x2(b_reg[2], b_reg[3], 0xEE); \ + a_reg[10] = _mm512_shuffle_i64x2(b_reg[4], b_reg[5], 0xEE); \ + a_reg[11] = _mm512_shuffle_i64x2(b_reg[6], b_reg[7], 0xEE); \ + a_reg[12] = _mm512_shuffle_i64x2(b_reg[8], b_reg[9], 0xEE); \ + a_reg[13] = _mm512_shuffle_i64x2(b_reg[10], b_reg[11], 0xEE); \ + a_reg[14] = _mm512_shuffle_i64x2(b_reg[12], b_reg[13], 0xEE); \ + a_reg[15] = _mm512_shuffle_i64x2(b_reg[14], b_reg[15], 0xEE); + + +void unpackb_nrlt16_bf16bf16f32of32_col_major + ( + const bfloat16* b, + bfloat16* unpack_b_buffer, + const dim_t KC, + dim_t ldb, + dim_t n0_partial_rem + ) +{ + dim_t NR = 16; + + // Used for permuting the mm512i elements for use in dpbf16_ps instruction. + __m512i selector1 = _mm512_setr_epi64( 0x0, 0x1, 0x8, 0x9, 0x4, 0x5, 0xC, 0xD ); + __m512i selector2 = _mm512_setr_epi64( 0x2, 0x3, 0xA, 0xB, 0x6, 0x7, 0xE, 0xF ); + + __m512i a_reg[16]; + __m512i b_reg[16]; + + // These registers are set with zeroes to avoid compiler warnings + // To-DO: TO be removed when pack code is optimized for fringe cases. + a_reg[0] = _mm512_setzero_si512(); + a_reg[1] = _mm512_setzero_si512(); + a_reg[2] = _mm512_setzero_si512(); + a_reg[3] = _mm512_setzero_si512(); + a_reg[4] = _mm512_setzero_si512(); + a_reg[5] = _mm512_setzero_si512(); + a_reg[6] = _mm512_setzero_si512(); + a_reg[7] = _mm512_setzero_si512(); + a_reg[8] = _mm512_setzero_si512(); + a_reg[9] = _mm512_setzero_si512(); + a_reg[10] = _mm512_setzero_si512(); + a_reg[11] = _mm512_setzero_si512(); + a_reg[12] = _mm512_setzero_si512(); + a_reg[13] = _mm512_setzero_si512(); + a_reg[14] = _mm512_setzero_si512(); + a_reg[15] = _mm512_setzero_si512(); + + dim_t kr = 0, jr = 0; + for ( kr = 0; ( kr + 31 ) < KC; kr += 32 ) + { + a_reg[0] = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + a_reg[1] = _mm512_loadu_si512( b + ( ( kr + 2 ) * NR ) ); + a_reg[2] = _mm512_loadu_si512( b + ( ( kr + 4 ) * NR ) ); + a_reg[3] = _mm512_loadu_si512( b + ( ( kr + 6 ) * NR ) ); + a_reg[4] = _mm512_loadu_si512( b + ( ( kr + 8 ) * NR ) ); + a_reg[5] = _mm512_loadu_si512( b + ( ( kr + 10 ) * NR ) ); + a_reg[6] = _mm512_loadu_si512( b + ( ( kr + 12 ) * NR ) ); + a_reg[7] = _mm512_loadu_si512( b + ( ( kr + 14 ) * NR ) ); + a_reg[8] = _mm512_loadu_si512( b + ( ( kr + 16 ) * NR ) ); + a_reg[9] = _mm512_loadu_si512( b + ( ( kr + 18 ) * NR ) ); + a_reg[10] = _mm512_loadu_si512( b + ( ( kr + 20 ) * NR ) ); + a_reg[11] = _mm512_loadu_si512( b + ( ( kr + 22 ) * NR ) ); + a_reg[12] = _mm512_loadu_si512( b + ( ( kr + 24 ) * NR ) ); + a_reg[13] = _mm512_loadu_si512( b + ( ( kr + 26 ) * NR ) ); + a_reg[14] = _mm512_loadu_si512( b + ( ( kr + 28 ) * NR ) ); + a_reg[15] = _mm512_loadu_si512( b + ( ( kr + 30 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + + for( jr = 0; jr < n0_partial_rem; jr += 1 ) + { + _mm512_storeu_si512( unpack_b_buffer + ( ldb * ( jr + 0 ) ) + kr, a_reg[jr] ); + } + } + + for ( ; ( kr + 15 ) < KC; kr += 16 ) + { + a_reg[0] = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + a_reg[1] = _mm512_loadu_si512( b + ( ( kr + 2 ) * NR ) ); + a_reg[2] = _mm512_loadu_si512( b + ( ( kr + 4 ) * NR ) ); + a_reg[3] = _mm512_loadu_si512( b + ( ( kr + 6 ) * NR ) ); + a_reg[4] = _mm512_loadu_si512( b + ( ( kr + 8 ) * NR ) ); + a_reg[5] = _mm512_loadu_si512( b + ( ( kr + 10 ) * NR ) ); + a_reg[6] = _mm512_loadu_si512( b + ( ( kr + 12 ) * NR ) ); + a_reg[7] = _mm512_loadu_si512( b + ( ( kr + 14 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + + for( jr = 0; jr < n0_partial_rem; jr += 1 ) + { + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 0 ) ) + kr, 0xFFFF, a_reg[jr] ); + } + } + + for ( ; ( kr + 7 ) < KC; kr += 8 ) + { + a_reg[0] = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + a_reg[1] = _mm512_loadu_si512( b + ( ( kr + 2 ) * NR ) ); + a_reg[2] = _mm512_loadu_si512( b + ( ( kr + 4 ) * NR ) ); + a_reg[3] = _mm512_loadu_si512( b + ( ( kr + 6 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + + for( jr = 0; jr < n0_partial_rem; jr += 1 ) + { + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 0 ) ) + kr, 0xFF, a_reg[jr] ); + } + } + + for ( ; (kr+3) < KC; kr += 4 ) + { + a_reg[0] = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + a_reg[1] = _mm512_loadu_si512( b + ( ( kr + 2 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + + for( jr = 0; jr < n0_partial_rem; jr += 1 ) + { + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 0 ) ) + kr, 0xF, a_reg[jr] ); + } + } + + for ( ; ( kr + 1 ) < KC; kr += 2 ) + { + a_reg[0] = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + + for( jr = 0; jr < n0_partial_rem; jr += 1 ) + { + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 0 ) ) + kr, 0x03, a_reg[jr] ); + } + } + + for ( ; kr < KC; kr += 1 ) + { + a_reg[0] = _mm512_loadu_si512( b + ( ( kr + 0 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + + for( jr = 0; jr < n0_partial_rem; jr += 1 ) + { + _mm512_mask_storeu_epi16( unpack_b_buffer + ( ldb * ( jr + 0 ) ) + kr, 0x01, a_reg[jr] ); + } + } +} + +void unpackb_nr_mult_16_bf16bf16f32of32_col_major + ( + const bfloat16* b, + bfloat16* unpack_b_buffer, + const dim_t NR, + const dim_t KC, + dim_t ldb + ) +{ + // Used for permuting the mm512i elements for use in dpbf16_ps instruction. + __m512i selector1 = _mm512_setr_epi64( 0x0, 0x1, 0x8, 0x9, 0x4, 0x5, 0xC, 0xD ); + __m512i selector2 = _mm512_setr_epi64( 0x2, 0x3, 0xA, 0xB, 0x6, 0x7, 0xE, 0xF ); + + __m512i a_reg[16]; + __m512i b_reg[16]; + + // These registers are set with zeroes to avoid compiler warnings + // To-DO: TO be removed when pack code is optimized for fringe cases. + a_reg[0] = _mm512_setzero_si512(); + a_reg[1] = _mm512_setzero_si512(); + a_reg[2] = _mm512_setzero_si512(); + a_reg[3] = _mm512_setzero_si512(); + a_reg[4] = _mm512_setzero_si512(); + a_reg[5] = _mm512_setzero_si512(); + a_reg[6] = _mm512_setzero_si512(); + a_reg[7] = _mm512_setzero_si512(); + a_reg[8] = _mm512_setzero_si512(); + a_reg[9] = _mm512_setzero_si512(); + a_reg[10] = _mm512_setzero_si512(); + a_reg[11] = _mm512_setzero_si512(); + a_reg[12] = _mm512_setzero_si512(); + a_reg[13] = _mm512_setzero_si512(); + a_reg[14] = _mm512_setzero_si512(); + a_reg[15] = _mm512_setzero_si512(); + + dim_t kr = 0; + for ( kr = 0; ( kr + 31 ) < KC; kr += 32 ) + { + for( dim_t jr = 0; jr < NR; jr += 16 ) + { + a_reg[0] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 0 ) * NR ) ); + a_reg[1] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 2 ) * NR ) ); + a_reg[2] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 4 ) * NR ) ); + a_reg[3] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 6 ) * NR ) ); + a_reg[4] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 8 ) * NR ) ); + a_reg[5] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 10 ) * NR ) ); + a_reg[6] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 12 ) * NR ) ); + a_reg[7] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 14 ) * NR ) ); + a_reg[8] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 16 ) * NR ) ); + a_reg[9] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 18 ) * NR ) ); + a_reg[10] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 20 ) * NR ) ); + a_reg[11] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 22 ) * NR ) ); + a_reg[12] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 24 ) * NR ) ); + a_reg[13] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 26 ) * NR ) ); + a_reg[14] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 28 ) * NR ) ); + a_reg[15] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 30 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + STORE_16_COLS_AVX512 + } + } + + for ( ; ( kr + 15 ) < KC; kr += 16 ) + { + for( dim_t jr = 0; jr < NR; jr += 16 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 64 elements in each row. + + a_reg[0] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 0 ) * NR ) ); + a_reg[1] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 2 ) * NR ) ); + a_reg[2] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 4 ) * NR ) ); + a_reg[3] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 6 ) * NR ) ); + a_reg[4] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 8 ) * NR ) ); + a_reg[5] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 10 ) * NR ) ); + a_reg[6] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 12 ) * NR ) ); + a_reg[7] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 14 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + MASK_STORE_16_COLS_AVX512( 0xFFFF ) + } + } + + for( ; ( kr +7 ) < KC; kr += 8 ) + { + for( dim_t jr = 0; jr < NR; jr += 16 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 64 elements in each row. + + a_reg[0] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 0 ) * NR ) ); + a_reg[1] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 2 ) * NR ) ); + a_reg[2] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 4 ) * NR ) ); + a_reg[3] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 6 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + MASK_STORE_16_COLS_AVX512( 0xFF ) + } + } + + for( ; ( kr +3 ) < KC; kr += 4 ) + { + for( dim_t jr = 0; jr < NR; jr += 16 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 64 elements in each row. + a_reg[0] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 0 ) * NR ) ); + a_reg[1] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 2 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + MASK_STORE_16_COLS_AVX512( 0xF ) + } + } + + for( ; ( kr +1 ) < KC; kr += 2 ) + { + for( dim_t jr = 0; jr < NR; jr += 16 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 64 elements in each row. + a_reg[0] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 0 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + MASK_STORE_16_COLS_AVX512( 0x3 ) + } + } + + for( ; kr < KC; kr += 1 ) + { + for( dim_t jr = 0; jr < NR; jr += 16 ) + { + // Rearrange for dpbf16_ps, read 2 rows from B with 64 elements in each row. + a_reg[0] = _mm512_loadu_si512( b + ( jr * 2 ) + ( ( kr + 0 ) * NR ) ); + + UNPACKHILO32_AVX512 + UNPACKHILO64_AVX512 + PERMUTEX2_VAR64_AVX512 + SHUFFLE64x2_AVX512 + MASK_STORE_16_COLS_AVX512( 0x1 ) + } + } +}; + +void unpackb_nr64_bf16bf16f32of32_col_major + ( + const bfloat16* b, + bfloat16* unpack_b_buffer, + const dim_t NC, + const dim_t KC, + dim_t ldb + ) +{ + dim_t NR = 64; + + dim_t n_full_pieces = NC / NR; + dim_t n_full_pieces_loop_limit = n_full_pieces * NR; + dim_t n_partial_pieces = NC % NR; + + dim_t k_partial_pieces = KC % 2; + + dim_t KC_updated = KC; + if ( k_partial_pieces > 0 ) + { + KC_updated += ( 2 - k_partial_pieces ); + } + + for ( dim_t jc = 0; jc < n_full_pieces_loop_limit; jc += NR ) + { + unpackb_nr_mult_16_bf16bf16f32of32_col_major + ( b + (jc * KC_updated), + unpack_b_buffer + (jc * ldb), 64, KC, ldb + ); + } + + if(n_partial_pieces > 0) + { + dim_t n0_partial_rem = n_partial_pieces % 16; + dim_t n0_partial_pack = 0; + + // Split into multiple smaller fringe kernels, so as to maximize + // vectorization after packing. Any n0 < NR(64) can be expressed + // as n0 = 48 + n` / n0 = 32 + n` / n0 = 16 + n`, where n` < 16. + dim_t n0_48 = n_partial_pieces / 48; + dim_t n0_32 = n_partial_pieces / 32; + dim_t n0_16 = n_partial_pieces / 16; + + if ( n0_48 == 1 ) + { + unpackb_nr_mult_16_bf16bf16f32of32_col_major + ( + ( b + ( n_full_pieces_loop_limit * KC_updated ) ), + ( unpack_b_buffer + n_full_pieces_loop_limit * ldb ), 48, KC, ldb + ); + + n0_partial_pack = 48; + } + else if ( n0_32 == 1 ) + { + unpackb_nr_mult_16_bf16bf16f32of32_col_major + ( + ( b + ( n_full_pieces_loop_limit * KC_updated ) ), + ( unpack_b_buffer + n_full_pieces_loop_limit * ldb ), 32, KC, ldb + ); + + n0_partial_pack = 32; + } + else if ( n0_16 == 1 ) + { + unpackb_nr_mult_16_bf16bf16f32of32_col_major + ( + ( b + ( n_full_pieces_loop_limit * KC_updated ) ), + ( unpack_b_buffer + n_full_pieces_loop_limit * ldb ), 16, KC, ldb + ); + + n0_partial_pack = 16; + } + + if ( n0_partial_rem > 0 ) + { + unpackb_nrlt16_bf16bf16f32of32_col_major + ( + ( b + ( n_full_pieces_loop_limit * KC_updated ) + + ( n0_partial_pack * KC_updated ) ), + ( unpack_b_buffer + ( n_full_pieces_loop_limit + n0_partial_pack ) * ldb ), KC, ldb, + n0_partial_rem + ); + } + } +}; + +void unpackb_nr64_bf16bf16f32of32 + ( + const bfloat16* b, + bfloat16* unpack_b_buffer_bf16bf16f32of32, + const dim_t NC, + const dim_t KC, + dim_t rs_b, + dim_t cs_b + ) +{ + if( cs_b == 1 ) + { + unpackb_nr64_bf16bf16f32of32_row_major( b, unpack_b_buffer_bf16bf16f32of32, NC, KC, rs_b ); + } + else + { + unpackb_nr64_bf16bf16f32of32_col_major( b, unpack_b_buffer_bf16bf16f32of32, NC, KC, cs_b ); + } +} +#endif // BLIS_ADDON_LPGEMM