PowerPC: Sgemm Optimization

This patch improves GEMM for FP32 Data Type on PowerPC

Implements GEMM on large blocks with configurable block size mc, nc, kc
(default: 256, 256, 256).
Packing Function optimized to access blocks as per memory layout.
GEMM Optimized to work on larger blocks.
Isolated Packing from GEMM Operations for better MMA utilization.

Verified functionality and correctness uing llama-cli and stand alone
test case (performs matmul and compares final mattrix C result with
base).

Performance Testing:

Observed 50% ~ 70% improvement in Prompt Processing Speed mesured using
llama-bench with Meta-Llama3-8B FP32 Model.  Similar gains observed with
Mistral-7b-Instruct-v0.3 Model.

 model                   Size	             Params	Backend	      Threads	Test	Patch	Base
 llama 8B all F32        29.92 GiB 	     8.03 B 	 CPU           20 	pp512	98.58	60.3
 llama 8B all F32        29.92 GiB 	     8.03 B 	 CPU           20 	pp1024	95.88	57.36
 llama 8B all F32        29.92 GiB 	     8.03 B 	 CPU           20 	pp2048	85.46	53.26
 llama 8B all F32        29.92 GiB 	     8.03 B 	 CPU           20 	pp4096	68.66	45.78
 llama 8B all F32        29.92 GiB 	     8.03 B 	 CPU           20 	pp6144	57.35	40.44

25 ~ 30% improvement in llama-batched-bench with Metla-Llama3-8B in
Prompt Processing Speed for large prompts (256, 512, 1024, 2048, 4096)tokens with various batch
sizes ( 1, 2, 4, 8, 16)

Signed-off-by: root <[email protected]>

Author: root

ggml/src/ggml-cpu/llamafile/sgemm.cpp

@@ -2175,14 +2175,29 @@ class tinyBLAS_PPC {
                 int ith, int nth)
         : A(A), B(B), C(C), k(k), lda(lda), ldb(ldb), ldc(ldc), ith(ith), nth(nth) {
     }
-
     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 ( m%mc == 0 && n%nc == 0) {
+          matmul_tiled(m, n, mc, nc, kc);
+       } else {
+          mnpack(0, m, 0, n);
+       }
     }
 
   private:
 
     void (tinyBLAS_PPC::*kernel)(int64_t, int64_t);
+    
+    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);// += *((float*)&vec_C[I]+J);
+             *c_ptr += *((float*)&vec_C[I]+J);
+          }
+       }
+    }
 
     inline void vector_permute_store_4(vector float *src, float *vecOffset) {
        vector float t1, t2, t3, t4, t5, t6, t7, t8;
@@ -2235,7 +2250,8 @@ class tinyBLAS_PPC {
           vec_xst(t8, 0, vecOffset + 28);
     }
 
- void packTranspose(const float* a, int64_t lda, int rows, int cols, float* vec) {
+    void packTranspose(const float* a, int64_t lda, int rows, int cols, float* vec) {
+        //printf("In packTranspose\n");
         int64_t i, j;
         float * aoffsets[8];
         float *aoffset = NULL, *boffset = NULL;
@@ -2249,13 +2265,18 @@ class tinyBLAS_PPC {
         if (j > 0) {
 
             do {
+                //printf("j=%d\n", j);
                 aoffsets[0] = aoffset;
                 for (int it = 1; it< 8; it++)
                     aoffsets[it] = aoffsets[it-1] + lda;
+                /*for (int x = 0; x< 8; x++)
+                        printf("x=%d val=%f\n", x, *(aoffsets[x]));*/
                 aoffset += 8 * lda;
+                //aoffset += 8;// * lda;
                 i = (cols >> 3);
                 if (i > 0) {
                     do {
+                        //printf("i=%d\n", i);
                         for (int it = 0; it< 8; it++) {
                             arr[it] = __builtin_vsx_lxvp(0, (__vector_pair*)aoffsets[it]);
                             __builtin_vsx_disassemble_pair(c[it], &arr[it]);
@@ -2265,10 +2286,16 @@ class tinyBLAS_PPC {
 
                         vector_permute_store_8(c1, boffset);
                         vector_permute_store_8(c2, boffset+32);
-                        for (int it = 0; it < 4; it++)
-                            aoffsets[it] = aoffsets[it] + 8*lda;
                         boffset += 64;
                         i--;
+                        if (i > 0) {
+                        for (int it = 0; it < 8; it++) {
+                            //aoffsets[it] = aoffsets[it] + 8*lda;
+                            aoffsets[it] = aoffsets[it] + 8;
+                        }
+                        /*for (int x = 0; x< 8; x++)
+                            printf("x=%d val=%f\n", x, *(aoffsets[x]));*/
+                        }
                     } while(i > 0);
                 }
                 if (cols & 4) {
@@ -2401,6 +2428,68 @@ class tinyBLAS_PPC {
         SAVE_ACC(&acc_3, ii+4, jj+4);
     }
 
+    inline void MMA_16x8(vec_t *vec_A, vec_t *vec_B, acc_t * acc) {
+        for (int x = 0; x < 16; x += 2) {
+            __builtin_mma_xvf32gerpp(&acc[0], vec_A[x + 0], vec_B[x]);
+            __builtin_mma_xvf32gerpp(&acc[1], vec_A[x + 0], vec_B[x + 1]);
+            __builtin_mma_xvf32gerpp(&acc[2], vec_A[x + 1], vec_B[x]);
+            __builtin_mma_xvf32gerpp(&acc[3], vec_A[x + 1], vec_B[x + 1]);
+            __builtin_mma_xvf32gerpp(&acc[4], vec_A[x + 16], vec_B[x]);
+            __builtin_mma_xvf32gerpp(&acc[5], vec_A[x + 16], vec_B[x + 1]);
+            __builtin_mma_xvf32gerpp(&acc[6], vec_A[x + 17], vec_B[x]);
+            __builtin_mma_xvf32gerpp(&acc[7], vec_A[x + 17], vec_B[x + 1]);
+        }
+
+   }
+
+   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) {
+            for (int64_t i = 0; i <mc;  i += 16) {
+                for (int64_t j = 0; j < nc; j += 8) {
+                    acc_t acc[8];
+                    for (int x = 0; x < 8; x++)
+                        __builtin_mma_xxsetaccz(&acc[x]);
+                    for (int64_t l = 0; l < kc; l+=8) {
+                        int A_block_idx = (l/8)*16+(i / 16)*32;
+                        int B_block_idx = (l/8)*16+ (j / 8)*16;
+                        vec_t* A_block = &vec_A[A_block_idx];
+                        vec_t* B_block = &vec_B[B_block_idx];
+                        MMA_16x8(A_block, B_block, acc);
+                    }
+                    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);
+            }
+        }
+    }
+
+    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/4];
+                 vec_t B_pack[kc*nc/4];
+                 packTranspose(A+(ii*lda)+kk, lda, kc, mc, (float*)A_pack);
+                 packTranspose(B+(jj*ldb)+kk, ldb, kc, nc, (float*)B_pack);
+                 KERNEL(ii, jj, mc, nc, kc, A_pack, B_pack);
+            }
+        }
+    }
+
     void mnpack(int64_t m0, int64_t m, int64_t n0, int64_t n) {
         int m_rem = MIN(m - m0, 8);
         int n_rem = MIN(n - n0, 8);