GEMM: bump MR to min(16,M) for skinny-N (n<=16) BF16 and F32 shapes (#524)

The default ZMM DE returns mr=6, nr=64, so for n<=16 shapes the JIT only reaches the lt16-mask kernel (bReg=1) and 6 of the 32 ZMMs hold C accumulators while the other ~25 sit idle. 

Overriding mr to min(16, m) lets each cached B line be consumed by up to 16 rows of A instead of 6, recovering the otherwise-wasted register file.

Change:

  * gemmBF16DEBackend / gemmF32DEBackend (ZMM fast path):
      For n<=16 and m>0, set mr = min(16, m). nr stays at nr_hint so
      the existing NR=64 packed-B layout, N-direction blocking, and
      rsB-divisor math are reused unchanged. F32 is additionally
      gated on !invokeRD and kc != 1.

Guards added so the bumped-MR path doesn't break the rest of the kernel set:

  * New `skinnyN` flag on kernel_frame::kernelInfo (threaded through
    the ctors, copy/move, operator== and
    gemmDEBackendUtils::checkPostOpsAndCreateKernelInfo). Set true
    only at the two ZMM override sites, when the MR bump actually
    fires; false everywhere else.

  * jitAmdZenFP32 / jitAmdZenBF16 generateAllKernels honor skinnyN by
    skipping nr>=2. Those wider NR variants (lt32 / 32 / lt48 / 48 /
    lt64 / 64) are unreachable for n<=16 and exceed the 32-ZMM budget
    at MR=16 (especially with post-ops and column-major beta scaling),
    so generating them only produced badKernelInfo aborts that took
    down the whole kernel set.

The default (n>16) path is untouched: skinnyN stays false and every
NR variant continues to be generated and dispatched as before.

[ AMD-Internal - SWLCSG-4250 ]

Author: BhaskarNallani

src/include/decision_engine/de_backend.hh

@@ -57,7 +57,8 @@ static const kernel_frame::kernelInfo INVALID_KERNEL_INFO{
     0,
     false,
     kernel_frame::kernelInstrPreference::none,
-    0
+    0,
+    false
 };
 
 class iDEBackend
@@ -387,11 +388,43 @@ class gemmF32DEBackend : public iDEBackend
             k_unroll = 4;
         }
 
+        // Increasing MR helps when n<=16 (mirrors the BF16 fix below).
+        //
+        // The default DE returns mr=6, nr=64. For shapes with n<=16 the
+        // dispatcher only ever reaches the NR=16 family of kernels (the
+        // lt16-mask kernel for n<16 and the NR=16 full kernel for n==16),
+        // so most of the 32 ZMMs sit idle as C accumulators. Bumping mr
+        // lets each cached B line be consumed by more rows of A.
+        //
+        // nr stays at nr_hint (=64) so the existing row-major NR=64
+        // packed-B layout and the framework's N-direction blocking are
+        // reused unchanged. The JIT generator below will see the bumped
+        // MR and skip the wider NR variants whose register budget would
+        // overflow at MR=16 (NR>=32 needs cReg>=32); those slots are
+        // never reached at runtime for n<=16 anyway.
+        //
+        // Override only when:
+        //  - n <= 16              (skinny-N: only NR=16 family reached)
+        //  - !invokeRD            (RD path has its own internal MR)
+        //  - kInstPref is ZMM     (32-ZMM budget is what the math relies on;
+        //                          AVX2 path has 16 regs, MR=16 won't fit)
+        //  - kc != 1              (the k=1 fused path is sized differently)
+        //
+        // For M < 16 we cap mr at m so the kernel uses an MR-partial
+        // kernel sized exactly to the input row count.
+        bool skinnyN = false;
+        if (!invokeRD && n <= 16 && m > 0 && kc != 1
+            && kInstPref
+                   == kernel_frame::kernelInstrPreference::avx512_zmm_favour) {
+            mr      = (m < 16) ? m : 16;
+            skinnyN = true;
+        }
+
         return gemmDEBackendUtils::checkPostOpsAndCreateKernelInfo(
             mr, nr, 0, k_unroll, kc, prefetch_c_dist, alphaScalingType,
             betaScalingType, mtag_a, mtag_b, allLtFringeKernels, invokeRD,
-            anyKOpsOrder, kInstPref, c_downscale, k_dtype, rs_c, cs_c,
-            metadata);
+            anyKOpsOrder, kInstPref, c_downscale, k_dtype, rs_c, cs_c, metadata,
+            skinnyN);
     }
 
     DLP_ALWAYS_INLINE
