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

da4ml/trace/pipeline.py

@@ -0,0 +1,155 @@
+from math import ceil, floor
+
+from ..cmvm.types import CascadedSolution, Op, Solution
+from .fixed_variable import FixedVariable, HWConfig
+from .tracer import comb_trace
+
+
+def retime_pipeline(csol: CascadedSolution, verbose=True):
+    n_stages = len(csol[0])
+    cutoff_high = ceil(max(max(sol.out_latency) / (i + 1) for i, sol in enumerate(csol[0])))
+    cutoff_low = 0
+    adder_size, carry_size = csol[0][0].adder_size, csol[0][0].carry_size
+    best = csol
+    while cutoff_high - cutoff_low > 1:
+        cutoff = (cutoff_high + cutoff_low) // 2
+        _hwconf = HWConfig(adder_size, carry_size, cutoff)
+        inp = [FixedVariable(*qint, hwconf=_hwconf) for qint in csol.inp_qint]
+        try:
+            out = list(csol(inp))
+        except AssertionError:
+            cutoff_low = cutoff
+            continue
+        _sol = to_pipeline(comb_trace(inp, out), cutoff, retiming=False)
+        if len(_sol[0]) > n_stages:
+            cutoff_low = cutoff
+        else:
+            cutoff_high = cutoff
+            best = _sol
+    if verbose:
+        print(f'actual cutoff: {cutoff_high}')
+    return best
+
+
+def to_pipeline(sol: Solution, latency_cutoff: int, retiming=True, verbose=True) -> CascadedSolution:
+    """Split the record into multiple stages based on the latency of the operations.
+    Only useful for HDL generation.
+
+    Parameters
+    ----------
+    sol : Solution
+        The solution to be split into multiple stages.
+    latency_cutoff : int
+        The latency cutoff for splitting the operations.
+    retiming : bool
+        Whether to retime the solution after splitting. Default is True.
+        If False, new stages are created when the propagation latency exceeds the cutoff.
+        If True, after the first round of splitting, the solution is retimed balance the delay within each stage.
+    verbose : bool
+        Whether to print the actual cutoff used for splitting. Only used if rebalance is True.
+        Default is True.
+
+    Returns
+    -------
+    CascadedSolution
+        The cascaded solution with multiple stages.
+    """
+    assert len(sol.ops) > 0, 'No operations in the record'
+    for i, op in enumerate(sol.ops):
+        if op.id1 != -1:
+            break
+
+    def get_stage(op: Op):
+        return floor(op.latency / (latency_cutoff + 1e-9)) if latency_cutoff > 0 else 0
+
+    opd: dict[int, list[Op]] = {}
+    out_idxd: dict[int, list[int]] = {}
+
+    locator: list[dict[int, int]] = []
+
+    ops = sol.ops.copy()
+    lat = max(ops[i].latency for i in sol.out_idxs)
+    for i in sol.out_idxs:
+        op_out = ops[i]
+        ops.append(Op(i, -1001, -1001, 0, op_out.qint, lat, 0.0))
+
+    for i, op in enumerate(ops):
+        stage = get_stage(op)
+        if op.opcode == -1:
+            # Copy from external buffer
+            opd.setdefault(stage, []).append(op)
+            locator.append({stage: len(opd[stage]) - 1})
+            continue
+        p0_stages = locator[op.id0].keys()
+        if stage not in p0_stages:
+            # Need to copy parent to later states
+            p0_stage = max(p0_stages)
+            p0_idx = locator[op.id0][p0_stage]
+            for j in range(p0_stage, stage):
+                op0 = ops[op.id0]
+                latency = float(latency_cutoff * (j + 1))
+                out_idxd.setdefault(j, []).append(locator[op.id0][j])
+                _copy_op = Op(len(out_idxd[j]) - 1, -1, -1, 0, op0.qint, latency, 0.0)
+                opd.setdefault(j + 1, []).append(_copy_op)
+                p0_idx = len(opd[j + 1]) - 1
+                locator[op.id0][j + 1] = p0_idx
+        else:
+            p0_idx = locator[op.id0][stage]
+
+        if op.opcode in (0, 1):
+            p1_stages = locator[op.id1].keys()
+            if stage not in p1_stages:
+                # Need to copy parent to later states
+                p1_stage = max(p1_stages)
+                p1_idx = locator[op.id1][p1_stage]
+                for j in range(p1_stage, stage):
+                    op1 = ops[op.id1]
+                    latency = float(latency_cutoff * (j + 1))
+                    out_idxd.setdefault(j, []).append(locator[op.id1][j])
+                    _copy_op = Op(len(out_idxd[j]) - 1, -1, -1, 0, op1.qint, latency, 0.0)
+                    opd.setdefault(j + 1, []).append(_copy_op)
+                    p1_idx = len(opd[j + 1]) - 1
+                    locator[op.id1][j + 1] = p1_idx
+            else:
+                p1_idx = locator[op.id1][stage]
+        else:
+            p1_idx = op.id1
+
+        if p1_idx == -1001:
+            # Output to external buffer
+            out_idxd.setdefault(stage, []).append(p0_idx)
+        else:
+            _Op = Op(p0_idx, p1_idx, op.opcode, op.data, op.qint, op.latency, op.cost)
+            opd.setdefault(stage, []).append(_Op)
+            locator.append({stage: len(opd[stage]) - 1})
+    sols = []
+    max_stage = max(opd.keys())
+    for i, stage in enumerate(opd.keys()):
+        _ops = opd[stage]
+        _out_idx = out_idxd[stage]
+        n_in = sum(op.opcode == -1 for op in _ops)
+        n_out = len(_out_idx)
+
+        if i == max_stage:
+            out_shifts = sol.out_shifts
+            out_negs = sol.out_negs
+        else:
+            out_shifts = [0] * len(_out_idx)
+            out_negs = [False] * len(_out_idx)
+
+        _sol = Solution(
+            shape=(n_in, n_out),
+            inp_shift=[0] * n_in,
+            out_idxs=_out_idx,
+            out_shifts=out_shifts,
+            out_negs=out_negs,
+            ops=_ops,
+            carry_size=sol.carry_size,
+            adder_size=sol.adder_size,
+        )
+        sols.append(_sol)
+    csol = CascadedSolution(tuple(sols))
+
+    if retiming:
+        csol = retime_pipeline(csol, verbose=verbose)
+    return csol

