ggml-cpu: Enable BF16 tiled gemm optimization on PowerPC (llama/26068)

This commit is contained in:
shalinib-ibm
2026-07-30 16:31:49 +03:00
committed by Georgi Gerganov
parent 2e5b78c1b0
commit 98a532e1da
+126 -20
View File
@@ -1797,14 +1797,6 @@ class tinyBLAS_Q0_AVX {
//PPC Implementation
#if defined(__MMA__)
#define SAVE_ACC(ACC, ii, jj) \
__builtin_mma_disassemble_acc(vec_C, ACC); \
for (int I = 0; I < 4; I++) { \
for (int J = 0; J < 4; J++) { \
*((float*)(C+ii+((jj+J)*ldc)+I)) = *((float*)&vec_C[I]+J); \
} \
} \
template<typename T>
struct mma_instr;
@@ -1834,10 +1826,49 @@ class tinyBLAS_HP16_PPC {
}
void matmul(int64_t m, int64_t n) {
mnpack(0, m, 0, n);
int64_t mc = 256;
int64_t nc = 256;
int64_t kc = 256;
#if defined(_AIX) || defined(__BIG_ENDIAN__)
mc = 128;
nc = 128;
kc = 128;
#endif
if (k < kc) {
kc = k;
}
bool can_use_tiled = (m % mc == 0) && (n % nc == 0) && (k % kc == 0);
if (can_use_tiled) {
matmul_tiled(m, n, mc, nc, kc);
} else {
mnpack(0, m, 0, n);
}
}
private:
__attribute__((always_inline))
inline void save_acc(acc_t * ACC, int64_t ii, int64_t jj) {
vec_t vec_C[4];
__builtin_mma_disassemble_acc(vec_C, ACC);
for (int I = 0; I < 4; I++) {
for (int J = 0; J < 4; J++) {
*((float *)(C+ii+((jj+J)*ldc)+I)) = *((float *)&vec_C[I]+J);
}
}
}
__attribute__((always_inline))
inline void add_save_acc(acc_t * ACC, int64_t ii, int64_t jj) {
vec_t vec_C[4];
__builtin_mma_disassemble_acc(vec_C, ACC);
for (int I = 0; I < 4; I++) {
for (int J = 0; J < 4; J++) {
float * c_ptr = (float *)(C+ii+((jj+J)*ldc)+I);
*c_ptr += *((float *)&vec_C[I]+J);
}
}
}
void vector_permute_store(vec_t *c, int numVec, unsigned char *vecOffset) {
vec_t t[8], s[8];
vec_t swiz1 = {0, 1, 2, 3, 16, 17, 18, 19, 4, 5, 6, 7, 20, 21, 22, 23};
@@ -1896,6 +1927,7 @@ class tinyBLAS_HP16_PPC {
j = (rows >> 3);
if (j > 0) {
do {
aoffsets[0] = aoffset;
if (cols == 4) {
aoffsets[0] = aoffset;
for (int it = 1; it < 4; ++it)
@@ -1910,17 +1942,17 @@ class tinyBLAS_HP16_PPC {
}
i = (cols >> 3);
if (i > 0) {
aoffsets[0] = aoffset;
for (int it = 1; it < 8; ++it) {
aoffsets[it] = aoffsets[it-1] + lda;
}
aoffset += 8 * lda;
do {
for (int it = 0; it < 8; ++it)
c_arr[it] = vec_xl(0, (vector unsigned char*)aoffsets[it]);
vector_permute_store(c_arr, 8, vecOffset);
for (int it = 0; it < 8; ++it)
aoffsets[it] = aoffsets[it] + 8*lda;
aoffsets[it] = aoffsets[it] + 8;
vecOffset += 128;
i--;
} while(i > 0);
@@ -2147,8 +2179,8 @@ class tinyBLAS_HP16_PPC {
mma_instr<TA>::outer_product(&acc_1, vec_A[x], vec_B[x+4]);
}
}
SAVE_ACC(&acc_0, ii, jj);
SAVE_ACC(&acc_1, ii, jj+4);
save_acc(&acc_0, ii, jj);
save_acc(&acc_1, ii, jj+4);
}
void KERNEL_8x4(int64_t ii, int64_t jj) {
@@ -2164,8 +2196,8 @@ class tinyBLAS_HP16_PPC {
mma_instr<TA>::outer_product(&acc_1, vec_A[x+4], vec_B[x]);
}
}
SAVE_ACC(&acc_0, ii, jj);
SAVE_ACC(&acc_1, ii+4, jj);
save_acc(&acc_0, ii, jj);
save_acc(&acc_1, ii+4, jj);
}
@@ -2186,13 +2218,64 @@ class tinyBLAS_HP16_PPC {
mma_instr<TA>::outer_product(&acc_3, vec_A[x+4], vec_B[x+4]);
}
}
SAVE_ACC(&acc_0, ii, jj);
SAVE_ACC(&acc_1, ii, jj+4);
SAVE_ACC(&acc_2, ii+4, jj);
SAVE_ACC(&acc_3, ii+4, jj+4);
save_acc(&acc_0, ii, jj);
save_acc(&acc_1, ii, jj+4);
save_acc(&acc_2, ii+4, jj);
save_acc(&acc_3, ii+4, jj+4);
}
inline void MMA_16x8(vec_t * vec_A0, vec_t * vec_A1, vec_t * vec_B, acc_t * acc) {
for (int x = 0; x < 4; x ++) {
mma_instr<TA>::outer_product(&acc[0], vec_A0[x], vec_B[x]);
mma_instr<TA>::outer_product(&acc[1], vec_A0[x], vec_B[x+4]);
mma_instr<TA>::outer_product(&acc[2], vec_A0[x+4], vec_B[x]);
mma_instr<TA>::outer_product(&acc[3], vec_A0[x+4], vec_B[x+4]);
mma_instr<TA>::outer_product(&acc[4], vec_A1[x], vec_B[x]);
mma_instr<TA>::outer_product(&acc[5], vec_A1[x], vec_B[x+4]);
mma_instr<TA>::outer_product(&acc[6], vec_A1[x+4], vec_B[x]);
mma_instr<TA>::outer_product(&acc[7], vec_A1[x+4], vec_B[x+4]);
}
}
void KERNEL(int64_t ii, int64_t jj, int64_t mc, int64_t nc, int64_t kc, vec_t * vec_A, vec_t * vec_B, int64_t kk) {
for (int64_t i = 0; i < mc; i += 16) {
int A_base_addr = (mc / 8) * (i / 8) * 8;
for (int64_t j = 0; j < nc; j += 8) {
int B_base_addr = (nc / 8) * (j / 8) * 8;
acc_t acc[8];
vec_t A0_block[8]; vec_t A1_block[8];
for (int x = 0; x < 8; x++)
__builtin_mma_xxsetaccz(&acc[x]);
for (int64_t l = 0; l < kc; l += 8) {
int A0_block_idx = A_base_addr + (l / 8) * 8;
int A1_block_idx = A0_block_idx + (mc / 8) * 8;
int B_block_idx = B_base_addr + (l / 8) * 8;
vec_t* A0_block = &vec_A[A0_block_idx];
vec_t* A1_block = &vec_A[A1_block_idx];
vec_t* B_block = &vec_B[B_block_idx];
MMA_16x8(A0_block, A1_block, B_block, acc);
}
if (kk == 0) {
save_acc(&acc[0], ii + i, jj + j);
save_acc(&acc[1], ii + i, jj + j + 4);
save_acc(&acc[2], ii + i + 4, jj + j);
save_acc(&acc[3], ii + i + 4, jj + j + 4);
save_acc(&acc[4], ii + i + 8, jj + j);
save_acc(&acc[5], ii + i + 8, jj + j + 4);
save_acc(&acc[6], ii + i + 12, jj + j);
save_acc(&acc[7], ii + i + 12, jj + j + 4);
} else {
add_save_acc(&acc[0], ii + i, jj + j);
add_save_acc(&acc[1], ii + i, jj + j + 4);
add_save_acc(&acc[2], ii + i + 4, jj + j);
add_save_acc(&acc[3], ii + i + 4, jj + j + 4);
add_save_acc(&acc[4], ii + i + 8, jj + j);
add_save_acc(&acc[5], ii + i + 8, jj + j + 4);
add_save_acc(&acc[6], ii + i + 12, jj + j);
add_save_acc(&acc[7], ii + i + 12, jj + j + 4);
}
}
}
}
template<int RM, int RN>
void gemm_small(int64_t m0, int64_t m, int64_t n0, int64_t n) {
int64_t ytiles = (m - m0) / RM;
@@ -2281,6 +2364,29 @@ class tinyBLAS_HP16_PPC {
}
}
void matmul_tiled(int64_t m, int64_t n, int64_t mc, int64_t nc, int64_t kc) {
int64_t ytiles = m / mc;
int64_t xtiles = n / nc;
int64_t tiles = xtiles * ytiles;
int64_t duty = (tiles + nth - 1) / nth;
int64_t start = duty * ith;
int64_t end = start + duty;
if (end > tiles) {
end = tiles;
}
for (int64_t job = start; job < end; ++job) {
int64_t ii = (job / xtiles) * mc;
int64_t jj = (job % xtiles) * nc;
for (int64_t kk = 0; kk < k; kk += kc) {
vec_t A_pack[kc * mc / 8];
vec_t B_pack[kc * nc / 8];
packNormal(A + (ii * lda) + kk, lda, kc, mc, (uint8_t *)A_pack);
packNormal(B + (jj * ldb) + kk, ldb, kc, nc, (uint8_t *)B_pack);
KERNEL(ii, jj, mc, nc, kc, A_pack, B_pack, kk);
}
}
}
template <int RM, int RN>
NOINLINE void gemm(int64_t m0, int64_t m, int64_t n0, int64_t n) {
int64_t ytiles = (m - m0) / RM;