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

da4ml/codegen/verilog/verilog_model.py

@@ -0,0 +1,265 @@
+import ctypes
+import os
+import re
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from uuid import uuid4
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ... import codegen
+from ...cmvm.types import CascadedSolution, Solution, _minimal_kif
+from ...trace.pipeline import to_pipeline
+from . import comb_binder_gen, comb_logic_gen, generate_io_wrapper, pipeline_binder_gen, pipeline_logic_gen
+
+
+def get_io_kifs(sol: Solution | CascadedSolution):
+    inp_kifs = tuple(zip(*map(_minimal_kif, sol.inp_qint)))
+    out_kifs = tuple(zip(*map(_minimal_kif, sol.out_qint)))
+    return np.array(inp_kifs, np.int8), np.array(out_kifs, np.int8)
+
+
+class VerilogModel:
+    def __init__(
+        self,
+        solution: Solution | CascadedSolution,
+        prj_name: str,
+        path: str | Path,
+        latency_cutoff: int = -1,
+        print_latency: bool = True,
+        part_name: str = 'xcvu13p-flga2577-2-e',
+        clock_period: int = 5,
+        clock_uncertainty: float = 0.1,
+        io_delay_minmax: tuple[float, float] = (0.2, 0.4),
+    ):
+        self._solution = solution
+        self._path = Path(path)
+        self._prj_name = prj_name
+        self._latency_cutoff = latency_cutoff
+        self._print_latency = print_latency
+        self.__src_root = Path(codegen.__file__).parent
+        self._part_name = part_name
+        self._clock_period = clock_period
+        self._clock_uncertainty = clock_uncertainty
+        self._io_delay_minmax = io_delay_minmax
+
+        self._pipe = solution if isinstance(solution, CascadedSolution) else None
+        if latency_cutoff > 0 and self._pipe is None:
+            assert isinstance(solution, Solution)
+            self._pipe = to_pipeline(solution, latency_cutoff, verbose=False)
+
+        if self._pipe is not None:
+            # get actual latency cutoff
+            latency_cutoff = int(max(max(st.latency) / (i + 1) for i, st in enumerate(self._pipe.solutions)))
+            self._latency_cutoff = latency_cutoff
+
+        self._lib = None
+
+    def write(self):
+        self._path.mkdir(parents=True, exist_ok=True)
+        if self._pipe is not None:  # Pipeline
+            # Main logic
+            codes = pipeline_logic_gen(self._pipe, self._prj_name, self._print_latency)
+            for k, v in codes.items():
+                with open(self._path / f'{k}.v', 'w') as f:
+                    f.write(v)
+
+            # Build script
+            with open(self.__src_root / 'verilog/source/build_prj.tcl') as f:
+                tcl = f.read()
+            tcl = tcl.replace('${DEVICE}', self._part_name)
+            tcl = tcl.replace('${PROJECT_NAME}', self._prj_name)
+            with open(self._path / 'build_prj.tcl', 'w') as f:
+                f.write(tcl)
+
+            # XDC
+            with open(self.__src_root / 'verilog/source/template.xdc') as f:
+                xdc = f.read()
+            xdc = xdc.replace('${CLOCK_PERIOD}', str(self._clock_period))
+            xdc = xdc.replace('${UNCERTAINITY_SETUP}', str(self._clock_uncertainty))
+            xdc = xdc.replace('${UNCERTAINITY_HOLD}', str(self._clock_uncertainty))
+            xdc = xdc.replace('${DELAY_MAX}', str(self._io_delay_minmax[0]))
+            xdc = xdc.replace('${DELAY_MIN}', str(self._io_delay_minmax[1]))
+            with open(self._path / f'{self._prj_name}.xdc', 'w') as f:
+                f.write(xdc)
+
+            # C++ binder w/
+            binder = pipeline_binder_gen(self._pipe, f'{self._prj_name}_wrapper', 1)
+
+            # Verilog IO wrapper (non-uniform bw to uniform one, clk passthrough)
+            io_wrapper = generate_io_wrapper(self._pipe, self._prj_name, True)
+
+            self._pipe.save(self._path / 'pipeline.json')
+        else:  # Comb
+            assert isinstance(self._solution, Solution)
+
+            # Main logic
+            code = comb_logic_gen(self._solution, self._prj_name, self._print_latency, '`timescale 1ns/1ps')
+            with open(self._path / f'{self._prj_name}.v', 'w') as f:
+                f.write(code)
+
+            # Verilog IO wrapper (non-uniform bw to uniform one, no clk)
+            io_wrapper = generate_io_wrapper(self._solution, self._prj_name, False)
+            binder = comb_binder_gen(self._solution, f'{self._prj_name}_wrapper')
+
+        with open(self._path / f'{self._prj_name}_wrapper.v', 'w') as f:
+            f.write(io_wrapper)
+        with open(self._path / f'{self._prj_name}_wrapper_binder.cc', 'w') as f:
+            f.write(binder)
+
+        # Common resource copy
+        shutil.copy(self.__src_root / 'verilog/source/shift_adder.v', self._path)
+        shutil.copy(self.__src_root / 'verilog/source/build_binder.mk', self._path)
+        shutil.copy(self.__src_root / 'verilog/source/ioutils.hh', self._path)
+        self._solution.save(self._path / 'model.json')
+        with open(self._path / 'misc.json', 'w') as f:
+            f.write(f'{{"cost": {self._solution.cost}}}')
+
+    def _compile(self, verbose=False, openmp=True, o3: bool = False, clean=True):
+        """Same as compile, but will not write to the library
+
+        Parameters
+        ----------
+        verbose : bool, optional
+            Verbose output, by default False
+        openmp : bool, optional
+            Enable openmp, by default True
+        o3 : bool | None, optional
+            Turn on -O3 flag, by default False
+        clean : bool, optional
+            Remove obsolete shared object files, by default True
+
+        Raises
+        ------
+        RuntimeError
+            If compilation fails
+        """
+
+        self._uuid = str(uuid4())
+        args = ['make', '-f', 'build_binder.mk']
+        env = os.environ.copy()
+        env['VM_PREFIX'] = f'{self._prj_name}_wrapper'
+        env['STAMP'] = self._uuid
+        env['EXTRA_CXXFLAGS'] = '-fopenmp' if openmp else ''
+        if o3:
+            args.append('fast')
+
+        if clean:
+            m = re.compile(r'^lib.*[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\.so$')
+            for p in self._path.iterdir():
+                if not p.is_dir() and m.match(p.name):
+                    p.unlink()
+
+        try:
+            r = subprocess.run(args, env=env, check=True, cwd=self._path, capture_output=not verbose)
+        except subprocess.CalledProcessError as e:
+            print(e.stderr.decode(), file=sys.stderr)
+            print(e.stdout.decode(), file=sys.stdout)
+            raise RuntimeError('Compilation failed!!') from e
+        if r.returncode != 0:
+            print(r.stderr.decode(), file=sys.stderr)
+            print(r.stdout.decode(), file=sys.stderr)
+            raise RuntimeError('Compilation failed!!')
+
+        self._load_lib(self._uuid)
+
+    def _load_lib(self, uuid: str | None = None):
+        uuid = uuid if uuid is not None else self._uuid
+        self._uuid = uuid
+        lib_path = self._path / f'lib{self._prj_name}_wrapper_{uuid}.so'
+        if not lib_path.exists():
+            raise RuntimeError(f'Library {lib_path} does not exist')
+        self._lib = ctypes.CDLL(str(lib_path))
+
+    def compile(self, verbose=False, openmp=True, o3: bool = False):
+        """Compile the generated code to a emulator for logic simulation.
+
+        Parameters
+        ----------
+        verbose : bool, optional
+            Verbose output, by default False
+        openmp : bool, optional
+            Enable openmp, by default True
+        o3 : bool | None, optional
+            Turn on -O3 flag, by default False
+
+        Raises
+        ------
+        RuntimeError
+            If compilation fails
+        """
+        self.write()
+        self._compile(verbose=verbose, openmp=openmp, o3=o3)
+        self._load_lib()
+
+    def predict(self, data: NDArray[np.floating]):
+        """Run the model on the input data.
+
+        Parameters
+        ----------
+        data : NDArray[np.floating]
+            Input data to the model. The shape is ignored, and the number of samples is
+            determined by the size of the data.
+
+        Returns
+        -------
+        NDArray[np.float64]
+            Output of the model in shape (n_samples, output_size).
+        """
+        assert self._lib is not None, 'Library not loaded, call .compile() first.'
+        inp_size, out_size = self._solution.shape
+
+        assert data.size % inp_size == 0, f'Input size {data.size} is not divisible by {inp_size}'
+        n_sample = data.size // inp_size
+
+        kifs_in, kifs_out = get_io_kifs(self._solution)
+        k_in, i_in, f_in = map(np.max, kifs_in)
+        k_out, i_out, f_out = map(np.max, kifs_out)
+        assert k_in + i_in + f_in <= 32, "Padded inp bw doesn't fit in int32. Emulation not supported"
+        assert k_out + i_out + f_out <= 32, "Padded out bw doesn't fit in int32. Emulation not supported"
+
+        inp_data = np.empty(n_sample * inp_size, dtype=np.int32)
+        out_data = np.empty(n_sample * out_size, dtype=np.int32)
+
+        # Convert to int32 matching the LSB position
+        inp_data[:] = data.ravel() * 2.0 ** np.max(f_in)
+
+        inp_buf = inp_data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
+        out_buf = out_data.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))
+        self._lib.inference(inp_buf, out_buf, n_sample)
+
+        # Unscale the output int32 to recover fp values
+        k, i, f = np.max(k_out), np.max(i_out), np.max(f_out)
+        a, b, c = 2.0 ** (k + i + f), 2.0 ** (i + f), 2.0**-f
+        return ((out_data.reshape(n_sample, out_size) + b) % a - b) * c
+
+    def __repr__(self):
+        inp_size, out_size = self._solution.shape
+        cost = round(self._solution.cost)
+        kifs_in, kifs_out = get_io_kifs(self._solution)
+        in_bits, out_bits = np.sum(kifs_in), np.sum(kifs_out)
+        if self._pipe is not None:
+            n_stage = len(self._pipe[0])
+            lat_cutoff = self._latency_cutoff
+            reg_bits = self._pipe.reg_bits
+            spec = f"""Top Module: {self._prj_name}\n====================
+{inp_size} ({in_bits} bits) -> {out_size} ({out_bits} bits)
+{n_stage} stages @ max_delay={lat_cutoff}
+Estimated cost: {cost} LUTs, {reg_bits} FFs"""
+
+        else:
+            spec = f"""Top Module: {self._prj_name}\n====================
+{inp_size} ({in_bits} bits) -> {out_size} ({out_bits} bits)
+combinational @ delay={self._solution.latency}
+Estimated cost: {cost} LUTs"""
+
+        is_compiled = self._lib is not None
+        if is_compiled:
+            openmp = 'with OpenMP' if self._lib.openmp_enabled else ''  # type: ignore
+            spec += f'\nEmulator is compiled {openmp} ({self._uuid[-12:]})'
+        else:
+            spec += '\nEmulator is **not compiled**'
+        return spec

