Xpu Kernels

# Files for xpu-kernels skill
SKILL.md
README.md
references/correctness.yaml
references/dtype_optimizations.yaml
references/fusion_patterns.yaml
references/huggingface-kernels-integration.md
references/implementation_reference.md
references/kernelbench-classification.md
references/memory_patterns.yaml
references/optimization_levels.yaml
references/optimization_strategies.md
references/persistent_kernel_patterns.yaml
references/workflow_details.md
references/xpu_optimizations.yaml
scripts/analyze_kernel.py
scripts/benchmark.py
scripts/benchmark_kernels.py
scripts/config.py
scripts/config.yaml
scripts/huggingface_kernels_example.py
scripts/requirements.txt
scripts/transformers_injection_example.py
scripts/trial_manager.py
scripts/validate_triton.py
scripts/xpu_profiler.py

constraints:
  - id: outputs_must_match
    name: "Outputs must match original"
    severity: info
    description: |
      The verification tool will check that outputs match the original.
      If it fails, try a different optimization approach.

  - id: streamk_output_must_be_prezeroed
    name: "Pre-zero output buffer when using atomic accumulation (Stream K)"
    severity: critical
    description: |
      When partial tiles use tl.atomic_add to accumulate results, the output
      tensor MUST be initialized to zero (torch.zeros, NOT torch.empty).
      Otherwise partial sums will include garbage values.

      WRONG:
      ```python
      c = torch.empty((M, N), device=a.device, dtype=torch.float32)
      first_wave[grid](a, b, c, ...)  # atomic_add onto garbage
      ```

      CORRECT:
      ```python
      c = torch.zeros((M, N), device=a.device, dtype=torch.float32)
      first_wave[grid](a, b, c, ...)  # atomic_add safely onto zeros
      ```

  - id: streamk_atomic_add_needs_mask
    name: "Atomic adds on partial tiles must be masked for boundary safety"
    severity: critical
    description: |
      When falling back to tl.atomic_add for partial tiles, you MUST apply
      boundary masks (rm < M, rn < N) to avoid writing out-of-bounds.

      ```python
      rm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
      rn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
      mask = (rm < M)[:, None] & (rn < N)[None, :]
      tl.atomic_add(c_ptr_, acc, mask=mask, sem='relaxed')
      ```

  - id: int64_cast_for_large_batch_offsets
    name: "Cast batch/stride products to int64 to prevent pointer overflow"
    severity: critical
    description: |
      When computing pointer offsets for batched operations, the product of
      a batch index and a stride can exceed int32 range for large tensors.
      Triton program_id returns int32 by default. You MUST cast to int64
      before multiplying by strides.

      WRONG (silent int32 overflow → wrong memory addresses):
      ```python
      bid = tl.program_id(axis=1)
      offset_a = bid * stride_az        # int32 * int32 → overflow for large tensors
      a_ptrs = a_ptr + offset_a + ...
      ```

      CORRECT:
      ```python
      bid = tl.program_id(axis=1)
      offset_a = bid.to(tl.int64) * stride_az   # safe for large tensors
      a_ptrs = a_ptr + offset_a + ...
      ```

      This applies whenever a program_id or loop index is multiplied by a
      stride that could produce values > 2^31 (≈2 billion elements). Common
      in batched GEMM, multi-head attention, and any kernel with a batch
      dimension over large tensors.

  - id: autotune_no_defaults
    name: "Do not put default values on @triton.autotune meta-parameters"
    severity: critical
    description: |
      When using @triton.autotune, the meta-parameters (BLOCK_M, BLOCK_N, etc.)
      must NOT have default values in the kernel signature. Default values cause
      a "Conflicting meta-parameters" error at runtime.

      WRONG:
      ```python
      @triton.autotune(configs=[...], key=['M', 'N', 'K'])
      @triton.jit
      def kernel(..., BLOCK_M: tl.constexpr = 128, ...):
          ...
      ```

      CORRECT:
      ```python
      @triton.autotune(configs=[...], key=['M', 'N', 'K'])
      @triton.jit
      def kernel(..., BLOCK_M: tl.constexpr, ...):
          ...
      ```

  - id: model_class_pattern
    name: "Model class must be compatible with ai-bench loading"
    severity: critical
    description: |
      ai-bench creates Model via direct `__init__()` and uses standard
      `load_state_dict()` for weight synchronization between reference
      and optimized models.

      The Model class should use standard nn.Module patterns:

      ```python
      class Model(nn.Module):
          def __init__(self, input_size, hidden_size, ...):
              super().__init__()
              self.gemm = nn.Linear(input_size, hidden_size)
              self._packed = False

          def _pack_weights(self):
              device = torch.device("xpu")
              w = self.gemm.weight.data.detach()
              b = self.gemm.bias.data.detach()
              self.weight_t = w.to(device, torch.float16).t().contiguous()
              self.bias_xpu = b.to(device, torch.float16).contiguous()
              self._packed = True

          def forward(self, x):
              if not self._packed:
                  self._pack_weights()
              # ... launch triton kernel ...
      ```

  - id: descriptor_no_boundary_check_arg
    name: "Tensor descriptor .load() does NOT accept boundary_check"
    severity: critical
    description: |
      Tensor descriptors are the preferred memory access API on XPU.
      Unlike block pointers which use tl.load(ptr, boundary_check=(0, 1)),
      tensor descriptors handle boundaries internally. The .load() method
      takes only a coordinate list.

      WRONG:
      ```python
      desc = tl.make_tensor_descriptor(base=ptr, shape=(M, K), ...)
      data = desc.load([row, col], boundary_check=(0, 1))
      ```

      CORRECT:
      ```python
      desc = tl.make_tensor_descriptor(base=ptr, shape=(M, K), ...)
      data = desc.load([row, col])
      ```

# Dtype Optimization Patterns for Intel XPU

patterns:
  - id: dtype_float64_to_float32
    name: "Float64 to Float32 Accumulator"
    stage: dtype_fix
    description: "Replace float64 accumulators with float32"
    rationale: |
      float64 throughput is 16-32x slower than float32 on GPUs/XPUs.
      This is the single biggest performance killer in many kernels.
      Using float64 alone can cap performance at around 2 TFLOPS on Intel XPU.
    pattern_before: |
      acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float64)
      a = a_fp32.to(tl.float64)
      b = b_fp32.to(tl.float64)
    pattern_after: |
      acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
      # No need to convert inputs - keep as float32
    expected_speedup: "5-10x"
    applies_to:
      - gemm
      - matmul
      - reduction
    examples:
      - before: |
          @triton.jit
          def kernel(a_ptr, b_ptr, c_ptr, ...):
              acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float64)
              for k in range(K):
                  a = tl.load(a_ptr + ...).to(tl.float64)
                  b = tl.load(b_ptr + ...).to(tl.float64)
                  acc = tl.dot(a, b, acc)
        after: |
          @triton.jit
          def kernel(a_ptr, b_ptr, c_ptr, ...):
              acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
              for k in range(K):
                  a = tl.load(a_ptr + ...)
                  b = tl.load(b_ptr + ...)
                  acc = tl.dot(a, b, acc)

  - id: dtype_input_conversion
    name: "Remove Unnecessary Type Conversions"
    stage: dtype_fix
    description: "Avoid converting inputs to higher precision unnecessarily"
    rationale: |
      Converting float16 inputs to float64 for computation wastes bandwidth
      and compute. Use float32 accumulators with float16 inputs for best
      performance on modern accelerators.
    pattern_before: |
      x = tl.load(x_ptr + offsets).to(tl.float64)
      result = x * x  # float64 computation
    pattern_after: |
      x = tl.load(x_ptr + offsets)  # Keep as float16
      x_fp32 = x.to(tl.float32)  # Upcast to float32 only if needed
      result = x_fp32 * x_fp32
    expected_speedup: "2-4x"
    applies_to:
      - elementwise
      - reduction

  - id: dtype_prepack_bf16
    name: "Pre-pack weights and inputs to bf16 before kernel launch"
    stage: dtype_fix
    description: |
      Convert weights to bf16 at _pack_weights() time and inputs to bf16
      before kernel launch, instead of loading fp32 and converting in-kernel.
    rationale: |
      Loading fp32 data and converting to bf16 inside the kernel wastes
      memory bandwidth:
        - fp32 load: 4 bytes per element from global memory
        - In-kernel .to(tl.bfloat16): discards half the loaded data
        - Net: 2x wasted bandwidth in the K-loop (the hottest path)

      Pre-packing to bf16 means the kernel loads 2 bytes per element directly.
      For a GEMM with K-loop iterations, this halves the memory traffic for
      both A and B tiles — often the difference between 2x and 4x+ speedup.
    pattern_before: |
      # In _pack_weights():
      self.weight_t = w.to(device).t().contiguous()       # stored as fp32

      # In forward():
      x = x.to(device).contiguous()                       # fp32 input

      # In kernel K-loop:
      a = tl.load(a_block_ptr, boundary_check=(0, 1))     # loads 4B per element
      a = a.to(tl.bfloat16)                               # converts to 2B — 2x waste
      b = tl.load(b_block_ptr, boundary_check=(0, 1))
      b = b.to(tl.bfloat16)
      acc += tl.dot(a, b)
    pattern_after: |
      # In _pack_weights():
      self.weight_t = w.to(device).t().contiguous().to(torch.bfloat16)  # bf16

      # In forward():
      x = x.to(device, torch.bfloat16).contiguous()       # bf16 input

      # In kernel K-loop (no conversion needed):
      a = tl.load(a_block_ptr, boundary_check=(0, 1))     # loads 2B directly
      b = tl.load(b_block_ptr, boundary_check=(0, 1))
      acc = tl.dot(a, b, acc=acc)                          # fused accumulate
    expected_speedup: "1.5-2x (halves K-loop memory traffic)"
    applies_to:
      - gemm
      - matmul
      - attention
      - inference
    notes: |
      - Keep bias and epilogue vectors in fp32 (small, precision-sensitive)
      - Combine with grf_mode='256' and tl.dot(a, b, acc=acc) for best results
      - Only for inference; training needs fp32 gradients
      - Works with both block pointers and tensor descriptors

