amd/blis
1
0
Fork 0
mirror of https://github.com/amd/blis.git synced 2026-07-14 11:07:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

Introduced un-reorder API for bf16bf16f32of32

Browse Source 
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
This commit is contained in:
Meghana Vankadari
2024-10-01 07:07:46 +05:30
committed by sireesha.sanga
parent bd8bc77381
commit aed0803606
14 changed files with 1495 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

123
addon/aocl_gemm/aocl_gemm_bf16_utils.c
View File

@@ -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))

16
addon/aocl_gemm/aocl_gemm_interface_apis.h
View File

@@ -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 \
( \

2
addon/aocl_gemm/config/lpgemm_config.c
View File

@@ -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

3
addon/aocl_gemm/config/lpgemm_func_map.h
View File

@@ -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) \

87
addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.c
View File

@@ -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,

10
addon/aocl_gemm/frame/bf16bf16f32/lpgemm_reorder_bf16.h
View File

@@ -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

1
addon/aocl_gemm/frame/lpgemm_types.h
View File

@@ -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;

21
addon/aocl_gemm/kernels/bf16bf16f32/lpgemm_pack_bf16.h
View File

@@ -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

3
bench/bench_aocl_gemm/Makefile
View File

@@ -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

35
bench/bench_aocl_gemm/bench_lpgemm.c
View File

@@ -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)

36
bench/bench_aocl_gemm/bench_lpgemm_helpers.h
View File

@@ -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

276
bench/bench_aocl_gemm/bench_lpgemm_unpack.c Normal file
View File

@@ -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;
}

6
bench/bench_aocl_gemm/bench_unpack_input.txt Normal file
View File

@@ -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

921
kernels/zen4/lpgemm/bf16bf16f32/lpgemm_unpackb_bf16_amd512vnni.c Normal file
View File

@@ -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 <immintrin.h>
#include <string.h>
#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