da4ml/trace/tracer.py

@@ -0,0 +1,120 @@
+from collections.abc import Sequence
+from decimal import Decimal
+from math import log2
+from typing import overload
+from uuid import UUID
+
+import numpy as np
+
+from ..cmvm.types import Op, QInterval, Solution
+from .fixed_variable import FixedVariable, _const_f
+from .fixed_variable_array import FixedVariableArray
+
+
+def _recursive_trace(v: FixedVariable, gathered: dict[UUID, FixedVariable]):
+    if v in gathered:
+        return
+    assert v._from is not None
+    for _v in v._from:
+        if _v.id not in gathered:
+            _recursive_trace(_v, gathered)
+    gathered[v.id] = v
+
+
+def gather_variables(inputs: Sequence[FixedVariable], outputs: Sequence[FixedVariable]):
+    gathered = {v.id: v for v in inputs}
+    for o in outputs:
+        _recursive_trace(o, gathered)
+
+    variables = list(gathered.values())
+
+    N = len(variables)
+    _index = sorted(list(range(N)), key=lambda i: variables[i].latency * N + i)
+    variables = [variables[i] for i in _index]
+    index = {variables[i].id: i for i in range(N)}
+
+    return variables, index
+
+
+def _comb_trace(inputs: Sequence[FixedVariable], outputs: Sequence[FixedVariable]):
+    variables, index = gather_variables(inputs, outputs)
+    ops: list[Op] = []
+    inp_uuids = {v.id: i for i, v in enumerate(inputs)}
+    for i, v in enumerate(variables):
+        if v.id in inp_uuids and v.opr != 'const':
+            id0 = inp_uuids[v.id]
+            ops.append(Op(id0, -1, -1, 0, v.unscaled.qint, v.latency, v.cost))
+            continue
+        if v.opr == 'new':
+            raise NotImplementedError('Operation "new" is only expected in the input list')
+        match v.opr:
+            case 'vadd':
+                v0, v1 = v._from
+                f0, f1 = v0._factor, v1._factor
+                id0, id1 = index[v0.id], index[v1.id]
+                sub = int(f1 < 0)
+                data = int(log2(abs(f1 / f0)))
+                assert id0 < i and id1 < i, f'{id0} {id1} {i} {v.id}'
+                ops.append(Op(id0, id1, sub, data, v.unscaled.qint, v.latency, v.cost))
+            case 'cadd':
+                v0 = v._from[0]
+                f0 = v0._factor
+                id0 = index[v0.id]
+                assert v._data is not None, 'cadd must have data'
+                qint = v.unscaled.qint
+                data = int(v._data / Decimal(qint.step))
+                assert id0 < i, f'{id0} {i} {v.id}'
+                ops.append(Op(id0, -1, 4, data, qint, v.latency, v.cost))
+            case 'wrap':
+                v0 = v._from[0]
+                id0 = index[v0.id]
+                assert id0 < i, f'{id0} {i} {v.id}'
+                opcode = -3 if v._from[0]._factor < 0 else 3
+                ops.append(Op(id0, -1, opcode, 0, v.unscaled.qint, v.latency, v.cost))
+            case 'relu':
+                v0 = v._from[0]
+                id0 = index[v0.id]
+                assert id0 < i, f'{id0} {i} {v.id}'
+                opcode = -2 if v._from[0]._factor < 0 else 2
+                ops.append(Op(id0, -1, opcode, 0, v.unscaled.qint, v.latency, v.cost))
+            case 'const':
+                qint = v.unscaled.qint
+                assert qint.min == qint.max, f'const {v.id} {qint.min} {qint.max}'
+                f = _const_f(qint.min)
+                step = 2.0**-f
+                qint = QInterval(qint.min, qint.min, step)
+                data = qint.min / step
+                ops.append(Op(-1, -1, 5, int(data), qint, v.latency, v.cost))
+            case _:
+                raise NotImplementedError(f'Operation "{v.opr}" is not supported in tracing')
+    out_index = [index[v.id] for v in outputs]
+    return ops, out_index
+
+
+@overload
+def comb_trace(inputs: Sequence[FixedVariable], outputs: Sequence[FixedVariable]) -> Solution: ...
+
+
+@overload
+def comb_trace(inputs: FixedVariableArray, outputs: FixedVariableArray) -> Solution: ...
+
+
+def comb_trace(inputs, outputs):
+    inputs, outputs = list(np.ravel(inputs)), list(np.ravel(outputs))
+    ops, out_index = _comb_trace(inputs, outputs)
+    shape = len(inputs), len(outputs)
+    inp_shift = [0] * shape[0]
+    out_sf = [v._factor for v in outputs]
+    out_shift = [int(log2(abs(sf))) for sf in out_sf]
+    out_neg = [sf < 0 for sf in out_sf]
+
+    return Solution(
+        shape,
+        inp_shift,
+        out_index,
+        out_shift,
+        out_neg,
+        ops,
+        outputs[0].hwconf.carry_size,
+        outputs[0].hwconf.adder_size,
+    )