# Kernel Fusion Patterns for Intel XPU
# Fusion can reduce memory traffic by eliminating intermediate writes/reads,
# but it can also hurt due to GRF/register pressure, reduced occupancy, and
# losing access to vendor-tuned primitives (e.g., GEMM).
#
# Guidance:
# - Fuse bandwidth-bound elementwise chains aggressively (usually a win).
# - Be cautious fusing into GEMM unless you already need a custom GEMM.
# - Do NOT fuse if the intermediate is not materialized, is dead/redundant,
#   or if fusion would replace a faster vendor primitive for a tiny epilogue.

constraints:
  - id: fuse_only_if_intermediate_is_materialized
    name: "Fuse only if it removes a materialized intermediate that is otherwise written/read"
    severity: critical
    description: |
      Fusion is beneficial only if the unfused baseline materializes an intermediate tensor
      (stores it to memory) and then reloads it in a separate kernel.

      If the intermediate is already:
      - handled by a fused library epilogue, OR
      - kept in registers within a single kernel, OR
      - dead / provably redundant for the workload,
      then fusion provides little/no benefit and may harm performance.

  - id: do_not_replace_vendor_gemm_for_tiny_epilogue
    name: "Do not replace vendor GEMM solely to fuse a tiny epilogue"
    severity: critical
    description: |
      If a vendor GEMM path exists and is faster, do not implement a custom Triton GEMM
      only to fuse a small epilogue (e.g., ReLU / min-sub / clamp).

      Prefer:
      - vendor GEMM + separate epilogue kernel, OR
      - vendor GEMM with native epilogue (if supported),
      unless the epilogue is substantial OR you must use a custom GEMM due to layout/math.

  - id: fusion_register_pressure_guard
    name: "Avoid fusion that causes GRF/register pressure collapse"
    severity: critical
    description: |
      Fusion increases live values and temporary tensors. If register pressure causes:
      - occupancy collapse, OR
      - GRF spills, OR
      - much lower parallelism,
      fusion can be slower than unfused.

      Red flags:
      - long activation chains with exp/log/tanh + multiple clamps
      - large tiles (e.g., 256x256) + high num_warps + many fused ops
      - reductions + heavy elementwise in same kernel on large blocks

      If unsure, fuse only a light epilogue (bias + simple activation) and keep the rest separate.

  - id: fuse_when_bandwidth_bound
    name: "Prefer fusion when the baseline is bandwidth-bound"
    severity: critical
    description: |
      Fusion helps most when the unfused path is bandwidth-bound:
      - elementwise chains (add/mul/clip/relu/etc.)
      - normalization epilogues (affine, bias, scale, clamp)
      - pointwise ops following reductions where intermediates are large

      Fusion helps least when:
      - compute-bound kernel dominates (large GEMM),
      - the intermediate is small,
      - or the baseline already uses a fused primitive.

  - id: skip_fusion_when_noop_or_redundant
    name: "Skip fusion when the operation is provably no-op or redundant"
    severity: critical
    description: |
      Do not fuse operations that are provably redundant for the workload:
      - multiply by 1, add 0
      - clamp/min/max where thresholds never trigger for the data range
      - dead outputs not consumed
      Extra fused instructions without reducing memory traffic can slow kernels down.

