da4ml 0.5.0__cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

da4ml/cmvm/api.py

@@ -0,0 +1,264 @@
+from math import ceil, log2
+from typing import TypedDict
+
+import numpy as np
+from numba import jit, prange
+
+from .core import _solve, create_state, to_solution
+from .types import Pipeline, QInterval
+from .util import kernel_decompose
+
+
+@jit(cache=True)
+def minimal_latency(
+    kernel: np.ndarray,
+    qintervals: list[QInterval],
+    latencies: list[float],
+    carry_size: int = -1,
+    adder_size: int = -1,
+):
+    """Fast latency calculation for a given kernel, QInterval, and input latencies.
+    When carry_size=-1, and the input latency is constant `l`:
+    this will be the same as `l + max(ceiling(log2(max(#CSD bits for each column, 1))))`.
+
+    Parameters
+    ----------
+    kernel : np.ndarray
+        The input kernel matrix.
+    qintervals : list[QInterval]
+        List of QIntervals for each input.
+    latencies : list[float]
+        List of latencies for each input
+    carry_size : int, optional
+        The size of the carry unit for latency computation, by default -1 (fixed latency for each addition operation)
+    adder_size : int, optional
+        The size of the adder unit for latency computation, by default -1 (fixed cost for each addition operation)
+
+    Returns
+    -------
+    float
+        The minimal latency for the given kernel, QInterval, and input latencies.
+    """
+
+    state = create_state(kernel, qintervals, latencies, no_stat_init=True)
+    sol = to_solution(state, adder_size=adder_size, carry_size=carry_size)
+    latencies = [sol.ops[i].latency if i >= 0 else 0.0 for i in sol.out_idxs]
+    return max(latencies)
+
+
+@jit(cache=True)
+def jit_solve(
+    kernel: np.ndarray,
+    method0: str = 'wmc',
+    method1: str = 'auto',
+    hard_dc: int = -1,
+    decompose_dc: int = -2,
+    qintervals: list[QInterval] | None = None,
+    latencies: list[float] | None = None,
+    adder_size: int = -1,
+    carry_size: int = -1,
+) -> Pipeline:
+    """Optimized implementation of a CMVM computation with cascaded two matrices.
+
+    Parameters
+    ----------
+    kernel : np.ndarray
+        The input kernel matrix to be implemented.
+    method0 : str, optional
+        Optimization method for the first stage. Must be one of [`wmc`, `wmc-dc`, `wmc-pdc`, `mc`, `mc-dc`, `mc-pdc`].
+    method1 : str, optional
+        Optimization method for the second stage. When 'auto', it will select based on hard_dc and method0, by default 'auto'
+    hard_dc : int, optional
+        Hard depth constraint (additional latency allowed beyond minimal latency), by default -1 (no constraint)
+    decompose_dc : int, optional
+        Decomposition depth constraint, by default -1 (no constraint, follows hard_dc)
+    qintervals : list[QInterval] | None, optional
+        List of quantization intervals for each input, by default None ([-128, 127, 1] for all inputs)
+    inp_latencies : list[float] | None, optional
+        List of input latencies, by default None (0. for all inputs)
+    adder_size : int, optional
+        Size of the adder unit for latency computation, by default -1 (fixed cost for each addition)
+    carry_size : int, optional
+        Size of the carry unit for latency computation, by default -1 (fixed latency for each addition)
+
+    Returns
+    -------
+    CascadedSolution
+        A solution containing the optimized implementation of the CMVM computation with cascaded stages.
+    """
+
+    if hard_dc < 0:
+        hard_dc = int(1e9)
+
+    if method1 == 'auto':
+        if hard_dc >= 6 or method0.endswith('dc'):
+            method1 = method0
+        else:
+            method1 = method0 + '-dc'
+    if hard_dc == 0 and not method0.endswith('dc'):
+        method0 = method0 + '-dc'
+
+    if qintervals is None:
+        _qintervals = [QInterval(-128.0, 127.0, 1.0)] * kernel.shape[0]
+    else:
+        _qintervals = list(qintervals)
+    if latencies is None:
+        _inp_latencies = [0.0] * kernel.shape[0]
+    else:
+        _inp_latencies = [float(lat) for lat in latencies]
+    assert len(_qintervals) == kernel.shape[0]
+    assert len(_inp_latencies) == kernel.shape[0]
+
+    min_lat = minimal_latency(kernel, _qintervals, _inp_latencies, carry_size=carry_size, adder_size=adder_size)
+    latency_allowed = hard_dc + min_lat
+    if decompose_dc == -2:
+        decompose_dc = min(hard_dc, ceil(log2(kernel.shape[0])))
+    else:
+        decompose_dc = min(hard_dc, decompose_dc, ceil(log2(kernel.shape[0])))
+
+    while True:
+        if decompose_dc < 0 and hard_dc >= 0:
+            if method0 != 'dummy':
+                method0, method1 = 'wmc-dc', 'wmc-dc'
+            else:
+                method0, method1 = 'dummy', 'dummy'
+        mat0, mat1 = kernel_decompose(kernel, dc=decompose_dc)
+        sol0 = _solve(
+            mat0, method=method0, qintervals=_qintervals, latencies=_inp_latencies, adder_size=adder_size, carry_size=carry_size
+        )
+        latencies0 = [sol0.ops[i].latency if i >= 0 else 0.0 for i in sol0.out_idxs]
+        qintervals0 = [sol0.ops[i].qint if i >= 0 else QInterval(0.0, 0.0, np.inf) for i in sol0.out_idxs]
+        if max(latencies0) > latency_allowed:
+            if not method0 == method1 == 'wmc-dc' or decompose_dc >= 0:
+                decompose_dc -= 1
+                continue
+        sol1 = _solve(
+            mat1, method=method1, qintervals=qintervals0, latencies=latencies0, adder_size=adder_size, carry_size=carry_size
+        )
+        latencies1 = [sol1.ops[i].latency if i >= 0 else 0.0 for i in sol1.out_idxs]
+        if max(latencies1) > latency_allowed:
+            # Prevent infinite loop, shouldn't happen though
+            if not method0 == method1 == 'wmc-dc' or decompose_dc >= 0:
+                decompose_dc -= 1
+                continue
+        break
+    if max(latencies1) > latency_allowed:
+        # When latency depends on the bw, may happen
+        print(f'Latency constraint not satisfied: {int(latency_allowed)} < {int(max(latencies1))}')
+    return Pipeline((sol0, sol1))
+
+
+@jit(cache=True, parallel=True)
+def solve(
+    kernel: np.ndarray,
+    method0: str = 'wmc',
+    method1: str = 'auto',
+    hard_dc: int = -1,
+    decompose_dc: int = -2,
+    qintervals: list[QInterval] | None = None,
+    latencies: list[float] | None = None,
+    adder_size: int = -1,
+    carry_size: int = -1,
+    search_all_decompose_dc: bool = True,
+) -> Pipeline:
+    """Solve the CMVM problem with cascaded two matrices.
+
+    Parameters
+    ----------
+    kernel : np.ndarray
+        The input kernel matrix to be implemented.
+    method0 : str, optional
+        Optimization method for the first stage. Must be one of [`wmc`, `wmc-dc`, `wmc-pdc`, `mc`, `mc-dc`, `mc-pdc`].
+    method1 : str, optional
+        Optimization method for the second stage. When 'auto', it will select based on hard_dc and method0, by default 'auto'
+    hard_dc : int, optional
+        Hard depth constraint (additional latency allowed beyond minimal latency), by default -1 (no constraint)
+    decompose_dc : int, optional
+        Decomposition depth constraint, by default -1 (no constraint, follows hard_dc)
+    qintervals : list[QInterval] | None, optional
+        List of quantization intervals for each input, by default None ([-128, 127, 1] for all inputs)
+    inp_latencies : list[float] | None, optional
+        List of input latencies, by default None (0. for all inputs)
+    adder_size : int, optional
+        Size of the adder unit for latency computation, by default -1 (fixed cost for each addition)
+    carry_size : int, optional
+        Size of the carry unit for latency computation, by default -1 (fixed latency for each addition)
+    search_all_decompose_dc : bool, optional
+        If True, search for all possible decomposition depth constraints. If False, use the provided decompose_dc value.
+        Default is True.
+
+    Returns
+    -------
+    CascadedSolution
+        A solution containing the optimized implementation of the CMVM computation with cascaded stages.
+    """
+
+    if qintervals is None:
+        _qintervals = [QInterval(-128.0, 127.0, 1.0)] * kernel.shape[0]
+    else:
+        _qintervals = list(qintervals)
+    if latencies is None:
+        _latencies = [0.0] * kernel.shape[0]
+    else:
+        _latencies = [float(lat) for lat in latencies]
+
+    if not search_all_decompose_dc:
+        return jit_solve(
+            kernel,
+            method0=method0,
+            method1=method1,
+            hard_dc=hard_dc,
+            decompose_dc=decompose_dc,
+            qintervals=_qintervals,
+            latencies=_latencies,
+            adder_size=adder_size,
+            carry_size=carry_size,
+        )
+
+    if hard_dc < 0:
+        hard_dc = int(1e9)
+
+    max_decompose_dc = min(hard_dc, ceil(log2(kernel.shape[0])))
+    try_decompose_dcs = list(range(-1, max_decompose_dc + 1))
+
+    costs = np.empty(len(try_decompose_dcs), dtype=np.float64)
+
+    for i in prange(len(try_decompose_dcs)):
+        decompose_dc = try_decompose_dcs[i]
+        _csol = jit_solve(
+            kernel,
+            method0=method0,
+            method1=method1,
+            hard_dc=hard_dc,
+            decompose_dc=decompose_dc,
+            qintervals=_qintervals,
+            latencies=_latencies,
+            adder_size=adder_size,
+            carry_size=carry_size,
+        )
+        _cost = sum([sum([op.cost for op in sol.ops]) for sol in _csol.solutions])
+        costs[i] = _cost
+
+    decompose_dc = try_decompose_dcs[np.argmin(costs)]
+    csol = jit_solve(
+        kernel,
+        method0=method0,
+        method1=method1,
+        hard_dc=hard_dc,
+        decompose_dc=decompose_dc,
+        qintervals=_qintervals,
+        latencies=_latencies,
+        adder_size=adder_size,
+        carry_size=carry_size,
+    )
+    return csol
+
+
+class solver_options_t(TypedDict, total=False):
+    method0: str
+    method1: str
+    hard_dc: int
+    decompose_dc: int
+    adder_size: int
+    carry_size: int
+    search_all_decompose_dc: bool

