src: gpu: intel: jit: gemm: limited a/b sum support for c repack

Author: hidefromkgb

src/gpu/intel/jit/gemm/generator/pieces/gemm_setup.cxx

@@ -1525,6 +1525,10 @@ bool BLASKernelGenerator<hw>::gemmAccumulateCSetup(GEMMProblem &problem, GEMMStr
                              state.repackA ? state.Ar_layout :
                                              state.A_layout;
         makeSumLayout(false, Ta, As_srcLayout, Tc, state.As_layout, strategy, state);
+        if (Tc != Tc_compute) {
+            std::swap(state.Asr_layout, state.As_layout);   /* TODO: trim down */
+            makeUnbackedRegLayout(Tc_compute, state.As_layout, unrollM, 1, true, 1);
+        }
         if (state.systolicSumA)
             setupTeardownAccumulateSumSystolic(true, Tb, problem, strategy, state);
     }
@@ -1538,6 +1542,10 @@ bool BLASKernelGenerator<hw>::gemmAccumulateCSetup(GEMMProblem &problem, GEMMStr
                              state.repackB ? state.Br_layout :
                                              state.B_layout;
         makeSumLayout(true,  Tb, Bs_srcLayout, Tc, state.Bs_layout, strategy, state);
+        if (Tc != Tc_compute) {
+            std::swap(state.Bsr_layout, state.Bs_layout);
+            makeUnbackedRegLayout(Tc_compute, state.Bs_layout, 1, unrollN, false, 1);
+        }
         if (state.systolicSumB)
             setupTeardownAccumulateSumSystolic(true, Ta, problem, strategy, state);
     }

src/gpu/intel/jit/gemm/generator/pieces/matrix_multiply.cxx

@@ -588,13 +588,19 @@ void BLASKernelGenerator<hw>::outerProductSystolic(int h, int ha, int hb, int op
                     B = state.sysSumAll1s[0];
                     nb = elementsPerGRF(hw, Tb);
                     B_block = &sumBlock;
-                    C = findBlockReg(Tc, state.As_layout, x, 0, state.As_regs, nc, C_block);
+                    if (repackC)
+                        C = findBlockReg(Tc, state.Asr_layout, x % Cr_unrollM, 0, state.Asr_regs, nc, C_block);
+                    else
+                        C = findBlockReg(Tc, state.As_layout, x, 0, state.As_regs, nc, C_block);
                 } else {
                     A = state.sysSumAll1s[0];
                     na = elementsPerGRF(hw, Ta);
                     A_block = &sumBlock;
                     B = findBlockReg(Tb, B_layout, hhb, x, B_regs, nb, B_block);
-                    C = findBlockReg(Tc, state.Bs_layout, 0, x, state.Bs_regs, nc, C_block);
+                    if (repackC)
+                        C = findBlockReg(Tc, state.Bsr_layout, 0, x % Cr_unrollN, state.Bsr_regs, nc, C_block);
+                    else
+                        C = findBlockReg(Tc, state.Bs_layout, 0, x, state.Bs_regs, nc, C_block);
                 }
 
                 int nv = globalCM ? na : nb;