patterns:
  - id: elementwise_chain_fusion
    name: "Elementwise Chain Fusion (safe default)"
    stage: fusion
    description: "Fuse multiple elementwise ops into one kernel to reduce memory traffic"
    rationale: |
      Elementwise chains are typically bandwidth-bound. Each unfused step reads+writes the full
      tensor again. Fusion usually wins unless the chain is very large or involves heavy
      transcendentals that blow up register pressure.
    pattern_before: |
      y = a + b
      y = y * scale
      y = clamp(y, lo, hi)
      y = relu(y)
    pattern_after: |
      @triton.jit
      def fused_elementwise(a_ptr, b_ptr, out_ptr, ..., lo, hi, scale):
          x = tl.load(a_ptr + offsets, mask=mask, other=0.0)
          y = tl.load(b_ptr + offsets, mask=mask, other=0.0)
          x = x + y
          x = x * scale
          x = tl.maximum(x, lo)
          x = tl.minimum(x, hi)
          x = tl.maximum(x, 0.0)
          tl.store(out_ptr + offsets, x, mask=mask)
    expected_speedup: "1.2-3x (bandwidth dependent)"
    applies_to:
      - all_elementwise
      - normalization_epilogues

  - id: gemm_activation_fusion
    name: "GEMM + Activation Fusion (conditional)"
    stage: fusion
    description: "Fuse activation into GEMM ONLY when you already need a custom GEMM"
    rationale: |
      Fusion removes an intermediate write+read of the GEMM output. This can be beneficial
      when you already use a custom GEMM kernel (special layout, custom math, no vendor path).
      However, replacing a vendor GEMM with a custom fused GEMM for a tiny activation can be slower.
    applies_when:
      - uses_custom_gemm: true
      - intermediate_materialized: true
      - replaces_vendor_gemm: false
    rejects_when:
      - replaces_vendor_gemm: true
      - epilogue_is_tiny: true
      - num_pid_m_le_1: true
    pattern_before: |
      # Kernel 1: GEMM (custom)
      @triton.jit
      def gemm_kernel(a_ptr, b_ptr, c_ptr, ...):
          acc = tl.dot(a, b)
          tl.store(c_ptr + offsets, acc, mask=mask)

      # Kernel 2: Activation (separate launch)
      @triton.jit
      def activation_kernel(c_ptr, out_ptr, ...):
          c = tl.load(c_ptr + offsets, mask=mask, other=0.0)
          out = tl.maximum(c, 0.0)  # ReLU
          tl.store(out_ptr + offsets, out, mask=mask)
    pattern_after: |
      # Single fused kernel (custom GEMM + activation epilogue)
      @triton.jit
      def fused_gemm_relu(a_ptr, b_ptr, out_ptr, ...):
          acc = tl.dot(a, b)
          acc = tl.maximum(acc, 0.0)  # ReLU
          tl.store(out_ptr + offsets, acc, mask=mask)
    expected_speedup: "workload-dependent (positive if avoiding a large materialized intermediate; negative if replacing vendor GEMM)"
    applies_to:
      - custom_gemm_only
      - mlp
      - transformer

  - id: gemm_bias_activation_fusion
    name: "GEMM + Bias + Activation Chain (partial fusion recommended)"
    stage: fusion
    description: "Fuse bias + light activation into GEMM epilogue; split heavy chains if needed"
    rationale: |
      Bias add + simple activations are often good epilogues. Very heavy activation chains
      (multiple exp/tanh/clamp) can increase register pressure and reduce occupancy.
      Prefer partial fusion when needed:
      - fuse bias + simple activation
      - keep heavy post-processing in a separate kernel
    applies_when:
      - uses_custom_gemm: true
      - intermediate_materialized: true
      - replaces_vendor_gemm: false
    rejects_when:
      - replaces_vendor_gemm: true
      - activation_chain_is_heavy: true
    pattern_before: |
      c = gemm(a, b)
      c = c + bias
      c = swish(c)
      c = clamp(c, -1, 1)
      c = tanh(c)
    pattern_after: |
      @triton.jit
      def fused_gemm_bias_activation(..., bias_ptr, out_ptr, ...):
          acc = tl.dot(a, b)
          bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0)
          acc = acc + bias[None, :]

          # Light activation epilogue (example: ReLU or SiLU)
          # ReLU:
          acc = tl.maximum(acc, 0.0)

          tl.store(out_ptr + offsets, acc, mask=mask)
      # If you truly need heavy chains (swish+clamp+tanh), consider splitting:
      # - fused GEMM+bias+swish
      # - separate clamp+tanh kernel
    expected_speedup: "workload-dependent"
    applies_to:
      - mlp
      - gelu
      - silu

  - id: reduction_elementwise_fusion
    name: "Reduction + Elementwise Fusion (usually good; watch pressure)"
    stage: fusion
    description: "Fuse reductions with subsequent broadcast elementwise ops"
    rationale: |
      Reductions followed by broadcast (softmax / layernorm / rmsnorm) often benefit from
      fusion because it avoids writing the reduction results (max/sum/mean/var) to memory.
      However, very large reductions or overly large blocks can increase register/shared
      usage. If performance regresses, consider a 2-stage approach.
    pattern_before: |
      # Compute max for numerical stability
      max_val = tl.max(x, axis=1)
      x_stable = x - max_val[:, None]
      # Compute exp
      exp_x = tl.exp(x_stable)
      # Sum for softmax denominator
      sum_exp = tl.sum(exp_x, axis=1)
      # Normalize
      softmax = exp_x / sum_exp[:, None]
    pattern_after: |
      max_val = tl.max(x, axis=1, keep_dims=True)
      x_stable = x - max_val
      exp_x = tl.exp(x_stable)
      sum_exp = tl.sum(exp_x, axis=1, keep_dims=True)
      softmax = exp_x / sum_exp
    expected_speedup: "1.2-3x (shape dependent)"
    applies_to:
      - softmax
      - layernorm
      - attention

  - id: fusion_skip_when_dead_or_constant
    name: "Skip fusion if the fused op is dead / constant / no-op"
    stage: fusion
    description: "Do not fuse operations that are provably redundant for the given workload"
    rationale: |
      If an op is effectively a no-op for the given constants/data distribution,
      fusion adds extra instructions without reducing memory traffic meaningfully.

      Examples:
      - clamp range is so wide it never triggers
      - min/max with a constant that is outside observed values
      - multiply by 1, add 0
    expected_speedup: "prevents regressions"
    applies_to:
      - all

  - id: gemm_matrix_add_epilogue
    name: "GEMM + matrix add (residual/bias matrix) epilogue fusion"
    stage: fusion
    description: "Load a full matrix D after the GEMM K-loop and add to accumulator before storing"
    rationale: |
      A common pattern is C = A @ B + D where D is a full [M, N] matrix (residual
      connection, bias broadcast, etc.). Fusing the matrix add into the GEMM epilogue
      avoids a separate kernel launch and an extra read+write of the [M, N] output.

      The epilogue is lightweight (one descriptor load + one add), so register pressure
      impact is minimal and this fusion is almost always a win when you already have a
      custom GEMM kernel.

      Tensor descriptors preferred on XPU; block pointers also supported.
    pattern_before: |
      # Separate: GEMM then add
      c = matmul(a, b)           # writes [M, N] to memory
      c = c + d                  # reads [M, N] + [M, N], writes [M, N]
    pattern_after: |
      # Fused: add D inside GEMM epilogue (tensor descriptor variant)
      accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
      off_k = 0
      for _ in range(0, K, BLOCK_SIZE_K):
          a = a_desc.load([pid_m * BLOCK_SIZE_M, off_k])
          b = b_desc.load([off_k, pid_n * BLOCK_SIZE_N])
          accumulator += tl.dot(a, b)
          off_k += BLOCK_SIZE_K

      # Epilogue: load D tile and add
      d_desc = tl.make_tensor_descriptor(base=d_ptr, shape=(M, N),
                                         strides=(stride_dm, stride_dn),
                                         block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
      d = d_desc.load([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N])
      c = accumulator + d

      c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N),
                                         strides=(stride_cm, stride_cn),
                                         block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
      c_desc.store([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N], c)
    expected_speedup: "eliminates one full [M, N] read+write pass"
    applies_to:
      - gemm
      - residual_add
      - bias_matrix
      - transformer
    notes: |
      - D can be the same dtype as the accumulator or a different dtype (Triton auto-casts)
      - For batched variants, create d_desc with base=d_ptr + batch_offset
      - This pattern extends to any lightweight elementwise epilogue (add, mul, clamp, relu)
        applied after the K-loop and before the store

  - id: algebraic_weight_folding
    name: "Algebraic Weight Folding (BN/scale/affine into GEMM weights)"
    stage: fusion
    description: |
      Fold per-channel linear transforms into GEMM weights at pack time,
      eliminating the epilogue entirely. This is the most impactful fusion
      for kernels with BatchNorm inference, per-channel scaling, or affine transforms.
    rationale: |
      BatchNorm inference is a per-channel linear operation:
        BN(y) = gamma * (y - mean) / sqrt(var + eps) + beta

      When GEMM output feeds through scaling then BN:
        y = (x @ W^T + bias) * scale
        output = BN(y) = gamma * (y - mean) * inv_std + beta

      This entire chain is linear in W and bias. We can pre-compute:
        alpha = gamma * scale * inv_std                           [N]
        W_fused[n, :] = alpha[n] * W[n, :]                       [N, K]
        b_fused[n] = alpha[n] * bias[n] - gamma[n] * inv_std[n] * mean[n] + beta[n]  [N]

      The kernel becomes a pure GEMM + bias — zero epilogue overhead, zero extra
      vector loads, maximum compute efficiency.
    applies_when:
      - BatchNorm or LayerNorm in inference mode (fixed running stats)
      - Per-channel scaling (multiply by [N] vector)
      - Any per-channel affine transform: y = a*x + b
    rejects_when:
      - Training mode (statistics computed per-batch, not fixed)
      - Non-linear epilogues (sigmoid, ReLU) — cannot be folded
      - Operations depending on the M dimension (row-wise reductions)
    pattern_before: |
      # _pack_weights: just transpose
      self.weight_t = w.to(device).t().contiguous()
      self.bias_xpu = b.to(device)
      self.scale_xpu = s.to(device)
      self.gamma_xpu = gamma.to(device)
      # ... 6 more buffers

      # kernel epilogue: 6 vector loads + arithmetic
      bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0)
      acc = acc + bias[None, :]
      scale = tl.load(scale_ptr + offs_n, mask=mask_n, other=1.0)
      acc = acc * scale[None, :]
      mean = tl.load(mean_ptr + offs_n, mask=mask_n, other=0.0)
      var = tl.load(var_ptr + offs_n, mask=mask_n, other=1.0)
      gamma = tl.load(gamma_ptr + offs_n, mask=mask_n, other=1.0)
      beta = tl.load(beta_ptr + offs_n, mask=mask_n, other=0.0)
      inv_std = 1.0 / tl.sqrt(var + eps)
      acc = (acc - mean[None, :]) * inv_std[None, :]
      acc = acc * gamma[None, :] + beta[None, :]
    pattern_after: |
      # _pack_weights: fold BN + scale into weights (one-time cost)
      inv_std = 1.0 / torch.sqrt(rv + self.eps)
      alpha = gamma * s * inv_std
      w_fused = alpha.unsqueeze(1) * w
      b_fused = alpha * b - gamma * inv_std * rm + beta
      self.weight_t = w_fused.to(device).t().contiguous().to(torch.bfloat16)
      self.bias_fused = b_fused.to(device)

      # kernel epilogue: just bias add (everything else is folded)
      bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0)
      acc = acc + bias[None, :]
    expected_speedup: "1.5-2x on top of Level 1 optimizations (eliminates epilogue entirely)"
    applies_to:
      - gemm
      - batchnorm
      - inference
      - mlp
    examples:
      - file: test_kernels/39_Gemm_Scale_BatchNorm_triton.py
        description: "Folds Linear + Scale + BatchNorm into pure GEMM. Level 1: 2.69x → Level 3: 5.28x"

# Intel XPU Memory Patterns Optimizations
# These patterns focus on memory access, layout, and avoiding hidden sync/overhead.
# They are intended to be used alongside your XPU-specific compute patterns.