da4ml/cmvm/core/__init__.py

@@ -0,0 +1,221 @@
+import heapq
+from math import log2
+
+import numpy as np
+from numba import jit
+
+from ..types import CombLogic, DAState, Op, QInterval
+from .indexers import (
+    idx_mc,
+    idx_mc_dc,
+    idx_wmc,
+    idx_wmc_dc,
+)
+from .state_opr import cost_add, create_state, qint_add, update_state
+
+
+@jit(cache=True)
+def cmvm(
+    kernel: np.ndarray,
+    method: str = 'wmc',
+    qintervals: list[QInterval] | None = None,
+    inp_latencies: list[float] | None = None,
+    adder_size: int = -1,
+    carry_size: int = -1,
+) -> DAState:
+    """Optimizes the kernel using the CMVM algorithm.
+
+    Parameters
+    ----------
+    kernel : np.ndarray
+        The kernel to optimize.
+    method : str, optional
+        Which indexing method to use, by default 'wmc' (weighted most common)
+        Must be one of [`mc`, `mc-dc`, `mc-pdc`, `wmc`, `wmc-dc`, `wmc-pdc`, `dummy`]
+    qintervals : list[QInterval] | None, optional
+        List of QIntervals for each input, by default None
+        If None, defaults to [-128., 127., 1.] for each input.
+    inp_latencies : list[float] | None, optional
+        List of latencies for each input, by default None
+        If None, defaults to 0. for each input.
+    adder_size : int, optional
+        The atomic size of the adder for cost computation, by default -1
+        if -1, each adder can be arbitrary large, and the cost will be the number of adders
+    carry_size : int, optional
+        The size of the carry unit for latency computation, by default -1
+        if -1, each carry unit can be arbitrary large, and the cost will be the depth of the adder tree
+
+    Returns
+    -------
+    DAState
+        The optimized kernel as a DAState object.
+    """
+
+    if qintervals is None:
+        _qintervals = [QInterval(-128.0, 127.0, 1.0)] * kernel.shape[0]
+    else:
+        _qintervals = [QInterval(*qi) for qi in qintervals]
+    if inp_latencies is None:
+        _inp_latencies = [0.0] * kernel.shape[0]
+    else:
+        _inp_latencies = [float(lat) for lat in inp_latencies]
+    assert len(_qintervals) == kernel.shape[0]
+    assert len(_inp_latencies) == kernel.shape[0]
+
+    state = create_state(kernel, _qintervals, _inp_latencies)
+    while True:
+        if len(state.freq_stat) == 0:
+            break
+        match method:
+            case 'mc':
+                pair_idx = idx_mc(state)
+            case 'mc-dc':
+                pair_idx = idx_mc_dc(state, absolute=True)
+            case 'mc-pdc':
+                pair_idx = idx_mc_dc(state, absolute=False)
+            case 'wmc':
+                pair_idx = idx_wmc(state)
+            case 'wmc-dc':
+                pair_idx = idx_wmc_dc(state, absolute=True)
+            case 'wmc-pdc':
+                pair_idx = idx_wmc_dc(state, absolute=False)
+            case 'dummy':
+                break
+            case _:
+                raise ValueError(f'Unknown method: {method}')
+        if pair_idx < 0:
+            break
+        pair_chosen = list(state.freq_stat.keys())[pair_idx]
+        state = update_state(state, pair_chosen, adder_size=adder_size, carry_size=carry_size)
+    return state
+
+
+@jit(cache=True)
+def to_solution(
+    state: DAState,
+    adder_size: int,
+    carry_size: int,
+):
+    """Converts the DAState to a Solution object with balanced tree reduction for the non-extracted bits in the kernel.
+
+    Parameters
+    ----------
+    state : DAState
+        The DAState to convert.
+    adder_size : int, optional
+        The atomic size of the adder for cost computation, by default -1
+        if -1, each adder can be arbitrary large, and the cost will be the number of adders
+    carry_size : int, optional
+        The size of the carry unit for latency computation, by default -1
+        if -1, each carry unit can be arbitrary large, and the cost will be the depth of the adder tree
+
+    Returns
+    -------
+    Solution
+        The Solution object with the optimized kernel.
+    """
+
+    ops = state.ops.copy()
+    n_out = state.kernel.shape[1]
+    expr = np.empty((len(state.expr), *state.expr[0].shape), dtype=np.int8)
+    for i, v in enumerate(state.expr):
+        expr[i] = v
+    in_shifts, out_shifts = state.shifts
+
+    out_qints = []
+    out_lats = []
+    out_idx = []
+    in_shift = in_shifts.copy()
+    out_shift = out_shifts.copy()
+    out_neg = []
+
+    _global_id = len(ops)
+    for i_out in range(n_out):
+        idx, shifts = np.where(expr[:, i_out] != 0)
+        sub = np.empty(len(idx), dtype=np.int64)
+        for i, (i_in, shift) in enumerate(zip(idx, shifts)):
+            sub[i] = expr[i_in, i_out, shift] == -1
+
+        qints: list[QInterval] = [state.ops[i].qint for i in idx]
+        lats: list[float] = [state.ops[i].latency for i in idx]
+
+        # No reduction required, dump the realized value directly
+        if len(sub) == 1:
+            out_shift[i_out] = out_shift[i_out] + shifts[0]
+            out_qints.append(qints[0])
+            out_lats.append(lats[0])
+            out_idx.append(idx[0])
+            out_neg.append(sub[0])
+            continue
+        # Output is zero
+        if len(sub) == 0:
+            out_idx.append(-1)  # -1 means output constant zero
+            out_qints.append(QInterval(0.0, 0.0, np.inf))
+            out_lats.append(0.0)
+            out_neg.append(False)
+            continue
+
+        # Sort by latency -> location of rightmost bit -> lower bound
+        left_align: list[int] = []
+        for i, qint in enumerate(qints):
+            n_int = int(log2(max(abs(qint.max + qint.step), abs(qint.min))))
+            left_align.append(n_int + shifts[i])
+        heap = list(zip(lats, sub, left_align, qints, idx, shifts))
+        heapq.heapify(heap)
+
+        while len(heap) > 1:
+            lat0, sub0, _, qint0, id0, shift0 = heapq.heappop(heap)
+            lat1, sub1, _, qint1, id1, shift1 = heapq.heappop(heap)
+
+            if sub0:
+                shift = shift0 - shift1
+                qint = qint_add(qint1, qint0, shift, sub1, sub0)
+                dlat, dcost = cost_add(qint1, qint0, shift=shift, sub=1 ^ sub1, adder_size=adder_size, carry_size=carry_size)
+                lat = max(lat0, lat1) + dlat
+                op = Op(id1, id0, 1 ^ sub1, shift, qint, lat, dcost)
+                shift = shift1
+            else:
+                shift = shift1 - shift0
+                qint = qint_add(qint0, qint1, shift, sub0, sub1)
+                dlat, dcost = cost_add(qint0, qint1, shift=shift, sub=sub1, adder_size=adder_size, carry_size=carry_size)
+                lat = max(lat0, lat1) + dlat
+                op = Op(id0, id1, sub1, shift, qint, lat, dcost)
+                shift = shift0
+
+            left_align = int(log2(max(abs(qint.max + qint.step), abs(qint.min)))) + shift
+            heapq.heappush(heap, (lat, sub0 & sub1, left_align, qint, _global_id, shift))
+            ops.append(op)
+            _global_id += 1
+
+        lat, sub, _, qint, id0, shift0 = heap[0]
+        out_idx.append(_global_id - 1)
+        out_qints.append(qint)
+        out_lats.append(lat)
+        out_neg.append(sub)
+        out_shift[i_out] = out_shift[i_out] + shift0
+
+    return CombLogic(
+        shape=state.kernel.shape,  # type: ignore
+        inp_shifts=list(in_shift),
+        out_idxs=out_idx,
+        out_shifts=list(out_shift),
+        out_negs=out_neg,
+        ops=ops,
+        carry_size=carry_size,
+        adder_size=adder_size,
+    )
+
+
+@jit
+def _solve(
+    kernel: np.ndarray,
+    method: str,
+    qintervals: list[QInterval],
+    latencies: list[float],
+    adder_size: int,
+    carry_size: int,
+):
+    state = cmvm(
+        kernel, method=method, qintervals=qintervals, inp_latencies=latencies, adder_size=adder_size, carry_size=carry_size
+    )
+    return to_solution(state, adder_size=adder_size, carry_size=carry_size)