@@ -567,8 +600,40 @@ class gemmBF16DEBackend : public iDEBackend
         std::tie(alphaScalingType, betaScalingType) =
             gemmDEBackendUtils::getScalingTypes<float>(alpha, beta, k, kc_hint);
 
-        md_t mr              = mr_hint;
-        md_t nr              = nr_hint;
+        md_t mr = mr_hint;
+        md_t nr = nr_hint;
+
+        // Increasing MR helps when n<=16.
+        //
+        // The default DE returns mr=6, nr=64. For shapes with n<=16 the
+        // dispatcher only ever reaches the NR=16 family of kernels (the
+        // lt16-mask kernel for n<16 and the NR=16 full kernel for n==16),
+        // so only 6 of the 32 ZMM registers are used as C accumulators --
+        // the other ~25 ZMMs sit idle.
+        //
+        // We bump mr to min(16, M) so each cached B line is now consumed
+        // by up to 16 rows of A instead of 6, raising B reuse and cutting
+        // the M-iteration count from ceil(M/6) to ceil(M/16). With
+        // bReg=1 (the only NR variant the skinny-N dispatch reaches),
+        // cReg=16 and aReg=15: well inside the 32-ZMM budget.
+        //
+        // nr stays at nr_hint (=64) so the existing row-major NR=64
+        // packed-B layout and the framework's N-direction blocking are
+        // reused unchanged. The JIT generator below skips the wider NR
+        // variants whose register budget would overflow at MR=16
+        // (NR>=32 needs cReg>=32); those slots are never reached at
+        // runtime for n<=16 anyway.
+        //
+        // For M < 16 we cap mr at m so the kernel uses an MR-partial
+        // kernel sized exactly to the input row count (single full
+        // panel, no fringe). This avoids leaving C ZMMs idle for tiny-M
+        // shapes.
+        bool skinnyN = false;
+        if (n <= 16 && m > 0) {
+            mr      = (m < 16) ? m : 16;
+            skinnyN = true;
+        }
+
         md_t k_unroll        = 1;
         md_t kc              = kc_hint;
         md_t prefetch_c_dist = getPrefetchDistance();
@@ -582,7 +647,7 @@ class gemmBF16DEBackend : public iDEBackend
         return gemmDEBackendUtils::checkPostOpsAndCreateKernelInfo(
             mr, nr, 0, k_unroll, kc, prefetch_c_dist, alphaScalingType,
             betaScalingType, mtag_a, mtag_b, false, false, anyKOpsOrder,
-            kInstPref, c_downscale, k_dtype, rs_c, cs_c, metadata);
+            kInstPref, c_downscale, k_dtype, rs_c, cs_c, metadata, skinnyN);
     }
 };
 

src/include/decision_engine/de_backend_utils.hh

@@ -156,7 +156,8 @@ class gemmDEBackendUtils
         dlp::kernel_frame::kernelDatatype   k_dtype,
         [[maybe_unused]] md_t               rs_c,
         [[maybe_unused]] md_t               cs_c,