constraints:
  - id: no_device_to_host_scalar_sync
    name: "Do NOT force device->host scalar sync in hot path"
    severity: critical
    description: |
      Avoid .item(), float(tensor), int(tensor), printing device tensors, or any device->host scalar
      extraction in forward()/kernel wrapper hot paths. This forces synchronization and kills perf.

      WRONG:
      ```python
      c = float(constant_tensor.item())  # syncs XPU -> host
      ```

      CORRECT:
      - Keep constants as Python floats / CPU tensors
      - Pass scalar to kernel as an argument

  - id: prefer_contiguous_inputs
    name: "Prefer contiguous tensors for Triton kernels"
    severity: critical
    description: |
      Triton kernels generally assume or strongly benefit from contiguous memory.
      If inputs are strided/non-contiguous, do an explicit .contiguous() in wrapper
      (outside the timed region if possible).

      WRONG:
      ```python
      # Non-contiguous input passed into kernel
      kernel[grid](x_transposed, ...)
      ```

      CORRECT:
      ```python
      if not x.is_contiguous():
        x = x.contiguous()
      kernel[grid](x, ...)
      ```

  - id: block_ptr_boundary_check_tuple
    name: "Block pointer boundary_check must be tuple of ints"
    severity: critical
    description: |
      tl.load(block_ptr) uses boundary_check=(dim0, dim1) where values are dimension indices.
      Use (0,1) not booleans.

      CORRECT:
      ```python
      tl.load(ptr, boundary_check=(0, 1))
      ```

  - id: no_tl_multiple_of_on_python_scalars
    name: "Do NOT call tl.multiple_of / tl.max_contiguous on Python scalars or constexpr-like values"
    severity: critical
    description: |
      tl.multiple_of / tl.max_contiguous are intended for Triton IR values (tensors/expressions),
      not Python integers/constexpr-like scalars such as stride arguments.

      WRONG (can error like: 'constexpr' object has no attribute 'shape'):
      ```python
      tl.multiple_of(stride_xk, 1)
      ```

      CORRECT (apply to tensor expressions):
      ```python
      offs_n = col_start + tl.arange(0, BLOCK_N)
      tl.multiple_of(offs_n, 8)  # only if you KNOW it is aligned/multiple
      ```