da4ml-0.2.0.dist-info/METADATA

@@ -0,0 +1,65 @@
+Metadata-Version: 2.4
+Name: da4ml
+Version: 0.2.0
+Summary: Digital Arithmetic for Machine Learning
+Author-email: Chang Sun <chsun@cern.ch>
+License: GNU Lesser General Public License v3 (LGPLv3)
+Project-URL: repository, https://github.com/calad0i/da4ml
+Keywords: CMVM,distributed arithmetic,hls4ml,MCM,subexpression elimination
+Classifier: Development Status :: 4 - Beta
+Classifier: License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: llvmlite>=0.44
+Requires-Dist: numba>=0.61
+Dynamic: license-file
+
+# da4ml: Distributed Arithmetic for Machine Learning
+
+This project performs Constant Matrix-Vector Multiplication (CMVM) with Distributed Arithmetic (DA) for Machine Learning (ML) on a Field Programmable Gate Arrays (FPGAs).
+
+CMVM optimization is done through greedy CSE of two-term subexpressions, with possible Delay Constraints (DC). The optimization is done in jitted Python (Numba), and a list of optimized operations is generated as traced Python code.
+
+At the moment, the project only generates Vitis HLS C++ code for the FPGA implementation of the optimized CMVM kernel. HDL code generation is planned for the future. Currently, the major use of this repository is through the `distributed_arithmetic` strategy in the [`hls4ml`](https://github.com/fastmachinelearning/hls4ml/) project.
+
+
+## Installation
+
+The project is available on PyPI and can be installed with pip:
+
+```bash
+pip install da4ml
+```
+
+Notice that `numba>=6.0.0` is required for the project to work. The project does not work with `python<3.10`. If the project fails to compile, try upgrading `numba` and `llvmlite` to the latest versions.
+
+## `hls4ml`
+
+The major use of this project is through the `distributed_arithmetic` strategy in the `hls4ml`:
+
+```python
+model_hls = hls4ml.converters.convert_from_keras_model(
+    model,
+    hls_config={
+        'Model': {
+            ...
+            'Strategy': 'distributed_arithmetic',
+        },
+        ...
+    },
+    ...
+)
+```
+
+Currently, `Dense/Conv1D/Conv2D` layers are supported for both `io_parallel` and `io_stream` dataflows. However, notice that distributed arithmetic implies `reuse_factor=1`, as the whole kernel is implemented in combinational logic.
+
+### Notice
+
+Currently, only the `da4ml-v3` branch of `hls4ml` supports the `distributed_arithmetic` strategy. The `da4ml-v3` branch is not yet merged into the `main` branch of `hls4ml`, so you need to install it from the GitHub repository.

da4ml-0.2.0.dist-info/RECORD

@@ -0,0 +1,39 @@
+da4ml/__init__.py,sha256=IETRRvzsJvPMLu1kzzi8UN5FYaM5MhNaXH2A_ZKr2_w,469
+da4ml/_version.py,sha256=iB5DfB5V6YB5Wo4JmvS-txT42QtmGaWcWp3udRT7zCI,511
+da4ml/cmvm/__init__.py,sha256=4Tbt913k9zP0w8R1p6Oss06v5jrManbUhskyHl6e-U0,154
+da4ml/cmvm/api.py,sha256=wwJoEnKtv86q2eyIV2yh_DXIpJPgIW826WYweU8YxBY,9891
+da4ml/cmvm/types.py,sha256=7Z6fSMZU3TcrivWCCWzRuygrSF2SIHydF1jR04-WWZc,17855
+da4ml/cmvm/core/__init__.py,sha256=FRqePEE3SXSSTSjSvByw6QDJkzHCeyvyeihxgxkXVJI,7721
+da4ml/cmvm/core/indexers.py,sha256=QjXgvExS-B2abHTJPDG4NufMdMEflo1i6cUhFOgJpH4,2945
+da4ml/cmvm/core/state_opr.py,sha256=wLqO8qVuM2-qCE5LDeYJDNkUruIPHy63obsv4-x-aR8,8661
+da4ml/cmvm/util/__init__.py,sha256=DkBlUEKA_Gu7n576ja_xZlAQfToWmNL9VXU-jmj6a-g,145
+da4ml/cmvm/util/bit_decompose.py,sha256=SUco70HRYf4r1JU6BXwcgabDrhm_yAmucae5FC67i4I,2216
+da4ml/cmvm/util/mat_decompose.py,sha256=eSJNlXwx_jxgqt5vLJrSLQaeq2ZXu8j9mC4d-eq883M,4094
+da4ml/codegen/__init__.py,sha256=g58EgubgPPoiwRTBduSzm6hAc-poPcK6egdoECfPx9o,329
+da4ml/codegen/cpp/__init__.py,sha256=Tw4XeU_oJsyUkTrsfEPuZ-r0rGAo8E2NX5wn_VTA7NM,90
+da4ml/codegen/cpp/cpp_codegen.py,sha256=gfnDKdmdrubLp5ANeRXvYGTezCMGMet-IQaDlPymaHM,4787
+da4ml/codegen/cpp/source/vitis.h,sha256=ovEefBOfW5-PXuDdRObPGNokGGFHiixDCpPWeTN6aTo,765
+da4ml/codegen/cpp/source/vitis_bridge.h,sha256=XvvGw3A4eAaXKi5jp50bMKUsNfd5iQ-HhUKtsty1uns,567
+da4ml/codegen/verilog/__init__.py,sha256=obRTdtMWhPHsxFHg2ADoPd3iDBEX8nk_6HuCet5EDz0,356
+da4ml/codegen/verilog/comb.py,sha256=EZONCceEvIKHHF8yLY-i2V_U_8THw_dJEQWujjCJ5iI,5592
+da4ml/codegen/verilog/io_wrapper.py,sha256=W7c9_jKRXqrcCm3M13NP7OJdbWtc3FSJgwjZlMEBDL4,7953
+da4ml/codegen/verilog/pipeline.py,sha256=oi8qaobHnTk7gpzpvctuLxKzRIaRJPRFb3BYnnVSSsw,1557
+da4ml/codegen/verilog/verilog_model.py,sha256=ewZuu653ZkhTmf0vRfn6QHYewLoEPdAl4sNwptlV0-M,10585
+da4ml/codegen/verilog/source/build_binder.mk,sha256=rQbI98itE_b1wIQ_0uCXfBzNmGK2XT4vWmRyCJNnPKk,960
+da4ml/codegen/verilog/source/build_prj.tcl,sha256=DMdIDBOF5stkm9Bknk6MiDpIK9-n85h6qwn4B6sV22A,3175
+da4ml/codegen/verilog/source/ioutils.hh,sha256=1o1-oIyQyYc9CU91bBxuitVzzcrNT8p4MTarFKiJoG4,3967
+da4ml/codegen/verilog/source/shift_adder.v,sha256=l2ofym56Y-_PeeY9fwkcZeW9MzrTL_WxvSTvoWERJrU,1885
+da4ml/codegen/verilog/source/template.xdc,sha256=ON8i-TK96Yo6FoZ66WzcVKELajTF5TBmbWFbEilna2U,1142
+da4ml/trace/__init__.py,sha256=1br9bWeFb33t69k6h1XQ50iJhLCqrRuEHtqEawELp-c,230
+da4ml/trace/fixed_variable.py,sha256=KPkYnJgk8bK3W_D6pwCHEndoIgyRn4wyPzTye5Pnh1g,11647
+da4ml/trace/fixed_variable_array.py,sha256=zYuYaXYK-LVcAEUDQ-TxKpR2a_B30RwBIpODBlp2Aq8,6400
+da4ml/trace/pipeline.py,sha256=dYduPBNUeyW2Ws392hZNGJEo0qI5ynpn-iC2n7UVahk,5687
+da4ml/trace/tracer.py,sha256=y0o_KeXCmlUwTLOxwmqZZsZm0oA8-kfu41J-I4-6LXU,4385
+da4ml/trace/ops/__init__.py,sha256=qz0DLPUyxBAu08RCN22kCkJj1EPKanC8ey8NB3_K8co,1640
+da4ml/trace/ops/conv_utils.py,sha256=LtgP3iSZ3fNV6QkEVBzT7ixt-7WTdmBDrFTtQ_9D5aE,3638
+da4ml/trace/ops/einsum_utils.py,sha256=miyMyzJwBLpLTEzXU4vErPE1Xk-ckZG0cjhd13MLAuA,11325
+da4ml-0.2.0.dist-info/licenses/LICENSE,sha256=46mU2C5kSwOnkqkw9XQAJlhBL2JAf1_uCD8lVcXyMRg,7652
+da4ml-0.2.0.dist-info/METADATA,sha256=LdaPbfaPUbMKwLjydibFiFT7qSTs6XuILpnEkN7qEgs,2849
+da4ml-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+da4ml-0.2.0.dist-info/top_level.txt,sha256=N0tnKVwRqFiffFdeAzCgFq71hUNySh5-ITbNd6-R58Q,6
+da4ml-0.2.0.dist-info/RECORD,,

{da4ml-0.1.2.dist-info → da4ml-0.2.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (77.0.3)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

da4ml/cmvm/balanced_reduction.py

@@ -1,46 +0,0 @@
-import numpy as np
-
-from .fixed_variable import FixedVariable
-from .scoring import py_scorer
-
-
-def _balanced_reduction(vars: list[FixedVariable]):
-    vars = vars.copy()
-
-    n = len(vars)
-    if n == 0:
-        return FixedVariable.from_const(0.0)
-    score_mat = np.full((len(vars), len(vars)), -np.inf, dtype=np.float32)
-    for i in range(len(vars)):
-        for j in range(i + 1, len(vars)):
-            score_mat[i, j] = py_scorer(vars[i], vars[j]) - 1000 * (vars[i]._depth + vars[j]._depth)
-
-    while n > 1:
-        idx = np.argmax(score_mat)
-        i, j = np.unravel_index(idx, score_mat.shape)
-        vars[i] = vars[i] + vars[j]
-        vars.pop(j)
-        score_mat[j : n - 1] = score_mat[j + 1 : n]
-        score_mat[:, j : n - 1] = score_mat[:, j + 1 : n]
-        score_mat = score_mat[: n - 1, : n - 1]
-        n -= 1
-        for k in range(n):
-            if k == i:
-                continue
-            if k < i:
-                score_mat[k, i] = py_scorer(vars[k], vars[i]) - 1000 * (vars[k]._depth + vars[i]._depth)
-            else:
-                score_mat[i, k] = py_scorer(vars[i], vars[k]) - 1000 * (vars[i]._depth + vars[k]._depth)
-
-    return vars[0]
-
-
-def balanced_reduction(vars: list[FixedVariable], signed=True):
-    if not signed:
-        return _balanced_reduction(vars)
-    pos_vars = [v for v in vars if v._factor > 0]
-    neg_vars = [v for v in vars if v._factor < 0]
-    for v in neg_vars:
-        v._factor = -v._factor
-
-    return _balanced_reduction(pos_vars) - _balanced_reduction(neg_vars)