conv2d_tiling_solver.py

"""
This module provides a hierarchical tiler for 2D convolution (NHWC layout):
    Y[Hout, Wout, Cout] = conv2d(X[H, W, Cin], W[KH, KW, Cin, Cout])

The tiler mirrors the same approach used in the GEMM tiler:
  - Hierarchical tiles across L1/L2/L3
  - Capacity constraints per level
  - Optional double buffering (default True)

---------------------------------------------------------------------------
Public API
---------------------------------------------------------------------------
  - MemoryBudget
  - Conv2DShape
  - Conv2DTile
  - get_conv2d_tiling(shape, budgets, dtype_bytes=4, double_buffering=True) -> Conv2DTile

---------------------------------------------------------------------------
Conv2D tiling model (hierarchical)
---------------------------------------------------------------------------
We tile the output space and the reduction (input channels) dimension.

Decision variables (each restricted to divisors of Hout/Wout/Cin/Cout by default):
  L1: TH, TW, TIC, TOC
  L2: TH_L2, TW_L2, TIC_L2, TOC_L2
  L3: TH_L3, TW_L3, TIC_L3, TOC_L3

Hierarchy constraints:
  • Non-decreasing across levels for each dimension
  • Divisibility across levels (e.g., TH_L2 % TH == 0, TH_L3 % TH_L2 == 0, ...)

Memory capacity model (elements, then bytes):
Let:
  OH, OW be output tile sizes
  IC be input-channel tile size
  OC be output-channel tile size
  KH, KW be kernel sizes
  stride_h, stride_w be strides (padding handled in Hout/Wout).

Input tile spatial footprint needed for an output tile (valid convolution model):
  IH = (OH - 1) * stride_h + KH
  IW = (OW - 1) * stride_w + KW
Input elements  = IH * IW * IC
Weight elements = KH * KW * IC * OC
Output elements = OH * OW * OC

When double_buffering=True (default):
  L1 (double-buffer X, W, and Y partial sums):
      L1 = 2*(X_L1 + W_L1 + Y_L1)*dtype_bytes + L1_overhead
  L2 (double-buffer X and W tiles; Y not stored in L2 by default):
      L2 = 2*(X_L2 + W_L2)*dtype_bytes + L2_overhead
  L3 (single-buffer X, W, and Y tiles):
      L3 = (X_L3 + W_L3 + Y_L3)*dtype_bytes + L3_overhead

When double_buffering=False:
  same expressions without factor 2.

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

Objective (lexicographic order):
  1) maximize work_L1 ~ TH*TW*TIC*TOC
  2) maximize work_L2 ~ TH_L2*TW_L2*TIC_L2*TOC_L2
  3) maximize work_L3 ~ TH_L3*TW_L3*TIC_L3*TOC_L3
  4) maximize out_area_L1 = TH*TW
  5) maximize TOC

---------------------------------------------------------------------------
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 Conv2DShape:
    # NHWC input, HWKCC output (but output is Hout x Wout x Cout)
    H: int
    W: int
    Cin: int
    Cout: int
    KH: int
    KW: int
    stride_h: int = 1
    stride_w: int = 1
    # Hout/Wout can be passed explicitly, otherwise derived as "valid" conv
    Hout: Optional[int] = None
    Wout: Optional[int] = None


@dataclass
class Conv2DTile:
    # L1 tile
    TH: int
    TW: int
    TIC: int
    TOC: int

    # L2 tile
    TH_L2: Optional[int] = None
    TW_L2: Optional[int] = None
    TIC_L2: Optional[int] = None
    TOC_L2: Optional[int] = None

    # L3 tile
    TH_L3: Optional[int] = None
    TW_L3: Optional[int] = None
    TIC_L3: Optional[int] = None
    TOC_L3: Optional[int] = None


# ==============================
# Conv2D adapter (hierarchical)
# ==============================

def _build_conv2d_problem_hierarchical(
    shape: Conv2DShape,
    budgets: MemoryBudget,
    *,
    double_buffering: bool,
    l1_overhead: int,
    l2_overhead: int,
    l3_overhead: int,
    l1_util: float,
    l2_util: float,
    l3_util: float,
) -> TilingProblem:
    H, W = int(shape.H), int(shape.W)
    Cin, Cout = int(shape.Cin), int(shape.Cout)
    KH, KW = int(shape.KH), int(shape.KW)
    sh, sw = int(shape.stride_h), int(shape.stride_w)

    if sh <= 0 or sw <= 0:
        raise ValueError("stride_h/stride_w must be positive")

    # Default to "valid" convolution output size if not provided.
    Hout = int(shape.Hout) if shape.Hout is not None else (H - KH) // sh + 1
    Wout = int(shape.Wout) if shape.Wout is not None else (W - KW) // sw + 1
    if Hout <= 0 or Wout <= 0:
        raise ValueError(f"Derived Hout/Wout invalid: Hout={Hout}, Wout={Wout} for H={H},W={W},KH={KH},KW={KW},stride=({sh},{sw})")

    # 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)

    TH_vals = divisors(Hout)
    TW_vals = divisors(Wout)
    TIC_vals = divisors(Cin)
    TOC_vals = divisors(Cout)

    var_domains = {
        # L1
        "TH": TH_vals, "TW": TW_vals, "TIC": TIC_vals, "TOC": TOC_vals,
        # L2
        "TH_L2": TH_vals, "TW_L2": TW_vals, "TIC_L2": TIC_vals, "TOC_L2": TOC_vals,
        # L3
        "TH_L3": TH_vals, "TW_L3": TW_vals, "TIC_L3": TIC_vals, "TOC_L3": TOC_vals,
    }

    # Upper bounds for products (safe but not too huge)
    max_in_l1 = (Hout - 1) * sh + KH
    max_in_w1 = (Wout - 1) * sw + KW

    def build_model(model, v, dtype_b: int, b: MemoryBudget) -> Dict[str, Any]:
        # --- Hierarchy constraints ---
        for dim in ("TH", "TW", "TIC", "TOC"):
            model.Add(v[dim] <= v[f"{dim}_L2"])
            model.Add(v[f"{dim}_L2"] <= v[f"{dim}_L3"])
            model.AddModuloEquality(0, v[f"{dim}_L2"], v[dim])
            model.AddModuloEquality(0, v[f"{dim}_L3"], v[f"{dim}_L2"])

        # --- Derived input spatial footprints per level ---
        # IH = (OH - 1)*stride + KH; IW = (OW - 1)*stride + KW
        # Since OH/OW are vars, use linear form: IH = sh*OH - sh + KH
        # All terms are integer.
        IH_L1 = model.NewIntVar(1, max_in_l1, "IH_L1")
        IW_L1 = model.NewIntVar(1, max_in_w1, "IW_L1")
        model.Add(IH_L1 == sh * v["TH"] - sh + KH)
        model.Add(IW_L1 == sw * v["TW"] - sw + KW)

        IH_L2 = model.NewIntVar(1, max_in_l1, "IH_L2")
        IW_L2 = model.NewIntVar(1, max_in_w1, "IW_L2")
        model.Add(IH_L2 == sh * v["TH_L2"] - sh + KH)
        model.Add(IW_L2 == sw * v["TW_L2"] - sw + KW)

        IH_L3 = model.NewIntVar(1, max_in_l1, "IH_L3")
        IW_L3 = model.NewIntVar(1, max_in_w1, "IW_L3")
        model.Add(IH_L3 == sh * v["TH_L3"] - sh + KH)
        model.Add(IW_L3 == sw * v["TW_L3"] - sw + KW)

        # --- Element counts per level ---
        # Input elements: IH * IW * IC
        # Weight elements: KH * KW * IC * OC
        # Output elements: OH * OW * OC

        # Helper products for L1
        IH_IW_L1 = model.NewIntVar(0, max_in_l1 * max_in_w1, "IH_IW_L1")
        model.AddMultiplicationEquality(IH_IW_L1, [IH_L1, IW_L1])

        X_L1 = model.NewIntVar(0, max_in_l1 * max_in_w1 * Cin, "X_L1_elems")
        model.AddMultiplicationEquality(X_L1, [IH_IW_L1, v["TIC"]])

        OH_OW_L1 = model.NewIntVar(0, Hout * Wout, "OH_OW_L1")
        model.AddMultiplicationEquality(OH_OW_L1, [v["TH"], v["TW"]])

        Y_L1 = model.NewIntVar(0, Hout * Wout * Cout, "Y_L1_elems")
        model.AddMultiplicationEquality(Y_L1, [OH_OW_L1, v["TOC"]])

        IC_OC_L1 = model.NewIntVar(0, Cin * Cout, "IC_OC_L1")
        model.AddMultiplicationEquality(IC_OC_L1, [v["TIC"], v["TOC"]])

        W_L1 = model.NewIntVar(0, KH * KW * Cin * Cout, "W_L1_elems")
        model.Add(W_L1 == (KH * KW) * IC_OC_L1)

        # Helper products for L2
        IH_IW_L2 = model.NewIntVar(0, max_in_l1 * max_in_w1, "IH_IW_L2")
        model.AddMultiplicationEquality(IH_IW_L2, [IH_L2, IW_L2])

        X_L2 = model.NewIntVar(0, max_in_l1 * max_in_w1 * Cin, "X_L2_elems")
        model.AddMultiplicationEquality(X_L2, [IH_IW_L2, v["TIC_L2"]])

        OH_OW_L2 = model.NewIntVar(0, Hout * Wout, "OH_OW_L2")
        model.AddMultiplicationEquality(OH_OW_L2, [v["TH_L2"], v["TW_L2"]])

        Y_L2 = model.NewIntVar(0, Hout * Wout * Cout, "Y_L2_elems")
        model.AddMultiplicationEquality(Y_L2, [OH_OW_L2, v["TOC_L2"]])

        IC_OC_L2 = model.NewIntVar(0, Cin * Cout, "IC_OC_L2")
        model.AddMultiplicationEquality(IC_OC_L2, [v["TIC_L2"], v["TOC_L2"]])

        W_L2 = model.NewIntVar(0, KH * KW * Cin * Cout, "W_L2_elems")
        model.Add(W_L2 == (KH * KW) * IC_OC_L2)

        # Helper products for L3
        IH_IW_L3 = model.NewIntVar(0, max_in_l1 * max_in_w1, "IH_IW_L3")
        model.AddMultiplicationEquality(IH_IW_L3, [IH_L3, IW_L3])

        X_L3 = model.NewIntVar(0, max_in_l1 * max_in_w1 * Cin, "X_L3_elems")
        model.AddMultiplicationEquality(X_L3, [IH_IW_L3, v["TIC_L3"]])

        OH_OW_L3 = model.NewIntVar(0, Hout * Wout, "OH_OW_L3")
        model.AddMultiplicationEquality(OH_OW_L3, [v["TH_L3"], v["TW_L3"]])

        Y_L3 = model.NewIntVar(0, Hout * Wout * Cout, "Y_L3_elems")
        model.AddMultiplicationEquality(Y_L3, [OH_OW_L3, v["TOC_L3"]])

        IC_OC_L3 = model.NewIntVar(0, Cin * Cout, "IC_OC_L3")
        model.AddMultiplicationEquality(IC_OC_L3, [v["TIC_L3"], v["TOC_L3"]])

        W_L3 = model.NewIntVar(0, KH * KW * Cin * Cout, "W_L3_elems")
        model.Add(W_L3 == (KH * KW) * IC_OC_L3)

        # --- Memory constraints (bytes) ---
        if L1_cap is not None:
            l1_bytes = model.NewIntVar(0, L1_cap, "l1_bytes")
            if double_buffering:
                model.Add(l1_bytes == (2 * (X_L1 + W_L1 + Y_L1) * dtype_b + l1_overhead))
            else:
                model.Add(l1_bytes == ((X_L1 + W_L1 + Y_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")
            # By default, L2 stores double-buffered input+weights only (no Y in L2).
            if double_buffering:
                model.Add(l2_bytes == (2 * (X_L2 + W_L2) * dtype_b + l2_overhead))
            else:
                model.Add(l2_bytes == ((X_L2 + W_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 == ((X_L3 + W_L3 + Y_L3) * dtype_b + l3_overhead))
            model.Add(l3_bytes <= L3_cap)

        # --- Objective vars ---
        # "Work" proxy: output area * IC * OC (kernel area is constant)
        out_area_L1 = OH_OW_L1
        work_L1 = model.NewIntVar(0, Hout * Wout * Cin * Cout, "work_L1")
        model.AddMultiplicationEquality(work_L1, [out_area_L1, IC_OC_L1])

        out_area_L2 = OH_OW_L2
        work_L2 = model.NewIntVar(0, Hout * Wout * Cin * Cout, "work_L2")
        model.AddMultiplicationEquality(work_L2, [out_area_L2, IC_OC_L2])

        out_area_L3 = OH_OW_L3
        work_L3 = model.NewIntVar(0, Hout * Wout * Cin * Cout, "work_L3")
        model.AddMultiplicationEquality(work_L3, [out_area_L3, IC_OC_L3])

        return {
            "out_area_L1": out_area_L1,
            "work_L1": work_L1,
            "work_L2": work_L2,
            "work_L3": work_L3,
        }

    objectives = ["work_L1", "work_L2", "work_L3", "out_area_L1", "TOC"]
    return TilingProblem(var_domains=var_domains, build_model=build_model, objectives=objectives)


def get_conv2d_tiling(
    shape: Conv2DShape,
    budgets: MemoryBudget,
    dtype_bytes: int = 4,
    *,
    double_buffering: bool = True,
) -> Conv2DTile:
    """
    Compute hierarchical tile sizes for Conv2D under memory constraints.

    Returned fields:
      - TH,TW,TIC,TOC are the L1 tile
      - TH_L2,TW_L2,TIC_L2,TOC_L2 are the L2 tile
      - TH_L3,TW_L3,TIC_L3,TOC_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_conv2d_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 Conv2DTile(
            TH=v["TH"], TW=v["TW"], TIC=v["TIC"], TOC=v["TOC"],
            TH_L2=v["TH_L2"], TW_L2=v["TW_L2"], TIC_L2=v["TIC_L2"], TOC_L2=v["TOC_L2"],
            TH_L3=v["TH_L3"], TW_L3=v["TW_L3"], TIC_L3=v["TIC_L3"], TOC_L3=v["TOC_L3"],
        )

    # Retry with full budgets (no utilization headroom)
    problem2 = _build_conv2d_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 Conv2DTile(
            TH=v["TH"], TW=v["TW"], TIC=v["TIC"], TOC=v["TOC"],
            TH_L2=v["TH_L2"], TW_L2=v["TW_L2"], TIC_L2=v["TIC_L2"], TOC_L2=v["TOC_L2"],
            TH_L3=v["TH_L3"], TW_L3=v["TW_L3"], TIC_L3=v["TIC_L3"], TOC_L3=v["TOC_L3"],
        )

    # Conservative fallback
    return Conv2DTile(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1)