patterns:
  - id: mem_block_pointers
    name: "Block pointers for structured tiles (fallback — prefer tensor descriptors on XPU)"
    stage: memory_access
    description: "Use tl.make_block_ptr + tl.load(boundary_check=...) for tiled loads/stores"
    rationale: |
      NOTE: On Intel XPU, tensor descriptors (tl.make_tensor_descriptor) are preferred
      over block pointers — they produce better address generation codegen. See
      xpu_optimizations.yaml (xpu_tensor_descriptor_note, xpu_descriptor_gemm_pattern).

      Block pointers improve addressing and bounds handling over manual pointer math.
      Use them as a fallback when tensor descriptors are not suitable (e.g., legacy code).
    pattern_before: |
      offs_m = pid_m * BM + tl.arange(0, BM)
      offs_k = k0 + tl.arange(0, BK)
      x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk
      x = tl.load(x_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K), other=0.0)
    pattern_after: |
      x_bp = tl.make_block_ptr(
        base=x_ptr,
        shape=(M, K),
        strides=(stride_xm, stride_xk),
        offsets=(pid_m * BM, k0),
        block_shape=(BM, BK),
        order=(1, 0),
      )
      x = tl.load(x_bp, boundary_check=(0, 1))
    expected_speedup: "workload-dependent (often positive when indexing is heavy)"
    applies_to: [gemm, conv, attention, all_memory_bound]
    examples:
      - before: |
          k0 += BK
        after: |
          x_bp = tl.advance(x_bp, (0, BK))

  - id: mem_skip_boundary_check_when_divisible
    name: "Specialize away boundary checks when shapes are divisible"
    stage: memory_access
    description: "Avoid boundary_check/masks when M/N/K are known multiples of tile sizes"
    rationale: |
      Masks and boundary checks add overhead. If your benchmark/problem guarantees divisibility
      (or you can specialize per shape), you can remove masks and use unconditional loads/stores.
    pattern_before: |
      x = tl.load(x_bp, boundary_check=(0, 1))
      tl.store(o_bp, acc, boundary_check=(0, 1))
    pattern_after: |
      # If M % BM == 0 and K % BK == 0 (specialized case):
      x = tl.load(x_bp)
      tl.store(o_bp, acc)
    expected_speedup: "1.05-1.20x"
    applies_to: [gemm, large_fixed_shapes]
    notes: |
      - Only safe when you can guarantee divisibility or guard with a specialized kernel variant.

  - id: mem_use_static_range_for_k_loop
    name: "Use tl.static_range for K loops (when K is known / bounded)"
    stage: memory_access
    description: "Enable better unrolling/pipelining by using tl.static_range"
    rationale: |
      tl.static_range can help Triton generate better scheduled loops and reduce loop overhead.
      Particularly useful in GEMM-like reductions and fixed-shape workloads.
    pattern_before: |
      for k in range(0, K, BK):
        ...
    pattern_after: |
      for k in tl.static_range(0, K, BK):
        ...
    expected_speedup: "1.05-1.15x"
    applies_to: [gemm, attention]
    notes: |
      Use only when K is a compile-time constant or you launch specialized kernels by K.

  - id: mem_pointer_arithmetic_coalescing
    name: "Prefer contiguous (coalesced) access in the innermost dimension"
    stage: memory_access
    description: "Make the fastest-changing index map to contiguous memory"
    rationale: |
      Coalesced loads/stores are critical. In most layouts, the last dimension is contiguous.
      Make tl.arange(...) along the contiguous stride dimension when possible.
    pattern_before: |
      # Strided / scattered pattern (example)
      offs = base + tl.arange(0, BN) * stride_n
      x = tl.load(x_ptr + offs)
    pattern_after: |
      # Coalesced: make arange hit contiguous dimension
      offs_n = col_start + tl.arange(0, BN)
      x = tl.load(x_ptr + offs_n)
    expected_speedup: "workload-dependent (can be large if fixing uncoalesced access)"
    applies_to: [all_memory_bound]

  - id: mem_cache_modifiers_optional
    name: "Optional cache modifiers (backend-dependent)"
    stage: memory_access
    description: "Use cache_modifier/eviction_policy carefully; keep it optional"
    rationale: |
      cache modifiers can help (or hurt) depending on backend and access pattern.
      Treat them as an autotune dimension rather than a hard rule.
    pattern_before: |
      x = tl.load(ptr)
    pattern_after: |
      # Example: try these as tunables if supported by your Triton backend
      x = tl.load(ptr, cache_modifier=".ca")
      # or:
      x = tl.load(ptr, cache_modifier=".cg")
    expected_speedup: "workload-dependent"
    applies_to: [streaming_reads, bandwidth_bound]
    notes: |
      - Do not assume Intel XPU backend interprets modifiers the same as CUDA.
      - Use only if you confirm correctness and performance.

  - id: mem_alignment_hints_on_offsets
    name: "Alignment hints on OFFSET tensors (not strides)"
    stage: memory_access
    description: "Use tl.multiple_of / tl.max_contiguous on tensor offset expressions"
    rationale: |
      Alignment hints can improve vectorization/coalescing when you truly have aligned access.
      Apply hints to offset tensors (e.g., tl.arange expressions) and only when guaranteed.
    pattern_before: |
      offs_n = col_start + tl.arange(0, BN)
      x = tl.load(x_ptr + offs_n)
    pattern_after: |
      offs_n = col_start + tl.arange(0, BN)
      tl.max_contiguous(offs_n, BN)
      # Only if you KNOW alignment/multiple is valid:
      # tl.multiple_of(offs_n, 8)
      x = tl.load(x_ptr + offs_n)
    expected_speedup: "small but sometimes measurable"
    applies_to: [gemm, bandwidth_bound]
    notes: |
      Do NOT apply tl.multiple_of blindly—incorrect assumptions can lead to wrong-code.

  - id: mem_avoid_redundant_reads
    name: "Load once, reuse in registers"
    stage: memory_access
    description: "Avoid re-loading the same values inside inner loops"
    rationale: |
      Redundant tl.load calls inside reduction loops can dominate memory traffic.
      If values are reused (e.g., bias, scale), load once per tile and broadcast.
    pattern_before: |
      for k in ...:
        b = tl.load(b_ptr + offs_n)  # re-loaded each iteration (waste)
        acc += ...
    pattern_after: |
      b = tl.load(b_ptr + offs_n)    # load once
      for k in ...:
        acc += ...
      acc += b[None, :]
    expected_speedup: "1.05-1.30x (depends on redundancy)"
    applies_to: [gemm_epilogues, fused_pointwise]

  - id: mem_wrapper_one_time_device_moves
    name: "One-time device moves in wrapper (avoid per-forward transfers)"
    stage: memory_access
    description: "Move parameters to XPU once, not every forward"
    rationale: |
      Per-forward .to('xpu') or .data = .to('xpu') costs time and can interfere with timing.
      Use a one-time guard in the module.
    pattern_before: |
      def forward(self, x):
        if self.weight.device.type != "xpu":
          self.weight.data = self.weight.data.to("xpu")
        ...
    pattern_after: |
      def __init__(...):
        self._moved_to_xpu = False

      def forward(self, x):
        if not self._moved_to_xpu:
          self.weight.data = self.weight.data.to("xpu")
          self._moved_to_xpu = True
    expected_speedup: "huge for microbenchmarks; correctness/measurement improvement"
    applies_to: [all]
    notes: |
      Ideally move model+inputs to XPU outside the timed region in the harness.

  - id: mem_layout_transform_prepack
    name: "Pre-pack / pre-transform weights for better access"
    stage: memory_access
    description: "Change weight layout to match access pattern (when allowed)"
    rationale: |
      If your kernel accesses W as W^T (K,N), storing W in a layout that makes K-contiguous
      for the inner loop can improve coalescing/cache behavior. This is a higher-level change.
    pattern_before: |
      # W stored as [N, K], kernel reads as [K, N] view
      w_ptr shape=(K, N), strides=(stride_wk, stride_wn)
    pattern_after: |
      # Prepack once (outside timed loop), store as [K, N] or blocked layout if acceptable
      # so inner K loads are contiguous in memory for each program.
      # (Exact layout depends on your kernel and framework constraints.)
    expected_speedup: "workload-dependent (can be large for bandwidth-limited kernels)"
    applies_to: [gemm, attention]
    notes: |
      - Only if you control weight storage and can amortize the prepack cost.
      - Ensure KernelBench correctness expectations still hold.

  - id: reduce_liveness_sink_load_and_prefetch
    name: "Reduce variable liveness: prefetch early, load late (sink load closer to dot)"
    stage: memory_access
    description: "Replace long-lived live-in operand loads with prefetch + per-iteration load near dot/use"
    rationale: |
      On Intel XPU, keeping a large operand tile live across a long loop can reserve many registers,
      increasing GRF pressure and causing spills. A better approach is:
      1) prefetch the data earlier to pull it into L1, then
      2) load the operand into registers only right before it is used (inside the loop),
      reducing liveness and helping the register allocator.

      This is especially relevant when a tensor is loaded outside a loop and used repeatedly
      inside the loop (classic FlashAttention Q tile scenario).
    pattern_before: |
      # Load once, keep live across loop (long liveness)
      q = tl.load(q_ptrs)  # q stays live for the whole loop
      for k0 in range(0, K, BLOCK_K):
          k = tl.load(k_ptrs)
          acc += tl.dot(q, k)   # q used repeatedly
          ...
    pattern_after: |
      # Prefetch outside loop (warm cache)
      # (pseudo; actual prefetch API may differ by backend/compiler lowering)
      tl.prefetch(q_ptrs)  # bring q closer (L1) without allocating regs long-term

      for k0 in range(0, K, BLOCK_K):
          # Load right before use (short liveness)
          q = tl.load(q_ptrs)   # q live only for a short window
          k = tl.load(k_ptrs)
          acc += tl.dot(q, k)
          ...
    expected_speedup: "1.05-1.15x (selective; can regress if cache misses dominate)"
    applies_to:
      - attention
      - flashattention
      - dot_in_loop
      - reduction_in_loop

  - id: reduce_liveness_duplicate_load_for_multi_use
    name: "Reduce liveness when operand has multiple uses (duplicate loads near each use)"
    stage: memory_access
    description: "If a large tensor has multiple distant uses, reload near each use to shorten live ranges"
    rationale: |
      If a large tensor is used in multiple regions far apart (e.g., two loops or two stages),
      keeping it live across both regions can reserve registers too long. On XPU, reloading
      from L1 can be cheaper than spilling registers to memory.
    pattern_before: |
      q = tl.load(q_ptrs)        # q stays live a long time
      for ...:
          acc += tl.dot(q, k)
          ...
      for ...:
          acc2 += tl.dot(q, k2)  # second use far away keeps q live across both loops
          ...
    pattern_after: |
      tl.prefetch(q_ptrs)
      for ...:
          q = tl.load(q_ptrs)    # short liveness
          acc += tl.dot(q, k)
          ...
      for ...:
          q = tl.load(q_ptrs)    # reload for second stage
          acc2 += tl.dot(q, k2)
          ...
    expected_speedup: "0.95-1.10x (only when it avoids spills; otherwise may regress)"
    applies_to:
      - attention
      - multi_stage_kernels
      - dot_in_loop

  - id: avoid_long_lived_large_tiles_across_control_flow
    name: "Avoid long-lived tiles across control-flow splits (two-loop / masked attention patterns)"
    stage: memory_access
    description: "Prefer loading tiles inside each loop when control-flow creates long liveness (off-band/on-band)"
    rationale: |
      Control-flow splits (e.g., two loops for masked/unmasked regions) often extend the
      lifetime of large tiles (Q) across both loops. On XPU, that increases GRF pressure.
      Loading inside each loop reduces the live range and can prevent spills.
    pattern_before: |
      q = tl.load(Q_block_ptr)    # q live across both loops
      for start_n in range(lo, hi, BLOCK_N):
          ... qk += tl.dot(q, k) ...
      for start_n in range(lo, hi, BLOCK_N):
          ... qk += tl.dot(q, k) ...  # second loop extends q lifetime
    pattern_after: |
      tl.prefetch(Q_block_ptr)
      for start_n in range(lo, hi, BLOCK_N):
          q = tl.load(Q_block_ptr)    # load per-iteration or per-loop (shorter liveness)
          ... qk += tl.dot(q, k) ...
      for start_n in range(lo, hi, BLOCK_N):
          q = tl.load(Q_block_ptr)    # reload for second loop
          ... qk += tl.dot(q, k) ...
    expected_speedup: "1.05-1.10x (most likely on masked/causal variants)"
    applies_to:
      - attention
      - flashattention
      - causal_mask

  - id: trade_bandwidth_for_regs_when_spilling
    name: "Prefer extra loads over GRF spills (XPU-first rule)"
    stage: memory_access
    description: "When register pressure is high, allow re-loads (from cache) to reduce spills"
    rationale: |
      When GRF pressure is high, the real cost is often spilling to memory.
      Reloading a tile (especially if it hits in L1) can be cheaper than spilling
      and enables other compiler optimizations (unrolling/scheduling).
    pattern_before: |
      # Keep many intermediates live to avoid re-loads
      a = tl.load(...)
      b = tl.load(...)
      c = tl.load(...)
      # long epilogue chain keeps many values live
      y = f(g(h(a,b,c,...)))
    pattern_after: |
      # Shorten live ranges: recompute/reload where cheap
      tl.prefetch(...)
      a = tl.load(...)   # load near use
      y = f(a)
      # if needed later, reload instead of keeping live
      a2 = tl.load(...)
      z = g(a2)
    expected_speedup: "workload-dependent; key benefit is avoiding catastrophic regressions"
    applies_to:
      - flashattention
      - long_epilogues
      - register_heavy_kernels

  - id: do_not_apply_if_cache_eviction_likely
    name: "Do NOT sink loads if cache eviction is likely (small L1 / high conflict risk)"
    stage: memory_access
    description: "Guard against turning cheap L1 loads into expensive global loads"
    rationale: |
      The sink-load approach assumes the data remains in L1 after prefetch.
      If cache conflicts/evictions are likely, sinking loads into the loop can increase
      global memory traffic and regress performance.
    pattern_before: |
      # Always sink loads
      tl.prefetch(ptrs)
      for ...:
          x = tl.load(ptrs)   # becomes expensive if evicted
          ...
    pattern_after: |
      # Keep original load placement OR stage to a closer explicitly-managed buffer (if available)
      x = tl.load(ptrs)       # load once if it avoids repeated global misses
      for ...:
          use(x)
    expected_speedup: "prevents regressions"
    applies_to:
      - cache_sensitive
      - small_cache_targets
      - very_large_working_sets

  - id: mem_atomic_relaxed_for_accumulation
    name: "Use sem='relaxed' for commutative atomic accumulation"
    stage: memory_access
    description: "Relaxed semantics for atomic_add when ordering doesn't matter"
    rationale: |
      When multiple programs accumulate partial sums into the same output
      (e.g., Stream K partial tiles), the order of additions doesn't matter.
      Using sem='relaxed' avoids unnecessary memory ordering overhead.
    pattern_after: |
      tl.atomic_add(c_ptr_, acc, mask=mask, sem='relaxed')
    expected_speedup: "avoids synchronization overhead vs stricter semantics"
    applies_to: [stream_k, split_accumulation, reduction]

  - id: mem_atomic_fallback_from_descriptors
    name: "Fall back from descriptors/block pointers to manual pointers for atomic stores"
    stage: memory_access
    description: "Use manual pointer arithmetic when atomic operations are needed"
    rationale: |
      Tensor descriptors (desc.store) and block pointer stores do not support
      atomic operations. When partial-tile accumulation requires tl.atomic_add,
      you must compute pointers manually. This is the standard pattern for
      Stream K partial tiles.
    pattern_before: |
      # Cannot do this:
      c_desc.atomic_add(...)   # NOT SUPPORTED
    pattern_after: |
      rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
      rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
      c_ptr_ = c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
      mask = (rm < M)[:, None] & (rn < N)[None, :]
      tl.atomic_add(c_ptr_, acc, mask=mask, sem='relaxed')
    expected_speedup: "N/A (required fallback when atomics are needed)"
    applies_to: [stream_k, partial_tiles, split_accumulation]