@@ -672,6 +678,10 @@ void BLASKernelGenerator<hw>::outerProductRepackC(int x0, int xr0, int nx, int h
     bool globalCM = isLayoutColMajor(C_layout);
     bool scaleA = state.lateScale2DA, scaleB = state.lateScale2DB;
 
+    bool sumA = problem.needsASums();
+    bool sumB = problem.needsBSums();
+    if (globalCM ? sumB : sumA) stub();
+
     if (Tc.size() != Tc_compute.size()) stub();
     if (state.C_buffers > 1) stub();
 
@@ -712,41 +722,56 @@ void BLASKernelGenerator<hw>::outerProductRepackC(int x0, int xr0, int nx, int h
     for (int x1 = 0; x1 < nx; x1 += 2 * nec) {
         int x = x0 + x1, xr = xr0 + x1;
         int xchunk = std::min(nx - x1, 2 * nec);
-        for (int y = 0; y < ny; y++) {
+        for (int y = 0; y < ny + sumA + sumB; y++) {
             auto i = globalCM ? x : y;
             auto j = globalCM ? y : x;
             auto ir = globalCM ? xr : y;
             auto jr = globalCM ? y : xr;
 
-            int ne, ner, nes[2];
-            const RegisterBlock *C_block, *Cr_block, *sblock;
-            auto C = findBlockReg(Tc, C_layout, i, j, C_regs, ne, C_block);
-            auto Cr = findBlockReg(Tc_compute, Cr_layout, ir, jr, Cr_regs, ner, Cr_block);
+            int ne = 0, ner = 0, nes[2] = {0, 0};
+            const RegisterBlock *C_block = nullptr, *Cr_block = nullptr;
+            const RegisterBlock *sblock = nullptr;
+            Subregister C, Cr;
+
+            bool doASum = sumA && y == ny;
+            bool doBSum = sumB && y == ny;
+
+            if (y < ny) {
+                C  = findBlockReg(Tc, C_layout, i, j, C_regs, ne, C_block);
+                Cr = findBlockReg(Tc_compute, Cr_layout, ir, jr, Cr_regs, ner, Cr_block);
+            } else if (doASum) {
+                C  = findBlockReg(Tc, state.As_layout, x, 0, state.As_regs, ne, C_block);
+                Cr = findBlockReg(Tc_compute, state.Asr_layout, xr, 0, state.Asr_regs, ner, Cr_block);
+            } else if (doBSum) {
+                C  = findBlockReg(Tc, state.Bs_layout, 0, x, state.Bs_regs, ne, C_block);
+                Cr = findBlockReg(Tc_compute, state.Bsr_layout, 0, xr, state.Bsr_regs, ner, Cr_block);
+            }
 
             std::array<Subregister, 2> scale;
             std::array<int, 2> scaleStride = {0, 0};
             int nscale = 0;
-            if (scaleA) {
+            if (scaleA && !doBSum) {
                 int js = ((jr + h) / problem.aqGroupK) % state.kaqLate;
                 scale[nscale] = findBlockReg(state.Ta_scaleInt, state.Ar_scaleLayout,
                                              i, js, state.Ar_scaleRegs, nes[0], sblock);
                 scaleStride[nscale] = globalCM ? 1 : 0;
                 nscale++;
             }
-            if (scaleB) {
+            if (scaleB && !doASum) {
                 int is = ((ir + h) / problem.bqGroupK) % state.kbqLate;
                 scale[nscale] = findBlockReg(state.Tb_scaleInt, state.Br_scaleLayout,
                                              is, j, state.Br_scaleRegs, nes[1], sblock);
                 scaleStride[nscale] = globalCM ? 0 : 1;
                 nscale++;
             }
 
-            ne = std::min(ne, ner);
+            ne = std::min({ne, ner, xchunk});
             if (scaleStride[0] == 1) ne = std::min(ne, nes[0]);
             if (scaleStride[1] == 1) ne = std::min(ne, nes[1]);
 
             if (ne < xchunk) stub();
-            if (C_block->crosspack != 1 || Cr_block->crosspack != 1) stub();
+            if ((C_block && C_block->crosspack != 1)
+                    || (Cr_block && Cr_block->crosspack != 1)) stub();
 
             WorkItem item = {C, Cr, ne, iacc, scale, scaleStride};
             bool coalesce = false;

src/gpu/intel/jit/gemm/generator/pieces/register_allocation.cxx

@@ -613,6 +613,8 @@ void BLASKernelGenerator<hw>::gemmAllocRegs(GEMMProblem &problem, GEMMStrategy &
         state.Bi_regs[q] = state.ra.alloc_range(state.Bi_regCount);
 
     // Allocate registers for A/B sums.
+    state.Asr_regs = state.ra.alloc_range(getRegCount(state.Asr_layout));
+    state.Bsr_regs = state.ra.alloc_range(getRegCount(state.Bsr_layout));
     state.As_regs = state.ra.alloc_range(getRegCount(state.As_layout));
     state.Bs_regs = state.ra.alloc_range(getRegCount(state.Bs_layout));
 

src/gpu/intel/jit/gemm/generator/pieces/state.hpp

@@ -261,6 +261,7 @@ struct GEMMState : public CommonState {
     GRFMultirange Ao_regs, Bo_regs;                         // Outgoing data to copy to SLM.
     GRFMultirange Ao_regsRem, Bo_regsRem;
     GRFMultirange As_regs, Bs_regs;                         // A row sums/B column sums.
+    GRFMultirange Asr_regs, Bsr_regs;                       // A row sums/B column sums to be repacked.
     GRFMultirange Ap_regs, Bp_regs, Cp_regs;                // A/B/C prefetch registers.
     GRFMultirange A_offsetRegs, B_offsetRegs;               // A/B offsets (grouped).
     GRFMultirange A_scaleRegs, B_scaleRegs;                 // A/B scales (grouped).
@@ -327,6 +328,7 @@ struct GEMMState : public CommonState {
     std::vector<RegisterBlock> Ai_layoutRem, Bi_layoutRem;
     std::vector<RegisterBlock> Ao_layout, Bo_layout;
     std::vector<RegisterBlock> As_layout, Bs_layout;
+    std::vector<RegisterBlock> Asr_layout, Bsr_layout;
     std::vector<RegisterBlock> Ap_layout, Bp_layout, Cp_layout;
     std::vector<RegisterBlock> Ap_layoutAlt, Bp_layoutAlt;
     std::vector<RegisterBlock> A_offsetLayout, B_offsetLayout;