-        dlp_gemm_post_op*                   metadata)
+        dlp_gemm_post_op*                   metadata,
+        bool                                skinnyN = false)
     {
         // Iterate over the post_ops list to get the number of post-ops.
         md_t              numPostOps    = 0;
@@ -184,7 +185,8 @@ class gemmDEBackendUtils
                                              0,
                                              anyKOpsOrder,
                                              kInstPref,
-                                             c_downscale };
+                                             c_downscale,
+                                             skinnyN };
         } else {
             kernel_frame::kernelInfo kI{ mr,
                                          nr,
@@ -202,7 +204,8 @@ class gemmDEBackendUtils
                                          0,
                                          anyKOpsOrder,
                                          kInstPref,
-                                         c_downscale };
+                                         c_downscale,
+                                         skinnyN };
             kI.kOpsArrSize = numPostOps;
             kI.kOpsArr =
                 kernel_frame::kernelInfo::allocateKernelOpsArray(numPostOps);

src/include/kernel_frame/kernel_frame_base.hh

@@ -304,6 +304,12 @@ struct kernelInfo
     bool                  anyKOpsOrder;
     kernelInstrPreference kInstPref;
     md_t                  c_downscale;
+    // True when the DE has applied the skinny-N (n<=16) override that
+    // bumps MR above the default. Signals to the JIT generator that
+    // only the lt16 (nr=0) and full-16 (nr=1) N-direction variants will
+    // ever be invoked at runtime, so the wider lt32/32/lt48/48/lt64/64
+    // slots can be skipped during generation.
+    bool skinnyN;
 
     // Empty constructor to create dummy kernelInfo.
     kernelInfo()
@@ -323,6 +329,7 @@ struct kernelInfo
         , anyKOpsOrder(false)
         , kInstPref(kernel_frame::kernelInstrPreference::none)
         , c_downscale(0)
+        , skinnyN(false)
     {
     }
 
@@ -342,7 +349,8 @@ struct kernelInfo
                std::size_t                        kOpsArrSize,
                bool                               anyKOpsOrder,
                kernelInstrPreference              instPref,
-               md_t                               c_downscale)
+               md_t                               c_downscale,
+               bool                               _skinnyN = false)
         : mr(mr)
         , nr(nr)
         , term_fringe_nr(_term_fringe_nr)
@@ -363,6 +371,7 @@ struct kernelInfo
         , anyKOpsOrder(anyKOpsOrder)
         , kInstPref(instPref)
         , c_downscale(c_downscale)
+        , skinnyN(_skinnyN)
     {
     }
 
@@ -386,6 +395,7 @@ struct kernelInfo
         , anyKOpsOrder(other.anyKOpsOrder)
         , kInstPref(other.kInstPref)
         , c_downscale(other.c_downscale)
+        , skinnyN(other.skinnyN)
     {
         if ((other.kOpsArr != nullptr) && (other.kOpsArrSize > 0)) {
             this->kOpsArr =
@@ -419,6 +429,7 @@ struct kernelInfo
         , anyKOpsOrder(other->anyKOpsOrder)
         , kInstPref(other->kInstPref)
         , c_downscale(other->c_downscale)
+        , skinnyN(other->skinnyN)
     {
         if ((other->kOpsArr != nullptr) && (other->kOpsArrSize > 0)) {
             other->kOpsArr     = nullptr;
@@ -448,6 +459,7 @@ struct kernelInfo
         , anyKOpsOrder(other.anyKOpsOrder)
         , kInstPref(other.kInstPref)
         , c_downscale(other.c_downscale)
+        , skinnyN(other.skinnyN)
     {
         if ((other.kOpsArr != nullptr) && (other.kOpsArrSize > 0)) {
             other.kOpsArr     = nullptr;
@@ -485,6 +497,7 @@ struct kernelInfo
             this->anyKOpsOrder = other.anyKOpsOrder;
             this->kInstPref    = other.kInstPref;
             this->c_downscale  = other.c_downscale;
+            this->skinnyN      = other.skinnyN;
         }
         return *this;
     }
@@ -517,6 +530,7 @@ struct kernelInfo
             this->anyKOpsOrder = other.anyKOpsOrder;
             this->kInstPref    = other.kInstPref;
             this->c_downscale  = other.c_downscale;
+            this->skinnyN      = other.skinnyN;
         }
         return *this;
     }
@@ -552,7 +566,8 @@ struct kernelInfo
                 && (this->kOpsArrSize == rhs.kOpsArrSize) && isKOpsArrEqual
                 && (this->anyKOpsOrder == rhs.anyKOpsOrder)
                 && (this->kInstPref == rhs.kInstPref)
-                && (this->c_downscale == rhs.c_downscale));
+                && (this->c_downscale == rhs.c_downscale)
+                && (this->skinnyN == rhs.skinnyN));
     }
 
     // TODO: Need to implement a subset function for kernelInfo

