Introduced OZAKI_SB

Author: hfp

samples/ozaki/kernels/ozaki1_int8.cl

@@ -27,6 +27,7 @@
  *   BK              - K-unroll factor for DPAS (32 for int8)
  *   KU              - K-loop unroll depth (2)
  *   NSLICES         - number of mantissa slices
+ *   OZAKI_SB        - slice-block width for the pair loop (1 = unblocked)
  *   MANT_BITS       - mantissa bit count (52 for fp64, 23 for fp32)
  *   BIAS_PLUS_MANT  - exponent bias + mantissa bits
  *   USE_DOUBLE      - if 1, fp64 accumulation; otherwise fp32
@@ -48,6 +49,17 @@
 #if !defined(NSLICES)
 # define NSLICES 8
 #endif
+/**
+ * Slice blocking for the pair loop.  The unblocked nest re-streams slice sa's
+ * A panel once per sb, so p slices cost p*p panel loads where p distinct
+ * panels exist.  Blocking sa and sb by OZAKI_SB hoists the fragment loads out
+ * of the pair dimension: loads drop by OZAKI_SB at the cost of OZAKI_SB^2
+ * accumulator sets held concurrently.  Requires the split load/compute macros,
+ * so it is Intel-only; other paths keep the unblocked nest.
+ */
+#if !defined(OZAKI_SB)
+# define OZAKI_SB 1
+#endif
 #if !defined(MANT_BITS)
 # define MANT_BITS 52
 #endif
@@ -78,6 +90,19 @@
 # define OZAKI_USE_OCL_KLOOP
 #endif
 
+/**
+ * Slice blocking needs the split load/compute macros (Intel DPAS path) and the
+ * square traversal: under OZAKI_SQ every ordered pair is its own block entry,
+ * whereas the triangular nest folds a transposed product into each off-diagonal
+ * pair, which a shared-fragment block cannot express.
+ */
+#if (1 < OZAKI_SB) && defined(INTEL) && (2 <= INTEL) && (RTM >= 2) && (RTN >= 2) && OZAKI_SQ
+# if 0 != (NSLICES % OZAKI_SB)
+#   error OZAKI_SB must divide NSLICES
+# endif
+# define OZAKI_SLICE_BLOCKED
+#endif
+
 /**
  * Bounds checks are OFF by default for performance.
  * Set -DOZAKI_BOUNDS=1 to enable per-element M/N guards
@@ -203,6 +228,49 @@
     } while (0)
 #endif
 
+#if defined(OZAKI_SLICE_BLOCKED)
+/**
+ * Slice-blocked K-loop: one K-step loads the A fragments of NA slices and the
+ * B fragments of NB slices, then issues every (ia, ib) sub-tile product from
+ * registers.  Load messages per K-step drop from 2*NA*NB (unblocked, one pair
+ * at a time) to NA+NB.
+ *
+ * AS_BASE/BS_BASE are the slice-0 panels, A_STRIDE/B_STRIDE the per-slice
+ * strides, SA0/SB0 the first slice of each block.  ACC holds NA*NB
+ * accumulator groups of RTM*RTN each, indexed (ia * NB + ib).
+ */
+# define OZAKI_KSTEP_BLOCKED(AS_BASE, BS_BASE, A_STRIDE, B_STRIDE, SA0, SB0, K_PAD_, N_PAD_, M_, MI, NJ, KOFF, ACC) \
+    do { \
+      ushort8 a_bk_[OZAKI_SB][RTM]; \
+      uint8 b_bk_[OZAKI_SB][RTN]; \
+      int ia_bk_, ib_bk_; \
+      UNROLL_FORCE(OZAKI_SB) for (ia_bk_ = 0; ia_bk_ < OZAKI_SB; ++ia_bk_) \
+      { \
+        OZAKI_LOAD_A_TILED((AS_BASE) + (long)((SA0) + ia_bk_) * (A_STRIDE), K_PAD_, M_, MI, KOFF, a_bk_[ia_bk_]); \
+      } \
+      UNROLL_FORCE(OZAKI_SB) for (ib_bk_ = 0; ib_bk_ < OZAKI_SB; ++ib_bk_) \
+      { \
+        OZAKI_LOAD_B_TILED((BS_BASE) + (long)((SB0) + ib_bk_) * (B_STRIDE), N_PAD_, K_PAD_, NJ, KOFF, b_bk_[ib_bk_]); \
+      } \
+      UNROLL_FORCE(OZAKI_SB) for (ia_bk_ = 0; ia_bk_ < OZAKI_SB; ++ia_bk_) \
+      { \
+        UNROLL_FORCE(OZAKI_SB) for (ib_bk_ = 0; ib_bk_ < OZAKI_SB; ++ib_bk_) \
+        { \
+          OZAKI_COMPUTE_TILED(a_bk_[ia_bk_], b_bk_[ib_bk_], (ACC) + (ia_bk_ * OZAKI_SB + ib_bk_) * RTM * RTN); \
+        } \
+      } \
+    } while (0)
+
+# define OZAKI_KLOOP_BLOCKED(AS_BASE, BS_BASE, A_STRIDE, B_STRIDE, SA0, SB0, K_PAD_, N_PAD_, M_, MI, NJ, ACC) \
+    do { \
+      int k_b_; \
+      for (k_b_ = 0; k_b_ < (K_PAD_); k_b_ += BK) { \
+        OZAKI_KSTEP_BLOCKED(AS_BASE, BS_BASE, A_STRIDE, B_STRIDE, SA0, SB0, \
+          K_PAD_, N_PAD_, M_, MI, NJ, k_b_, ACC); \
+      } \
+    } while (0)
+#endif /* OZAKI_SLICE_BLOCKED */
+
 /* K-loop prefetch: opt-in via OZAKI_PREFETCH=1 (default off on PVC). */
 #if defined(OZAKI_PREFETCH) && (0 < OZAKI_PREFETCH)
 # define OZAKI_KLOOP_PREFETCH(AS, BS, K, N, M, KOFF, MI, NJ) OZAKI_PREFETCH_TILED(AS, BS, K, N, M, KOFF, MI, NJ)