da4ml/trace/__init__.py

@@ -0,0 +1,6 @@
+from .fixed_variable import HWConfig
+from .fixed_variable_array import FixedVariableArray
+from .pipeline import to_pipeline
+from .tracer import comb_trace
+
+__all__ = ['to_pipeline', 'comb_trace', 'FixedVariableArray', 'HWConfig']

da4ml/trace/fixed_variable.py

@@ -0,0 +1,358 @@
+from decimal import Decimal
+from math import ceil, floor, log2
+from typing import NamedTuple
+from uuid import UUID, uuid4
+
+from ..cmvm.core import cost_add
+from ..cmvm.types import QInterval
+
+
+class HWConfig(NamedTuple):
+    adder_size: int
+    carry_size: int
+    latency_cutoff: float
+
+
+def _const_f(const: float | Decimal):
+    const = float(const)
+    _low, _high = -32, 32
+    while _high - _low > 1:
+        _mid = (_high + _low) // 2
+        _value = const * (2.0**_mid)
+        if _value == int(_value):
+            _high = _mid
+        else:
+            _low = _mid
+    return _high
+
+
+class FixedVariable:
+    def __init__(
+        self,
+        low: float | Decimal,
+        high: float | Decimal,
+        step: float | Decimal,
+        latency: float | None = None,
+        hwconf=HWConfig(-1, -1, -1),
+        opr: str = 'new',
+        cost: float | None = None,
+        _from: tuple['FixedVariable', ...] = (),
+        _factor: float | Decimal = 1.0,
+        _data: Decimal | None = None,
+        _id: UUID | None = None,
+    ) -> None:
+        assert low <= high, f'low {low} must be less than high {high}'
+
+        if low == high:
+            opr = 'const'
+            _factor = 1.0
+            _from = ()
+
+        low, high, step = Decimal(low), Decimal(high), Decimal(step)
+        low, high = floor(low / step) * step, ceil(high / step) * step
+        self.low = low
+        self.high = high
+        self.step = step
+        self._factor = Decimal(_factor)
+        self._from: tuple[FixedVariable, ...] = _from
+        opr = opr
+        self.opr = opr
+        self._data = _data
+        self.id = _id or uuid4()
+        self.hwconf = hwconf
+
+        if opr == 'cadd':
+            assert _data is not None, 'cadd must have data'
+
+        if cost is None or latency is None:
+            _cost, _latency = self.get_cost_and_latency()
+        else:
+            _cost, _latency = cost, latency
+
+        self.latency = _latency
+        self.cost = _cost
+
+    def get_cost_and_latency(self):
+        if self.opr == 'const':
+            return 0.0, 0.0
+        if self.opr in ('vadd', 'cadd'):
+            adder_size = self.hwconf.adder_size
+            carry_size = self.hwconf.carry_size
+            latency_cutoff = self.hwconf.latency_cutoff
+
+            if self.opr == 'vadd':
+                assert len(self._from) == 2
+                v0, v1 = self._from
+                int0, int1 = v0.qint, v1.qint
+                base_latency = max(v0.latency, v1.latency)
+                dlat, _cost = cost_add(int0, int1, 0, False, adder_size, carry_size)
+            else:
+                assert len(self._from) == 1
+                assert self._data is not None, 'cadd must have data'
+                # int0 = self._from[0].qint
+                # int1 = QInterval(float(self._data), float(self._data), float(self.step))
+                _f = _const_f(self._data)
+                _cost = float(ceil(log2(abs(self._data) + Decimal(2) ** -_f))) + _f
+                base_latency = self._from[0].latency
+                dlat = 0.0
+
+            _latency = dlat + base_latency
+            if latency_cutoff > 0 and ceil(_latency / latency_cutoff) > ceil(base_latency / latency_cutoff):
+                # Crossed the latency cutoff boundry
+                assert (
+                    dlat <= latency_cutoff
+                ), f'Latency of an atomic operation {dlat} is larger than the pipelining latency cutoff {latency_cutoff}'
+                _latency = ceil(base_latency / latency_cutoff) * latency_cutoff + dlat
+        elif self.opr in ('relu', 'wrap'):
+            assert len(self._from) == 1
+            _latency = self._from[0].latency
+            _cost = 0.0
+            # Assume LUT5 used here (2 fan-out per LUT6, thus *1/2)
+            if self._from[0]._factor < 0:
+                _cost += sum(self.kif) / 2
+            if self.opr == 'relu':
+                _cost += sum(self.kif) / 2
+
+        elif self.opr == 'new':
+            # new variable, no cost
+            _latency = 0.0
+            _cost = 0.0
+        else:
+            raise NotImplementedError(f'Operation {self.opr} is unknown')
+        return _cost, _latency
+
+    @property
+    def unscaled(self):
+        return self * (1 / self._factor)
+
+    @property
+    def qint(self) -> QInterval:
+        return QInterval(float(self.low), float(self.high), float(self.step))
+
+    @property
+    def kif(self) -> tuple[bool, int, int]:
+        if self.step == 0:
+            return False, 0, 0
+        f = -int(log2(self.step))
+        i = ceil(log2(max(-self.low, self.high + self.step)))
+        k = self.low < 0
+        return k, i, f
+
+    def __repr__(self) -> str:
+        if self._factor == 1:
+            return f'FixedVariable({self.low}, {self.high}, {self.step})'
+        return f'({self._factor}) FixedVariable({self.low}, {self.high}, {self.step})'
+
+    def __neg__(self):
+        return FixedVariable(
+            -self.high,
+            -self.low,
+            self.step,
+            _from=self._from,
+            _factor=-self._factor,
+            latency=self.latency,
+            cost=self.cost,
+            opr=self.opr,
+            _id=self.id,
+            _data=self._data,
+            hwconf=self.hwconf,
+        )
+
+    def __add__(self, other: 'FixedVariable|float|Decimal|int'):
+        if not isinstance(other, FixedVariable):
+            return self._const_add(other)
+        if other.high == other.low:
+            return self._const_add(other.low)
+        if self.high == self.low:
+            return other._const_add(self.low)
+
+        assert self.hwconf == other.hwconf, 'FixedVariable must have the same hwconf'
+
+        f0, f1 = self._factor, other._factor
+        if f0 < 0:
+            if f1 > 0:
+                return other + self
+            else:
+                return -((-self) + (-other))
+
+        return FixedVariable(
+            self.low + other.low,
+            self.high + other.high,
+            min(self.step, other.step),
+            _from=(self, other),
+            _factor=f0,
+            opr='vadd',
+            hwconf=self.hwconf,
+        )
+
+    def _const_add(self, other: float | Decimal):
+        if not isinstance(other, (int, float, Decimal)):
+            other = float(other)  # direct numpy to decimal raises error
+        other = Decimal(other)
+        if other == 0:
+            return self
+
+        if self.opr != 'cadd':
+            cstep = Decimal(2.0 ** -_const_f(other))
+
+            return FixedVariable(
+                self.low + other,
+                self.high + other,
+                min(self.step, cstep),
+                _from=(self,),
+                _factor=self._factor,
+                _data=other / self._factor,
+                opr='cadd',
+                hwconf=self.hwconf,
+            )
+
+        # cadd, combine the constant
+        assert len(self._from) == 1
+        parent = self._from[0]
+        assert self._data is not None, 'cadd must have data'
+        sf = self._factor / parent._factor
+        other1 = (self._data * parent._factor) + other / sf
+        return (parent + other1) * sf
+
+    def __sub__(self, other: 'FixedVariable|int|float|Decimal'):
+        return self + (-other)
+
+    def __mul__(
+        self,
+        other: 'float|Decimal',
+    ):
+        if other == 0:
+            return FixedVariable(0, 0, 1, hwconf=self.hwconf)
+
+        assert log2(abs(other)) % 1 == 0, 'Only support pow2 multiplication'
+
+        other = Decimal(other)
+
+        low = min(self.low * other, self.high * other)
+        high = max(self.low * other, self.high * other)
+        step = abs(self.step * other)
+        _factor = self._factor * other
+
+        return FixedVariable(
+            low,
+            high,
+            step,
+            _from=self._from,
+            _factor=_factor,
+            opr=self.opr,
+            latency=self.latency,
+            cost=self.cost,
+            _id=self.id,
+            _data=self._data,
+            hwconf=self.hwconf,
+        )
+
+    def __radd__(self, other: 'float|Decimal|int|FixedVariable'):
+        return self + other
+
+    def __rsub__(self, other: 'float|Decimal|int|FixedVariable'):
+        return (-self) + other
+
+    def __rmul__(self, other: 'float|Decimal|int|FixedVariable'):
+        return self * other
+
+    def relu(self, i: int | None = None, f: int | None = None, round_mode: str = 'TRN'):
+        round_mode = round_mode.upper()
+        assert round_mode in ('TRN', 'RND')
+
+        if self.opr == 'const':
+            val = self.low * (self.low > 0)
+            f = _const_f(val) if not f else f
+            step = Decimal(2) ** -f
+            i = ceil(log2(val + step)) if not i else i
+            eps = step / 2 if round_mode == 'RND' else 0
+            val = floor(val / step + eps) % Decimal(2) ** i * step
+            return FixedVariable(val, val, step, hwconf=self.hwconf)
+
+        step = max(Decimal(2) ** -f, self.step) if f is not None else self.step
+        if step > self.step and round_mode == 'RND':
+            return (self + step / 2).relu(i, f, 'TRN')
+        low = max(Decimal(0), self.low)
+        high = max(Decimal(0), self.high)
+        if i is not None:
+            _high = Decimal(2) ** i - step
+            if _high < high:
+                # overflows
+                low = Decimal(0)
+                high = _high
+        _factor = self._factor
+        return FixedVariable(
+            low,
+            high,
+            step,
+            _from=(self,),
+            _factor=abs(_factor),
+            opr='relu',
+            hwconf=self.hwconf,
+            cost=sum(self.kif) * (1 if _factor > 0 else 2),
+        )
+
+    def quantize(
+        self,
+        k: int | bool,
+        i: int,
+        f: int,
+        overflow_mode: str = 'WRAP',
+        round_mode: str = 'TRN',
+    ):
+        overflow_mode, round_mode = overflow_mode.upper(), round_mode.upper()
+        assert overflow_mode in ('WRAP', 'SAT')
+        assert round_mode in ('TRN', 'RND')
+
+        _k, _i, _f = self.kif
+
+        if k >= _k and i >= _i and f >= _f:
+            return self
+
+        if f < _f and round_mode == 'RND':
+            return (self + 2.0 ** (-f - 1)).quantize(k, i, f, overflow_mode, 'TRN')
+
+        if self.low == self.high:
+            val = self.low
+            step = Decimal(2) ** -f
+            _high = Decimal(2) ** i
+            high, low = _high - step, -_high * k
+            val = (floor(val / step) * step - low) % (2 * _high) + low
+            return FixedVariable(val, val, step, hwconf=self.hwconf)
+
+        # TODO: corner cases exists (e.g., overflow to negative, or negative overflow to high value)
+        # bit-exactness will be lost in these cases, but they should never happen (quantizers are used in a weird way)
+        # Keeping this for now; change if absolutely necessary
+        f = min(f, _f)
+        k = min(k, _k)
+        i = min(i, _i)
+
+        step = max(Decimal(2) ** -f, self.step)
+
+        low = -k * Decimal(2) ** i
+        high = Decimal(2) ** i - step
+        _low, _high = self.low, self.high
+
+        if _low >= low and _high <= high:
+            low, high = _low, _high
+
+        if low > high:
+            return FixedVariable(0, 0, 1, hwconf=self.hwconf)
+
+        return FixedVariable(
+            low,
+            high,
+            step,
+            _from=(self,),
+            _factor=abs(self._factor),
+            opr='wrap' if overflow_mode == 'WRAP' else 'sat',
+            latency=self.latency,
+            hwconf=self.hwconf,
+        )
+
+    @classmethod
+    def from_kif(cls, k: int | bool, i: int, f: int, **kwargs):
+        step = Decimal(2) ** -f
+        _high = Decimal(2) ** i
+        low, high = k * _high, _high - step
+        return cls(low, high, step, **kwargs)