da4ml 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

da4ml/cmvm/types.py

@@ -0,0 +1,569 @@
+import json
+from decimal import Decimal
+from functools import reduce, singledispatch
+from math import ceil, floor, log2
+from pathlib import Path
+from typing import TYPE_CHECKING, NamedTuple, TypeVar
+
+import numpy as np
+from numba import jit
+from numpy import float32, int8
+from numpy.typing import NDArray
+
+if TYPE_CHECKING:
+    from ..trace.tracer import FixedVariable
+
+
+class QInterval(NamedTuple):
+    """A class representing a quantized interval: [min, max] with a step size."""
+
+    min: float
+    max: float
+    step: float
+
+    @classmethod
+    def from_kif(cls, k: int | bool, i: int, f: int):
+        _high = 2.0**i
+        step = 2.0**-f
+        low, high = -k * step, _high - step
+        return cls(low, high, step)
+
+    @classmethod
+    def from_precision(cls, prec: 'Precision'):
+        return cls.from_kif(*prec)
+
+    @property
+    def precision(self):
+        return Precision.from_qint(self)
+
+    def __repr__(self):
+        return f'[{self.min}, {self.max}, {self.step}]'
+
+
+class Precision(NamedTuple):
+    """A class representing the precision of a quantized interval."""
+
+    keep_negative: bool
+    integers: int
+    fractional: int
+
+    def __str__(self):
+        k, i, f = self.keep_negative, self.integers, self.fractional
+        k, B, I = k, i + f + k, i + k
+        return f'fixed({k}, {B}, {I})'
+
+    def __repr__(self):
+        return str(self)
+
+    @classmethod
+    def from_qint(cls, qint: QInterval, symmetric: bool = False):
+        return _minimal_kif(qint, symmetric=symmetric)
+
+    @property
+    def qint(self):
+        return QInterval.from_kif(*self)
+
+
+class Op(NamedTuple):
+    """One single operation on the data buffer.
+
+    Parameters
+    ----------
+    id0: int
+        index of the first operand
+    id1: int
+        index of the second operand, or special opcode if negative
+    opcode: int
+        0: addition, 1: subtraction, 2: relu, 3: quantize, 4: const addition
+    data: int
+        Data to be used in the operation
+    qint: QInterval
+        Quantization interval of the resultant buffer
+    latency: float
+        Latency of the data generated by this operation (t_available)
+    cost: float
+        Cost of the operation
+    """
+
+    id0: int
+    id1: int
+    opcode: int
+    data: int
+    qint: QInterval
+    latency: float
+    cost: float
+
+
+class Pair(NamedTuple):
+    """An operation representing data[id0] +/- data[id1] * 2**shift."""
+
+    id0: int
+    id1: int
+    sub: bool
+    shift: int
+
+
+class DAState(NamedTuple):
+    """Internal state of the DA algorithm."""
+
+    shifts: tuple[NDArray[int8], NDArray[int8]]
+    expr: list[NDArray[int8]]
+    ops: list[Op]
+    freq_stat: dict[Pair, int]
+    kernel: NDArray[float32]
+
+
+def _minimal_kif(qi: QInterval, symmetric: bool = False) -> Precision:
+    """Calculate the minimal KIF for a given QInterval.
+
+    Parameters
+    ----------
+    qi : QInterval
+        The QInterval to calculate the KIF for.
+    symmetric : bool
+        Only relevant if qi may be negative. If True, -2**i will be regarded as forbidden.
+        May be useful in special cases only.
+        Default is False.
+
+    Returns
+    -------
+    Precision
+        A named tuple with the KIF values.
+    """
+
+    if qi.min == qi.max == 0:
+        return Precision(keep_negative=False, integers=0, fractional=0)
+    keep_negative = qi.min < 0
+    fractional = int(-log2(qi.step))
+    int_min, int_max = round(qi.min / qi.step), round(qi.max / qi.step)
+    if symmetric:
+        bits = int(ceil(log2(max(abs(int_min), int_max) + 1)))
+    else:
+        bits = int(ceil(log2(max(abs(int_min), int_max + 1))))
+    integers = bits - fractional
+    return Precision(keep_negative=keep_negative, integers=integers, fractional=fractional)
+
+
+if TYPE_CHECKING:
+
+    def minimal_kif(qi: QInterval, symmetric: bool = False) -> Precision: ...
+else:
+    minimal_kif = jit(_minimal_kif)
+
+
+T = TypeVar('T', 'FixedVariable', float, int, np.float32, np.float64, Decimal)
+
+
+@singledispatch
+def _relu(v: 'T', i: int | None = None, f: int | None = None, inv: bool = False, round_mode: str = 'TRN') -> 'T':
+    from ..trace.fixed_variable import FixedVariable
+
+    assert isinstance(v, FixedVariable), f'Unknown type {type(v)} for symbolic relu'
+    return v.relu(i, f, round_mode=round_mode)
+
+
+@_relu.register(float)
+@_relu.register(int)
+@_relu.register(np.float32)
+@_relu.register(np.float64)
+def _(v, i: int | None = None, f: int | None = None, inv: bool = False, round_mode: str = 'TRN'):
+    if inv:
+        v = -v
+    v = max(0, v)
+    if f is not None:
+        if round_mode.upper() == 'RND':
+            v += 2.0 ** (-f - 1)
+        sf = 2.0**f
+        v = floor(v * sf) / sf
+    if i is not None:
+        v = v % 2.0**i
+    return v
+
+
+@_relu.register
+def _(v: Decimal, i: int | None = None, f: int | None = None, inv: bool = False, round_mode: str = 'TRN'):
+    if inv:
+        v = -v
+    v = max(Decimal(0), v)
+    if f is not None:
+        if round_mode.upper() == 'RND':
+            v += Decimal(2) ** (-f - 1)
+        sf = Decimal(2) ** f
+        v = floor(v * sf) / sf
+    if i is not None:
+        v = v % Decimal(2) ** i
+    return v
+
+
+@singledispatch
+def _quantize(v: 'T', k: int | bool, i: int, f: int, round_mode: str = 'TRN') -> 'T':
+    from ..trace.fixed_variable import FixedVariable
+
+    assert isinstance(v, FixedVariable), f'Unknown type {type(v)} for symbolic quantization'
+    return v.quantize(k, i, f, round_mode=round_mode)
+
+
+@_quantize.register(float)
+@_quantize.register(int)
+@_quantize.register(np.float32)
+@_quantize.register(np.float64)
+def _(v, k: int | bool, i: int, f: int, round_mode: str = 'TRN'):
+    if round_mode.upper() == 'RND':
+        v += 2.0 ** (-f - 1)
+    b = k + i + f
+    bias = 2.0 ** (b - 1) * k
+    eps = 2.0**-f
+    return eps * ((np.floor(v / eps) + bias) % 2**b - bias)
+
+
+@_quantize.register
+def _(v: Decimal, k: int | bool, i: int, f: int, round_mode: str = 'TRN'):
+    if round_mode.upper() == 'RND':
+        v += Decimal(2) ** (-f - 1)
+    b = k + i + f
+    bias = Decimal(2) ** (b - 1) * k
+    eps = Decimal(2) ** -f
+    return eps * ((floor(v / eps) + bias) % Decimal(2) ** b - bias)
+
+
+class Solution(NamedTuple):
+    """Represents a series of operations that can be applied to a vector of data.
+    May represent a CMVM solution or a general neural network
+
+    Attributes
+    ----------
+    shape: tuple[int, int]
+        #input, #output
+    inp_shift: list[int]
+        The shifts that should be applied to the input data.
+    out_idxs: list[int]
+        The indices of the output data in the buffer.
+    out_shifts: list[int]
+        The shifts that should be applied to the output data.
+    out_negs: list[bool]
+        The signs of the output data.
+    ops: list[Op]
+        Core list of operations for generating each buffer element.
+    carry_size: int
+        Size of the carrier for the adder.
+    adder_size: int
+        Elementary size of the adder.
+
+
+    The core part of the solution is the operations in the ops list.
+    For the exact operations executed with Op, refer to the Op class.
+    After all operations are executed, the output data is read from data[op.out_idx] and multiplied by 2**out_shift.
+
+    """
+
+    shape: tuple[int, int]
+    inp_shift: list[int]
+    out_idxs: list[int]
+    out_shifts: list[int]
+    out_negs: list[bool]
+    ops: list[Op]
+    carry_size: int
+    adder_size: int
+
+    def __call__(self, inp: list | np.ndarray | tuple, quantize=False, debug=False, dump=False):
+        """Executes the solution on the input data.
+
+        Parameters
+        ----------
+        inp : list | np.ndarray | tuple
+            Input data to be processed. The input data should be a list or numpy array of objects.
+        quantize : bool
+            If True, the input data will be quantized to the output quantization intervals.
+            Only floating point data types are supported when quantize is True.
+            Default is False.
+        debug : bool
+            If True, the function will print debug information about the operations being performed.
+            Default is False.
+        dump : bool
+            If True, the return the whole buffer, without applying the output shifts and signs.
+            Default is False.
+
+        Returns
+        -------
+        np.ndarray
+            The output data after applying the operations defined in the solution.
+
+        """
+        buf = np.empty(len(self.ops), dtype=object)
+        inp = np.asarray(inp)
+
+        inp_qint = [op.qint for op in self.ops if op.opcode == -1]
+        if quantize:  # TRN and WRAP
+            k, i, f = map(np.array, zip(*map(minimal_kif, inp_qint)))
+            eps = 2.0**-f
+            _low, _high = -(2.0 ** (i + f)) * k, 2.0 ** (i + f) - 1
+            inp = eps * ((np.floor(inp / eps) - _low) % 2.0 ** (k + i + f) + _low)
+
+        inp = inp * (2.0 ** np.array(self.inp_shift))
+        for i, op in enumerate(self.ops):
+            match op.opcode:
+                case -1:  # copy form external buffer
+                    buf[i] = inp[op.id0]
+                case 0 | 1:  # addition
+                    v0, v1 = buf[op.id0], 2.0**op.data * buf[op.id1]
+                    buf[i] = v0 + v1 if op.opcode == 0 else v0 - v1
+                case 2 | -2:  # relu(+/-x)
+                    v = buf[op.id0]
+                    _, _i, _f = _minimal_kif(op.qint)
+                    buf[i] = _relu(v, _i, _f, inv=op.opcode == -2, round_mode='TRN')
+                case 3 | -3:  # quantize(+/-x)
+                    v = buf[op.id0] if op.opcode == 3 else -buf[op.id0]
+                    _k, _i, _f = _minimal_kif(op.qint)
+                    buf[i] = _quantize(v, _k, _i, _f, round_mode='TRN')
+                case 4:  # const addition
+                    bias = op.data * op.qint.step
+                    buf[i] = buf[op.id0] + bias
+                case 5:
+                    buf[i] = op.data * op.qint.step  # const definition
+                case _:
+                    raise ValueError(f'Unknown opcode {op.opcode} in {op}')
+
+        sf = 2.0 ** np.array(self.out_shifts)
+        sign = np.where(self.out_negs, -1, 1)
+        out_idx = np.array(self.out_idxs)
+        mask = np.where(out_idx < 0, 0, 1)
+        if debug:
+            for i, v in enumerate(buf):
+                op = self.ops[i]
+                match op.opcode:
+                    case -1:
+                        op_str = 'inp'
+                    case 0:
+                        op_str = f'buf[{op.id0}] + buf[{op.id1}]<<{op.data}'
+                    case 1:
+                        op_str = f'buf[{op.id0}] - buf[{op.id1}]<<{op.data}'
+                    case 2:
+                        op_str = f'relu(buf[{op.id0}])'
+                    case -2:
+                        op_str = f'relu(-buf[{op.id0}])'
+                    case 3:
+                        op_str = f'quantize(buf[{op.id0}])'
+                    case -3:
+                        op_str = f'quantize(-buf[{op.id0}])'
+                    case 4:
+                        op_str = f'buf[{op.id0}] + {op.data * op.qint.step}'
+                    case 5:
+                        op_str = f'const {op.data * op.qint.step}'
+                    case _:
+                        raise ValueError(f'Unknown opcode {op.opcode} in {op}')
+
+                print(f'{op_str:24} |-> buf[{i}] = {v}')
+
+        if dump:
+            return buf
+        return buf[out_idx] * sf * sign * mask
+
+    @property
+    def kernel(self):
+        """the kernel represented by the solution, when applicable."""
+        kernel = np.empty(self.shape, dtype=np.float32)
+        for i, one_hot in enumerate(np.identity(self.shape[0])):
+            kernel[i] = self(one_hot)
+        return kernel
+
+    @property
+    def cost(self):
+        """Total cost of the solution."""
+        return float(sum(op.cost for op in self.ops))
+
+    @property
+    def latency(self):
+        """Minimum and maximum latency of the solution."""
+        latency = [self.ops[i].latency for i in self.out_idxs]
+        if len(latency) == 0:
+            return 0.0, 0.0
+        return min(latency), max(latency)
+
+    def __repr__(self):
+        n_in, n_out = self.shape
+        cost = self.cost
+        lat_min, lat_max = self.latency
+        return f'Solution([{n_in} -> {n_out}], cost={cost}, latency={lat_min}-{lat_max})'
+
+    @property
+    def out_latency(self):
+        """Latencies of all output elements of the solution."""
+        return [self.ops[i].latency if i >= 0 else 0.0 for i in self.out_idxs]
+
+    @property
+    def out_qint(self):
+        """Quantization intervals of the output elements."""
+        buf = []
+        for i, idx in enumerate(self.out_idxs):
+            _min, _max, _step = self.ops[idx].qint
+            sf = 2.0 ** self.out_shifts[i]
+            _min, _max, _step = _min * sf, _max * sf, _step * sf
+            if self.out_negs[i]:
+                _min, _max = -_max, -_min
+            buf.append(QInterval(_min, _max, _step))
+        return buf
+
+    @property
+    def inp_latency(self):
+        """Latencies of all input elements of the solution."""
+        return [op.latency for op in self.ops if op.opcode == -1]
+
+    @property
+    def inp_qint(self):
+        """Quantization intervals of the input elements."""
+        return [op.qint for op in self.ops if op.opcode == -1]
+
+    def save(self, path: str | Path):
+        """Save the solution to a file."""
+        with open(path, 'w') as f:
+            json.dump(self, f)
+
+    @classmethod
+    def deserialize(cls, data: dict):
+        """Load the solution from a file."""
+        ops = []
+        for _op in data[5]:
+            op = Op(*_op[:4], QInterval(*_op[4]), *_op[5:])  # type: ignore
+            ops.append(op)
+        return cls(
+            shape=tuple(data[0]),
+            inp_shift=data[1],
+            out_idxs=data[2],
+            out_shifts=data[3],
+            out_negs=data[4],
+            ops=ops,
+            carry_size=data[6],
+            adder_size=data[7],
+        )
+
+    @classmethod
+    def load(cls, path: str | Path):
+        """Load the solution from a file."""
+        with open(path) as f:
+            data = json.load(f)
+        return cls.deserialize(data)
+
+
+class CascadedSolution(NamedTuple):
+    """A solution that implements cascaded matrix-vector multiplications through multiple CMVM stages.
+
+    CascadedSolution represents a sequence of Solution objects where the output of each stage
+    is fed as input to the next stage.
+
+    Attributes
+    ----------
+    solutions: tuple[Solution, ...]
+        A tuple containing the individual Solution objects for each stage of the cascade.
+
+    Properties
+    ----------
+    kernel: NDArray[float32]
+        The overall kernel matrix which the cascaded solution implements: vec @ kernel = solution(vec).
+        This is calculated as the matrix product of all individual solution kernels.
+    cost: float
+        The total cost of the cascaded solution, computed as the sum of the costs of all stages.
+    latency: tuple[float, float]
+        The minimum and maximum latency of the cascaded solution.
+    inp_qint: list[QInterval]
+        Input quantization intervals
+    inp_lat: list[float]
+        Input latencies
+    in_shift: list[int]
+        Input shifts
+    out_qint: list[QInterval]
+        Output quantization intervals
+    out_lat: list[float]
+        Output latencies
+    out_shift: list[int]
+        Output shifts
+    out_neg: list[bool]
+        Output signs
+    shape: tuple[int, int]
+        The shape of the corresponding kernel matrix.
+    """
+
+    solutions: tuple[Solution, ...]
+
+    def __call__(self, inp: list | np.ndarray | tuple, quantize=False, debug=False):
+        out = np.asarray(inp)
+        for sol in self.solutions:
+            out = sol(out, quantize=quantize, debug=debug)
+        return out
+
+    @property
+    def kernel(self):
+        return reduce(lambda x, y: x @ y, [sol.kernel for sol in self.solutions])
+
+    @property
+    def cost(self):
+        return sum(sol.cost for sol in self.solutions)
+
+    @property
+    def latency(self):
+        return self.solutions[-1].latency
+
+    @property
+    def inp_qint(self):
+        return self.solutions[0].inp_qint
+
+    @property
+    def inp_latency(self):
+        return self.solutions[0].inp_latency
+
+    @property
+    def out_qint(self):
+        return self.solutions[-1].out_qint
+
+    @property
+    def out_latencies(self):
+        return self.solutions[-1].out_latency
+
+    @property
+    def shape(self):
+        return self.solutions[0].shape[0], self.solutions[-1].shape[1]
+
+    @property
+    def inp_shift(self):
+        return self.solutions[0].inp_shift
+
+    @property
+    def out_shift(self):
+        return self.solutions[-1].out_shifts
+
+    @property
+    def out_neg(self):
+        return self.solutions[-1].out_negs
+
+    def __repr__(self) -> str:
+        n_ins = [sol.shape[0] for sol in self.solutions] + [self.shape[1]]
+        shape_str = ' -> '.join(map(str, n_ins))
+        _cost = self.cost
+        lat_min, lat_max = self.latency
+        return f'CascatedSolution([{shape_str}], cost={_cost}, latency={lat_min}-{lat_max})'
+
+    def save(self, path: str | Path):
+        """Save the solution to a file."""
+        with open(path, 'w') as f:
+            json.dump(self, f)
+
+    @classmethod
+    def deserialize(cls, data: dict):
+        """Load the solution from a file."""
+        return cls(solutions=tuple(Solution.deserialize(sol) for sol in data[0]))
+
+    @classmethod
+    def load(cls, path: str):
+        """Load the solution from a file."""
+        with open(path) as f:
+            data = json.load(f)
+        return cls.deserialize(data)
+
+    @property
+    def reg_bits(self):
+        """The number of bits used for the register in the solution."""
+        bits = 0
+        for _sol in self.solutions:
+            kifs = [_minimal_kif(qint) for qint in _sol.out_qint]
+            _bits = sum(map(sum, kifs))
+            bits += _bits
+        return bits

