From 0d0254754ee00273dca95a99833c38cd119074a8 Mon Sep 17 00:00:00 2001
From: "Wang, Yi A" <yi.a.wang@intel.com>
Date: Mon, 6 Jul 2026 05:28:19 +0000
Subject: [PATCH] Add XPU guidance for register-spill-aware autotune pruning to
 the `xpu-kernels` skill.

This PR extends the XPU kernel optimization skill with a new pattern for pruning Triton autotune configurations that are likely to spill registers on Intel XPU. The existing Xe-Forge workflow already detects register spills through profiling and recommends GRF/liveness mitigations; this change adds a proactive pruning strategy that can reduce wasted compile and benchmark time for obviously spill-prone configs.

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
---
 kernel-builder/skills/xpu-kernels/SKILL.md    | 24 +++++++++
 .../references/xpu_optimizations.yaml         | 49 +++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/kernel-builder/skills/xpu-kernels/SKILL.md b/kernel-builder/skills/xpu-kernels/SKILL.md
index e62aeaa51..df0d24b0d 100644
--- a/kernel-builder/skills/xpu-kernels/SKILL.md
+++ b/kernel-builder/skills/xpu-kernels/SKILL.md
@@ -176,6 +176,30 @@ def kernel(...):
     ...
 ```
 
+### Register-Spill-Aware Autotune Pruning
+
+For autotuned XPU kernels, prune configs that are likely to spill even after large-GRF recompilation. This avoids slow double-compile/benchmark cycles where the Intel backend reports spills, recompiles in large GRF mode, and still spills.
+
+Use an `early_config_prune` that estimates per-thread GRF pressure from the live accumulator and operand tiles:
+
+```text
+regs ~= (accumulators * BM * BN + stages * (BM * BK + weight_tiles * BN * BK))
+        / (num_warps * warp_size)
+```
+
+Guidelines:
+- Apply this only on XPU; leave other backends' configs unchanged.
+- Read `BM`, `BN`, and `BK` from either autotune config kwargs or launch args.
+- Count accumulator tiles explicitly: gate/up fused kernels often have 2 accumulators, down/simple GEMM usually has 1.
+- Count weight tiles explicitly: gate/up fused kernels often load 2 weight tiles, down/simple GEMM usually loads 1.
+- Use `num_stages` for pipelined MMA paths; use 1 for scalar/non-pipelined paths.
+- Keep a never-empty fallback, such as the config with the lowest estimated register pressure.
+- Treat the model as a coarse filter for catastrophic spillers, not a perfect predictor; leave borderline configs for real benchmarking.
+
+This complements `grf_mode='256'`: large GRF increases the register budget, while pruning removes configs that are still too register-heavy to be useful.
+
+For an implementation pattern, see `references/xpu_optimizations.yaml (xpu_register_spill_autotune_pruning)`.
+
 ### Tile Swizzling
 
 Use 1D grid with GROUP_SIZE_M for L2 locality:
diff --git a/kernel-builder/skills/xpu-kernels/references/xpu_optimizations.yaml b/kernel-builder/skills/xpu-kernels/references/xpu_optimizations.yaml
index 3fe22d2b5..b85782e35 100644
--- a/kernel-builder/skills/xpu-kernels/references/xpu_optimizations.yaml
+++ b/kernel-builder/skills/xpu-kernels/references/xpu_optimizations.yaml
@@ -272,6 +272,55 @@ patterns:
     expected_speedup: "workload-dependent"
     applies_to: [gemm, compute_bound]
 
+  - id: xpu_register_spill_autotune_pruning
+    name: "Register-spill-aware autotune pruning"
+    stage: xpu_specific
+    description: "Use early_config_prune to skip autotune configs that are very likely to spill on XPU"
+    rationale: |
+      XPU kernels with large accumulator and operand live sets can spend substantial autotune time
+      compiling and benchmarking configurations that spill even after large-GRF recompilation.
+      Add a coarse register-pressure estimator as an early_config_prune filter, then keep borderline
+      configs for real benchmarking.
+    pattern_after: |
+      def grf_config_pruner(
+          configs,
+          named_args,
+          *,
+          n_accumulators=1,
+          n_weight_tiles=1,
+          grf_regs=256,
+          headroom=0.82,
+        ):
+          if not is_xpu():
+              return configs
+
+          def meta(config, name, default):
+              return config.kwargs.get(name, named_args.get(name, default))
+
+          def estimate(config):
+              bm = meta(config, "BLOCK_M", 1)
+              bn = meta(config, "BLOCK_N", 1)
+              bk = meta(config, "BLOCK_K", 1)
+              stages = getattr(config, "num_stages", 1) or 1
+              warps = getattr(config, "num_warps", 1) or 1
+              warp_size = 32
+
+              live_values = (
+                  n_accumulators * bm * bn
+                  + stages * (bm * bk + n_weight_tiles * bn * bk)
+              )
+              return live_values / (warps * warp_size)
+
+          budget = grf_regs * headroom
+          kept = [config for config in configs if estimate(config) <= budget]
+          return kept or [min(configs, key=estimate)]
+    pitfalls:
+      - "Use this as a coarse pre-filter, not a replacement for benchmarking."
+      - "Keep a never-empty fallback so autotune always has at least one config."
+      - "Treat headroom as a workload- and backend-version-dependent tuning knob; calibrate it with profiler spill counters."
+      - "Do not apply globally to non-XPU backends."
+    applies_to: [gemm, grouped_gemm, moe, compute_bound]
+
   - id: xpu_warp_count
     name: "Warp count (tune; avoid hard-coded '32 warps')"
     stage: xpu_specific