gemm_tiling_solver.py

"""
This module provides a hierarchical tiler for GEMM:
    C[M, N] = A[M, K] @ B[K, N]

and is structured so the same solver can be reused for other kernels.

---------------------------------------------------------------------------
Public API
---------------------------------------------------------------------------
  - MemoryBudget
  - GemmShape
  - GemmTile
  - get_gemm_tiling(shape, budgets, dtype_bytes=4, double_buffering=True)
        -> GemmTile

---------------------------------------------------------------------------
Minimal generic layer
---------------------------------------------------------------------------
  - TilingProblem: finite-domain variables + model builder + lexicographic objectives
  - TilingResult:  solution values + objective values
  - solve_tiling(problem, budgets, dtype_bytes=..., double_buffering=...)
        -> TilingResult

`get_gemm_tiling(...)` is a wrapper that:
  1) builds a GEMM-specific TilingProblem (hierarchical L1/L2/L3)
  2) calls solve_tiling(...)
  3) maps the result back to GemmTile

---------------------------------------------------------------------------
GEMM tiling model (hierarchical)
---------------------------------------------------------------------------
Decision variables (each restricted to divisors of M/N/K by default):
  L1: TM, TN, TK
  L2: TM_L2, TN_L2, TK_L2
  L3: TM_L3, TN_L3, TK_L3

Hierarchy constraints:
  • Non-decreasing across levels:
      TM <= TM_L2 <= TM_L3   (and similarly for TN, TK)
  • Divisibility:
      TM_L2 % TM == 0,  TM_L3 % TM_L2 == 0
      TN_L2 % TN == 0,  TN_L3 % TN_L2 == 0
      TK_L2 % TK == 0,  TK_L3 % TK_L2 == 0

---------------------------------------------------------------------------
Memory capacity model (bytes)
---------------------------------------------------------------------------
The memory model depends on the `double_buffering` flag.

When double_buffering=True (default):

  L1 (double-buffered A, B, and C partial sums):
      L1 = 2*(TM*TK + TK*TN + TM*TN)*dtype_bytes + L1_overhead

  L2 (double-buffered A and B tiles):
      L2 = 2*(TM_L2*TK_L2 + TK_L2*TN_L2)*dtype_bytes + L2_overhead

  L3 (single-buffer A, B, and C tiles):
      L3 = (TM_L3*TK_L3 + TK_L3*TN_L3 + TM_L3*TN_L3)*dtype_bytes
           + L3_overhead

When double_buffering=False:

  L1:
      L1 = (TM*TK + TK*TN + TM*TN)*dtype_bytes + L1_overhead

  L2:
      L2 = (TM_L2*TK_L2 + TK_L2*TN_L2)*dtype_bytes + L2_overhead

  L3:
      L3 = (TM_L3*TK_L3 + TK_L3*TN_L3 + TM_L3*TN_L3)*dtype_bytes
           + L3_overhead

If a budget level is None, that constraint is disabled.

---------------------------------------------------------------------------
Objective (lexicographic order)
---------------------------------------------------------------------------
  1) maximize vol_L1 = TM * TN * TK
  2) maximize vol_L2 = TM_L2 * TN_L2 * TK_L2
  3) maximize vol_L3 = TM_L3 * TN_L3 * TK_L3
  4) maximize area_L1 = TM * TN
  5) maximize TK

---------------------------------------------------------------------------
Dependencies
---------------------------------------------------------------------------
    pip install ortools
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, List, Tuple, Dict, Callable, Any
import math
from generic_tiling_solver import *

# ===========
# Public API 
# ===========

@dataclass
class GemmShape:
    M: int
    N: int
    K: int

@dataclass
class GemmTile:
    # Dimensions of L1 buffers 
    TM: int
    TN: int
    TK: int

    # Optional hierarchical fields
    TM_L2: Optional[int] = None
    TN_L2: Optional[int] = None
    TK_L2: Optional[int] = None

    TM_L3: Optional[int] = None
    TN_L3: Optional[int] = None
    TK_L3: Optional[int] = None


# ============================
# GEMM adapter (hierarchical)
# ============================

def _build_gemm_problem_hierarchical(
    shape: GemmShape,
    budgets: MemoryBudget,
    *,
    double_buffering: bool,
    l1_overhead: int,
    l2_overhead: int,
    l3_overhead: int,
    l1_util: float,
    l2_util: float,
    l3_util: float,
) -> TilingProblem:
    M, N, K = int(shape.M), int(shape.N), int(shape.K)

    TM_vals = divisors(M)
    TN_vals = divisors(N)
    TK_vals = divisors(K)

    # Effective caps (utilization headroom)
    L1_cap = None if budgets.L1_bytes is None else int(budgets.L1_bytes * l1_util)
    L2_cap = None if budgets.L2_bytes is None else int(budgets.L2_bytes * l2_util)
    L3_cap = None if budgets.L3_bytes is None else int(budgets.L3_bytes * l3_util)

    var_domains = {
        # L1
        "TM": TM_vals, "TN": TN_vals, "TK": TK_vals,
        # L2
        "TM_L2": TM_vals, "TN_L2": TN_vals, "TK_L2": TK_vals,
        # L3
        "TM_L3": TM_vals, "TN_L3": TN_vals, "TK_L3": TK_vals,
    }

    def build_model(model, v, dtype_b: int, b: MemoryBudget) -> Dict[str, Any]:
        # --- Hierarchy constraints ---
        model.Add(v["TM"] <= v["TM_L2"])
        model.Add(v["TN"] <= v["TN_L2"])
        model.Add(v["TK"] <= v["TK_L2"])

        model.Add(v["TM_L2"] <= v["TM_L3"])
        model.Add(v["TN_L2"] <= v["TN_L3"])
        model.Add(v["TK_L2"] <= v["TK_L3"])

        # Divisibility
        model.AddModuloEquality(0, v["TM_L2"], v["TM"])
        model.AddModuloEquality(0, v["TN_L2"], v["TN"])
        model.AddModuloEquality(0, v["TK_L2"], v["TK"])

        model.AddModuloEquality(0, v["TM_L3"], v["TM_L2"])
        model.AddModuloEquality(0, v["TN_L3"], v["TN_L2"])
        model.AddModuloEquality(0, v["TK_L3"], v["TK_L2"])

        # --- Products for L1 ---
        TM_TK_L1 = model.NewIntVar(0, M * K, "TM_TK_L1")
        TK_TN_L1 = model.NewIntVar(0, K * N, "TK_TN_L1")
        TM_TN_L1 = model.NewIntVar(0, M * N, "TM_TN_L1")
        model.AddMultiplicationEquality(TM_TK_L1, [v["TM"], v["TK"]])
        model.AddMultiplicationEquality(TK_TN_L1, [v["TK"], v["TN"]])
        model.AddMultiplicationEquality(TM_TN_L1, [v["TM"], v["TN"]])

        # --- Products for L2 ---
        TM_TK_L2 = model.NewIntVar(0, M * K, "TM_TK_L2_prod")
        TK_TN_L2 = model.NewIntVar(0, K * N, "TK_TN_L2_prod")
        TM_TN_L2 = model.NewIntVar(0, M * N, "TM_TN_L2_prod")
        model.AddMultiplicationEquality(TM_TK_L2, [v["TM_L2"], v["TK_L2"]])
        model.AddMultiplicationEquality(TK_TN_L2, [v["TK_L2"], v["TN_L2"]])
        model.AddMultiplicationEquality(TM_TN_L2, [v["TM_L2"], v["TN_L2"]])

        # --- Products for L3 ---
        TM_TK_L3 = model.NewIntVar(0, M * K, "TM_TK_L3_prod")
        TK_TN_L3 = model.NewIntVar(0, K * N, "TK_TN_L3_prod")
        TM_TN_L3 = model.NewIntVar(0, M * N, "TM_TN_L3_prod")
        model.AddMultiplicationEquality(TM_TK_L3, [v["TM_L3"], v["TK_L3"]])
        model.AddMultiplicationEquality(TK_TN_L3, [v["TK_L3"], v["TN_L3"]])
        model.AddMultiplicationEquality(TM_TN_L3, [v["TM_L3"], v["TN_L3"]])

        # --- Memory constraints ---
        if L1_cap is not None:
            l1_bytes = model.NewIntVar(0, L1_cap, "l1_bytes")
            if double_buffering:
                model.Add(l1_bytes == (2 * (TM_TK_L1 + TK_TN_L1 + TM_TN_L1) * dtype_b + l1_overhead))
            else:
                model.Add(l1_bytes == ((TM_TK_L1 + TK_TN_L1 + TM_TN_L1) * dtype_b + l1_overhead))
            model.Add(l1_bytes <= L1_cap)

        if L2_cap is not None:
            l2_bytes = model.NewIntVar(0, L2_cap, "l2_bytes")
            if double_buffering:
                model.Add(l2_bytes == (((2 * (TM_TK_L2 + TK_TN_L2))) * dtype_b + l2_overhead))
            else:
                model.Add(l2_bytes == ((TM_TK_L2 + TK_TN_L2) * dtype_b + l2_overhead))
            model.Add(l2_bytes <= L2_cap)

        if L3_cap is not None:
            l3_bytes = model.NewIntVar(0, L3_cap, "l3_bytes")
            model.Add(l3_bytes == ((TM_TK_L3 + TK_TN_L3 + TM_TN_L3) * dtype_b + l3_overhead))
            model.Add(l3_bytes <= L3_cap)

        # --- Objective vars ---
        vol_L1 = model.NewIntVar(0, M * N * K, "vol_L1")
        model.AddMultiplicationEquality(vol_L1, [TM_TN_L1, v["TK"]])

        vol_L2 = model.NewIntVar(0, M * N * K, "vol_L2")
        model.AddMultiplicationEquality(vol_L2, [TM_TN_L2, v["TK_L2"]])

        vol_L3 = model.NewIntVar(0, M * N * K, "vol_L3")
        model.AddMultiplicationEquality(vol_L3, [TM_TN_L3, v["TK_L3"]])

        return {
            "area_L1": TM_TN_L1,
            "vol_L1": vol_L1,
            "vol_L2": vol_L2,
            "vol_L3": vol_L3,
        }

    objectives = ["vol_L1", "vol_L2", "vol_L3", "area_L1", "TK"]
    return TilingProblem(var_domains=var_domains, build_model=build_model, objectives=objectives)


# ---------------- API entry point ----------------

def get_gemm_tiling(shape: GemmShape, budgets: MemoryBudget, dtype_bytes: int = 4, *, double_buffering: bool = True) -> GemmTile:
    """
    Compute hierarchical tile sizes for GEMM under memory constraints.

    Returned fields:
      - TM,TN,TK are the L1 tile
      - TM_L2,TN_L2,TK_L2 are the L2 tile
      - TM_L3,TN_L3,TK_L3 are the L3 tile

    The hierarchy is consistent (non-decreasing and divisible across levels).
    """
    if dtype_bytes <= 0:
        raise ValueError(f"dtype_bytes must be > 0, got {dtype_bytes}")

    problem = _build_gemm_problem_hierarchical(
        shape,
        budgets,
        double_buffering=bool(double_buffering),
        l1_overhead=DEFAULT_L1_OVERHEAD,
        l2_overhead=DEFAULT_L2_OVERHEAD,
        l3_overhead=DEFAULT_L3_OVERHEAD,
        l1_util=DEFAULT_L1_UTIL,
        l2_util=DEFAULT_L2_UTIL,
        l3_util=DEFAULT_L3_UTIL,
    )
    res = solve_tiling(problem, budgets, dtype_bytes=int(dtype_bytes),
                       max_time_sec=DEFAULT_MAX_TIME_SEC, num_workers=DEFAULT_NUM_WORKERS)

    if res is not None:
        v = res.values
        return GemmTile(
            TM=v["TM"], TN=v["TN"], TK=v["TK"],
            TM_L2=v["TM_L2"], TN_L2=v["TN_L2"], TK_L2=v["TK_L2"],
            TM_L3=v["TM_L3"], TN_L3=v["TN_L3"], TK_L3=v["TK_L3"],
        )

    # Retry with full budgets (no utilization headroom)
    problem2 = _build_gemm_problem_hierarchical(
        shape,
        budgets,
        double_buffering=bool(double_buffering),
        l1_overhead=DEFAULT_L1_OVERHEAD,
        l2_overhead=DEFAULT_L2_OVERHEAD,
        l3_overhead=DEFAULT_L3_OVERHEAD,
        l1_util=1.0,
        l2_util=1.0,
        l3_util=1.0,
    )
    res2 = solve_tiling(problem2, budgets, dtype_bytes=int(dtype_bytes),
                        max_time_sec=max(1.0, DEFAULT_MAX_TIME_SEC), num_workers=DEFAULT_NUM_WORKERS)
    if res2 is not None:
        v = res2.values
        return GemmTile(
            TM=v["TM"], TN=v["TN"], TK=v["TK"],
            TM_L2=v["TM_L2"], TN_L2=v["TN_L2"], TK_L2=v["TK_L2"],
            TM_L3=v["TM_L3"], TN_L3=v["TN_L3"], TK_L3=v["TK_L3"],
        )

    # Conservative fallback (if infeasible)
    return GemmTile(1, 1, 1, 1, 1, 1, 1, 1, 1)