@@ -599,6 +667,49 @@ kernel void gemm_fused(
   }
 #endif
 
+#if defined(OZAKI_SLICE_BLOCKED)
+  /**
+   * Slice-blocked pair loop.  Each (sa0, sb0) block covers up to OZAKI_SB
+   * consecutive slices per side and shares the loaded fragments across all
+   * pairs inside it.  The block bounds are compile-time constants, so the
+   * cutoff predication below folds away entirely; only the tail blocks (where
+   * NSLICES or the cutoff truncates the block) retain a narrower extent.
+   *
+   * Square traversal only: with OZAKI_SQ the mirror pair is a separate (sa, sb)
+   * block entry, so no in-block transpose term is needed.
+   */
+  for (sa = 0; sa < (SINT)NSLICES && (int)sa <= OZAKI_CUTOFF; sa += OZAKI_SB) {
+    SINT sb0;
+    for (sb0 = 0; sb0 < (SINT)NSLICES; sb0 += OZAKI_SB) {
+      OZAKI_ACC_T c_blk[OZAKI_SB * OZAKI_SB * RTM * RTN];
+      int ia, ib;
+      UNROLL_FORCE(OZAKI_SB * OZAKI_SB * RTM * RTN)
+      for (ia = 0; ia < OZAKI_SB * OZAKI_SB * RTM * RTN; ++ia) c_blk[ia] = OZAKI_ACC_ZERO;
+      OZAKI_KLOOP_BLOCKED(as_base, bs_base, a_stride, b_stride, sa, sb0, K_pad, N_pad, M, mi_base, nj_base, c_blk);
+      /**
+       * Every slice of the block is loaded and multiplied unconditionally --
+       * OZAKI_SB divides NSLICES, so the indices stay in range -- and pairs
+       * past the cutoff are simply not flushed.  Their products are dead work
+       * on the tail blocks only, which is why OZAKI_SB stays small.
+       */
+      UNROLL_FORCE(OZAKI_SB) for (ia = 0; ia < OZAKI_SB; ++ia)
+      {
+        const int high_a = MANT_BITS - (7 * ((int)sa + ia));
+        const int low_a = MAX(0, high_a - 6);
+        UNROLL_FORCE(OZAKI_SB) for (ib = 0; ib < OZAKI_SB; ++ib)
+        {
+          const int high_b = MANT_BITS - (7 * ((int)sb0 + ib));
+          const int low_b = MAX(0, high_b - 6);
+          if ((int)sa + ia + (int)sb0 + ib <= OZAKI_CUTOFF) {
+            const real_t pscale = OZAKI_ALPHA_MUL(alpha, EXP2I(low_a + low_b - 2 * MANT_BITS));
+            OZAKI_SCALE_FLUSH(c_blk + (ia * OZAKI_SB + ib) * RTM * RTN, c, ldc, ea_cache, eb_cache, mi_base, nj_base, sg_lid,
+              M, N, pscale);
+          }
+        }
+      }
+    }
+  }
+#else
     for (sa = 0; sa < (SINT)NSLICES && (int)sa <= OZAKI_CUTOFF; ++sa) {
       const int high_sa = MANT_BITS - (7 * (int)sa);
       const int low_bit_sa = MAX(0, high_sa - 6);
@@ -760,6 +871,7 @@ kernel void gemm_fused(
 #endif
       }
     }