da4ml/cmvm/util/__init__.py

@@ -0,0 +1,7 @@
+from .bit_decompose import csd_decompose
+from .mat_decompose import kernel_decompose
+
+__all__ = [
+    'csd_decompose',
+    'kernel_decompose',
+]

da4ml/cmvm/util/bit_decompose.py

@@ -0,0 +1,86 @@
+import numpy as np
+from numba import jit
+from numpy.typing import NDArray
+
+
+@jit
+def _volatile_int_arr_to_csd(x: NDArray) -> NDArray[np.int8]:
+    x = x
+    N = np.max(np.ceil(np.log2(np.abs(x) * 1.5 + 1e-19)))
+    N = int(max(N, 1))
+    buf = np.zeros((*np.shape(x), N), dtype=np.int8)
+
+    for n in range(N - 1, -1, -1):
+        _2pn = 2**n
+        thres = _2pn / 1.5
+        bit = (x > thres).astype(np.int8)
+        bit -= (x < -thres).astype(np.int8)
+        x -= _2pn * bit
+        buf[..., n] = bit
+    return buf
+
+
+@jit(error_model='numpy')
+def _shift_centering(arr: NDArray):
+    low, high = -64, 64
+    if np.all(arr == 0):
+        high = low = 0
+    while high - low > 1:
+        mid = (high + low) // 2
+        xs = arr * (2.0**mid)
+        if np.all(xs == np.floor(xs)):
+            high = mid
+        else:
+            low = mid
+    return -high
+
+
+@jit(error_model='numpy')
+def shift_centering(arr: NDArray, axis: int):
+    n = arr.shape[axis]
+    shifts = np.empty(n, dtype=np.int8)
+    for i in range(n):
+        shifts[i] = _shift_centering(arr.take(i, axis=axis))
+    return shifts
+
+
+@jit
+def _center(arr: NDArray):
+    shift1 = shift_centering(arr, 1)  # d_out
+    arr = arr * (2.0**-shift1)
+    shift0 = shift_centering(arr, 0)  # d_in
+    arr = arr * (2.0 ** -shift0[:, None])
+    return arr, shift0, shift1
+
+
+@jit
+def csd_decompose(arr: NDArray, center=True):
+    """
+    Convert an 2D array to CSD representation.
+
+    Parameters
+    ----------
+    arr : ndarray
+        Input array to be converted.
+    center : bool, optional
+        If True, the array is centered before conversion. Default is True.
+        If False, the function may accept non-2D arrays.
+
+    Returns
+    -------
+    csd : ndarray
+        CSD representation of the input array after centering, if center is True.
+    shift0 : ndarray
+        Shift values for the first axis.
+    shift1 : ndarray
+        Shift values for the second axis.
+    """
+
+    if center:
+        arr, shift0, shift1 = _center(arr)
+    else:
+        shift0 = np.zeros(arr.shape[0], dtype=np.int8)
+        shift1 = np.zeros(arr.shape[1], dtype=np.int8)
+        arr = arr.copy()
+    csd = _volatile_int_arr_to_csd(arr)
+    return csd, shift0, shift1