src/jit/amdzen/amdzen_generator.cc

@@ -602,7 +602,22 @@ jitAmdZenFP32::generateAllKernels(const dlp::jit::jitGeneratorContext& jI)
             params.kernelOps.push_back((jI.kI).kOpsArr[ii]);
         }
 
-        // Generate all kernels for the given MR and NR
+        // Generate all kernels for the given MR and NR. Any per-variant
+        // generator failure (after passing the feasibility filter below)
+        // is fatal (goto cleanup): we want the existing fail-fast contract
+        // to hold, so we never call generateKernel() for an (mr, nr) pair
+        // we already know cannot fit.
+        //
+        // Feasibility filter: when the DE bumps MR (e.g. MR=16 for the
+        // skinny-N n<=16 override), the wider-NR variants (NR>=32) would
+        // exceed the register budget (cReg=MR*bReg, aReg = numRegs -
+        // cReg - bReg - maskVecReg, must have aReg >= 1 -- mirrors
+        // jitGEMMF32::allocateReg()). Skip those slots up front instead
+        // of relying on the per-variant generator to return badKernelInfo.
+        // The dispatcher only reaches the lt-mask kernel and the NR=16
+        // full kernel for n<=16, both of which always pass the filter.
+        const int kNumRegs =
+            (kType == utils::kernelInstrType::avx2_ymm_16_reg) ? 16 : 32;
         for (iter_t mr = 0; mr < numMRVariants; mr++) {
             for (iter_t nr = 0; nr < numNRVariants; nr++) {
                 params.MR    = mr == 0 ? MR : mr;
@@ -611,6 +626,42 @@ jitAmdZenFP32::generateAllKernels(const dlp::jit::jitGeneratorContext& jI)
                 int correspondingMainFringe = 0;
                 deriveGEMMNRAndMaskUse(nr, params, correspondingMainFringe);
 
+                // Skinny-N override: when the DE has bumped MR via the
+                // n<=16 override (jI.kI.skinnyN), only the lt-numElems
+                // (nr=0) and the full numElems (nr=1) variants are ever
+                // dispatched at runtime. The wider NR slots (nr>=2:
+                // lt2x/2x/lt3x/3x/lt4x/4x) are unreachable AND would
+                // exceed the register budget at bumped MR -- skip them
+                // entirely so we don't waste codegen on dead kernels
+                // (and don't trigger badKernelInfo on infeasible ones).
+                if (jI.kI.skinnyN && nr >= 2) {
+                    continue;
+                }
+
+                // Pre-filter register-infeasible (MR, NR) variants. This
+                // mirrors jitGEMMF32<KType>::allocateReg(): bFullReg =
+                // NR / numElemsPerReg, bMaskReg = useMask ? numMaskRegs
+                // : 0, bReg = bFullReg + bMaskReg, cReg = MR * bReg.
+                // For AVX2 ymm, the mask consumes vector registers
+                // (maskVecReg = numMaskRegs); for AVX-512 the mask is in
+                // Opmask regs and does not draw from the vector budget.
+                {
+                    int bFullReg = params.NR / numElemsPerReg;
+                    int bMaskReg = params.useMask ? params.numMaskRegs : 0;
+                    int bReg     = bFullReg + bMaskReg;
+                    int cReg     = params.MR * bReg;
+                    int maskVecReg =
+                        (kType == utils::kernelInstrType::avx2_ymm_16_reg)
+                            ? bMaskReg
+                            : 0;
+                    if (kNumRegs - cReg - bReg - maskVecReg < 1) {
+                        // Slot stays nullptr (zero-initialized by
+                        // resize). The dispatcher never reaches it for
+                        // any DE-blessed shape that bumped MR.
+                        continue;
+                    }
+                }
+
                 std::unique_ptr<Xbyak::CodeGenerator> gen;
                 switch (kType) {
                     case utils::kernelInstrType::avx512_zmm_32_reg: {
@@ -1414,14 +1465,54 @@ jitAmdZenBF16::generateAllKernels(const dlp::jit::jitGeneratorContext& jI)
             params.kernelOps.push_back((jI.kI).kOpsArr[ii]);
         }
 
-        // Generate all kernels for the given MR and NR
+        // Generate all kernels for the given MR and NR. Any per-variant
+        // generator failure (after passing the feasibility filter below)
+        // is fatal (goto cleanup): we want the existing fail-fast contract
+        // to hold, so we never call generateKernel() for an (mr, nr) pair
+        // we already know cannot fit.
+        //
+        // Feasibility filter: when the DE bumps MR (e.g. MR=16 for the
+        // skinny-N n<=16 override), the wider-NR variants (NR>=32) would
+        // exceed the 32-ZMM budget (cReg=MR*bReg, must have aReg = 32 -
+        // cReg - bReg >= aRegMin -- mirrors jitGEMMBF16::allocateReg()).
+        // Skip those slots up front instead of relying on the per-variant
+        // generator to return badKernelInfo. The dispatcher only reaches
+        // the lt16-mask kernel and the NR=16 full kernel for n<=16, both
+        // of which always pass the filter.
+        constexpr int kZmmRegs = 32;
+        const int     aRegMin  = ((jI.kI).c_downscale < DLP_F32) ? 2 : 1;
         for (iter_t mr = 0; mr < numMRVariants; mr++) {
             for (iter_t nr = 0; nr < numNRVariants; nr++) {
                 params.MR          = (mr == 0) ? MR : mr;
                 params.mLoop       = (mr == 0);
                 params.NR          = (nr * numElemsPerReg);
                 params.useMask     = (nr == 0);
                 params.numMaskRegs = (params.useMask) ? 1 : 0;
+
+                // Skinny-N override: when the DE has bumped MR via the
+                // n<=16 override (jI.kI.skinnyN), only the lt16 (nr=0)
+                // and full-16 (nr=1) variants are ever dispatched. The
+                // wider NR slots (nr>=2: lt32/32/lt48/48/lt64/64) are
+                // unreachable AND would exceed the 32-ZMM budget at
+                // bumped MR -- skip them entirely so we don't waste
+                // codegen on dead kernels (and don't trigger
+                // badKernelInfo on infeasible ones).
+                if (jI.kI.skinnyN && nr >= 2) {
+                    continue;
+                }
+
+                // For BF16 ZMM: bFullReg = (2*NR)/nBF16ElemsPerReg
+                // = NR/16 = nr (with numElemsPerReg=16). bMaskReg=1 for
+                // useMask, else 0. So bReg = max(1, nr).
+                int bReg = (nr == 0) ? 1 : static_cast<int>(nr);
+                int cReg = params.MR * bReg;
+                if (kZmmRegs - cReg - bReg < aRegMin) {
+                    // Slot stays nullptr (zero-initialized by resize).
+                    // The dispatcher never reaches it for any DE-blessed
+                    // shape that bumped MR.
+                    continue;
+                }
+
                 auto gen = std::make_unique<GEMMcodeGenerator::jitGEMMBF16<
                     utils::kernelInstrType::avx512_zmm_32_reg>>(
                     utils::JIT_KERNEL_SIZE);