mirror of
https://github.com/amd/blis.git
synced 2026-05-04 06:21:12 +00:00
bfc512d3e1
Details: - The batch matmul performs a series of matmuls, processing more than one GEMM problem at once. - Introduced a new parameter called batch_size for the user to indicate number of GEMM problems in a batch/group. - This operation supports processing GEMM problems with different parameters including dims,post-ops,stor-schemes etc., - This operation is optimized for problems where all the GEMMs in a batch are of same size and shape. - For now, the threads are distributed among different GEMM problems equally irrespective of their dimensions which leads to better performance for batches with identical GEMMs but performs sub-optimally for batches with non-identical GEMMs. - Optimizations for batches with non-identical GEMMs is in progress. - Added bench and input files for batch_matmul. AMD-Internal: [SWLCSG-2944] Change-Id: Idc59db5b8c5794bf19f6f86bcb8455cd2599c155
480 lines
13 KiB
C
480 lines
13 KiB
C
/*
|
|
|
|
BLIS
|
|
An object-based framework for developing high-performance BLAS-like
|
|
libraries.
|
|
|
|
Copyright (C) 2025, 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 "blis.h"
|
|
#include "aocl_gemm_interface_apis.h"
|
|
#include "aocl_gemm_check.h"
|
|
#include "lpgemm_types.h"
|
|
#include "lpgemm_post_ops.h"
|
|
#include "lpgemm_thread_decor_openmp.h"
|
|
#include "lpgemm_5loop_interface_apis.h"
|
|
#include "lpgemm_config.h"
|
|
#include "lpgemm_utils.h"
|
|
|
|
AOCL_BGEMM_MATMUL(bfloat16,bfloat16,float,float,bf16bf16f32of32)
|
|
{
|
|
// Check if avx512_vnni 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();
|
|
|
|
#ifdef LPGEMM_BF16_JIT
|
|
if( jit_kernels_generated == FALSE )
|
|
{
|
|
bli_print_msg(" Could not generate bf16bf16f32of32 "
|
|
" kernels using JIT.", __FILE__, __LINE__ );
|
|
return;
|
|
}
|
|
#endif
|
|
|
|
trans_t blis_transa;
|
|
trans_t blis_transb;
|
|
|
|
inc_t rs_a[batch_size];
|
|
inc_t cs_a[batch_size];
|
|
|
|
inc_t rs_b[batch_size];
|
|
inc_t cs_b[batch_size];
|
|
|
|
inc_t rs_c[batch_size];
|
|
inc_t cs_c[batch_size];
|
|
|
|
AOCL_MEMORY_TAG mtag_a[batch_size];
|
|
AOCL_MEMORY_TAG mtag_b[batch_size];
|
|
|
|
bfloat16 *a_local[batch_size], *b_local[batch_size];
|
|
dim_t m_local[batch_size], n_local[batch_size];
|
|
|
|
// Convert post op struct to post op linked list format.
|
|
lpgemm_post_op post_op_list[batch_size][AOCL_MAX_POST_OPS];
|
|
|
|
for( dim_t bs_i = 0; bs_i < batch_size; bs_i++ )
|
|
{
|
|
// check for validity of params.
|
|
AOCL_BATCH_GEMM_CHECK
|
|
(
|
|
"batch_bf16bf16f32of32",
|
|
order[bs_i], transa[bs_i], transb[bs_i],
|
|
bs_i,
|
|
m[bs_i], n[bs_i], k[bs_i],
|
|
a[bs_i], lda[bs_i], mem_format_a[bs_i],
|
|
b[bs_i], ldb[bs_i], mem_format_b[bs_i],
|
|
c[bs_i], ldc[bs_i]
|
|
);
|
|
|
|
/* Map BLAS chars to their corresponding BLIS enumerated type value. */
|
|
bli_param_map_netlib_to_blis_trans( transa[bs_i], &blis_transa );
|
|
bli_param_map_netlib_to_blis_trans( transb[bs_i], &blis_transb );
|
|
|
|
bool is_column_major = ( ( order[bs_i] == 'c' ) || ( order[bs_i] == 'C' ) );
|
|
|
|
if( is_column_major == TRUE )
|
|
{
|
|
rs_a[bs_i] = ldb[bs_i];
|
|
cs_a[bs_i] = 1;
|
|
|
|
if( bli_is_trans( blis_transb ) )
|
|
{
|
|
rs_a[bs_i] = 1;
|
|
cs_a[bs_i] = ldb[bs_i];
|
|
}
|
|
|
|
rs_b[bs_i] = lda[bs_i];
|
|
cs_b[bs_i] = 1;
|
|
|
|
if( bli_is_trans( blis_transa ) )
|
|
{
|
|
rs_b[bs_i] = 1;
|
|
cs_b[bs_i] = lda[bs_i];
|
|
}
|
|
|
|
bli_param_map_char_to_lpmtag( mem_format_a[bs_i], &(mtag_b[bs_i]) );
|
|
bli_param_map_char_to_lpmtag( mem_format_b[bs_i], &(mtag_a[bs_i]) );
|
|
|
|
// Inputs swapped in column major, A becomes B from kernel point of view.
|
|
// Reorder is not supported for column major matrices.
|
|
if ( ( ( mtag_b[bs_i] == REORDERED ) || ( mtag_a[bs_i] == REORDERED ) ) )
|
|
{
|
|
bli_print_msg(" Reordering of column major matrices is not supported.", __FILE__, __LINE__ );
|
|
return;
|
|
}
|
|
// From 5-loop function point of view,
|
|
// A matrix when in column major storage needs to be packed to row-major
|
|
// storage as kernel expects A matrix to be in row-major format.
|
|
// Inputs swapped in column major, A becomes B from kernel point of view.
|
|
if ( bli_is_trans(blis_transb ) )
|
|
{
|
|
mtag_a[bs_i] = PACK;
|
|
}
|
|
|
|
// swap m & n in case of col-major matrices
|
|
m_local[bs_i] = n[bs_i];
|
|
n_local[bs_i] = m[bs_i];
|
|
|
|
// swap a & b pointers in case of col-major matrices
|
|
a_local[bs_i] = (bfloat16*)(b[bs_i]);
|
|
b_local[bs_i] = (bfloat16*)(a[bs_i]);
|
|
}
|
|
else // row-major
|
|
{
|
|
rs_a[bs_i] = lda[bs_i];
|
|
cs_a[bs_i] = 1;
|
|
|
|
if( bli_is_trans( blis_transa ) )
|
|
{
|
|
rs_a[bs_i] = 1;
|
|
cs_a[bs_i] = lda[bs_i];
|
|
}
|
|
|
|
rs_b[bs_i] = ldb[bs_i];
|
|
cs_b[bs_i] = 1;
|
|
|
|
if( bli_is_trans( blis_transb ) )
|
|
{
|
|
rs_b[bs_i] = 1;
|
|
cs_b[bs_i] = ldb[bs_i];
|
|
}
|
|
|
|
bli_param_map_char_to_lpmtag( mem_format_a[bs_i], &(mtag_a[bs_i]) );
|
|
bli_param_map_char_to_lpmtag( mem_format_b[bs_i], &(mtag_b[bs_i]) );
|
|
|
|
// Reorder is not supported for A matrix
|
|
if( mtag_a[bs_i] == REORDERED )
|
|
{
|
|
bli_print_msg(" Reordering of A matrix is not supported in row major case.", __FILE__, __LINE__ );
|
|
return;
|
|
}
|
|
// From 5-loop function point of view,
|
|
// A matrix when in column major storage needs to be packed to row-major
|
|
// storage as kernel expects A matrix to be in row-major format.
|
|
if( bli_is_trans(blis_transa ) )
|
|
{
|
|
mtag_a[bs_i] = PACK;
|
|
}
|
|
|
|
// copy the values of m & n
|
|
m_local[bs_i] = m[bs_i];
|
|
n_local[bs_i] = n[bs_i];
|
|
|
|
// copy the values of a & b pointers
|
|
a_local[bs_i] = (bfloat16*)(a[bs_i]);
|
|
b_local[bs_i] = (bfloat16*)(b[bs_i]);
|
|
}
|
|
|
|
rs_c[bs_i] = ldc[bs_i];
|
|
cs_c[bs_i] = 1;
|
|
|
|
// From 5-loop function point of view
|
|
// B matrix needs to be packed in a certain format in order to be loaded
|
|
// and used in bf16 instrution. As such the mtag_b always needs to be either
|
|
// packed or reordered. B matrix as it is (unpacked) cannot be used, and
|
|
// the mtag_b is set to packed to enable runtime packing.
|
|
if ( mtag_b[bs_i] == UNPACKED )
|
|
{
|
|
mtag_b[bs_i] = PACK;
|
|
}
|
|
|
|
err_t err = lpgemm_translate_to_post_ops_list
|
|
(
|
|
post_op_unparsed[bs_i], post_op_list[bs_i],
|
|
( void* )c[bs_i], ( void* )( (order + bs_i) ),
|
|
m[bs_i], n[bs_i]
|
|
);
|
|
|
|
if( err != BLIS_SUCCESS ) return;
|
|
|
|
}
|
|
|
|
// 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 );
|
|
|
|
#ifdef BLIS_ENABLE_OPENMP
|
|
batch_lpgemm_bf16bf16f32of32_openmp_thread_decorator
|
|
(
|
|
batch_size, m_local, n_local, k,
|
|
(const bfloat16**)a_local, rs_a, cs_a, mtag_a,
|
|
(const bfloat16**)b_local, rs_b, cs_b, mtag_b,
|
|
c, rs_c, cs_c,
|
|
alpha, beta,
|
|
&rntm_g, lcntx_g,
|
|
post_op_list, F32
|
|
);
|
|
|
|
|
|
#else
|
|
batch_lpgemm_bf16bf16f32of32_thread_decorator
|
|
(
|
|
batch_size, m_local, n_local, k,
|
|
(const bfloat16**)a_local, rs_a, cs_a, mtag_a,
|
|
(const bfloat16**)b_local, rs_b, cs_b, mtag_b,
|
|
c, rs_c, cs_c,
|
|
alpha, beta,
|
|
&rntm_g, lcntx_g,
|
|
post_op_list, F32
|
|
);
|
|
#endif
|
|
}
|
|
|
|
AOCL_BGEMM_MATMUL(bfloat16,bfloat16,bfloat16,float,bf16bf16f32obf16)
|
|
{
|
|
// Check if avx512_vnni 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();
|
|
|
|
#ifdef LPGEMM_BF16_JIT
|
|
if( jit_kernels_generated == FALSE )
|
|
{
|
|
bli_print_msg(" Could not generate bf16bf16f32of32 "
|
|
" kernels using JIT.", __FILE__, __LINE__ );
|
|
return;
|
|
}
|
|
#endif
|
|
|
|
trans_t blis_transa;
|
|
trans_t blis_transb;
|
|
|
|
inc_t rs_a[batch_size];
|
|
inc_t cs_a[batch_size];
|
|
|
|
inc_t rs_b[batch_size];
|
|
inc_t cs_b[batch_size];
|
|
|
|
inc_t rs_c[batch_size];
|
|
inc_t cs_c[batch_size];
|
|
|
|
AOCL_MEMORY_TAG mtag_a[batch_size];
|
|
AOCL_MEMORY_TAG mtag_b[batch_size];
|
|
|
|
bfloat16 *a_local[batch_size], *b_local[batch_size];
|
|
dim_t m_local[batch_size], n_local[batch_size];
|
|
|
|
// Convert post op struct to post op linked list format.
|
|
lpgemm_post_op post_op_list[batch_size][AOCL_MAX_POST_OPS];
|
|
|
|
for( dim_t bs_i = 0; bs_i < batch_size; bs_i++ )
|
|
{
|
|
// check for validity of params.
|
|
AOCL_BATCH_GEMM_CHECK
|
|
(
|
|
"batch_bf16bf16f32obf16",
|
|
order[bs_i], transa[bs_i], transb[bs_i],
|
|
bs_i,
|
|
m[bs_i], n[bs_i], k[bs_i],
|
|
a[bs_i], lda[bs_i], mem_format_a[bs_i],
|
|
b[bs_i], ldb[bs_i], mem_format_b[bs_i],
|
|
c[bs_i], ldc[bs_i]
|
|
);
|
|
|
|
/* Map BLAS chars to their corresponding BLIS enumerated type value. */
|
|
bli_param_map_netlib_to_blis_trans( transa[bs_i], &blis_transa );
|
|
bli_param_map_netlib_to_blis_trans( transb[bs_i], &blis_transb );
|
|
|
|
bool is_column_major = ( ( order[bs_i] == 'c' ) || ( order[bs_i] == 'C' ) );
|
|
|
|
if( is_column_major == TRUE )
|
|
{
|
|
rs_a[bs_i] = ldb[bs_i];
|
|
cs_a[bs_i] = 1;
|
|
|
|
if( bli_is_trans( blis_transb ) )
|
|
{
|
|
rs_a[bs_i] = 1;
|
|
cs_a[bs_i] = ldb[bs_i];
|
|
}
|
|
|
|
rs_b[bs_i] = lda[bs_i];
|
|
cs_b[bs_i] = 1;
|
|
|
|
if( bli_is_trans( blis_transa ) )
|
|
{
|
|
rs_b[bs_i] = 1;
|
|
cs_b[bs_i] = lda[bs_i];
|
|
}
|
|
|
|
bli_param_map_char_to_lpmtag( mem_format_a[bs_i], &(mtag_b[bs_i]) );
|
|
bli_param_map_char_to_lpmtag( mem_format_b[bs_i], &(mtag_a[bs_i]) );
|
|
|
|
// Inputs swapped in column major, A becomes B from kernel point of view.
|
|
// Reorder is not supported for column major matrices.
|
|
if ( ( ( mtag_b[bs_i] == REORDERED ) || ( mtag_a[bs_i] == REORDERED ) ) )
|
|
{
|
|
bli_print_msg(" Reordering of column major matrices is not supported.", __FILE__, __LINE__ );
|
|
return;
|
|
}
|
|
// From 5-loop function point of view,
|
|
// A matrix when in column major storage needs to be packed to row-major
|
|
// storage as kernel expects A matrix to be in row-major format.
|
|
// Inputs swapped in column major, A becomes B from kernel point of view.
|
|
if ( bli_is_trans(blis_transb ) )
|
|
{
|
|
mtag_a[bs_i] = PACK;
|
|
}
|
|
|
|
// swap m & n in case of col-major matrices
|
|
m_local[bs_i] = n[bs_i];
|
|
n_local[bs_i] = m[bs_i];
|
|
|
|
// swap a & b pointers in case of col-major matrices
|
|
a_local[bs_i] = (bfloat16*)(b[bs_i]);
|
|
b_local[bs_i] = (bfloat16*)(a[bs_i]);
|
|
}
|
|
else // row-major
|
|
{
|
|
rs_a[bs_i] = lda[bs_i];
|
|
cs_a[bs_i] = 1;
|
|
|
|
if( bli_is_trans( blis_transa ) )
|
|
{
|
|
rs_a[bs_i] = 1;
|
|
cs_a[bs_i] = lda[bs_i];
|
|
}
|
|
|
|
rs_b[bs_i] = ldb[bs_i];
|
|
cs_b[bs_i] = 1;
|
|
|
|
if( bli_is_trans( blis_transb ) )
|
|
{
|
|
rs_b[bs_i] = 1;
|
|
cs_b[bs_i] = ldb[bs_i];
|
|
}
|
|
|
|
bli_param_map_char_to_lpmtag( mem_format_a[bs_i], &(mtag_a[bs_i]) );
|
|
bli_param_map_char_to_lpmtag( mem_format_b[bs_i], &(mtag_b[bs_i]) );
|
|
|
|
// Reorder is not supported for A matrix
|
|
if( mtag_a[bs_i] == REORDERED )
|
|
{
|
|
bli_print_msg(" Reordering of A matrix is not supported in row major case.", __FILE__, __LINE__ );
|
|
return;
|
|
}
|
|
// From 5-loop function point of view,
|
|
// A matrix when in column major storage needs to be packed to row-major
|
|
// storage as kernel expects A matrix to be in row-major format.
|
|
if( bli_is_trans(blis_transa ) )
|
|
{
|
|
mtag_a[bs_i] = PACK;
|
|
}
|
|
|
|
// copy the values of m & n
|
|
m_local[bs_i] = m[bs_i];
|
|
n_local[bs_i] = n[bs_i];
|
|
|
|
// copy the values of a & b pointers
|
|
a_local[bs_i] = (bfloat16*)(a[bs_i]);
|
|
b_local[bs_i] = (bfloat16*)(b[bs_i]);
|
|
}
|
|
|
|
rs_c[bs_i] = ldc[bs_i];
|
|
cs_c[bs_i] = 1;
|
|
|
|
// From 5-loop function point of view
|
|
// B matrix needs to be packed in a certain format in order to be loaded
|
|
// and used in bf16 instrution. As such the mtag_b always needs to be either
|
|
// packed or reordered. B matrix as it is (unpacked) cannot be used, and
|
|
// the mtag_b is set to packed to enable runtime packing.
|
|
if ( mtag_b[bs_i] == UNPACKED )
|
|
{
|
|
mtag_b[bs_i] = PACK;
|
|
}
|
|
|
|
err_t err = lpgemm_translate_to_post_ops_list
|
|
(
|
|
post_op_unparsed[bs_i], post_op_list[bs_i],
|
|
( void* )c[bs_i], ( void* )( (order + bs_i) ),
|
|
m[bs_i], n[bs_i]
|
|
);
|
|
|
|
if( err != BLIS_SUCCESS ) return;
|
|
|
|
}
|
|
|
|
// 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 );
|
|
|
|
#ifdef BLIS_ENABLE_OPENMP
|
|
batch_lpgemm_bf16bf16f32of32_openmp_thread_decorator
|
|
(
|
|
batch_size, m_local, n_local, k,
|
|
(const bfloat16**)a_local, rs_a, cs_a, mtag_a,
|
|
(const bfloat16**)b_local, rs_b, cs_b, mtag_b,
|
|
(float**)c, rs_c, cs_c,
|
|
alpha, beta,
|
|
&rntm_g, lcntx_g,
|
|
post_op_list, BF16
|
|
);
|
|
|
|
|
|
#else
|
|
batch_lpgemm_bf16bf16f32of32_thread_decorator
|
|
(
|
|
batch_size, m_local, n_local, k,
|
|
(const bfloat16**)a_local, rs_a, cs_a, mtag_a,
|
|
(const bfloat16**)b_local, rs_b, cs_b, mtag_b,
|
|
(float**)c, rs_c, cs_c,
|
|
alpha, beta,
|
|
&rntm_g, lcntx_g,
|
|
post_op_list, BF16
|
|
);
|
|
#endif
|
|
}
|