da4ml/cmvm/core/indexers.py

@@ -0,0 +1,83 @@
+from math import ceil, log2
+
+import numpy as np
+from numba import jit
+
+from ..types import DAState, QInterval
+
+
+@jit
+def idx_mc(state: DAState):
+    """Choose the pair with highest frequency."""
+    freqs = list(state.freq_stat.values())
+    max_freq = max(freqs)
+    pair_idx = freqs.index(max_freq)
+    return pair_idx
+
+
+@jit
+def idx_mc_dc(state: DAState, absolute: bool = False):
+    """Choose the pair with highest frequency with latency penalty.
+    If absolute is True, return -1 if any latency overhead may present."""
+    freqs = list(state.freq_stat.values())
+    factor = max(freqs) + 1
+    ops = state.ops
+    lat_penalty = [abs(ops[pair.id1].latency - ops[pair.id0].latency) * factor for pair in state.freq_stat.keys()]
+    score = [freq - lat_penalty[i] for i, freq in enumerate(freqs)]
+    max_score = max(score)
+    if absolute and max_score < 0:
+        return -1
+    pair_idx = score.index(max_score)
+    return pair_idx
+
+
+@jit
+def overlap_and_accum(qint0: QInterval, qint1: QInterval):
+    """Calculate the overlap and total number of bits for two QIntervals, when represented in fixed-point format."""
+
+    min0, max0, step0 = qint0
+    min1, max1, step1 = qint1
+    max0, max1 = max0 + step0, max1 + step1
+
+    f = -log2(max(step0, step1))
+    i_high = ceil(log2(max(abs(min0), abs(min1), abs(max0), abs(max1))))
+    i_low = ceil(log2(min(max(abs(min0), abs(max0)), max(abs(min1), abs(max1)))))
+    k = int(qint0.min < 0 or qint1.min < 0)
+    n_accum = k + i_high + f
+    n_overlap = k + i_low + f
+    return n_overlap, n_accum
+
+
+@jit
+def idx_wmc(state: DAState):
+    """Choose the pair with the highest weighted most common subexpression (WMC) score."""
+    freqs = list(state.freq_stat.values())
+    keys = list(state.freq_stat.keys())
+    score = np.empty(len(freqs), dtype=np.float32)
+    for i, (k, v) in enumerate(zip(keys, freqs)):
+        id0, id1 = k.id0, k.id1
+        qint0, qint1 = state.ops[id0].qint, state.ops[id1].qint
+        n_overlap, _ = overlap_and_accum(qint0, qint1)
+        score[i] = v * n_overlap
+    max_score = np.max(score)
+    if max_score < 0:
+        return -1
+    return np.argmax(score)
+
+
+@jit
+def idx_wmc_dc(state: DAState, absolute: bool = False):
+    """Choose the pair with the highest weighted most common subexpression (WMC) score with latency and cost penalty.
+    When absolute is True, return -1 if any latency overhead may present."""
+    freqs = list(state.freq_stat.values())
+    keys = list(state.freq_stat.keys())
+    score = np.empty(len(freqs), dtype=np.float32)
+    for i, (k, v) in enumerate(zip(keys, freqs)):
+        id0, id1 = k.id0, k.id1
+        qint0, qint1 = state.ops[id0].qint, state.ops[id1].qint
+        lat0, lat1 = state.ops[id0].latency, state.ops[id1].latency
+        n_overlap, _ = overlap_and_accum(qint0, qint1)
+        score[i] = v * n_overlap - 256 * abs(lat0 - lat1)
+    if absolute and np.max(score) < 0:
+        return -1
+    return np.argmax(score)