OSchip
/
OpenBLAS
Not watched
1
0
Fork 0
Code
Releases 66 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: cd8ac192a9
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'cd8ac192a9'
${ noResults }
OpenBLAS/kernel/x86_64/sbgemm_microk_cooperlake_te...

1836 lines
94 kB
Raw Blame History 

  1. #include "bf16_common_macros.h"

  2. #include <immintrin.h>

  3. 

  4. #define BF16_BLOCK_STEP_N 8

  5. #define BF16_BLOCK_THRES_K 1024

  6. #define BF16_BLOCK_THRES_M 32

  7. #define BF16_BLOCK_THRES_N 1024

  8. 

  9. #define A(i,j) A[(i)*lda+(j)]

  10. #define B(i,j) B[(i)*ldb+(j)]

  11. #define C(i,j) C[(i)*ldc+(j)]

  12. 

  13. #define ONE  1.e0f

  14. #define ZERO  0.e0f

  15. 

  16. #define SHUFFLE_MAGIC_NO (const int) 0x39

  17. 

  18. #undef STORE16_COMPLETE_RESULT

  19. #undef STORE16_MASK_COMPLETE_RESULT

  20. #undef SBGEMM_BLOCK_KERNEL_NN_32x8xK

  21. #undef SBGEMM_BLOCK_KERNEL_NN_16x8xK

  22. #undef SBGEMM_BLOCK_KERNEL_NN_32xNx32

  23. #undef SBGEMM_BLOCK_KERNEL_NN_16xNx32

  24. #undef SBGEMM_BLOCK_KERNEL_NT_32x8xK

  25. #undef SBGEMM_BLOCK_KERNEL_NT_16x8xK

  26. #undef SBGEMM_BLOCK_KERNEL_NT_32xNxK

  27. #undef SBGEMM_BLOCK_KERNEL_NT_16xNxK

  28. #undef SBGEMM_BLOCK_KERNEL_TN_32x8xK

  29. #undef SBGEMM_BLOCK_KERNEL_TN_16x8xK

  30. #undef SBGEMM_BLOCK_KERNEL_TN_32xNx32

  31. #undef SBGEMM_BLOCK_KERNEL_TN_16xNx32

  32. #undef SBGEMM_BLOCK_KERNEL_TT_32x8xK

  33. #undef SBGEMM_BLOCK_KERNEL_TT_16x8xK

  34. #undef SBGEMM_BLOCK_KERNEL_TT_32xNxK

  35. #undef SBGEMM_BLOCK_KERNEL_TT_16xNxK

  36. #undef SBGEMM_BLOCKING_KERNEL_NN

  37. #undef SBGEMM_BLOCKING_KERNEL_NT

  38. #undef SBGEMM_BLOCKING_KERNEL_TN

  39. #undef SBGEMM_BLOCKING_KERNEL_TT

  40. 

  41. #ifndef ONE_ALPHA      // ALPHA is not ONE

  42.     #define STORE16_COMPLETE_RESULT          STORE16_COMPLETE_RESULT_ALPHA_ONE

  43.     #define STORE16_MASK_COMPLETE_RESULT     STORE16_MASK_COMPLETE_RESULT_ALPHA_ONE

  44. 

  45.     #define SBGEMM_BLOCK_KERNEL_NN_32x8xK    sbgemm_block_kernel_nn_32x8xK_alpha

  46.     #define SBGEMM_BLOCK_KERNEL_NN_16x8xK    sbgemm_block_kernel_nn_16x8xK_alpha

  47.     #define SBGEMM_BLOCK_KERNEL_NN_32xNx32   sbgemm_block_kernel_nn_32xNx32_alpha

  48.     #define SBGEMM_BLOCK_KERNEL_NN_16xNx32   sbgemm_block_kernel_nn_16xNx32_alpha

  49. 

  50.     #define SBGEMM_BLOCK_KERNEL_NT_32x8xK    SBGEMM_BLOCK_KERNEL_NN_32x8xK

  51.     #define SBGEMM_BLOCK_KERNEL_NT_16x8xK    SBGEMM_BLOCK_KERNEL_NN_16x8xK

  52.     #define SBGEMM_BLOCK_KERNEL_NT_32xNxK    sbgemm_block_kernel_nt_32xNxK_alpha

  53.     #define SBGEMM_BLOCK_KERNEL_NT_16xNxK    sbgemm_block_kernel_nt_16xNxK_alpha

  54. 

  55.     #define SBGEMM_BLOCK_KERNEL_TN_32x8xK    sbgemm_block_kernel_tn_32x8xK_alpha

  56.     #define SBGEMM_BLOCK_KERNEL_TN_16x8xK    sbgemm_block_kernel_tn_16x8xK_alpha

  57.     #define SBGEMM_BLOCK_KERNEL_TN_32xNx32   sbgemm_block_kernel_tn_32xNx32_alpha

  58.     #define SBGEMM_BLOCK_KERNEL_TN_16xNx32   sbgemm_block_kernel_tn_16xNx32_alpha

  59. 

  60.     #define SBGEMM_BLOCK_KERNEL_TT_32x8xK    SBGEMM_BLOCK_KERNEL_TN_32x8xK

  61.     #define SBGEMM_BLOCK_KERNEL_TT_16x8xK    SBGEMM_BLOCK_KERNEL_TN_16x8xK

  62.     #define SBGEMM_BLOCK_KERNEL_TT_32xNxK    sbgemm_block_kernel_tt_32xNxK_alpha

  63.     #define SBGEMM_BLOCK_KERNEL_TT_16xNxK    sbgemm_block_kernel_tt_16xNxK_alpha

  64. 

  65.     #define SBGEMM_BLOCKING_KERNEL_NN        sbgemm_blocking_kernel_nn_alpha

  66.     #define SBGEMM_BLOCKING_KERNEL_NT        sbgemm_blocking_kernel_nt_alpha

  67.     #define SBGEMM_BLOCKING_KERNEL_TN        sbgemm_blocking_kernel_tn_alpha

  68.     #define SBGEMM_BLOCKING_KERNEL_TT        sbgemm_blocking_kernel_tt_alpha

  69. #else                  // ALPHA is ONE

  70.     #define STORE16_COMPLETE_RESULT          STORE16_COMPLETE_RESULT_ONE_ONE

  71.     #define STORE16_MASK_COMPLETE_RESULT     STORE16_MASK_COMPLETE_RESULT_ONE_ONE

  72. 

  73.     #define SBGEMM_BLOCK_KERNEL_NN_32x8xK    sbgemm_block_kernel_nn_32x8xK_one

  74.     #define SBGEMM_BLOCK_KERNEL_NN_16x8xK    sbgemm_block_kernel_nn_16x8xK_one

  75.     #define SBGEMM_BLOCK_KERNEL_NN_32xNx32   sbgemm_block_kernel_nn_32xNx32_one

  76.     #define SBGEMM_BLOCK_KERNEL_NN_16xNx32   sbgemm_block_kernel_nn_16xNx32_one

  77. 

  78.     #define SBGEMM_BLOCK_KERNEL_NT_32x8xK    SBGEMM_BLOCK_KERNEL_NN_32x8xK

  79.     #define SBGEMM_BLOCK_KERNEL_NT_16x8xK    SBGEMM_BLOCK_KERNEL_NN_16x8xK

  80.     #define SBGEMM_BLOCK_KERNEL_NT_32xNxK    sbgemm_block_kernel_nt_32xNxK_one

  81.     #define SBGEMM_BLOCK_KERNEL_NT_16xNxK    sbgemm_block_kernel_nt_16xNxK_one

  82. 

  83.     #define SBGEMM_BLOCK_KERNEL_TN_32x8xK    sbgemm_block_kernel_tn_32x8xK_one

  84.     #define SBGEMM_BLOCK_KERNEL_TN_16x8xK    sbgemm_block_kernel_tn_16x8xK_one

  85.     #define SBGEMM_BLOCK_KERNEL_TN_32xNx32   sbgemm_block_kernel_tn_32xNx32_one

  86.     #define SBGEMM_BLOCK_KERNEL_TN_16xNx32   sbgemm_block_kernel_tn_16xNx32_one

  87. 

  88.     #define SBGEMM_BLOCK_KERNEL_TT_32x8xK    SBGEMM_BLOCK_KERNEL_TN_32x8xK

  89.     #define SBGEMM_BLOCK_KERNEL_TT_16x8xK    SBGEMM_BLOCK_KERNEL_TN_16x8xK

  90.     #define SBGEMM_BLOCK_KERNEL_TT_32xNxK    sbgemm_block_kernel_tt_32xNxK_one

  91.     #define SBGEMM_BLOCK_KERNEL_TT_16xNxK    sbgemm_block_kernel_tt_16xNxK_one

  92. 

  93.     #define SBGEMM_BLOCKING_KERNEL_NN        sbgemm_blocking_kernel_nn_one

  94.     #define SBGEMM_BLOCKING_KERNEL_NT        sbgemm_blocking_kernel_nt_one

  95.     #define SBGEMM_BLOCKING_KERNEL_TN        sbgemm_blocking_kernel_tn_one

  96.     #define SBGEMM_BLOCKING_KERNEL_TT        sbgemm_blocking_kernel_tt_one

  97. #endif

  98. 

  99. extern bfloat16 * block_A;

  100. extern bfloat16 * block_B;

  101. 

  102. /* --------------------------------------------- NN kernels ------------------------------------------ */

  103. // SBGEMM Kernel for 16<M<=32, N=8, K can be any number, but the processing will take 32 as a base

  104. #ifndef ONE_ALPHA      // ALPHA is not ONE

  105. void sbgemm_block_kernel_nn_32x8xK_alpha(BLASLONG m, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  106. #else                  // ALPHA is ONE

  107. void sbgemm_block_kernel_nn_32x8xK_one(BLASLONG m, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  108. #endif

  109. {

  110.     bfloat16 * A_addr = A;

  111.     bfloat16 * B_addr = B;

  112.     float    * C_addr = C;

  113. 

  114. #ifndef ONE_ALPHA

  115.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  116. #endif

  117. 

  118.     __m512i arrayA_512_0, arrayA_512_1;

  119.     __m512i arrayB_512_0, arrayB_512_1, arrayB_512_2, arrayB_512_3, arrayB_512_4, arrayB_512_5, arrayB_512_6, arrayB_512_7;

  120.     __m512  result_512_0, result_512_1, result_512_2, result_512_3, result_512_4, result_512_5, result_512_6, result_512_7,

  121.             result_512_8, result_512_9, result_512_10, result_512_11, result_512_12, result_512_13, result_512_14, result_512_15;

  122.     __m512  result_512_tmp_0, result_512_tmp_1, result_512_tmp_2, result_512_tmp_3;

  123. 

  124.     __m512i M512_EPI32_8      = _mm512_set1_epi32(8);

  125.     __m512i shuffle_idx_base0 = _mm512_set_epi32(23, 22, 21, 20, 7, 6, 5, 4, 19, 18, 17, 16, 3, 2, 1, 0);

  126.     __m512i shuffle_idx_base1 = _mm512_add_epi32(shuffle_idx_base0, M512_EPI32_8);

  127. 

  128.     result_512_0  = _mm512_setzero_ps();

  129.     result_512_1  = _mm512_setzero_ps();

  130.     result_512_2  = _mm512_setzero_ps();

  131.     result_512_3  = _mm512_setzero_ps();

  132.     result_512_4  = _mm512_setzero_ps();

  133.     result_512_5  = _mm512_setzero_ps();

  134.     result_512_6  = _mm512_setzero_ps();

  135.     result_512_7  = _mm512_setzero_ps();

  136.     result_512_8  = _mm512_setzero_ps();

  137.     result_512_9  = _mm512_setzero_ps();

  138.     result_512_10 = _mm512_setzero_ps();

  139.     result_512_11 = _mm512_setzero_ps();

  140.     result_512_12 = _mm512_setzero_ps();

  141.     result_512_13 = _mm512_setzero_ps();

  142.     result_512_14 = _mm512_setzero_ps();

  143.     result_512_15 = _mm512_setzero_ps();

  144. 

  145.     for (BLASLONG idx_k = 0; idx_k < k; idx_k += 2) {

  146.         // Each two rows are a group for 32-pair bf16 elements

  147.         arrayA_512_0 = _mm512_loadu_si512(A_addr);

  148.         arrayA_512_1 = _mm512_loadu_si512(A_addr + 32);

  149. 

  150.         _MM512_BROADCASTD_EPI32(B_addr + 0,  arrayB_512_0);

  151.         _MM512_BROADCASTD_EPI32(B_addr + 2,  arrayB_512_1);

  152.         _MM512_BROADCASTD_EPI32(B_addr + 4,  arrayB_512_2);

  153.         _MM512_BROADCASTD_EPI32(B_addr + 6,  arrayB_512_3);

  154.         _MM512_BROADCASTD_EPI32(B_addr + 8,  arrayB_512_4);

  155.         _MM512_BROADCASTD_EPI32(B_addr + 10, arrayB_512_5);

  156.         _MM512_BROADCASTD_EPI32(B_addr + 12, arrayB_512_6);

  157.         _MM512_BROADCASTD_EPI32(B_addr + 14, arrayB_512_7);

  158. 

  159.         result_512_0  = _mm512_dpbf16_ps(result_512_0,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_0);

  160.         result_512_1  = _mm512_dpbf16_ps(result_512_1,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_1);

  161.         result_512_2  = _mm512_dpbf16_ps(result_512_2,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_2);

  162.         result_512_3  = _mm512_dpbf16_ps(result_512_3,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_3);

  163.         result_512_4  = _mm512_dpbf16_ps(result_512_4,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_4);

  164.         result_512_5  = _mm512_dpbf16_ps(result_512_5,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_5);

  165.         result_512_6  = _mm512_dpbf16_ps(result_512_6,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_6);

  166.         result_512_7  = _mm512_dpbf16_ps(result_512_7,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_7);

  167. 

  168.         result_512_8  = _mm512_dpbf16_ps(result_512_8,  (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_0);

  169.         result_512_9  = _mm512_dpbf16_ps(result_512_9,  (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_1);

  170.         result_512_10 = _mm512_dpbf16_ps(result_512_10, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_2);

  171.         result_512_11 = _mm512_dpbf16_ps(result_512_11, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_3);

  172.         result_512_12 = _mm512_dpbf16_ps(result_512_12, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_4);

  173.         result_512_13 = _mm512_dpbf16_ps(result_512_13, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_5);

  174.         result_512_14 = _mm512_dpbf16_ps(result_512_14, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_6);

  175.         result_512_15 = _mm512_dpbf16_ps(result_512_15, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_7);

  176. 

  177.         // Load B with unroll 8

  178.         B_addr += 16;

  179.         // Load A with unroll 64

  180.         A_addr += 64;

  181.     }

  182. 

  183.     if (m != 32) {

  184.         unsigned short tail_mask_value = (((unsigned short)0xffff) >> (32-m));

  185.         __mmask16 tail_mask = *((__mmask16*) &tail_mask_value);

  186.         result_512_tmp_0 = _mm512_permutex2var_ps(result_512_0, shuffle_idx_base0, result_512_8);

  187.         result_512_tmp_1 = _mm512_permutex2var_ps(result_512_0, shuffle_idx_base1, result_512_8);

  188.         result_512_tmp_2 = _mm512_permutex2var_ps(result_512_1, shuffle_idx_base0, result_512_9);

  189.         result_512_tmp_3 = _mm512_permutex2var_ps(result_512_1, shuffle_idx_base1, result_512_9);

  190.         STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr))

  191.         STORE16_MASK_COMPLETE_RESULT(result_512_tmp_1, (C_addr + 16), tail_mask)

  192.         STORE16_COMPLETE_RESULT(result_512_tmp_2, (C_addr + ldc*1))

  193.         STORE16_MASK_COMPLETE_RESULT(result_512_tmp_3, (C_addr + ldc*1 + 16), tail_mask)

  194.         result_512_tmp_0 = _mm512_permutex2var_ps(result_512_2, shuffle_idx_base0, result_512_10);

  195.         result_512_tmp_1 = _mm512_permutex2var_ps(result_512_2, shuffle_idx_base1, result_512_10);

  196.         result_512_tmp_2 = _mm512_permutex2var_ps(result_512_3, shuffle_idx_base0, result_512_11);

  197.         result_512_tmp_3 = _mm512_permutex2var_ps(result_512_3, shuffle_idx_base1, result_512_11);

  198.         STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*2))

  199.         STORE16_MASK_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*2 + 16), tail_mask)

  200.         STORE16_COMPLETE_RESULT(result_512_tmp_2, (C_addr + ldc*3))

  201.         STORE16_MASK_COMPLETE_RESULT(result_512_tmp_3, (C_addr + ldc*3 + 16), tail_mask)

  202.         result_512_tmp_0 = _mm512_permutex2var_ps(result_512_4, shuffle_idx_base0, result_512_12);

  203.         result_512_tmp_1 = _mm512_permutex2var_ps(result_512_4, shuffle_idx_base1, result_512_12);

  204.         result_512_tmp_2 = _mm512_permutex2var_ps(result_512_5, shuffle_idx_base0, result_512_13);

  205.         result_512_tmp_3 = _mm512_permutex2var_ps(result_512_5, shuffle_idx_base1, result_512_13);

  206.         STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*4))

  207.         STORE16_MASK_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*4 + 16), tail_mask)

  208.         STORE16_COMPLETE_RESULT(result_512_tmp_2, (C_addr + ldc*5))

  209.         STORE16_MASK_COMPLETE_RESULT(result_512_tmp_3, (C_addr + ldc*5 + 16), tail_mask)

  210.         result_512_tmp_0 = _mm512_permutex2var_ps(result_512_6, shuffle_idx_base0, result_512_14);

  211.         result_512_tmp_1 = _mm512_permutex2var_ps(result_512_6, shuffle_idx_base1, result_512_14);

  212.         result_512_tmp_2 = _mm512_permutex2var_ps(result_512_7, shuffle_idx_base0, result_512_15);

  213.         result_512_tmp_3 = _mm512_permutex2var_ps(result_512_7, shuffle_idx_base1, result_512_15);

  214.         STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*6))

  215.         STORE16_MASK_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*6 + 16), tail_mask)

  216.         STORE16_COMPLETE_RESULT(result_512_tmp_2, (C_addr + ldc*7))

  217.         STORE16_MASK_COMPLETE_RESULT(result_512_tmp_3, (C_addr + ldc*7 + 16), tail_mask)

  218.     } else {

  219.         result_512_tmp_0 = _mm512_permutex2var_ps(result_512_0, shuffle_idx_base0, result_512_8);

  220.         result_512_tmp_1 = _mm512_permutex2var_ps(result_512_0, shuffle_idx_base1, result_512_8);

  221.         result_512_tmp_2 = _mm512_permutex2var_ps(result_512_1, shuffle_idx_base0, result_512_9);

  222.         result_512_tmp_3 = _mm512_permutex2var_ps(result_512_1, shuffle_idx_base1, result_512_9);

  223.         STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr))

  224.         STORE16_COMPLETE_RESULT(result_512_tmp_1, (C_addr + 16))

  225.         STORE16_COMPLETE_RESULT(result_512_tmp_2, (C_addr + ldc*1))

  226.         STORE16_COMPLETE_RESULT(result_512_tmp_3, (C_addr + ldc*1 + 16))

  227.         result_512_tmp_0 = _mm512_permutex2var_ps(result_512_2, shuffle_idx_base0, result_512_10);

  228.         result_512_tmp_1 = _mm512_permutex2var_ps(result_512_2, shuffle_idx_base1, result_512_10);

  229.         result_512_tmp_2 = _mm512_permutex2var_ps(result_512_3, shuffle_idx_base0, result_512_11);

  230.         result_512_tmp_3 = _mm512_permutex2var_ps(result_512_3, shuffle_idx_base1, result_512_11);

  231.         STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*2))

  232.         STORE16_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*2 + 16))

  233.         STORE16_COMPLETE_RESULT(result_512_tmp_2, (C_addr + ldc*3))

  234.         STORE16_COMPLETE_RESULT(result_512_tmp_3, (C_addr + ldc*3 + 16))

  235.         result_512_tmp_0 = _mm512_permutex2var_ps(result_512_4, shuffle_idx_base0, result_512_12);

  236.         result_512_tmp_1 = _mm512_permutex2var_ps(result_512_4, shuffle_idx_base1, result_512_12);

  237.         result_512_tmp_2 = _mm512_permutex2var_ps(result_512_5, shuffle_idx_base0, result_512_13);

  238.         result_512_tmp_3 = _mm512_permutex2var_ps(result_512_5, shuffle_idx_base1, result_512_13);

  239.         STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*4))

  240.         STORE16_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*4 + 16))

  241.         STORE16_COMPLETE_RESULT(result_512_tmp_2, (C_addr + ldc*5))

  242.         STORE16_COMPLETE_RESULT(result_512_tmp_3, (C_addr + ldc*5 + 16))

  243.         result_512_tmp_0 = _mm512_permutex2var_ps(result_512_6, shuffle_idx_base0, result_512_14);

  244.         result_512_tmp_1 = _mm512_permutex2var_ps(result_512_6, shuffle_idx_base1, result_512_14);

  245.         result_512_tmp_2 = _mm512_permutex2var_ps(result_512_7, shuffle_idx_base0, result_512_15);

  246.         result_512_tmp_3 = _mm512_permutex2var_ps(result_512_7, shuffle_idx_base1, result_512_15);

  247.         STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*6))

  248.         STORE16_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*6 + 16))

  249.         STORE16_COMPLETE_RESULT(result_512_tmp_2, (C_addr + ldc*7))

  250.         STORE16_COMPLETE_RESULT(result_512_tmp_3, (C_addr + ldc*7 + 16))

  251.     }

  252. }

  253. 

  254. // SBGEMM Kernel for M<=16, N=8, K can be any number

  255. #ifndef ONE_ALPHA      // ALPHA is not ONE

  256. void sbgemm_block_kernel_nn_16x8xK_alpha(BLASLONG m, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  257. #else                  // ALPHA is ONE

  258. void sbgemm_block_kernel_nn_16x8xK_one(BLASLONG m, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  259. #endif

  260. {

  261.     bfloat16 * A_addr = A;

  262.     bfloat16 * B_addr = B;

  263.     float    * C_addr = C;

  264. 

  265. #ifndef ONE_ALPHA

  266.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  267. #endif

  268. 

  269.     __m512i arrayA_512_0;

  270.     __m512i arrayB_512_0, arrayB_512_1, arrayB_512_2, arrayB_512_3, arrayB_512_4, arrayB_512_5, arrayB_512_6, arrayB_512_7;

  271.     __m512  result_512_0, result_512_1, result_512_2, result_512_3, result_512_4, result_512_5, result_512_6, result_512_7;

  272. 

  273.     result_512_0  = _mm512_setzero_ps();

  274.     result_512_1  = _mm512_setzero_ps();

  275.     result_512_2  = _mm512_setzero_ps();

  276.     result_512_3  = _mm512_setzero_ps();

  277.     result_512_4  = _mm512_setzero_ps();

  278.     result_512_5  = _mm512_setzero_ps();

  279.     result_512_6  = _mm512_setzero_ps();

  280.     result_512_7  = _mm512_setzero_ps();

  281. 

  282.     for (BLASLONG idx_k = 0; idx_k < k; idx_k += 2) {

  283.         // Each two rows are a group for 32-pair bf16 elements

  284.         // Load two rows into a 512 register

  285.         arrayA_512_0 = _mm512_loadu_si512(A_addr);

  286. 

  287.         _MM512_BROADCASTD_EPI32(B_addr + 0,  arrayB_512_0);

  288.         _MM512_BROADCASTD_EPI32(B_addr + 2,  arrayB_512_1);

  289.         _MM512_BROADCASTD_EPI32(B_addr + 4,  arrayB_512_2);

  290.         _MM512_BROADCASTD_EPI32(B_addr + 6,  arrayB_512_3);

  291.         _MM512_BROADCASTD_EPI32(B_addr + 8,  arrayB_512_4);

  292.         _MM512_BROADCASTD_EPI32(B_addr + 10, arrayB_512_5);

  293.         _MM512_BROADCASTD_EPI32(B_addr + 12, arrayB_512_6);

  294.         _MM512_BROADCASTD_EPI32(B_addr + 14, arrayB_512_7);

  295. 

  296.         result_512_0  = _mm512_dpbf16_ps(result_512_0,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_0);

  297.         result_512_1  = _mm512_dpbf16_ps(result_512_1,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_1);

  298.         result_512_2  = _mm512_dpbf16_ps(result_512_2,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_2);

  299.         result_512_3  = _mm512_dpbf16_ps(result_512_3,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_3);

  300.         result_512_4  = _mm512_dpbf16_ps(result_512_4,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_4);

  301.         result_512_5  = _mm512_dpbf16_ps(result_512_5,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_5);

  302.         result_512_6  = _mm512_dpbf16_ps(result_512_6,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_6);

  303.         result_512_7  = _mm512_dpbf16_ps(result_512_7,  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_7);

  304. 

  305.         // Load B with unroll 8

  306.         B_addr += 16;

  307.         // Load A with unroll 16

  308.         A_addr += 32;

  309.     }

  310. 

  311.     if (m != 16) {

  312.         unsigned short tail_mask = (((unsigned short)0xffff) >> (16-m));

  313. 

  314.         result_512_0 = _mm512_shuffle_f32x4(result_512_0, result_512_0, 0xd8);

  315.         result_512_1 = _mm512_shuffle_f32x4(result_512_1, result_512_1, 0xd8);

  316.         result_512_2 = _mm512_shuffle_f32x4(result_512_2, result_512_2, 0xd8);

  317.         result_512_3 = _mm512_shuffle_f32x4(result_512_3, result_512_3, 0xd8);

  318.         STORE16_MASK_COMPLETE_RESULT(result_512_0, (C_addr), tail_mask)

  319.         STORE16_MASK_COMPLETE_RESULT(result_512_1, (C_addr + ldc*1), tail_mask)

  320.         STORE16_MASK_COMPLETE_RESULT(result_512_2, (C_addr + ldc*2), tail_mask)

  321.         STORE16_MASK_COMPLETE_RESULT(result_512_3, (C_addr + ldc*3), tail_mask)

  322.         result_512_4 = _mm512_shuffle_f32x4(result_512_4, result_512_4, 0xd8);

  323.         result_512_5 = _mm512_shuffle_f32x4(result_512_5, result_512_5, 0xd8);

  324.         result_512_6 = _mm512_shuffle_f32x4(result_512_6, result_512_6, 0xd8);

  325.         result_512_7 = _mm512_shuffle_f32x4(result_512_7, result_512_7, 0xd8);

  326.         STORE16_MASK_COMPLETE_RESULT(result_512_4, (C_addr + ldc*4), tail_mask)

  327.         STORE16_MASK_COMPLETE_RESULT(result_512_5, (C_addr + ldc*5), tail_mask)

  328.         STORE16_MASK_COMPLETE_RESULT(result_512_6, (C_addr + ldc*6), tail_mask)

  329.         STORE16_MASK_COMPLETE_RESULT(result_512_7, (C_addr + ldc*7), tail_mask)

  330.     } else {

  331.         result_512_0 = _mm512_shuffle_f32x4(result_512_0, result_512_0, 0xd8);

  332.         result_512_1 = _mm512_shuffle_f32x4(result_512_1, result_512_1, 0xd8);

  333.         result_512_2 = _mm512_shuffle_f32x4(result_512_2, result_512_2, 0xd8);

  334.         result_512_3 = _mm512_shuffle_f32x4(result_512_3, result_512_3, 0xd8);

  335.         STORE16_COMPLETE_RESULT(result_512_0, (C_addr))

  336.         STORE16_COMPLETE_RESULT(result_512_1, (C_addr + ldc*1))

  337.         STORE16_COMPLETE_RESULT(result_512_2, (C_addr + ldc*2))

  338.         STORE16_COMPLETE_RESULT(result_512_3, (C_addr + ldc*3))

  339.         result_512_4 = _mm512_shuffle_f32x4(result_512_4, result_512_4, 0xd8);

  340.         result_512_5 = _mm512_shuffle_f32x4(result_512_5, result_512_5, 0xd8);

  341.         result_512_6 = _mm512_shuffle_f32x4(result_512_6, result_512_6, 0xd8);

  342.         result_512_7 = _mm512_shuffle_f32x4(result_512_7, result_512_7, 0xd8);

  343.         STORE16_COMPLETE_RESULT(result_512_4, (C_addr + ldc*4))

  344.         STORE16_COMPLETE_RESULT(result_512_5, (C_addr + ldc*5))

  345.         STORE16_COMPLETE_RESULT(result_512_6, (C_addr + ldc*6))

  346.         STORE16_COMPLETE_RESULT(result_512_7, (C_addr + ldc*7))

  347.     }

  348. }

  349. 

  350. // SBGEMM Kernel for 16<M<=32, N<8, K can be any number, but the processing will take 32 as a base

  351. #ifndef ONE_ALPHA      // ALPHA is not ONE

  352. void sbgemm_block_kernel_nn_32xNx32_alpha(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  353. #else                  // ALPHA is ONE

  354. void sbgemm_block_kernel_nn_32xNx32_one(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  355. #endif

  356. {

  357.     bfloat16 * A_addr = A;

  358.     bfloat16 * B_addr = B;

  359.     float    * C_addr = C;

  360. 

  361.     BLASLONG tag_k_32x = k & (~31);

  362. 

  363. #ifndef ONE_ALPHA

  364.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  365. #endif

  366. 

  367.     __m512i arrayA_512[2];

  368.     __m512i arrayB_512[8];

  369.     __m512  result_512[16];

  370.     __m512  result_512_tmp_0, result_512_tmp_1;

  371. 

  372.     __m512i M512_EPI32_8      = _mm512_set1_epi32(8);

  373.     __m512i shuffle_idx_base0 = _mm512_set_epi32(23, 22, 21, 20, 7, 6, 5, 4, 19, 18, 17, 16, 3, 2, 1, 0);

  374.     __m512i shuffle_idx_base1 = _mm512_add_epi32(shuffle_idx_base0, M512_EPI32_8);

  375. 

  376.     for (int i = 0; i < 15; i += 2) {

  377.         result_512[i]    = _mm512_setzero_ps();

  378.         result_512[i+1]  = _mm512_setzero_ps();

  379.     }

  380. 

  381.     for (BLASLONG idx_k = 0; idx_k < tag_k_32x; idx_k += 32) {

  382.         // Load B with unroll n

  383.         for (int i = 0; i < n; i ++) {

  384.             arrayB_512[i] = _mm512_loadu_si512(B_addr);

  385.             B_addr += 32;

  386.         }

  387. 

  388.         for (BLASLONG idx = 0; idx < 32;) {

  389.             // Each two rows are a group for 32-pair bf16 elements

  390.             arrayA_512[0] = _mm512_loadu_si512(A_addr);

  391.             arrayA_512[1] = _mm512_loadu_si512(A_addr + 32);

  392.             A_addr += 64;

  393. 

  394.             for (int i = 0; i < n; i++) {

  395.                 result_512[i]   = _mm512_dpbf16_ps(result_512[i]  , (__m512bh) arrayA_512[0], (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  396.                 result_512[i+8] = _mm512_dpbf16_ps(result_512[i+8], (__m512bh) arrayA_512[1], (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  397.                 arrayB_512[i]   = _mm512_shuffle_epi32(arrayB_512[i], SHUFFLE_MAGIC_NO);

  398.             }

  399. 

  400.             idx += 2;

  401.             // Every 4 loops we need to switch to next 128 bits of arrayB registers

  402.             if ((idx & (~7)) == idx) {

  403.                 for (int i = 0; i < n; i++) {

  404.                     arrayB_512[i] = _mm512_shuffle_i32x4(arrayB_512[i], arrayB_512[i], SHUFFLE_MAGIC_NO);

  405.                 }

  406.             }

  407.         }

  408.     }

  409. 

  410.     if (tag_k_32x != k) {

  411.         // Load B with unroll n

  412.         for (int i = 0; i < n; i ++) {

  413.             arrayB_512[i] = _mm512_loadu_si512(B_addr);

  414.             B_addr += 32;

  415.         }

  416. 

  417.         BLASLONG width = k - tag_k_32x;

  418.         for (BLASLONG idx = 0; idx < width;) {

  419.             // Each two rows are a group for 32-pair bf16 elements

  420.             arrayA_512[0] = _mm512_loadu_si512(A_addr);

  421.             arrayA_512[1] = _mm512_loadu_si512(A_addr + 32);

  422.             A_addr += 64;

  423. 

  424.             for (int i = 0; i < n; i++) {

  425.                 result_512[i]   = _mm512_dpbf16_ps(result_512[i]  , (__m512bh) arrayA_512[0], (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  426.                 result_512[i+8] = _mm512_dpbf16_ps(result_512[i+8], (__m512bh) arrayA_512[1], (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  427.                 arrayB_512[i]   = _mm512_shuffle_epi32(arrayB_512[i], SHUFFLE_MAGIC_NO);

  428.             }

  429. 

  430.             idx += 2;

  431.             // Every 4 loops we need to switch to next 128 bits of arrayB registers

  432.             if ((idx & (~7)) == idx) {

  433.                 for (int i = 0; i < n; i++) {

  434.                     arrayB_512[i] = _mm512_shuffle_i32x4(arrayB_512[i], arrayB_512[i], SHUFFLE_MAGIC_NO);

  435.                 }

  436.             }

  437.         }

  438.     }

  439. 

  440.     if (m != 32) {

  441.         unsigned short tail_mask = (((unsigned short)0xffff) >> (32-m));

  442.         for (int i = 0; i < n; i++) {

  443.             result_512_tmp_0 = _mm512_permutex2var_ps(result_512[i], shuffle_idx_base0, result_512[i+8]);

  444.             result_512_tmp_1 = _mm512_permutex2var_ps(result_512[i], shuffle_idx_base1, result_512[i+8]);

  445.             STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*i))

  446.             STORE16_MASK_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*i + 16), tail_mask)

  447.         }

  448.     } else {

  449.         for (int i = 0; i < n; i++) {

  450.             result_512_tmp_0 = _mm512_permutex2var_ps(result_512[i], shuffle_idx_base0, result_512[i+8]);

  451.             result_512_tmp_1 = _mm512_permutex2var_ps(result_512[i], shuffle_idx_base1, result_512[i+8]);

  452.             STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*i))

  453.             STORE16_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*i + 16))

  454.         }

  455.     }

  456. }

  457. 

  458. // SBGEMM Kernel for 16<=M, N<8, K can be any number, but the processing will take 32 as a base

  459. #ifndef ONE_ALPHA      // ALPHA is not ONE

  460. void sbgemm_block_kernel_nn_16xNx32_alpha(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  461. #else                  // ALPHA is ONE

  462. void sbgemm_block_kernel_nn_16xNx32_one(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  463. #endif

  464. {

  465.     bfloat16 * A_addr = A;

  466.     bfloat16 * B_addr = B;

  467.     float    * C_addr = C;

  468. 

  469.     BLASLONG tag_k_32x = k & (~31);

  470. 

  471. #ifndef ONE_ALPHA

  472.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  473. #endif

  474. 

  475.     __m512i arrayA_512;

  476.     __m512i arrayB_512[8];

  477.     __m512  result_512[8];

  478. 

  479.     for (int i = 0; i < 8; i += 2) {

  480.         result_512[i]    = _mm512_setzero_ps();

  481.         result_512[i+1]  = _mm512_setzero_ps();

  482.     }

  483. 

  484.     for (BLASLONG idx_k = 0; idx_k < tag_k_32x; idx_k += 32) {

  485.         // Load B with unroll n

  486.         for (int i = 0; i < n; i++) {

  487.             arrayB_512[i] = _mm512_loadu_si512(B_addr);

  488.             B_addr += 32;

  489.         }

  490. 

  491.         for (BLASLONG idx = 0; idx < 32;) {

  492.             // Each two rows are a group for 32-pair bf16 elements

  493.             // Load two rows into a 512 register

  494.             arrayA_512 = _mm512_loadu_si512(A_addr);

  495.             A_addr += 32;

  496. 

  497.             for (int i = 0; i < n; i ++) {

  498.                 result_512[i]  = _mm512_dpbf16_ps(result_512[i],  (__m512bh) arrayA_512, (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  499.                 arrayB_512[i] = _mm512_shuffle_epi32(arrayB_512[i], SHUFFLE_MAGIC_NO);

  500.             }

  501. 

  502.             idx += 2;

  503.             // Every 4 loops we need to switch to next 128 bits of arrayB registers

  504.             if ((idx & (~7)) == idx) {

  505.                 for (int i = 0; i < n; i++) {

  506.                     arrayB_512[i] = _mm512_shuffle_i32x4(arrayB_512[i], arrayB_512[i], SHUFFLE_MAGIC_NO);

  507.                 }

  508.             }

  509.         }

  510.     }

  511. 

  512.     if (tag_k_32x != k) {

  513.         // Load B with unroll n

  514.         for (int i = 0; i < n; i++) {

  515.             arrayB_512[i] = _mm512_loadu_si512(B_addr);

  516.             B_addr += 32;

  517.         }

  518. 

  519.         BLASLONG width = k - tag_k_32x;

  520.         for (BLASLONG idx = 0; idx < width;) {

  521.             // Each two rows are a group for 32-pair bf16 elements

  522.             // Load two rows into a 512 register

  523.             arrayA_512 = _mm512_loadu_si512(A_addr);

  524.             A_addr += 32;

  525. 

  526.             for (int i = 0; i < n; i++) {

  527.                 result_512[i]  = _mm512_dpbf16_ps(result_512[i],  (__m512bh) arrayA_512, (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  528.                 arrayB_512[i] = _mm512_shuffle_epi32(arrayB_512[i], SHUFFLE_MAGIC_NO);

  529.             }

  530. 

  531.             idx += 2;

  532.             // Every 4 loops we need to switch to next 128 bits of arrayB registers

  533.             if ((idx & (~7)) == idx) {

  534.                 for (int i = 0; i < n; i++) {

  535.                     arrayB_512[i] = _mm512_shuffle_i32x4(arrayB_512[i], arrayB_512[i], SHUFFLE_MAGIC_NO);

  536.                 }

  537.             }

  538.         }

  539.     }

  540. 

  541.     if (m != 16) {

  542.         unsigned short tail_mask = (((unsigned short)0xffff) >> (16-m));

  543.         for (int i = 0; i < n; i++) {

  544.             result_512[i] = _mm512_shuffle_f32x4(result_512[i], result_512[i], 0xd8);

  545.             STORE16_MASK_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i), tail_mask)

  546.         }

  547.     } else {

  548.         for (int i = 0; i < n; i++) {

  549.             result_512[i] = _mm512_shuffle_f32x4(result_512[i], result_512[i], 0xd8);

  550.             STORE16_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i))

  551.         }

  552.     }

  553. }

  554. 

  555. 

  556. #ifndef ONE_ALPHA      // ALPHA is not ONE

  557. void sbgemm_blocking_kernel_nn_alpha(blasint M, blasint N, blasint K, float alpha, bfloat16 *A, blasint lda, bfloat16 *B, blasint ldb, float *C, blasint ldc, bfloat16 * block_A, bfloat16 * block_B)

  558. #else                  // ALPHA is ONE

  559. void sbgemm_blocking_kernel_nn_one(blasint M, blasint N, blasint K, float alpha, bfloat16 *A, blasint lda, bfloat16 *B, blasint ldb, float *C, blasint ldc, bfloat16 * block_A, bfloat16 * block_B)

  560. #endif

  561. {

  562.     BLASLONG m_step, n_step, k_step, k_step_round32;

  563.     BLASLONG tag_m_Nx = M & (~(BF16_BLOCK_THRES_M-1));

  564. 

  565.     BLASLONG n_from, n_to;

  566.     BLASLONG tag_n_Nx;

  567. 

  568.     n_from = 0;

  569.     n_to = (BF16_BLOCK_THRES_N > N) ? N : BF16_BLOCK_THRES_N;

  570.     tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  571. 

  572.     k_step = (K > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : K;

  573.     k_step_round32 = k_step & (~31);

  574.     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  575. 

  576.     if (M >= BF16_BLOCK_THRES_M) {

  577.         while (n_from < N) {

  578.             for (BLASLONG idx_k = 0; idx_k < K;) {

  579.                 // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  580.                 COL_MAJOR_INCOPY_KERNEL_Kx32(k_step, 32, &A(idx_k, 0), lda, block_A);

  581.                 for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  582.                     // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  583.                     COL_MAJOR_ONCOPY_KERNEL_8x32(k_step, &B(idx_n, idx_k), ldb, block_B + (idx_n-n_from)*k_step_round32);

  584.                     SBGEMM_BLOCK_KERNEL_NN_32x8xK(32, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  585.                 }

  586. 

  587.                 if (tag_n_Nx != n_to) {

  588.                     n_step = n_to - tag_n_Nx;

  589.                     COL_MAJOR_ONCOPY_KERNEL_Nx32(n_step, k_step, &B(tag_n_Nx, idx_k), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  590.                     SBGEMM_BLOCK_KERNEL_NN_32xNx32(32, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  591.                 }

  592. 

  593.                 for (BLASLONG idx_m = BF16_BLOCK_THRES_M; idx_m < tag_m_Nx; idx_m += BF16_BLOCK_THRES_M) {

  594.                     COL_MAJOR_INCOPY_KERNEL_Kx32(k_step, 32, &A(idx_k, idx_m), lda, block_A);

  595.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  596.                         SBGEMM_BLOCK_KERNEL_NN_32x8xK(32, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, idx_m), ldc);

  597.                     }

  598. 

  599.                     if (tag_n_Nx != n_to) {

  600.                         n_step = n_to - tag_n_Nx;

  601.                         SBGEMM_BLOCK_KERNEL_NN_32xNx32(32, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, idx_m), ldc);

  602.                     }

  603.                 }

  604. 

  605.                 if (tag_m_Nx != M) {

  606.                     m_step = M - tag_m_Nx;

  607.                     if (m_step > 16) {

  608.                         COL_MAJOR_INCOPY_KERNEL_Kx32(k_step, m_step, &A(idx_k, tag_m_Nx), lda, block_A);

  609.                         for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  610.                             SBGEMM_BLOCK_KERNEL_NN_32x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, tag_m_Nx), ldc);

  611.                         }

  612. 

  613.                         if (tag_n_Nx != n_to) {

  614.                             n_step = n_to - tag_n_Nx;

  615.                             SBGEMM_BLOCK_KERNEL_NN_32xNx32(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, tag_m_Nx), ldc);

  616.                         }

  617.                     } else {

  618.                         COL_MAJOR_INCOPY_KERNEL_Kx16(k_step, m_step, &A(idx_k, tag_m_Nx), lda, block_A);

  619.                         for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  620.                             SBGEMM_BLOCK_KERNEL_NN_16x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, tag_m_Nx), ldc);

  621.                         }

  622. 

  623.                         if (tag_n_Nx != n_to) {

  624.                             n_step = n_to - tag_n_Nx;

  625.                             SBGEMM_BLOCK_KERNEL_NN_16xNx32(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, tag_m_Nx), ldc);

  626.                         }

  627.                     }

  628.                 }

  629. 

  630.                 idx_k += k_step;

  631.                 k_step = K - idx_k;

  632.                 k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  633.                 k_step_round32 = k_step & (~31);

  634.                 k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  635.             }

  636. 

  637.             n_from = n_to;

  638.             n_to += BF16_BLOCK_THRES_N;

  639.             n_to = (n_to > N) ? N : n_to;

  640.             tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));           

  641.         }

  642.     } else {

  643.         m_step = M;

  644.         if (m_step > 16) {

  645.             while (n_from < N) {

  646.                 for (BLASLONG idx_k = 0; idx_k < K;) {

  647.                     // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  648.                     COL_MAJOR_INCOPY_KERNEL_Kx32(k_step, m_step, &A(idx_k, 0), lda, block_A);

  649.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  650.                         // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  651.                         COL_MAJOR_ONCOPY_KERNEL_8x32(k_step, &B(idx_n, idx_k), ldb, block_B + (idx_n-n_from)*k_step_round32);

  652.                         SBGEMM_BLOCK_KERNEL_NN_32x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  653.                     }

  654. 

  655.                     if (tag_n_Nx != n_to) {

  656.                         n_step = n_to - tag_n_Nx;

  657.                         COL_MAJOR_ONCOPY_KERNEL_Nx32(n_step, k_step, &B(tag_n_Nx, idx_k), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  658.                         SBGEMM_BLOCK_KERNEL_NN_32xNx32(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  659.                     }

  660. 

  661.                     idx_k += k_step;

  662.                     k_step = K - idx_k;

  663.                     k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  664.                     k_step_round32 = k_step & (~31);

  665.                     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  666.                 }

  667.                 n_from = n_to;

  668.                 n_to += BF16_BLOCK_THRES_N;

  669.                 n_to = (n_to > N) ? N : n_to;

  670.                 tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  671.             }

  672.         } else {

  673.             while (n_from < N) {

  674.                 for (BLASLONG idx_k = 0; idx_k < K;) {

  675.                     COL_MAJOR_INCOPY_KERNEL_Kx16(k_step, m_step, &A(idx_k, 0), lda, block_A);

  676.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  677.                         // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  678.                         COL_MAJOR_ONCOPY_KERNEL_8x32(k_step, &B(idx_n, idx_k), ldb, block_B + (idx_n-n_from)*k_step_round32);

  679.                         SBGEMM_BLOCK_KERNEL_NN_16x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  680.                     }

  681. 

  682.                     if (tag_n_Nx != n_to) {

  683.                         n_step = n_to - tag_n_Nx;

  684.                         COL_MAJOR_ONCOPY_KERNEL_Nx32(n_step, k_step, &B(tag_n_Nx, idx_k), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  685.                         SBGEMM_BLOCK_KERNEL_NN_16xNx32(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  686.                     }

  687. 

  688.                     idx_k += k_step;

  689.                     k_step = K - idx_k;

  690.                     k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  691.                     k_step_round32 = k_step & (~31);

  692.                     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  693.                 }

  694.                 n_from = n_to;

  695.                 n_to += BF16_BLOCK_THRES_N;

  696.                 n_to = (n_to > N) ? N : n_to;

  697.                 tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  698.             }

  699.         }

  700.     }

  701. }

  702. /* ----------------------------------------- End of NN kernels --------------------------------------- */

  703. 

  704. /* --------------------------------------------- NT kernels ------------------------------------------ */

  705. // SBGEMM Kernel for 16<M<=32, N<8, K can be any number

  706. #ifndef ONE_ALPHA      // ALPHA is not ONE

  707. void sbgemm_block_kernel_nt_32xNxK_alpha(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  708. #else                  // ALPHA is ONE

  709. void sbgemm_block_kernel_nt_32xNxK_one(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  710. #endif

  711. {

  712.     bfloat16 * A_addr = A;

  713.     bfloat16 * B_addr = B;

  714.     float    * C_addr = C;

  715. 

  716. #ifndef ONE_ALPHA

  717.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  718. #endif

  719. 

  720.     __m512i arrayA_512_0, arrayA_512_1;

  721.     __m512i arrayB_512[8];

  722.     __m512  result_512[16];

  723.     __m512  result_512_tmp_0, result_512_tmp_1;

  724. 

  725.     __m512i M512_EPI32_8      = _mm512_set1_epi32(8);

  726.     __m512i shuffle_idx_base0 = _mm512_set_epi32(23, 22, 21, 20, 7, 6, 5, 4, 19, 18, 17, 16, 3, 2, 1, 0);

  727.     __m512i shuffle_idx_base1 = _mm512_add_epi32(shuffle_idx_base0, M512_EPI32_8);

  728. 

  729.     result_512[0]  = _mm512_setzero_ps();

  730.     result_512[1]  = _mm512_setzero_ps();

  731.     result_512[2]  = _mm512_setzero_ps();

  732.     result_512[3]  = _mm512_setzero_ps();

  733.     result_512[4]  = _mm512_setzero_ps();

  734.     result_512[5]  = _mm512_setzero_ps();

  735.     result_512[6]  = _mm512_setzero_ps();

  736.     result_512[7]  = _mm512_setzero_ps();

  737.     result_512[8]  = _mm512_setzero_ps();

  738.     result_512[9]  = _mm512_setzero_ps();

  739.     result_512[10] = _mm512_setzero_ps();

  740.     result_512[11] = _mm512_setzero_ps();

  741.     result_512[12] = _mm512_setzero_ps();

  742.     result_512[13] = _mm512_setzero_ps();

  743.     result_512[14] = _mm512_setzero_ps();

  744.     result_512[15] = _mm512_setzero_ps();

  745. 

  746.     for (BLASLONG idx_k = 0; idx_k < k; idx_k += 2) {

  747.         // Each two rows are a group for 32-pair bf16 elements

  748.         arrayA_512_0 = _mm512_loadu_si512(A_addr);

  749.         arrayA_512_1 = _mm512_loadu_si512(A_addr + 32);

  750.         A_addr += 64;

  751. 

  752.         for (int i = 0; i < n; i ++) {

  753.             _MM512_BROADCASTD_EPI32(B_addr + i*2,  arrayB_512[i]);

  754.         }

  755.         B_addr += 16;

  756. 

  757.         for (int i = 0; i < n; i ++) {

  758.             result_512[i] = _mm512_dpbf16_ps(result_512[i],  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512[i]);

  759.             result_512[i+8] = _mm512_dpbf16_ps(result_512[i+8],  (__m512bh) arrayA_512_1, (__m512bh) arrayB_512[i]);

  760.         }

  761.     }

  762. 

  763.     if (m != 32) {

  764.         unsigned short tail_mask = (((unsigned short)0xffff) >> (32-m));

  765.         for (int i = 0; i < n; i ++) {

  766.             result_512_tmp_0 = _mm512_permutex2var_ps(result_512[i], shuffle_idx_base0, result_512[i+8]);

  767.             result_512_tmp_1 = _mm512_permutex2var_ps(result_512[i], shuffle_idx_base1, result_512[i+8]);

  768.             STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*i))

  769.             STORE16_MASK_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*i + 16), tail_mask)

  770.         }

  771.     } else {

  772.         for (int i = 0; i < n; i ++) {

  773.             result_512_tmp_0 = _mm512_permutex2var_ps(result_512[i], shuffle_idx_base0, result_512[i+8]);

  774.             result_512_tmp_1 = _mm512_permutex2var_ps(result_512[i], shuffle_idx_base1, result_512[i+8]);

  775.             STORE16_COMPLETE_RESULT(result_512_tmp_0, (C_addr + ldc*i))

  776.             STORE16_COMPLETE_RESULT(result_512_tmp_1, (C_addr + ldc*i + 16))

  777.         }

  778.     }

  779. }

  780. 

  781. // SBGEMM Kernel for M<=16, N<8, K can be any number

  782. #ifndef ONE_ALPHA      // ALPHA is not ONE

  783. void sbgemm_block_kernel_nt_16xNxK_alpha(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  784. #else                  // ALPHA is ONE

  785. void sbgemm_block_kernel_nt_16xNxK_one(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  786. #endif

  787. {

  788.     bfloat16 * A_addr = A;

  789.     bfloat16 * B_addr = B;

  790.     float    * C_addr = C;

  791. 

  792. #ifndef ONE_ALPHA

  793.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  794. #endif

  795. 

  796.     __m512i arrayA_512_0;

  797.     __m512i arrayB_512[8];

  798.     __m512  result_512[8];

  799. 

  800.     result_512[0]  = _mm512_setzero_ps();

  801.     result_512[1]  = _mm512_setzero_ps();

  802.     result_512[2]  = _mm512_setzero_ps();

  803.     result_512[3]  = _mm512_setzero_ps();

  804.     result_512[4]  = _mm512_setzero_ps();

  805.     result_512[5]  = _mm512_setzero_ps();

  806.     result_512[6]  = _mm512_setzero_ps();

  807.     result_512[7]  = _mm512_setzero_ps();

  808. 

  809.     for (BLASLONG idx_k = 0; idx_k < k; idx_k += 2) {

  810.         // Each two rows are a group for 16-pair bf16 elements

  811.         // Load two rows into a 512 register

  812.         arrayA_512_0 = _mm512_loadu_si512(A_addr);

  813.         A_addr += 32;

  814. 

  815.         for (int i = 0; i < n; i ++) {

  816.             _MM512_BROADCASTD_EPI32(B_addr + i*2,  arrayB_512[i]);

  817.         }

  818.         B_addr += 16;

  819. 

  820.         for (int i = 0; i < n; i ++) {

  821.             result_512[i] = _mm512_dpbf16_ps(result_512[i],  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512[i]);

  822.         }

  823.     }

  824. 

  825.     if (m != 16) {

  826.         unsigned short tail_mask = (((unsigned short)0xffff) >> (16-m));

  827.         for (int i = 0; i < n; i++) {

  828.             result_512[i] = _mm512_shuffle_f32x4(result_512[i], result_512[i], 0xd8);

  829.             STORE16_MASK_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i), tail_mask)

  830.         }

  831.     } else {

  832.         for (int i = 0; i < n; i++) {

  833.             result_512[i] = _mm512_shuffle_f32x4(result_512[i], result_512[i], 0xd8);

  834.             STORE16_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i))

  835.         }

  836.     }

  837. }

  838. 

  839. #ifndef ONE_ALPHA      // ALPHA is not ONE

  840. void sbgemm_blocking_kernel_nt_alpha(blasint M, blasint N, blasint K, float alpha, bfloat16 *A, blasint lda, bfloat16 *B, blasint ldb, float *C, blasint ldc, bfloat16 * block_A, bfloat16 * block_B)

  841. #else                  // ALPHA is ONE

  842. void sbgemm_blocking_kernel_nt_one(blasint M, blasint N, blasint K, float alpha, bfloat16 *A, blasint lda, bfloat16 *B, blasint ldb, float *C, blasint ldc, bfloat16 * block_A, bfloat16 * block_B)

  843. #endif

  844. {

  845.     BLASLONG m_step, n_step, k_step, k_step_round32;

  846.     BLASLONG tag_m_Nx = M & (~(BF16_BLOCK_THRES_M-1));

  847. 

  848.     BLASLONG n_from, n_to;

  849.     BLASLONG tag_n_Nx;

  850. 

  851.     n_from = 0;

  852.     n_to = (BF16_BLOCK_THRES_N > N) ? N : BF16_BLOCK_THRES_N;

  853.     tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  854. 

  855.     k_step = (K > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : K;

  856.     k_step_round32 = k_step & (~31);

  857.     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  858. 

  859.     if (M >= BF16_BLOCK_THRES_M) {

  860.         while (n_from < N) {

  861.             for (BLASLONG idx_k = 0; idx_k < K;) {

  862.                 // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  863.                 COL_MAJOR_INCOPY_KERNEL_Kx32(k_step, 32, &A(idx_k, 0), lda, block_A);

  864.                 for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  865.                     // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  866.                     COL_MAJOR_OTCOPY_KERNEL_Kx8(k_step, &B(idx_k, idx_n), ldb, block_B + (idx_n-n_from)*k_step_round32);

  867.                     SBGEMM_BLOCK_KERNEL_NT_32x8xK(32, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  868.                 }

  869. 

  870.                 if (tag_n_Nx != n_to) {

  871.                     n_step = n_to - tag_n_Nx;

  872.                     COL_MAJOR_OTCOPY_KERNEL_Kx8m(k_step, n_step, &B(idx_k, tag_n_Nx), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  873.                     SBGEMM_BLOCK_KERNEL_NT_32xNxK(32, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  874.                 }

  875. 

  876.                 for (BLASLONG idx_m = BF16_BLOCK_THRES_M; idx_m < tag_m_Nx; idx_m += BF16_BLOCK_THRES_M) {

  877.                     COL_MAJOR_INCOPY_KERNEL_Kx32(k_step, 32, &A(idx_k, idx_m), lda, block_A);

  878.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  879.                         SBGEMM_BLOCK_KERNEL_NT_32x8xK(32, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, idx_m), ldc);

  880.                     }

  881. 

  882.                     if (tag_n_Nx != n_to) {

  883.                         n_step = n_to - tag_n_Nx;

  884.                         SBGEMM_BLOCK_KERNEL_NT_32xNxK(32, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, idx_m), ldc);

  885.                     }

  886.                 }

  887. 

  888.                 if (tag_m_Nx != M) {

  889.                     m_step = M - tag_m_Nx;

  890.                     if (m_step > 16) {

  891.                         COL_MAJOR_INCOPY_KERNEL_Kx32(k_step, m_step, &A(idx_k, tag_m_Nx), lda, block_A);

  892.                         for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  893.                             SBGEMM_BLOCK_KERNEL_NT_32x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, tag_m_Nx), ldc);

  894.                         }

  895. 

  896.                         if (tag_n_Nx != n_to) {

  897.                             n_step = n_to - tag_n_Nx;

  898.                             SBGEMM_BLOCK_KERNEL_NT_32xNxK(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, tag_m_Nx), ldc);

  899.                         }

  900.                     } else {

  901.                         COL_MAJOR_INCOPY_KERNEL_Kx16(k_step, m_step, &A(idx_k, tag_m_Nx), lda, block_A);

  902.                         for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  903.                             SBGEMM_BLOCK_KERNEL_NT_16x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, tag_m_Nx), ldc);

  904.                         }

  905. 

  906.                         if (tag_n_Nx != n_to) {

  907.                             n_step = n_to - tag_n_Nx;

  908.                             SBGEMM_BLOCK_KERNEL_NT_16xNxK(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, tag_m_Nx), ldc);

  909.                         }

  910.                     }

  911.                 }

  912. 

  913.                 idx_k += k_step;

  914.                 k_step = K - idx_k;

  915.                 k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  916.                 k_step_round32 = k_step & (~31);

  917.                 k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  918.             }

  919. 

  920.             n_from = n_to;

  921.             n_to += BF16_BLOCK_THRES_N;

  922.             n_to = (n_to > N) ? N : n_to;

  923.             tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  924.         }

  925.     } else {

  926.         m_step = M;

  927.         if (m_step > 16) {

  928.             while (n_from < N) {

  929.                 for (BLASLONG idx_k = 0; idx_k < K;) {

  930.                     // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  931.                     COL_MAJOR_INCOPY_KERNEL_Kx32(k_step, m_step, &A(idx_k, 0), lda, block_A);

  932.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  933.                         // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  934.                         COL_MAJOR_OTCOPY_KERNEL_Kx8(k_step, &B(idx_k, idx_n), ldb, block_B + (idx_n-n_from)*k_step_round32);

  935.                         SBGEMM_BLOCK_KERNEL_NT_32x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  936.                     }

  937. 

  938.                     if (tag_n_Nx != n_to) {

  939.                         n_step = n_to - tag_n_Nx;

  940.                         COL_MAJOR_OTCOPY_KERNEL_Kx8m(k_step, n_step, &B(idx_k, tag_n_Nx), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  941.                         SBGEMM_BLOCK_KERNEL_NT_32xNxK(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  942.                     }

  943. 

  944.                     idx_k += k_step;

  945.                     k_step = K - idx_k;

  946.                     k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  947.                     k_step_round32 = k_step & (~31);

  948.                     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  949.                 }

  950.                 n_from = n_to;

  951.                 n_to += BF16_BLOCK_THRES_N;

  952.                 n_to = (n_to > N) ? N : n_to;

  953.                 tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  954.             }

  955.         } else {

  956.             while (n_from < N) {

  957.                 for (BLASLONG idx_k = 0; idx_k < K;) {

  958.                     // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  959.                     COL_MAJOR_INCOPY_KERNEL_Kx16(k_step, m_step, &A(idx_k, 0), lda, block_A);

  960.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  961.                         // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  962.                         COL_MAJOR_OTCOPY_KERNEL_Kx8(k_step, &B(idx_k, idx_n), ldb, block_B + (idx_n-n_from)*k_step_round32);

  963.                         SBGEMM_BLOCK_KERNEL_NT_16x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  964.                     }

  965. 

  966.                     if (tag_n_Nx != n_to) {

  967.                         n_step = n_to - tag_n_Nx;

  968.                         COL_MAJOR_OTCOPY_KERNEL_Kx8m(k_step, n_step, &B(idx_k, tag_n_Nx), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  969.                         SBGEMM_BLOCK_KERNEL_NT_16xNxK(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  970.                     }

  971. 

  972.                     idx_k += k_step;

  973.                     k_step = K - idx_k;

  974.                     k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  975.                     k_step_round32 = k_step & (~31);

  976.                     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  977.                 }

  978.                 n_from = n_to;

  979.                 n_to += BF16_BLOCK_THRES_N;

  980.                 n_to = (n_to > N) ? N : n_to;

  981.                 tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  982.             }

  983.         }

  984.     }

  985. }

  986. /* ----------------------------------------- End of NT kernels --------------------------------------- */

  987. 

  988. /* --------------------------------------------- TN kernels ------------------------------------------ */

  989. // SBGEMM Kernel for 16<M<=32, N=8, K=Any number

  990. #ifndef ONE_ALPHA      // ALPHA is not ONE

  991. void sbgemm_block_kernel_tn_32x8xK_alpha(BLASLONG m, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  992. #else                  // ALPHA is ONE

  993. void sbgemm_block_kernel_tn_32x8xK_one(BLASLONG m, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  994. #endif

  995. {

  996.     bfloat16 * A_addr = A;

  997.     bfloat16 * B_addr = B;

  998.     float    * C_addr = C;

  999. 

  1000. #ifndef ONE_ALPHA

  1001.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  1002. #endif

  1003. 

  1004.     __m512i arrayA_512_0, arrayA_512_1;

  1005.     __m512i arrayB_512_0, arrayB_512_1, arrayB_512_2, arrayB_512_3, arrayB_512_4, arrayB_512_5, arrayB_512_6, arrayB_512_7;

  1006.     __m512  result_512_0, result_512_1, result_512_2, result_512_3, result_512_4, result_512_5, result_512_6, result_512_7,

  1007.             result_512_8, result_512_9, result_512_10, result_512_11, result_512_12, result_512_13, result_512_14, result_512_15;

  1008. 

  1009.     result_512_0  = _mm512_setzero_ps();

  1010.     result_512_1  = _mm512_setzero_ps();

  1011.     result_512_2  = _mm512_setzero_ps();

  1012.     result_512_3  = _mm512_setzero_ps();

  1013.     result_512_4  = _mm512_setzero_ps();

  1014.     result_512_5  = _mm512_setzero_ps();

  1015.     result_512_6  = _mm512_setzero_ps();

  1016.     result_512_7  = _mm512_setzero_ps();

  1017.     result_512_8  = _mm512_setzero_ps();

  1018.     result_512_9  = _mm512_setzero_ps();

  1019.     result_512_10 = _mm512_setzero_ps();

  1020.     result_512_11 = _mm512_setzero_ps();

  1021.     result_512_12 = _mm512_setzero_ps();

  1022.     result_512_13 = _mm512_setzero_ps();

  1023.     result_512_14 = _mm512_setzero_ps();

  1024.     result_512_15 = _mm512_setzero_ps();

  1025. 

  1026.     for (BLASLONG idx_k = 0; idx_k < k; idx_k += 2) {

  1027.         // Load 32 pair of BF16 elements from A (32 rows)

  1028.         arrayA_512_0 = _mm512_loadu_si512(A_addr);

  1029.         arrayA_512_1 = _mm512_loadu_si512(A_addr + 32);

  1030. 

  1031.         // Load 8 rows of B

  1032.         _MM512_BROADCASTD_EPI32(B_addr + 0,  arrayB_512_0);

  1033.         _MM512_BROADCASTD_EPI32(B_addr + 2,  arrayB_512_1);

  1034.         _MM512_BROADCASTD_EPI32(B_addr + 4,  arrayB_512_2);

  1035.         _MM512_BROADCASTD_EPI32(B_addr + 6,  arrayB_512_3);

  1036.         _MM512_BROADCASTD_EPI32(B_addr + 8,  arrayB_512_4);

  1037.         _MM512_BROADCASTD_EPI32(B_addr + 10, arrayB_512_5);

  1038.         _MM512_BROADCASTD_EPI32(B_addr + 12, arrayB_512_6);

  1039.         _MM512_BROADCASTD_EPI32(B_addr + 14, arrayB_512_7);

  1040. 

  1041.         result_512_0  = _mm512_dpbf16_ps(result_512_0, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_0);

  1042.         result_512_1  = _mm512_dpbf16_ps(result_512_1, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_1);

  1043.         result_512_2  = _mm512_dpbf16_ps(result_512_2, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_2);

  1044.         result_512_3  = _mm512_dpbf16_ps(result_512_3, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_3);

  1045.         result_512_4  = _mm512_dpbf16_ps(result_512_4, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_4);

  1046.         result_512_5  = _mm512_dpbf16_ps(result_512_5, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_5);

  1047.         result_512_6  = _mm512_dpbf16_ps(result_512_6, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_6);

  1048.         result_512_7  = _mm512_dpbf16_ps(result_512_7, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_7);

  1049. 

  1050.         result_512_8  = _mm512_dpbf16_ps(result_512_8,  (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_0);

  1051.         result_512_9  = _mm512_dpbf16_ps(result_512_9,  (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_1);

  1052.         result_512_10 = _mm512_dpbf16_ps(result_512_10, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_2);

  1053.         result_512_11 = _mm512_dpbf16_ps(result_512_11, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_3);

  1054.         result_512_12 = _mm512_dpbf16_ps(result_512_12, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_4);

  1055.         result_512_13 = _mm512_dpbf16_ps(result_512_13, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_5);

  1056.         result_512_14 = _mm512_dpbf16_ps(result_512_14, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_6);

  1057.         result_512_15 = _mm512_dpbf16_ps(result_512_15, (__m512bh) arrayA_512_1, (__m512bh) arrayB_512_7);

  1058. 

  1059.         // Load B with unroll 8

  1060.         B_addr += 16;

  1061.         // Load A with unroll 64

  1062.         A_addr += 64;

  1063.     }

  1064. 

  1065.     if (m != 32) {

  1066.         unsigned short tail_mask_value = (((unsigned short)0xffff) >> (32-m));

  1067.         __mmask16 tail_mask = *((__mmask16*) &tail_mask_value);

  1068.         STORE16_COMPLETE_RESULT(result_512_0,  (C_addr))

  1069.         STORE16_MASK_COMPLETE_RESULT(result_512_8,  (C_addr + 16), tail_mask)

  1070.         STORE16_COMPLETE_RESULT(result_512_1,  (C_addr + ldc))

  1071.         STORE16_MASK_COMPLETE_RESULT(result_512_9,  (C_addr + ldc + 16), tail_mask)

  1072.         STORE16_COMPLETE_RESULT(result_512_2,  (C_addr + ldc*2))

  1073.         STORE16_MASK_COMPLETE_RESULT(result_512_10, (C_addr + ldc*2 + 16), tail_mask)

  1074.         STORE16_COMPLETE_RESULT(result_512_3,  (C_addr + ldc*3))

  1075.         STORE16_MASK_COMPLETE_RESULT(result_512_11, (C_addr + ldc*3 + 16), tail_mask)

  1076.         STORE16_COMPLETE_RESULT(result_512_4,  (C_addr + ldc*4))

  1077.         STORE16_MASK_COMPLETE_RESULT(result_512_12, (C_addr + ldc*4 + 16), tail_mask)

  1078.         STORE16_COMPLETE_RESULT(result_512_5,  (C_addr + ldc*5))

  1079.         STORE16_MASK_COMPLETE_RESULT(result_512_13, (C_addr + ldc*5 + 16), tail_mask)

  1080.         STORE16_COMPLETE_RESULT(result_512_6,  (C_addr + ldc*6))

  1081.         STORE16_MASK_COMPLETE_RESULT(result_512_14, (C_addr + ldc*6 + 16), tail_mask)

  1082.         STORE16_COMPLETE_RESULT(result_512_7,  (C_addr + ldc*7))

  1083.         STORE16_MASK_COMPLETE_RESULT(result_512_15, (C_addr + ldc*7 + 16), tail_mask)

  1084.     } else {

  1085.         STORE16_COMPLETE_RESULT(result_512_0,  (C_addr))

  1086.         STORE16_COMPLETE_RESULT(result_512_8,  (C_addr + 16))

  1087.         STORE16_COMPLETE_RESULT(result_512_1,  (C_addr + ldc))

  1088.         STORE16_COMPLETE_RESULT(result_512_9,  (C_addr + ldc + 16))

  1089.         STORE16_COMPLETE_RESULT(result_512_2,  (C_addr + ldc*2))

  1090.         STORE16_COMPLETE_RESULT(result_512_10, (C_addr + ldc*2 + 16))

  1091.         STORE16_COMPLETE_RESULT(result_512_3,  (C_addr + ldc*3))

  1092.         STORE16_COMPLETE_RESULT(result_512_11, (C_addr + ldc*3 + 16))

  1093.         STORE16_COMPLETE_RESULT(result_512_4,  (C_addr + ldc*4))

  1094.         STORE16_COMPLETE_RESULT(result_512_12, (C_addr + ldc*4 + 16))

  1095.         STORE16_COMPLETE_RESULT(result_512_5,  (C_addr + ldc*5))

  1096.         STORE16_COMPLETE_RESULT(result_512_13, (C_addr + ldc*5 + 16))

  1097.         STORE16_COMPLETE_RESULT(result_512_6,  (C_addr + ldc*6))

  1098.         STORE16_COMPLETE_RESULT(result_512_14, (C_addr + ldc*6 + 16))

  1099.         STORE16_COMPLETE_RESULT(result_512_7,  (C_addr + ldc*7))

  1100.         STORE16_COMPLETE_RESULT(result_512_15, (C_addr + ldc*7 + 16))

  1101.     }

  1102. }

  1103. 

  1104. // SBGEMM Kernel for M=16, N=8, K=Any number

  1105. #ifndef ONE_ALPHA      // ALPHA is not ONE

  1106. void sbgemm_block_kernel_tn_16x8xK_alpha(BLASLONG m, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1107. #else                  // ALPHA is ONE

  1108. void sbgemm_block_kernel_tn_16x8xK_one(BLASLONG m, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1109. #endif

  1110. {

  1111.     bfloat16 * A_addr = A;

  1112.     bfloat16 * B_addr = B;

  1113.     float    * C_addr = C;

  1114. 

  1115. #ifndef ONE_ALPHA

  1116.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  1117. #endif

  1118. 

  1119.     __m512i arrayA_512_0;

  1120.     __m512i arrayB_512_0, arrayB_512_1, arrayB_512_2, arrayB_512_3, arrayB_512_4, arrayB_512_5, arrayB_512_6, arrayB_512_7;

  1121.     __m512  result_512_0, result_512_1, result_512_2, result_512_3, result_512_4, result_512_5, result_512_6, result_512_7;

  1122. 

  1123.     result_512_0  = _mm512_setzero_ps();

  1124.     result_512_1  = _mm512_setzero_ps();

  1125.     result_512_2  = _mm512_setzero_ps();

  1126.     result_512_3  = _mm512_setzero_ps();

  1127.     result_512_4  = _mm512_setzero_ps();

  1128.     result_512_5  = _mm512_setzero_ps();

  1129.     result_512_6  = _mm512_setzero_ps();

  1130.     result_512_7  = _mm512_setzero_ps();

  1131. 

  1132.     for (BLASLONG idx_k = 0; idx_k < k; idx_k += 2) {

  1133.         // Load 16 pair of BF16 elements from A (16 rows)

  1134.         arrayA_512_0 = _mm512_loadu_si512(A_addr + 0);

  1135. 

  1136.         // Load 8 rows of B

  1137.         _MM512_BROADCASTD_EPI32(B_addr + 0,  arrayB_512_0);

  1138.         _MM512_BROADCASTD_EPI32(B_addr + 2,  arrayB_512_1);

  1139.         _MM512_BROADCASTD_EPI32(B_addr + 4,  arrayB_512_2);

  1140.         _MM512_BROADCASTD_EPI32(B_addr + 6,  arrayB_512_3);

  1141.         _MM512_BROADCASTD_EPI32(B_addr + 8,  arrayB_512_4);

  1142.         _MM512_BROADCASTD_EPI32(B_addr + 10, arrayB_512_5);

  1143.         _MM512_BROADCASTD_EPI32(B_addr + 12, arrayB_512_6);

  1144.         _MM512_BROADCASTD_EPI32(B_addr + 14, arrayB_512_7);

  1145. 

  1146.         result_512_0  = _mm512_dpbf16_ps(result_512_0, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_0);

  1147.         result_512_1  = _mm512_dpbf16_ps(result_512_1, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_1);

  1148.         result_512_2  = _mm512_dpbf16_ps(result_512_2, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_2);

  1149.         result_512_3  = _mm512_dpbf16_ps(result_512_3, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_3);

  1150.         result_512_4  = _mm512_dpbf16_ps(result_512_4, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_4);

  1151.         result_512_5  = _mm512_dpbf16_ps(result_512_5, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_5);

  1152.         result_512_6  = _mm512_dpbf16_ps(result_512_6, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_6);

  1153.         result_512_7  = _mm512_dpbf16_ps(result_512_7, (__m512bh) arrayA_512_0, (__m512bh) arrayB_512_7);

  1154. 

  1155.         // Load B with unroll 8

  1156.         B_addr += 16;

  1157.         // Load A with unroll 32

  1158.         A_addr += 32;

  1159.     }

  1160. 

  1161.     if (m != 16) {

  1162.         unsigned short tail_mask_value = (((unsigned short)0xffff) >> (16-m));

  1163.         __mmask16 tail_mask = *((__mmask16*) &tail_mask_value);

  1164.         STORE16_MASK_COMPLETE_RESULT(result_512_0,  (C_addr), tail_mask)

  1165.         STORE16_MASK_COMPLETE_RESULT(result_512_1,  (C_addr + ldc), tail_mask)

  1166.         STORE16_MASK_COMPLETE_RESULT(result_512_2,  (C_addr + ldc*2), tail_mask)

  1167.         STORE16_MASK_COMPLETE_RESULT(result_512_3,  (C_addr + ldc*3), tail_mask)

  1168.         STORE16_MASK_COMPLETE_RESULT(result_512_4,  (C_addr + ldc*4), tail_mask)

  1169.         STORE16_MASK_COMPLETE_RESULT(result_512_5,  (C_addr + ldc*5), tail_mask)

  1170.         STORE16_MASK_COMPLETE_RESULT(result_512_6,  (C_addr + ldc*6), tail_mask)

  1171.         STORE16_MASK_COMPLETE_RESULT(result_512_7,  (C_addr + ldc*7), tail_mask)

  1172.     } else {

  1173.         STORE16_COMPLETE_RESULT(result_512_0,  (C_addr))

  1174.         STORE16_COMPLETE_RESULT(result_512_1,  (C_addr + ldc))

  1175.         STORE16_COMPLETE_RESULT(result_512_2,  (C_addr + ldc*2))

  1176.         STORE16_COMPLETE_RESULT(result_512_3,  (C_addr + ldc*3))

  1177.         STORE16_COMPLETE_RESULT(result_512_4,  (C_addr + ldc*4))

  1178.         STORE16_COMPLETE_RESULT(result_512_5,  (C_addr + ldc*5))

  1179.         STORE16_COMPLETE_RESULT(result_512_6,  (C_addr + ldc*6))

  1180.         STORE16_COMPLETE_RESULT(result_512_7,  (C_addr + ldc*7))

  1181.     }

  1182. }

  1183. 

  1184. // SBGEMM Kernel for 16<M<=32, N<8, K=Any number but will be processed based on 32

  1185. #ifndef ONE_ALPHA      // ALPHA is not ONE

  1186. void sbgemm_block_kernel_tn_32xNx32_alpha(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1187. #else                  // ALPHA is ONE

  1188. void sbgemm_block_kernel_tn_32xNx32_one(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1189. #endif

  1190. {

  1191.     bfloat16 * A_addr = A;

  1192.     bfloat16 * B_addr = B;

  1193.     float    * C_addr = C;

  1194. 

  1195.     BLASLONG tag_k_32x = k & (~31);

  1196. 

  1197. #ifndef ONE_ALPHA

  1198.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  1199. #endif

  1200. 

  1201.     __m512i arrayA_512[2];

  1202.     __m512i arrayB_512[8];

  1203.     __m512  result_512[16];

  1204. 

  1205.     for (int i = 0; i < 15; i++) {

  1206.         result_512[i] = _mm512_setzero_ps();

  1207.     }

  1208. 

  1209.     for (BLASLONG idx_k = 0; idx_k < tag_k_32x; idx_k += 32) {

  1210.         // Load B with unroll n

  1211.         for (int i = 0; i < n; i ++) {

  1212.             arrayB_512[i] = _mm512_loadu_si512(B_addr);

  1213.             B_addr += 32;

  1214.         }

  1215. 

  1216.         for (BLASLONG idx = 0; idx < 32;) {

  1217.             // Each two rows are a group for 32-pair bf16 elements

  1218.             arrayA_512[0] = _mm512_loadu_si512(A_addr);

  1219.             arrayA_512[1] = _mm512_loadu_si512(A_addr + 32);

  1220.             A_addr += 64;

  1221. 

  1222.             for (int i = 0; i < n; i++) {

  1223.                 result_512[i]   = _mm512_dpbf16_ps(result_512[i]  , (__m512bh) arrayA_512[0], (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  1224.                 result_512[i+8] = _mm512_dpbf16_ps(result_512[i+8], (__m512bh) arrayA_512[1], (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  1225.                 arrayB_512[i]   = _mm512_shuffle_epi32(arrayB_512[i], SHUFFLE_MAGIC_NO);

  1226.             }

  1227. 

  1228.             idx += 2;

  1229.             // Every 4 loops we need to switch to next 128 bits of arrayB registers

  1230.             if ((idx & (~7)) == idx) {

  1231.                 for (int i = 0; i < n; i++) {

  1232.                     arrayB_512[i] = _mm512_shuffle_i32x4(arrayB_512[i], arrayB_512[i], SHUFFLE_MAGIC_NO);

  1233.                 }

  1234.             }

  1235.         }

  1236.     }

  1237. 

  1238.     if (tag_k_32x != k) {

  1239.             // Load B with unroll n

  1240.         for (int i = 0; i < n; i ++) {

  1241.             arrayB_512[i] = _mm512_loadu_si512(B_addr);

  1242.             B_addr += 32;

  1243.         }

  1244. 

  1245.         BLASLONG width = k - tag_k_32x;

  1246.         for (BLASLONG idx = 0; idx < width;) {

  1247.             // Each two rows are a group for 32-pair bf16 elements

  1248.             arrayA_512[0] = _mm512_loadu_si512(A_addr);

  1249.             arrayA_512[1] = _mm512_loadu_si512(A_addr + 32);

  1250.             A_addr += 64;

  1251. 

  1252.             for (int i = 0; i < n; i++) {

  1253.                 result_512[i]   = _mm512_dpbf16_ps(result_512[i]  , (__m512bh) arrayA_512[0], (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  1254.                 result_512[i+8] = _mm512_dpbf16_ps(result_512[i+8], (__m512bh) arrayA_512[1], (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  1255.                 arrayB_512[i]   = _mm512_shuffle_epi32(arrayB_512[i], SHUFFLE_MAGIC_NO);

  1256.             }

  1257. 

  1258.             idx += 2;

  1259.             // Every 4 loops we need to switch to next 128 bits of arrayB registers

  1260.             if ((idx & (~7)) == idx) {

  1261.                 for (int i = 0; i < n; i++) {

  1262.                     arrayB_512[i] = _mm512_shuffle_i32x4(arrayB_512[i], arrayB_512[i], SHUFFLE_MAGIC_NO);

  1263.                 }

  1264.             }

  1265.         }

  1266.     }

  1267. 

  1268.     if (m != 32) {

  1269.         unsigned short tail_mask = (((unsigned short)0xffff) >> (32-m));

  1270.         for (int i = 0; i < n; i++) {

  1271.             STORE16_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i))

  1272.             STORE16_MASK_COMPLETE_RESULT(result_512[i+8], (C_addr + ldc*i + 16), tail_mask)

  1273.         }

  1274.     } else {

  1275.         for (int i = 0; i < n; i++) {

  1276.             STORE16_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i))

  1277.             STORE16_COMPLETE_RESULT(result_512[i+8], (C_addr + ldc*i + 16))

  1278.         }

  1279.     }

  1280. }

  1281. 

  1282. // SBGEMM Kernel for M<=16, N<8, K=Any number but will be processed based on 32

  1283. #ifndef ONE_ALPHA      // ALPHA is not ONE

  1284. void sbgemm_block_kernel_tn_16xNx32_alpha(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1285. #else                  // ALPHA is ONE

  1286. void sbgemm_block_kernel_tn_16xNx32_one(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1287. #endif

  1288. {

  1289.     bfloat16 * A_addr = A;

  1290.     bfloat16 * B_addr = B;

  1291.     float    * C_addr = C;

  1292. 

  1293.     BLASLONG tag_k_32x = k & (~31);

  1294. 

  1295. #ifndef ONE_ALPHA

  1296.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  1297. #endif

  1298. 

  1299.     __m512i arrayA_512;

  1300.     __m512i arrayB_512[8];

  1301.     __m512  result_512[8];

  1302. 

  1303.     for (int i = 0; i < 8; i++) {

  1304.         result_512[i]    = _mm512_setzero_ps();

  1305.     }

  1306. 

  1307.     for (BLASLONG idx_k = 0; idx_k < tag_k_32x; idx_k += 32) {

  1308.         // Load B with unroll n

  1309.         for (int i = 0; i < n; i ++) {

  1310.             arrayB_512[i] = _mm512_loadu_si512(B_addr);

  1311.             B_addr += 32;

  1312.         }

  1313. 

  1314.         for (BLASLONG idx = 0; idx < 32;) {

  1315.             // Each two rows are a group for 32-pair bf16 elements

  1316.             arrayA_512 = _mm512_loadu_si512(A_addr);

  1317.             A_addr += 32;

  1318. 

  1319.             for (int i = 0; i < n; i++) {

  1320.                 result_512[i]   = _mm512_dpbf16_ps(result_512[i], (__m512bh) arrayA_512, (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  1321.                 arrayB_512[i]   = _mm512_shuffle_epi32(arrayB_512[i], SHUFFLE_MAGIC_NO);

  1322.             }

  1323. 

  1324.             idx += 2;

  1325.             // Every 4 loops we need to switch to next 128 bits of arrayB registers

  1326.             if ((idx & (~7)) == idx) {

  1327.                 for (int i = 0; i < n; i++) {

  1328.                     arrayB_512[i] = _mm512_shuffle_i32x4(arrayB_512[i], arrayB_512[i], SHUFFLE_MAGIC_NO);

  1329.                 }

  1330.             }

  1331.         }

  1332.     }

  1333. 

  1334.     if (tag_k_32x != k) {

  1335.         // Load B with unroll n

  1336.         for (int i = 0; i < n; i ++) {

  1337.             arrayB_512[i] = _mm512_loadu_si512(B_addr);

  1338.             B_addr += 32;

  1339.         }

  1340. 

  1341.         BLASLONG width = k - tag_k_32x;

  1342.         for (BLASLONG idx = 0; idx < width;) {

  1343.             // Each two rows are a group for 32-pair bf16 elements

  1344.             arrayA_512 = _mm512_loadu_si512(A_addr);

  1345.             A_addr += 32;

  1346. 

  1347.             for (int i = 0; i < n; i++) {

  1348.                 result_512[i]   = _mm512_dpbf16_ps(result_512[i], (__m512bh) arrayA_512, (__m512bh) _mm512_broadcastd_epi32(_mm512_castsi512_si128(arrayB_512[i])));

  1349.                 arrayB_512[i]   = _mm512_shuffle_epi32(arrayB_512[i], SHUFFLE_MAGIC_NO);

  1350.             }

  1351. 

  1352.             idx += 2;

  1353.             // Every 4 loops we need to switch to next 128 bits of arrayB registers

  1354.             if ((idx & (~7)) == idx) {

  1355.                 for (int i = 0; i < n; i++) {

  1356.                     arrayB_512[i] = _mm512_shuffle_i32x4(arrayB_512[i], arrayB_512[i], SHUFFLE_MAGIC_NO);

  1357.                 }

  1358.             }

  1359.         }

  1360.     }

  1361. 

  1362.     if (m != 16) {

  1363.         unsigned short tail_mask = (((unsigned short)0xffff) >> (16-m));

  1364.         for (int i = 0; i < n; i++) {

  1365.             STORE16_MASK_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i), tail_mask)

  1366.         }

  1367.     } else {

  1368.         for (int i = 0; i < n; i++) {

  1369.             STORE16_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i))

  1370.         }

  1371.     }

  1372. }

  1373. 

  1374. #ifndef ONE_ALPHA      // ALPHA is not ONE

  1375. void sbgemm_blocking_kernel_tn_alpha(blasint M, blasint N, blasint K, float alpha, bfloat16 *A, blasint lda, bfloat16 *B, blasint ldb, float *C, blasint ldc, bfloat16 * block_A, bfloat16 * block_B)

  1376. #else                  // ALPHA is ONE

  1377. void sbgemm_blocking_kernel_tn_one(blasint M, blasint N, blasint K, float alpha, bfloat16 *A, blasint lda, bfloat16 *B, blasint ldb, float *C, blasint ldc, bfloat16 * block_A, bfloat16 * block_B)

  1378. #endif

  1379. {

  1380.     BLASLONG m_step, n_step, k_step, k_step_round32;

  1381.     BLASLONG tag_m_Nx = M & (~(BF16_BLOCK_THRES_M-1));

  1382. 

  1383.     BLASLONG n_from, n_to;

  1384.     BLASLONG tag_n_Nx;

  1385. 

  1386.     n_from = 0;

  1387.     n_to = (BF16_BLOCK_THRES_N > N) ? N : BF16_BLOCK_THRES_N;

  1388.     tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  1389. 

  1390.     k_step = (K > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : K;

  1391.     k_step_round32 = k_step & (~31);

  1392.     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  1393. 

  1394.     if (M >= BF16_BLOCK_THRES_M) {

  1395.         while (n_from < N) {

  1396.             for (BLASLONG idx_k = 0; idx_k < K;) {

  1397.                 // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  1398.                 COL_MAJOR_ITCOPY_KERNEL_Kx32(k_step, &A(0, idx_k), lda, block_A);

  1399.                 for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1400.                     // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  1401.                     COL_MAJOR_ONCOPY_KERNEL_8x32(k_step, &B(idx_n, idx_k), ldb, block_B + (idx_n-n_from)*k_step_round32);

  1402.                     SBGEMM_BLOCK_KERNEL_TN_32x8xK(32, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc); // TODO how to process m

  1403.                 }

  1404. 

  1405.                 if (tag_n_Nx != n_to) {

  1406.                     n_step = n_to - tag_n_Nx;

  1407.                     COL_MAJOR_ONCOPY_KERNEL_Nx32(n_step, k_step, &B(tag_n_Nx, idx_k), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  1408.                     SBGEMM_BLOCK_KERNEL_TN_32xNx32(32, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  1409.                 }

  1410. 

  1411.                 for (BLASLONG idx_m = BF16_BLOCK_THRES_M; idx_m < tag_m_Nx; idx_m += BF16_BLOCK_THRES_M) {

  1412.                     COL_MAJOR_ITCOPY_KERNEL_Kx32(k_step, &A(idx_m, idx_k), lda, block_A);

  1413.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1414.                         SBGEMM_BLOCK_KERNEL_TN_32x8xK(32, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, idx_m), ldc);

  1415.                     }

  1416. 

  1417.                     if (tag_n_Nx != n_to) {

  1418.                         n_step = n_to - tag_n_Nx;

  1419.                         SBGEMM_BLOCK_KERNEL_TN_32xNx32(32, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, idx_m), ldc);

  1420.                     }

  1421.                 }

  1422. 

  1423.                 if (tag_m_Nx != M) {

  1424.                     m_step = M - tag_m_Nx;

  1425.                     if (m_step > 16) {

  1426.                         COL_MAJOR_ITCOPY_KERNEL_Kx32m(m_step, k_step, &A(tag_m_Nx, idx_k), lda, block_A);

  1427.                         for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1428.                             SBGEMM_BLOCK_KERNEL_TN_32x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, tag_m_Nx), ldc);

  1429.                         }

  1430. 

  1431.                         if (tag_n_Nx != n_to) {

  1432.                             n_step = n_to - tag_n_Nx;

  1433.                             SBGEMM_BLOCK_KERNEL_TN_32xNx32(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, tag_m_Nx), ldc);

  1434.                         }

  1435.                     } else {

  1436.                         COL_MAJOR_ITCOPY_KERNEL_Kx16m(m_step, k_step, &A(tag_m_Nx, idx_k), lda, block_A);

  1437.                         for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1438.                             SBGEMM_BLOCK_KERNEL_TN_16x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, tag_m_Nx), ldc);

  1439.                         }

  1440. 

  1441.                         if (tag_n_Nx != n_to) {

  1442.                             n_step = n_to - tag_n_Nx;

  1443.                             SBGEMM_BLOCK_KERNEL_TN_16xNx32(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, tag_m_Nx), ldc);

  1444.                         }

  1445.                     }

  1446.                 }

  1447. 

  1448.                 idx_k += k_step;

  1449.                 k_step = K - idx_k;

  1450.                 k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  1451.                 k_step_round32 = k_step & (~31);

  1452.                 k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  1453.             }

  1454. 

  1455.             n_from = n_to;

  1456.             n_to += BF16_BLOCK_THRES_N;

  1457.             n_to = (n_to > N) ? N : n_to;

  1458.             tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  1459.         }

  1460.     } else {

  1461.         m_step = M;

  1462.         if (m_step > 16) {

  1463.             while (n_from < N) {

  1464.                 for (BLASLONG idx_k = 0; idx_k < K;) {

  1465.                     // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  1466.                     COL_MAJOR_ITCOPY_KERNEL_Kx32m(m_step, k_step, &A(0, idx_k), lda, block_A);

  1467.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1468.                         // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  1469.                         COL_MAJOR_ONCOPY_KERNEL_8x32(k_step, &B(idx_n, idx_k), ldb, block_B + (idx_n-n_from)*k_step_round32);

  1470.                         SBGEMM_BLOCK_KERNEL_TN_32x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  1471.                     }

  1472. 

  1473.                     if (tag_n_Nx != n_to) {

  1474.                         n_step = n_to - tag_n_Nx;

  1475.                         COL_MAJOR_ONCOPY_KERNEL_Nx32(n_step, k_step, &B(tag_n_Nx, idx_k), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  1476.                         SBGEMM_BLOCK_KERNEL_TN_32xNx32(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  1477.                     }

  1478. 

  1479.                     idx_k += k_step;

  1480.                     k_step = K - idx_k;

  1481.                     k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  1482.                     k_step_round32 = k_step & (~31);

  1483.                     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  1484.                 }

  1485.                 n_from = n_to;

  1486.                 n_to += BF16_BLOCK_THRES_N;

  1487.                 n_to = (n_to > N) ? N : n_to;

  1488.                 tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  1489.             }

  1490.         } else {

  1491.             while (n_from < N) {

  1492.                 for (BLASLONG idx_k = 0; idx_k < K;) {

  1493.                     // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  1494.                     COL_MAJOR_ITCOPY_KERNEL_Kx16m(m_step, k_step, &A(0, idx_k), lda, block_A);

  1495.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1496.                         // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  1497.                         COL_MAJOR_ONCOPY_KERNEL_8x32(k_step, &B(idx_n, idx_k), ldb, block_B + (idx_n-n_from)*k_step_round32);

  1498.                         SBGEMM_BLOCK_KERNEL_TN_16x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  1499.                     }

  1500. 

  1501.                     if (tag_n_Nx != n_to) {

  1502.                         n_step = n_to - tag_n_Nx;

  1503.                         COL_MAJOR_ONCOPY_KERNEL_Nx32(n_step, k_step, &B(tag_n_Nx, idx_k), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  1504.                         SBGEMM_BLOCK_KERNEL_TN_16xNx32(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  1505.                     }

  1506. 

  1507.                     idx_k += k_step;

  1508.                     k_step = K - idx_k;

  1509.                     k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  1510.                     k_step_round32 = k_step & (~31);

  1511.                     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  1512.                 }

  1513.                 n_from = n_to;

  1514.                 n_to += BF16_BLOCK_THRES_N;

  1515.                 n_to = (n_to > N) ? N : n_to;

  1516.                 tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  1517.             }

  1518.         }

  1519.     }

  1520. }

  1521. /* ----------------------------------------- End of TN kernels --------------------------------------- */

  1522. 

  1523. /* --------------------------------------------- TT kernels ------------------------------------------ */

  1524. // SBGEMM Kernel for 16<M<=32, N<8, K can be any number

  1525. #ifndef ONE_ALPHA      // ALPHA is not ONE

  1526. void sbgemm_block_kernel_tt_32xNxK_alpha(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1527. #else                  // ALPHA is ONE

  1528. void sbgemm_block_kernel_tt_32xNxK_one(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1529. #endif

  1530. {

  1531.     bfloat16 * A_addr = A;

  1532.     bfloat16 * B_addr = B;

  1533.     float    * C_addr = C;

  1534. 

  1535. #ifndef ONE_ALPHA

  1536.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  1537. #endif

  1538. 

  1539.     __m512i arrayA_512_0, arrayA_512_1;

  1540.     __m512i arrayB_512[8];

  1541.     __m512  result_512[16];

  1542. 

  1543.     result_512[0]  = _mm512_setzero_ps();

  1544.     result_512[1]  = _mm512_setzero_ps();

  1545.     result_512[2]  = _mm512_setzero_ps();

  1546.     result_512[3]  = _mm512_setzero_ps();

  1547.     result_512[4]  = _mm512_setzero_ps();

  1548.     result_512[5]  = _mm512_setzero_ps();

  1549.     result_512[6]  = _mm512_setzero_ps();

  1550.     result_512[7]  = _mm512_setzero_ps();

  1551.     result_512[8]  = _mm512_setzero_ps();

  1552.     result_512[9]  = _mm512_setzero_ps();

  1553.     result_512[10] = _mm512_setzero_ps();

  1554.     result_512[11] = _mm512_setzero_ps();

  1555.     result_512[12] = _mm512_setzero_ps();

  1556.     result_512[13] = _mm512_setzero_ps();

  1557.     result_512[14] = _mm512_setzero_ps();

  1558.     result_512[15] = _mm512_setzero_ps();

  1559. 

  1560.     for (BLASLONG idx_k = 0; idx_k < k; idx_k += 2) {

  1561.         // Each two rows are a group for 32-pair bf16 elements

  1562.         arrayA_512_0 = _mm512_loadu_si512(A_addr);

  1563.         arrayA_512_1 = _mm512_loadu_si512(A_addr + 32);

  1564.         A_addr += 64;

  1565. 

  1566.         for (int i = 0; i < n; i ++) {

  1567.             _MM512_BROADCASTD_EPI32(B_addr + i*2,  arrayB_512[i]);

  1568.         }

  1569.         B_addr += 16;

  1570. 

  1571.         for (int i = 0; i < n; i ++) {

  1572.             result_512[i] = _mm512_dpbf16_ps(result_512[i],  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512[i]);

  1573.             result_512[i+8] = _mm512_dpbf16_ps(result_512[i+8],  (__m512bh) arrayA_512_1, (__m512bh) arrayB_512[i]);

  1574.         }

  1575.     }

  1576. 

  1577.     if (m != 32) {

  1578.         unsigned short tail_mask = (((unsigned short)0xffff) >> (32-m));

  1579.         for (int i = 0; i < n; i ++) {

  1580.             STORE16_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i))

  1581.             STORE16_MASK_COMPLETE_RESULT(result_512[i+8], (C_addr + ldc*i + 16), tail_mask)

  1582.         }

  1583.     } else {

  1584.         for (int i = 0; i < n; i ++) {

  1585.             STORE16_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i))

  1586.             STORE16_COMPLETE_RESULT(result_512[i+8], (C_addr + ldc*i + 16))

  1587.         }

  1588.     }

  1589. }

  1590. 

  1591. // SBGEMM Kernel for M<=16, N<8, K can be any number

  1592. #ifndef ONE_ALPHA      // ALPHA is not ONE

  1593. void sbgemm_block_kernel_tt_16xNxK_alpha(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1594. #else                  // ALPHA is ONE

  1595. void sbgemm_block_kernel_tt_16xNxK_one(BLASLONG m, BLASLONG n, BLASLONG k, float alpha, bfloat16 *A, bfloat16 *B, float *C, int ldc)

  1596. #endif

  1597. {

  1598.     bfloat16 * A_addr = A;

  1599.     bfloat16 * B_addr = B;

  1600.     float    * C_addr = C;

  1601. 

  1602. #ifndef ONE_ALPHA

  1603.     __m512  ALPHAVECTOR = _mm512_set1_ps(alpha);

  1604. #endif

  1605. 

  1606.     __m512i arrayA_512_0;

  1607.     __m512i arrayB_512[8];

  1608.     __m512  result_512[8];

  1609. 

  1610.     result_512[0]  = _mm512_setzero_ps();

  1611.     result_512[1]  = _mm512_setzero_ps();

  1612.     result_512[2]  = _mm512_setzero_ps();

  1613.     result_512[3]  = _mm512_setzero_ps();

  1614.     result_512[4]  = _mm512_setzero_ps();

  1615.     result_512[5]  = _mm512_setzero_ps();

  1616.     result_512[6]  = _mm512_setzero_ps();

  1617.     result_512[7]  = _mm512_setzero_ps();

  1618. 

  1619.     for (BLASLONG idx_k = 0; idx_k < k; idx_k += 2) {

  1620.         // Each two rows are a group for 16-pair bf16 elements

  1621.         // Load two rows into a 512 register

  1622.         arrayA_512_0 = _mm512_loadu_si512(A_addr);

  1623.         A_addr += 32;

  1624. 

  1625.         for (int i = 0; i < n; i ++) {

  1626.             _MM512_BROADCASTD_EPI32(B_addr + i*2,  arrayB_512[i]);

  1627.         }

  1628.         B_addr += 16;

  1629. 

  1630.         for (int i = 0; i < n; i ++) {

  1631.             result_512[i] = _mm512_dpbf16_ps(result_512[i],  (__m512bh) arrayA_512_0, (__m512bh) arrayB_512[i]);

  1632.         }

  1633.     }

  1634. 

  1635.     if (m != 16) {

  1636.         unsigned short tail_mask = (((unsigned short)0xffff) >> (16-m));

  1637.         for (int i = 0; i < n; i++) {

  1638.             STORE16_MASK_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i), tail_mask)

  1639.         }

  1640.     } else {

  1641.         for (int i = 0; i < n; i++) {

  1642.             STORE16_COMPLETE_RESULT(result_512[i], (C_addr + ldc*i))

  1643.         }

  1644.     }

  1645. }

  1646. 

  1647. #ifndef ONE_ALPHA      // ALPHA is not ONE

  1648. void sbgemm_blocking_kernel_tt_alpha(blasint M, blasint N, blasint K, float alpha, bfloat16 *A, blasint lda, bfloat16 *B, blasint ldb, float *C, blasint ldc, bfloat16 * block_A, bfloat16 * block_B)

  1649. #else                  // ALPHA is ONE

  1650. void sbgemm_blocking_kernel_tt_one(blasint M, blasint N, blasint K, float alpha, bfloat16 *A, blasint lda, bfloat16 *B, blasint ldb, float *C, blasint ldc, bfloat16 * block_A, bfloat16 * block_B)

  1651. #endif

  1652. {

  1653.     BLASLONG m_step, n_step, k_step, k_step_round32;

  1654.     BLASLONG tag_m_Nx = M & (~(BF16_BLOCK_THRES_M-1));

  1655. 

  1656.     BLASLONG n_from, n_to;

  1657.     BLASLONG tag_n_Nx;

  1658. 

  1659.     n_from = 0;

  1660.     n_to = (BF16_BLOCK_THRES_N > N) ? N : BF16_BLOCK_THRES_N;

  1661.     tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  1662. 

  1663.     k_step = (K > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : K;

  1664.     k_step_round32 = k_step & (~31);

  1665.     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  1666. 

  1667.     if (M >= BF16_BLOCK_THRES_M) {

  1668.         while (n_from < N) {

  1669.             for (BLASLONG idx_k = 0; idx_k < K;) {

  1670.                 // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  1671.                 COL_MAJOR_ITCOPY_KERNEL_Kx32(k_step, &A(0, idx_k), lda, block_A);

  1672.                 for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1673.                     // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  1674.                     COL_MAJOR_OTCOPY_KERNEL_Kx8(k_step, &B(idx_k, idx_n), ldb, block_B + (idx_n-n_from)*k_step_round32);

  1675.                     SBGEMM_BLOCK_KERNEL_TT_32x8xK(32, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  1676.                 }

  1677. 

  1678.                 if (tag_n_Nx != n_to) {

  1679.                     n_step = n_to - tag_n_Nx;

  1680.                     COL_MAJOR_OTCOPY_KERNEL_Kx8m(k_step, n_step, &B(idx_k, tag_n_Nx), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  1681.                     SBGEMM_BLOCK_KERNEL_TT_32xNxK(32, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  1682.                 }

  1683. 

  1684.                 for (BLASLONG idx_m = BF16_BLOCK_THRES_M; idx_m < tag_m_Nx; idx_m += BF16_BLOCK_THRES_M) {

  1685.                     COL_MAJOR_ITCOPY_KERNEL_Kx32(k_step, &A(idx_m, idx_k), lda, block_A);

  1686.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1687.                         SBGEMM_BLOCK_KERNEL_TT_32x8xK(32, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, idx_m), ldc);

  1688.                     }

  1689. 

  1690.                     if (tag_n_Nx != n_to) {

  1691.                         n_step = n_to - tag_n_Nx;

  1692.                         SBGEMM_BLOCK_KERNEL_TT_32xNxK(32, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, idx_m), ldc);

  1693.                     }

  1694.                 }

  1695. 

  1696.                 if (tag_m_Nx != M) {

  1697.                     m_step = M - tag_m_Nx;

  1698.                     if (m_step > 16) {

  1699.                         COL_MAJOR_ITCOPY_KERNEL_Kx32m(m_step, k_step, &A(tag_m_Nx, idx_k), lda, block_A);

  1700.                         for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1701.                             SBGEMM_BLOCK_KERNEL_TT_32x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, tag_m_Nx), ldc);

  1702.                         }

  1703. 

  1704.                         if (tag_n_Nx != n_to) {

  1705.                             n_step = n_to - tag_n_Nx;

  1706.                             SBGEMM_BLOCK_KERNEL_TT_32xNxK(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, tag_m_Nx), ldc);

  1707.                         }

  1708.                     } else {

  1709.                         COL_MAJOR_ITCOPY_KERNEL_Kx16m(m_step, k_step, &A(tag_m_Nx, idx_k), lda, block_A);

  1710.                         for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1711.                             SBGEMM_BLOCK_KERNEL_TT_16x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, tag_m_Nx), ldc);

  1712.                         }

  1713. 

  1714.                         if (tag_n_Nx != n_to) {

  1715.                             n_step = n_to - tag_n_Nx;

  1716.                             SBGEMM_BLOCK_KERNEL_TT_16xNxK(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, tag_m_Nx), ldc);

  1717.                         }

  1718.                     }

  1719.                 }

  1720. 

  1721.                 idx_k += k_step;

  1722.                 k_step = K - idx_k;

  1723.                 k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  1724.                 k_step_round32 = k_step & (~31);

  1725.                 k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  1726.             }

  1727. 

  1728.             n_from = n_to;

  1729.             n_to += BF16_BLOCK_THRES_N;

  1730.             n_to = (n_to > N) ? N : n_to;

  1731.             tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  1732.         }

  1733.     } else {

  1734.         m_step = M;

  1735.         if (m_step > 16) {

  1736.             while (n_from < N) {

  1737.                 for (BLASLONG idx_k = 0; idx_k < K;) {

  1738.                     // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  1739.                     COL_MAJOR_ITCOPY_KERNEL_Kx32m(m_step, k_step, &A(0, idx_k), lda, block_A);

  1740.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1741.                         // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  1742.                         COL_MAJOR_OTCOPY_KERNEL_Kx8(k_step, &B(idx_k, idx_n), ldb, block_B + (idx_n-n_from)*k_step_round32);

  1743.                         SBGEMM_BLOCK_KERNEL_TT_32x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  1744.                     }

  1745. 

  1746.                     if (tag_n_Nx != n_to) {

  1747.                         n_step = n_to - tag_n_Nx;

  1748.                         COL_MAJOR_OTCOPY_KERNEL_Kx8m(k_step, n_step, &B(idx_k, tag_n_Nx), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  1749.                         SBGEMM_BLOCK_KERNEL_TT_32xNxK(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  1750.                     }

  1751. 

  1752.                     idx_k += k_step;

  1753.                     k_step = K - idx_k;

  1754.                     k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  1755.                     k_step_round32 = k_step & (~31);

  1756.                     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  1757.                 }

  1758.                 n_from = n_to;

  1759.                 n_to += BF16_BLOCK_THRES_N;

  1760.                 n_to = (n_to > N) ? N : n_to;

  1761.                 tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  1762.             }

  1763.         } else {

  1764.             while (n_from < N) {

  1765.                 for (BLASLONG idx_k = 0; idx_k < K;) {

  1766.                     // Use Kx32 kernel when BF16_BLOCK_THRES_M==32, Kx16 kernel when BF16_BLOCK_THRES_M==16, ...

  1767.                     COL_MAJOR_ITCOPY_KERNEL_Kx16m(m_step, k_step, &A(0, idx_k), lda, block_A);

  1768.                     for (BLASLONG idx_n = n_from; idx_n < tag_n_Nx; idx_n += BF16_BLOCK_STEP_N) {

  1769.                         // Use 8x32 kernel when BF16_BLOCK_THRES_N==8, 4x32 kernel when BF16_BLOCK_THRES_N==4, ...

  1770.                         COL_MAJOR_OTCOPY_KERNEL_Kx8(k_step, &B(idx_k, idx_n), ldb, block_B + (idx_n-n_from)*k_step_round32);

  1771.                         SBGEMM_BLOCK_KERNEL_TT_16x8xK(m_step, k_step, alpha, block_A, block_B + (idx_n-n_from)*k_step_round32, &C(idx_n, 0), ldc);

  1772.                     }

  1773. 

  1774.                     if (tag_n_Nx != n_to) {

  1775.                         n_step = n_to - tag_n_Nx;

  1776.                         COL_MAJOR_OTCOPY_KERNEL_Kx8m(k_step, n_step, &B(idx_k, tag_n_Nx), ldb, block_B + (tag_n_Nx-n_from)*k_step_round32);

  1777.                         SBGEMM_BLOCK_KERNEL_TT_16xNxK(m_step, n_step, k_step, alpha, block_A, block_B + (tag_n_Nx-n_from)*k_step_round32, &C(tag_n_Nx, 0), ldc);

  1778.                     }

  1779. 

  1780.                     idx_k += k_step;

  1781.                     k_step = K - idx_k;

  1782.                     k_step = (k_step > BF16_BLOCK_THRES_K) ? BF16_BLOCK_THRES_K : k_step;

  1783.                     k_step_round32 = k_step & (~31);

  1784.                     k_step_round32 = (k_step > k_step_round32) ? (k_step_round32 + 32) : k_step_round32;

  1785.                 }

  1786.                 n_from = n_to;

  1787.                 n_to += BF16_BLOCK_THRES_N;

  1788.                 n_to = (n_to > N) ? N : n_to;

  1789.                 tag_n_Nx = n_to & (~(BF16_BLOCK_STEP_N-1));

  1790.             }

  1791.         }

  1792.     }

  1793. }

  1794. /* ----------------------------------------- End of TT kernels --------------------------------------- */

  1795. 

  1796. /*

  1797. #ifndef ONE_ALPHA      // ALPHA is not ONE

  1798. void sbgemm_internal_kernel_alpha(OPENBLAS_CONST enum CBLAS_ORDER Order, OPENBLAS_CONST enum CBLAS_TRANSPOSE TransA, OPENBLAS_CONST enum CBLAS_TRANSPOSE TransB, OPENBLAS_CONST blasint M, OPENBLAS_CONST blasint N, OPENBLAS_CONST blasint K,

  1799. 		 OPENBLAS_CONST float alpha, OPENBLAS_CONST bfloat16 *A, OPENBLAS_CONST blasint lda, OPENBLAS_CONST bfloat16 *B, OPENBLAS_CONST blasint ldb, float *C, OPENBLAS_CONST blasint ldc)

  1800. #else                  // ALPHA is ONE

  1801. void sbgemm_internal_kernel_one(OPENBLAS_CONST enum CBLAS_ORDER Order, OPENBLAS_CONST enum CBLAS_TRANSPOSE TransA, OPENBLAS_CONST enum CBLAS_TRANSPOSE TransB, OPENBLAS_CONST blasint M, OPENBLAS_CONST blasint N, OPENBLAS_CONST blasint K,

  1802. 		 OPENBLAS_CONST float alpha, OPENBLAS_CONST bfloat16 *A, OPENBLAS_CONST blasint lda, OPENBLAS_CONST bfloat16 *B, OPENBLAS_CONST blasint ldb, float *C, OPENBLAS_CONST blasint ldc)

  1803. #endif

  1804. {

  1805.     if (Order == CblasColMajor) {

  1806.         if (TransA == CblasNoTrans) {

  1807.             if (TransB == CblasNoTrans) {

  1808.                 SBGEMM_BLOCKING_KERNEL_NN(M, N, K, alpha, A, lda, B, ldb, C, ldc, block_A, block_B);

  1809.             } else if (TransB == CblasTrans) {

  1810.                 SBGEMM_BLOCKING_KERNEL_NT(M, N, K, alpha, A, lda, B, ldb, C, ldc, block_A, block_B);

  1811.             }

  1812.         } else {

  1813.             if (TransB == CblasNoTrans) {

  1814.                 SBGEMM_BLOCKING_KERNEL_TN(M, N, K, alpha, A, lda, B, ldb, C, ldc, block_A, block_B);

  1815.             } else if (TransB == CblasTrans) {

  1816.                 SBGEMM_BLOCKING_KERNEL_TT(M, N, K, alpha, A, lda, B, ldb, C, ldc, block_A, block_B);

  1817.             }

  1818.         }

  1819.     } else {

  1820.         if (TransA == CblasNoTrans) {

  1821.             if (TransB == CblasNoTrans) {

  1822.                 SBGEMM_BLOCKING_KERNEL_NN(N, M, K, alpha, B, ldb, A, lda, C, ldc, block_A, block_B);

  1823.             } else if (TransB == CblasTrans) {

  1824.                 SBGEMM_BLOCKING_KERNEL_TN(N, M, K, alpha, B, ldb, A, lda, C, ldc, block_A, block_B);

  1825.             }

  1826.         } else {

  1827.             if (TransB == CblasNoTrans) {

  1828.                 SBGEMM_BLOCKING_KERNEL_NT(N, M, K, alpha, B, ldb, A, lda, C, ldc, block_A, block_B);

  1829.             } else if (TransB == CblasTrans) {

  1830.                 SBGEMM_BLOCKING_KERNEL_TT(N, M, K, alpha, B, ldb, A, lda, C, ldc, block_A, block_B);

  1831.             }

  1832.         }

  1833.     }

  1834. }

  1835. */