# Optimization Levels for Iterative Kernel Development
#
# Use this framework when optimizing a kernel. Don't stop at Level 1 —
# check the "try harder" decision tree to decide whether deeper optimization
# is warranted. Most production kernels should reach at least Level 2.
#
# Case study: kernel #39 (Gemm_Scale_BatchNorm)
#   Level 1 alone: 2.69x speedup (tensor descriptors, tile swizzling, fused epilogue)
#   Level 2 + 3:   5.28x speedup (bf16 pre-pack, BN folded into weights)
#   See: test_kernels/39_Gemm_Scale_BatchNorm_triton.py

levels:
  - id: level_1_baseline_xpu
    name: "Level 1: Baseline XPU Optimizations"
    description: |
      Get the kernel working with standard XPU patterns. This is the minimum
      viable optimization — necessary but rarely sufficient for production.
    checklist:
      - Tensor descriptors (preferred on XPU) or block pointers (not manual pointer arithmetic)
      - bf16/fp16 dot inputs, fp32 accumulator
      - 1D grid with tile swizzling (GROUP_SIZE_M) for GEMM kernels
      - "@triton.autotune with XPU-optimized configs"
      - Pre-packed weight transpose cached in _pack_weights()
      - Fused light epilogues (bias, simple activations)
      - Correct boundary_check=(0, 1) for block pointers (tensor descriptors handle boundaries internally); no default block params
    typical_speedup: "1.5-3x vs PyTorch"
    when_done: "Move to Level 2 — there is almost always more to gain."

  - id: level_2_bandwidth_reduction
    name: "Level 2: Bandwidth Reduction"
    description: |
      Reduce memory bandwidth consumption — the most common bottleneck after
      Level 1. Loading fp32 data and converting in-kernel wastes 2x bandwidth.
    checklist:
      - "Pre-pack weights to bf16 at pack time (not in-kernel conversion)"
      - "Pre-convert inputs to bf16 before kernel launch: x.to(device, torch.bfloat16)"
      - "Use grf_mode='256' for large register file (XPU-specific)"
      - "Use tl.dot(a, b, acc=acc) fused accumulate pattern"
      - "Keep bias in fp32 (small, precision-sensitive)"
    code_example: |
      # BEFORE (Level 1 — wastes 2x bandwidth):
      self.weight_t = w.to(device).t().contiguous()           # fp32 on device
      # kernel: a = tl.load(ptr).to(tl.bfloat16)              # loads 4B, uses 2B

      # AFTER (Level 2 — halves K-loop memory traffic):
      self.weight_t = w.to(device).t().contiguous().to(torch.bfloat16)  # bf16
      # kernel: a = tl.load(ptr)                               # loads 2B directly
    typical_speedup: "2-4x vs PyTorch"
    when_done: |
      Check Level 3 — if the epilogue contains linear transforms (BN, scale,
      affine), they can often be eliminated entirely.

  - id: level_3_algebraic_fusion
    name: "Level 3: Algebraic Fusion"
    description: |
      Eliminate epilogue work by algebraically folding linear transforms into
      GEMM weights at pack time. This is the highest-impact optimization for
      kernels with BatchNorm, per-channel scaling, or affine transforms.

      Key insight: BatchNorm inference is a per-channel linear operation:
        BN(y) = gamma * (y - mean) / sqrt(var + eps) + beta

      When the GEMM output feeds into BN (possibly through scaling), the entire
      chain can be reduced to a single GEMM + bias:
        alpha_n = gamma_n * scale_n / sqrt(var_n + eps)
        W_fused[n,:] = alpha_n * W[n,:]
        b_fused[n]   = alpha_n * bias_n - gamma_n * mean_n / sqrt(var_n + eps) + beta_n
    checklist:
      - Identify linear per-channel transforms in the epilogue
        (BatchNorm inference, scaling, affine, LayerNorm with fixed stats)
      - Derive fused weight and bias formulas algebraically (on paper first)
      - Implement folding in _pack_weights() (one-time cost, not in hot path)
      - Verify numerical equivalence (bf16 folding introduces small diffs)
      - The kernel becomes a pure GEMM + bias — maximum compute efficiency
    applies_when:
      - "BatchNorm in inference mode (running_mean/var are fixed)"
      - "Per-channel scaling (multiply by [N] vector)"
      - "Any per-channel affine transform (y = a*x + b where a,b are [N] vectors)"
    does_not_apply_when:
      - "BatchNorm in training mode (statistics are computed per-batch)"
      - "Non-linear epilogues (sigmoid, ReLU, etc.) — these cannot be folded"
      - "Operations that depend on the M dimension (row-wise reductions)"
    typical_speedup: "3-6x vs PyTorch"
    when_done: |
      For most kernels, this is sufficient. Move to Level 4 only for
      critical-path kernels where profiling shows remaining bottlenecks.
    example:
      file: test_kernels/39_Gemm_Scale_BatchNorm_triton.py
      description: |
        Kernel #39 folds Linear + Scale + BatchNorm into a pure GEMM.
        Before (Level 1): 6 vector loads + arithmetic in epilogue → 2.69x
        After  (Level 3): zero epilogue, pure GEMM + bias → 5.28x

  - id: level_4_expert
    name: "Level 4: Expert Techniques"
    description: |
      Advanced patterns for squeezing the last 10-30% of performance.
      Only pursue these for critical-path kernels after profiling confirms
      the bottleneck. These techniques increase code complexity significantly.
    checklist:
      - Stream K decomposition (for non-square or non-tile-divisible GEMMs)
      - Persistent kernels (fixed program count, iterate over tiles)
      - Hardware capability queries (gpu_subslice_count for program count)
      - Warp size sweeping (16 vs 32 on XPU)
      - Shape-specific specialization (skip boundary checks when M % BLOCK_M == 0)
      - Atomic partial tile accumulation
      - Adaptive grid layouts (small vs large workloads)
    typical_speedup: "5-10x+ vs PyTorch (shape-dependent)"
    when_done: "Benchmark shows no further improvement — try a fundamentally different approach."
    reference_files:
      - kb/persistent_kernel_patterns.yaml
      - kb/examples/stream_k_gemm_descriptors.py


# ============================================================================
# "Try Harder" Decision Tree
# ============================================================================
# Use this after completing Level 1 to decide whether to keep optimizing.

try_harder:
  description: |
    After each optimization level, measure speedup and consult this tree.
    The biggest gains typically come from Level 2 (bandwidth) and Level 3
    (algebraic fusion). Level 4 has diminishing returns for most workloads.

  decisions:
    - condition: "Speedup < 2x after Level 1"
      diagnosis: "Likely bandwidth-bound — loading fp32 and converting wastes 2x BW"
      action: "Apply Level 2: pre-pack to bf16, add grf_mode='large'. Run 'python skills/xpu_profiler.py' to confirm bandwidth bottleneck (look for high XVE Stalled %)."
      priority: high

    - condition: "Speedup 2-3x after Level 2"
      diagnosis: "Epilogue may be adding unnecessary work"
      action: |
        Check: does the epilogue contain only linear per-channel transforms?
        If yes → Apply Level 3: fold into weights algebraically.
        If no (non-linear ops like sigmoid, tanh) → the epilogue is already minimal.
        Consider Level 4 techniques or accept current speedup.
      priority: high

    - condition: "Speedup 3-5x after Level 3"
      diagnosis: "Good performance — likely near hardware limits for this shape"
      action: |
        Run 'python skills/xpu_profiler.py' to identify remaining bottleneck.
        If compute-bound (XVE Active > Stalled) → try larger tiles or Stream K.
        If memory-bound (XVE Stalled > Active) → check for unnecessary copies or layout transforms.
        Keep going — try a different approach if current strategy has plateaued.
      priority: medium

    - condition: "Speedup > 5x"
      diagnosis: "Excellent — diminishing returns ahead"
      action: "Stop unless this kernel is on the critical path. Focus effort elsewhere."
      priority: low

    - condition: "Speedup < 1x (slower than PyTorch)"
      diagnosis: "Something is fundamentally wrong"
      action: |
        Check for: fp64 usage, N-loop serialization, missing tile swizzling,
        very small tiles, or replacing a fast vendor GEMM with a slow custom one.
        Re-read kb/correctness.yaml and kb/fusion_patterns.yaml constraints.
      priority: critical

# Intel XPU Persistent Kernel & Stream K Patterns
# Stage: persistent_kernel, stream_k
#
# Persistent kernels are ADVANCED and should NOT be auto-applied by default.
# They help mainly when:
# - there are many tiles (launch/tail overhead),
# - the workload is bandwidth/reduction bound,
# - cache reuse across tiles is valuable,
# and they do NOT replace a faster vendor primitive (e.g., vendor GEMM).
#
# Stream K is a related advanced scheduling strategy that splits tiles'
# K-loop iterations evenly across programs to fix quantization inefficiency.