+#endif /* OZAKI_SLICE_BLOCKED */
 
 #if !defined(OZAKI_USE_OCL_KLOOP) && !(defined(NV_MMA) && (NV_MMA))
   /* Final write: register C -> global C */

samples/ozaki/ozaki_opencl.c

@@ -432,7 +432,7 @@ int ozaki_init(ozaki_context_t* ctx, int tm, int tn, int use_double, int kind, i
     char build_options[128];
     const int mant_bits = use_double ? 52 : 23;
     const int bias_plus_mant = use_double ? 1075 : 150;
-    int rtm = 0, rtn = 0, rtn_req = 0, ku_req, biggrf, hier;
+    int rtm = 0, rtn = 0, rtm_req = 0, rtn_req = 0, ku_req, biggrf, hier;
     size_t max_wgs;
     int v;
     {
@@ -461,7 +461,10 @@ int ozaki_init(ozaki_context_t* ctx, int tm, int tn, int use_double, int kind, i
     max_wgs = (0 != biggrf) ? devinfo->wgsize[0] / 2 : devinfo->wgsize[0];
     /* Read optional user overrides for register tiling factors. */
     env = getenv("OZAKI_RTM");
-    if (NULL != env && 0 < atoi(env)) rtm = atoi(env);
+    if (NULL != env && 0 < atoi(env)) {
+      rtm = atoi(env);
+      rtm_req = rtm; /* explicit request applies to both schemes */
+    }
     env = getenv("OZAKI_RTN");
     if (NULL != env && 0 < atoi(env)) {
       rtn = atoi(env);
@@ -503,9 +506,31 @@ int ozaki_init(ozaki_context_t* ctx, int tm, int tn, int use_double, int kind, i
      */
     env = getenv("OZAKI_NPANEL");
     ctx->npanel = (NULL != env) ? atoi(env) : 0;
+    /**
+     * Scheme-1 slice blocking (1 = unblocked).  The kernel shares each loaded
+     * A/B fragment across all pairs of an OZAKI_SB-wide slice block, trading
+     * OZAKI_SB^2 concurrent accumulator sets for an OZAKI_SB-fold cut in load
+     * messages.
+     *
+     * Opt-in: blocking only pays where the halved RTM it requires still leaves
+     * registers to spare, which on PVC means 256-GRF (SB=2 measured -18% at
+     * n=1024 and -5% at 6144 there).  Under the default GRF128 the row tiling
+     * is already 2, so the doubled accumulator set spills instead -- 3.3x
+     * slower at n=6144 -- and fp32 lacks the pair redundancy to recover even
+     * with registers available (+1.5% at n=4096, nslices=4).
+     */
+    env = getenv("OZAKI_SB");
+    { const int sb = (NULL != env && 0 < atoi(env)) ? atoi(env) : 1;
+      /* The kernel indexes whole blocks, so OZAKI_SB must divide NSLICES. */
+      ctx->sb = (1 < sb && sb <= nslices && 0 == (nslices % sb)) ? sb : 1;
+      if (1 < sb && ctx->sb != sb && 0 != verbosity) {
+        fprintf(stderr, "INFO OZAKI: OZAKI_SB=%d does not divide nslices=%d -- ignored\n", sb, nslices);
+      }
+    }
     if (0 == rtm) {
       if (0 != devinfo->intel && 0 != gpu) {
-        rtm = (0 != biggrf) ? 4 : 2;
+        /* Slice blocking holds SB^2 accumulator sets, so it halves RTM. */
+        rtm = (0 != biggrf) ? (1 < ctx->sb ? 2 : 4) : 2;
       }
       else if (0 != ctx->nv_mma && 0 != gpu) {
         rtm = 2;
@@ -650,10 +675,10 @@ int ozaki_init(ozaki_context_t* ctx, int tm, int tn, int use_double, int kind, i
         " -DNSLICES=%d -DUSE_DOUBLE=%d"
         " -DMANT_BITS=%d -DBIAS_PLUS_MANT=%d"
         " -DBM_PRE=%d -DBN_PRE=%d -DBK_PRE=%d"
-        " -DRTM=%d -DRTN=%d"
+        " -DRTM=%d -DRTN=%d -DOZAKI_SB=%d"
         " -DOZAKI_SQ=%d -DCONSTANT=global",
         bk_pre, ctx->ku, ctx->rc, sg, (int)devinfo->intel, nv,
-        nslices, use_double, mant_bits, bias_plus_mant, bm_pre, bn_pre, bk_pre, rtm, rtn, sq_jit);
+        nslices, use_double, mant_bits, bias_plus_mant, bm_pre, bn_pre, bk_pre, rtm, rtn, ctx->sb, sq_jit);
       if (0 != ctx->nv_mma) {
         goff += (size_t)LIBXS_SNPRINTF(build_params + goff, sizeof(build_params) - goff, " -DNV_MMA=1");
       }
@@ -737,7 +762,14 @@ int ozaki_init(ozaki_context_t* ctx, int tm, int tn, int use_double, int kind, i
       /* Fractional CRT is double-only: without fp64 fall back to Garner. */
       const int fraccrt = (0 == has_fp64) ? 0 : ((1 == fraccrt_req || 2 == fraccrt_req) ? fraccrt_req : 0);
       const int crt_hier = (1 == fraccrt) ? 0 : (0 != ctx->hier || 3 == kind || 2 == fraccrt);
-      const int crt_rtm = (0 != crt_hier && 0 != biggrf && 0 == ctx->hier) ? LIBXS_MAX(rtm / 2, 1) : rtm;
+      /**
+       * Scheme 2 has no pair loop, so slice blocking never applies to it and it
+       * must not inherit the halved RTM that blocking imposes on Scheme 1.
+       * rtm_crt_base is the row tiling Scheme 1 would have used unblocked.
+       */
+      const int rtm_crt_base = (1 < ctx->sb && 0 == rtm_req) ? rtm * 2 : rtm;
+      const int crt_rtm =
+        (0 != crt_hier && 0 != biggrf && 0 == ctx->hier) ? LIBXS_MAX(rtm_crt_base / 2, 1) : rtm_crt_base;
       /**
        * MMA gives a sub-tile 16 rows but only 8 columns, so reaching a square
        * register tile needs twice the column tiling. Scheme 2 measured +36% at
@@ -1009,6 +1041,7 @@ int ozaki_init(ozaki_context_t* ctx, int tm, int tn, int use_double, int kind, i
     if (ctx->crt_rtm != ctx->rtm) ozaki_print_opt(stderr, "crt_rtm", ctx->crt_rtm);
     ozaki_print_opt(stderr, "rtn", ctx->rtn);
     if (ctx->crt_rtn != ctx->rtn) ozaki_print_opt(stderr, "crt_rtn", ctx->crt_rtn);
+    if (1 < ctx->sb) ozaki_print_opt(stderr, "sb", ctx->sb);
     if (0 != devinfo->intel) {
       const int crt_grf128 = (0 != ctx->crt_rtm && ctx->crt_rtm < ctx->rtm);
       ozaki_print_opt(stderr, "grf", ctx->biggrf ? 256 : 128);

samples/ozaki/ozaki_opencl.h

@@ -187,6 +187,7 @@ typedef struct ozaki_context_t {
    */
   int rtm, rtn, crt_rtm, crt_rtn;
   int ku; /* K-loop unroll factor (compiled into kernel) */
+  int sb; /* Scheme-1 slice-block width for the pair loop (1 = unblocked) */
   int rc; /* DPAS repeat count: 8 (default) or 4 (split) */
   int nv_mma; /* NV MMA path enabled (m16n8k32, SG=32) */
   int pb; /* CRT prime batching factor (compiled into kernel) */