patterns:
  - id: persistent_kernel_basic_tile_loop
    name: "Persistent kernel: fixed program count loops over tiles"
    stage: persistent_kernel
    description: "Launch a fixed number of programs and iterate over all tiles with stride = NUM_PROGS"
    rationale: |
      Persistent kernels reduce kernel launch overhead and can improve cache reuse by having
      fewer programs loop over many tiles. This can help Intel XPU when the baseline launches
      huge grids (many tiles) and suffers from tail effects or overhead.

      WARNING:
      - Persistent kernels can increase GRF/register pressure and reduce occupancy.
      - Do NOT use by default. Use only when gated conditions indicate it.
    pattern_before: |
      # Standard per-tile launch
      grid = (tl.cdiv(M, BLOCK_M), tl.cdiv(N, BLOCK_N))

      @triton.jit
      def kernel(...):
          pid_m = tl.program_id(0)
          pid_n = tl.program_id(1)
          # Process one tile
          ...
    pattern_after: |
      # Persistent kernel launch (1D grid)
      # NUM_PROGS is a tuned constant (do NOT assume it equals SM/EU count)
      grid = (NUM_PROGS,)

      @triton.jit
      def kernel(..., NUM_PROGS: tl.constexpr):
          pid = tl.program_id(0)
          tiles_m = tl.cdiv(M, BLOCK_M)
          tiles_n = tl.cdiv(N, BLOCK_N)
          num_tiles = tiles_m * tiles_n

          tile = pid
          while tile < num_tiles:
              pid_m = tile // tiles_n
              pid_n = tile % tiles_n

              row_start = pid_m * BLOCK_M
              col_start = pid_n * BLOCK_N

              # Process tile (GEMM / reduction / etc.)
              ...

              tile += NUM_PROGS
    expected_speedup: "0.8-2.0x (highly shape/workload dependent; can regress)"
    applies_to:
      - gemm
      - attention
      - softmax
      - layernorm
      - large_workloads

  - id: persistent_kernel_autotuned_num_progs
    name: "Persistent kernel: autotune NUM_PROGS (required)"
    stage: persistent_kernel
    description: "Treat NUM_PROGS as a tunable meta-parameter instead of guessing hardware counts"
    rationale: |
      Intel XPU does not map cleanly to CUDA's SM model. Do not set NUM_PROGS = 'num_sm'.
      Instead, autotune NUM_PROGS across a small set (e.g., 32/64/128/256) to find the best.
    pattern_before: |
      NUM_PROGS = gpu_sm_count  # WRONG: not portable / not reliable on XPU
      grid = (NUM_PROGS,)
    pattern_after: |
      @triton.autotune(
          configs=[
              triton.Config({"NUM_PROGS": 32,  "grf_mode": "256"}, num_warps=4, num_stages=2),
              triton.Config({"NUM_PROGS": 64,  "grf_mode": "256"}, num_warps=4, num_stages=2),
              triton.Config({"NUM_PROGS": 128, "grf_mode": "256"}, num_warps=4, num_stages=2),
              triton.Config({"NUM_PROGS": 256, "grf_mode": "256"}, num_warps=4, num_stages=2),
          ],
          key=["M", "N", "K"],
      )
      @triton.jit
      def kernel(..., NUM_PROGS: tl.constexpr):
          ...
      # Launch with 1D grid:
      grid = lambda META: (META["NUM_PROGS"],)
    expected_speedup: "prevents regressions; enables win when persistent is appropriate"
    applies_to:
      - gemm
      - attention
      - large_workloads

  - id: streamk_two_wave_decomposition
    name: "Stream K: two-wave SK + data-parallel decomposition"
    stage: stream_k
    description: "Split tiles into stream-K (wave 1) and standard blocking (wave 2)"
    rationale: |
      Pure stream-K distributes ALL tiles via iteration splitting, maximizing
      atomic contention. A hybrid two-wave approach is more efficient:

      Wave 1 (first_wave): Only "remainder" tiles that cause quantization
        inefficiency. Programs share K-loop iterations, using atomic_add
        for partial results. Program count = num_xe_core (subslice count).

      Wave 2 (full_tiles): Remaining tiles are 1:1 (standard GEMM tiling),
        no atomics needed. Grid size = blocking_tiles.

      This minimizes atomic traffic while fixing utilization gaps.
    pattern_after: |
      num_xe_core = torch.xpu.get_device_capability(0)['gpu_subslice_count']
      streamk_programs = num_xe_core

      total_tiles = num_block_m * num_block_n
      iters_per_tile = triton.cdiv(K, BLOCK_SIZE_K)

      # Two-tile SK + DP heuristic
      streamk_tiles = total_tiles % streamk_programs
      if total_tiles - streamk_tiles > streamk_programs:
          streamk_tiles += streamk_programs

      blocking_tiles = total_tiles - streamk_tiles
      streamk_iters = streamk_tiles * iters_per_tile

      streamk_full_tiles = streamk_iters // streamk_programs
      streamk_partial_tiles = streamk_iters % streamk_programs

      # Wave 1: stream-K (fixed program count, iteration-split)
      first_wave[(streamk_programs,)](a, b, c, ...,
          streamk_full_tiles, streamk_partial_tiles, iters_per_tile)

      # Wave 2: standard 1:1 tiling (offset by streamk_tiles)
      full_tiles[(blocking_tiles,)](a, b, c, ..., streamk_tiles)
    expected_speedup: "significant for non-divisible tile counts; minimal for well-fitting shapes"
    applies_to: [gemm, large_workloads, variable_shapes]
    notes: |
      - Requires output to be pre-zeroed (torch.zeros) — see correctness.yaml
      - Wave 1 uses atomic_add for partial tiles only; full tiles use direct store
      - Wave 2 tile_ids offset by streamk_tiles to avoid overlap with wave 1

  - id: streamk_mac_loop_partial_tile
    name: "Stream K: MAC loop with arbitrary iteration range"
    stage: stream_k
    description: "Inner multiply-accumulate loop that handles [start_iter, end_iter) within a tile"
    rationale: |
      In stream-K, a program may own only a subset of a tile's K iterations.
      The MAC loop must:
      1. Derive tile_id from start_iter (tile_id = start_iter // iters_per_tile)
      2. Compute K offset (remain_iters * BLOCK_SIZE_K)
      3. Accumulate over the assigned range
      4. Store via descriptor (full tile) or atomic_add (partial tile)
    pattern_after: |
      tile_id = start_iter // iters_per_tile
      remain_iters = start_iter % iters_per_tile
      pid_m, pid_n = swizzle_tile(tile_id, ...)

      acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
      off_k = remain_iters * BLOCK_SIZE_K

      for _ in range(start_iter, end_iter):
          a = a_desc.load([pid_m * BLOCK_SIZE_M, off_k])
          b = b_desc.load([off_k, pid_n * BLOCK_SIZE_N])
          acc += tl.dot(a, b)
          off_k += BLOCK_SIZE_K

      # Full tile → direct store; partial tile → atomic add
      if remain_iters == 0 and end_iter % iters_per_tile == 0:
          c_desc.store([pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N], acc)
      else:
          rm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
          rn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
          c_ptr_ = c_ptr + rm[:, None] * stride_cm + rn[None, :] * stride_cn
          mask = (rm < M)[:, None] & (rn < N)[None, :]
          tl.atomic_add(c_ptr_, acc, mask=mask, sem='relaxed')
    applies_to: [gemm, stream_k]

  - id: streamk_iteration_distribution
    name: "Stream K: even iteration distribution (floor + remainder)"
    stage: stream_k
    description: "Distribute total stream-K iterations across programs with even split"
    rationale: |
      Each program gets floor(total_iters / num_programs) iterations, plus
      one extra if pid < remainder. Programs iterate through their range,
      snapping to tile boundaries so each mac_loop stays within one tile.
    pattern_after: |
      pid = tl.program_id(axis=0)
      start_iter = pid * full_tiles + tl.minimum(pid, partial_tiles)
      last_iter = (pid + 1) * full_tiles + tl.minimum(pid + 1, partial_tiles)

      while start_iter < last_iter:
          end_iter = start_iter + (iters_per_tile - start_iter % iters_per_tile)
          end_iter = tl.minimum(end_iter, last_iter)
          mac_loop(..., start_iter, end_iter, ...)
          start_iter = end_iter
    applies_to: [stream_k]
    notes: |
      - The while loop handles programs that span multiple tiles
      - end_iter snaps to tile boundaries within the assigned range

#!/usr/bin/env python3
"""
Analyze PyTorch kernel to extract structure and guide Triton optimization.

Usage:
    python scripts/analyze_kernel.py <pytorch_file>
"""

import ast
import re
import sys
from pathlib import Path
from typing import Dict, List, Set, Tuple


class KernelAnalyzer(ast.NodeVisitor):
    """AST visitor to analyze PyTorch model operations."""

    def __init__(self):
        self.operations = []
        self.shapes = {}
        self.dtypes = set()
        self.has_matmul = False
        self.has_linear = False
        self.activations = []
        self.reductions = []
        self.elementwise = []

    def visit_Call(self, node):
        """Visit function calls to identify operations."""
        # torch.matmul
        if isinstance(node.func, ast.Attribute):
            if hasattr(node.func.value, "id") and node.func.value.id == "torch":
                op_name = node.func.attr
                self.operations.append(op_name)

                if op_name == "matmul":
                    self.has_matmul = True
                elif op_name in ["sum", "mean", "max", "min"]:
                    self.reductions.append(op_name)
                elif op_name in ["sigmoid", "tanh", "relu", "gelu", "silu"]:
                    self.activations.append(op_name)
                elif op_name == "clamp":
                    self.elementwise.append("clamp")

            # torch.nn.functional
            elif hasattr(node.func.value, "attr"):
                if node.func.value.attr == "functional":
                    op_name = node.func.attr
                    self.operations.append(f"F.{op_name}")
                    if op_name in ["gelu", "relu", "silu", "softmax", "sigmoid"]:
                        self.activations.append(op_name)

        self.generic_visit(node)

    def visit_BinOp(self, node):
        """Visit binary operations (*, /, +, -)."""
        op_map = {
            ast.Mult: "multiply",
            ast.Div: "divide",
            ast.Add: "add",
            ast.Sub: "subtract",
        }
        op_type = type(node.op)
        if op_type in op_map:
            self.elementwise.append(op_map[op_type])
        self.generic_visit(node)

    def visit_Assign(self, node):
        """Visit assignments to track nn.Linear."""
        if isinstance(node.value, ast.Call):
            if hasattr(node.value.func, "attr") and node.value.func.attr == "Linear":
                self.has_linear = True
        self.generic_visit(node)


def analyze_pytorch_kernel(filepath: Path) -> Dict:
    """Analyze PyTorch kernel file and extract optimization hints."""

    with open(filepath, "r") as f:
        source = f.read()

    tree = ast.parse(source)
    analyzer = KernelAnalyzer()
    analyzer.visit(tree)

    # Extract shape information from module-level variables
    shapes = {}
    for line in source.split("\n"):
        if "=" in line and any(
            dim in line
            for dim in ["batch_size", "in_features", "out_features", "hidden_size", "input_size"]
        ):
            match = re.match(r"(\w+)\s*=\s*(\d+)", line.strip())
            if match:
                shapes[match.group(1)] = int(match.group(2))

    # Determine kernel type
    kernel_type = "unknown"
    if analyzer.has_matmul or analyzer.has_linear:
        if analyzer.activations or analyzer.elementwise:
            kernel_type = "gemm_epilogue"
        elif analyzer.reductions:
            kernel_type = "gemm_reduction"
        else:
            kernel_type = "gemm"
    elif analyzer.reductions:
        kernel_type = "reduction"
    elif analyzer.elementwise:
        kernel_type = "elementwise"

    # Fusion analysis
    fusion_opportunities = []
    if analyzer.has_matmul or analyzer.has_linear:
        if len(analyzer.activations) <= 2 and len(analyzer.elementwise) <= 3:
            fusion_opportunities.append("Light epilogue fusion (GEMM + activation + elementwise)")
        else:
            fusion_opportunities.append("Heavy epilogue - consider partial fusion or split")

        if analyzer.reductions:
            fusion_opportunities.append(
                "WARNING: GEMM + reduction - use 2D GEMM then separate reduction kernel"
            )

    # Memory pattern recommendation
    memory_pattern = "block_pointers"  # default
    if "Stream K" in str(analyzer.operations) or len(analyzer.reductions) > 1:
        memory_pattern = "tensor_descriptors"

    return {
        "kernel_type": kernel_type,
        "operations": analyzer.operations,
        "activations": analyzer.activations,
        "reductions": analyzer.reductions,
        "elementwise": analyzer.elementwise,
        "shapes": shapes,
        "fusion_opportunities": fusion_opportunities,
        "memory_pattern": memory_pattern,
        "has_gemm": analyzer.has_matmul or analyzer.has_linear,
    }


def print_analysis(analysis: Dict, filepath: Path):
    """Pretty print the analysis results."""

    print(f"\n{'=' * 70}")
    print(f"Analysis: {filepath.name}")
    print(f"{'=' * 70}\n")

    print(f"Kernel Type: {analysis['kernel_type'].upper()}")
    print(f"Memory Pattern: {analysis['memory_pattern']}")
    print()

    if analysis["shapes"]:
        print("Shapes:")
        for key, val in analysis["shapes"].items():
            print(f"  {key}: {val}")
        print()

    print("Operations:")
    print(f"  Total: {len(analysis['operations'])}")
    if analysis["has_gemm"]:
        print(f"  ✓ GEMM/Linear")
    if analysis["activations"]:
        print(f"  ✓ Activations: {', '.join(set(analysis['activations']))}")
    if analysis["reductions"]:
        print(f"  ✓ Reductions: {', '.join(set(analysis['reductions']))}")
    if analysis["elementwise"]:
        print(f"  ✓ Elementwise: {', '.join(set(analysis['elementwise']))}")
    print()

    if analysis["fusion_opportunities"]:
        print("Fusion Opportunities:")
        for opp in analysis["fusion_opportunities"]:
            if "WARNING" in opp:
                print(f"  ⚠️  {opp}")
            else:
                print(f"  → {opp}")
        print()

    # Recommendations
    print("Recommended Optimizations:")

    if analysis["has_gemm"]:
        print("  1. Use tensor descriptors (preferred on XPU) or block pointers")
        print("  2. Apply tile swizzling (GROUP_SIZE_M)")

        # Tile size recommendations based on shape
        batch_size = analysis["shapes"].get("batch_size", 0)
        if batch_size and batch_size < 256:
            print("  3. Use smaller BLOCK_M (32-64) for skinny M")
        else:
            print("  3. Try large tiles (256x256) with autotune")

        print("  4. Autotune: num_warps={4,8,16,32}, grf_mode='256'")
        print("  5. Mixed precision: bf16/fp16 inputs, fp32 accumulator")
        print("  6. Pre-pack weight transpose: weight_t = weight.t().contiguous()")

    if "sigmoid" in analysis["activations"]:
        print("  → Use exp2-based sigmoid (faster on XPU)")

    if "tanh" in analysis["activations"]:
        print("  → Implement tanh via sigmoid: tanh(x) = 2*sigmoid(2x) - 1")

    if "gelu" in analysis["activations"]:
        print("  → Use tanh-approximation GeLU with JIT helper")

    if analysis["reductions"]:
        if analysis["has_gemm"]:
            print("  ⚠️  Split GEMM and reduction into separate kernels")
            print("     (Don't serialize over N tiles inside one program)")
        else:
            print("  → Use multi-row tiling for reductions")
            print("  → Query max_work_group_size for BLOCK_SIZE_Y")

    print()

    print("Relevant Reference Files:")
    print("  • references/xpu_optimizations.yaml - Core XPU patterns")
    if analysis["fusion_opportunities"]:
        print("  • references/fusion_patterns.yaml - Fusion guidelines")
    print("  • references/memory_patterns.yaml - Memory access patterns")
    print("  • references/correctness.yaml - Critical constraints")
    print()

    # Template suggestion
    if analysis["kernel_type"] == "gemm":
        print("Suggested Template: See GEMM pattern in references/implementation_reference.md")
    elif analysis["kernel_type"] == "gemm_epilogue":
        print("Suggested Template: See GEMM with epilogue pattern in references/implementation_reference.md")
    elif analysis["kernel_type"] in ["reduction", "gemm_reduction"]:
        print("Suggested Template: See reduction pattern in references/implementation_reference.md")
    print()


def main():
    if len(sys.argv) != 2:
        print("Usage: python scripts/analyze_kernel.py <pytorch_file>")
        sys.exit(1)

    filepath = Path(sys.argv[1])
    if not filepath.exists():
        print(f"Error: File not found: {filepath}")
        sys.exit(1)

    analysis = analyze_pytorch_kernel(filepath)
    print_analysis(analysis, filepath)


if __name__ == "__main__":
    main()

"""Shared config loader — reads config.yaml from project root."""

from pathlib import Path

import yaml

_CONFIG_DIR = Path(__file__).resolve().parent
_DEFAULTS = {
    "max_trials": 10,
    "vtune_enabled": True,
    "vtune_bin": "/bin64/vtune",
}


def load_config() -> dict:
    """Load config.yaml, falling back to defaults for missing keys."""
    config_path = _CONFIG_DIR / "config.yaml"
    cfg = {}
    if config_path.exists():
        with open(config_path) as f:
            cfg = yaml.safe_load(f) or {}
    return {**_DEFAULTS, **cfg}

# Project configuration — edit these values to control optimization sessions.
max_trials: 10             # Maximum number of optimization trials (3-20)
vtune_enabled: true        # Set to false to skip VTune profiling entirely
vtune_bin: "/bin64/vtune"  # Path to VTune binary

transformers
safetensors
huggingface-hub
kernels
ai-bench[xpu] @ git+https://github.com/libxsmm/AI-bench.git