da4ml 0.5.0__cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl

da4ml/trace/ops/einsum_utils.py

@@ -0,0 +1,305 @@
+from math import prod
+from typing import TYPE_CHECKING, TypedDict, overload
+
+import numpy as np
+from numpy.typing import NDArray
+
+if TYPE_CHECKING:
+    from ..fixed_variable_array import FixedVariableArray
+
+
+class EinsumRecipe(TypedDict):
+    direct_sum_axis: tuple[tuple[int, ...], tuple[int, ...]]
+    in_transpose_idxs: tuple[tuple[int, ...], tuple[int, ...]]
+    L0: int
+    L1: int
+    I: int
+    C: int
+    out_interpert_shape: tuple[int, ...]
+    out_transpose_idxs: tuple[int, ...]
+
+
+def _validate_einsum_expr(fn: str, shape0: tuple[int, ...], shape1: tuple[int, ...]):
+    """Validate, resolve broadcasting, and compute output shape for einsum string
+
+    Parameters
+    ----------
+    fn : str
+        einsum string, e.g. 'ij,jk->ik'
+    shape0 : tuple[int,...]
+        shape of input0
+    shape1 : tuple[int,...]
+        shape of input1
+
+    Returns
+    -------
+    tuple[str, tuple[int,...]]
+        einsum string w/o broadcasting, and output shape
+
+    Raises
+    ------
+    ValueError
+        If the einsum string is invalid, or if it is incompatible with the input shapes
+    """
+    inp, out = map(str.strip, fn.split('->'))
+    in0, in1 = map(str.strip, inp.split(','))
+    alphabets = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
+    s_alphabets = set(alphabets)
+
+    # Invalid characters
+    if not (s_alphabets >= set(in0.replace('...', '') + in1.replace('...', '') + out.replace('...', ''))):
+        raise ValueError(f"einsum string {fn} is invalid: subscripts should be in [a-zA-Z] and '...' only")
+
+    in0 = in0.replace('...', '0')
+    in1 = in1.replace('...', '0')
+    out = out.replace('...', '0')
+    ax_in0, ax_in1, ax_out = list(in0), list(in1), list(out)
+    sax_in0, sax_in1, sax_out = set(ax_in0), set(ax_in1), set(ax_out)
+    free_indices = ''.join(sorted(s_alphabets - sax_in0 - sax_in1 - sax_out))
+
+    # Repeated indices
+    if len(sax_in0) != len(ax_in0):
+        for a in in0:
+            if in0.count(a) == 1:
+                continue
+            a = a if a != '0' else '...'
+            raise ValueError(f"einsum string {fn} is invalid: input0 subscripts includes '{a}' multiple times")
+    if len(sax_in1) != len(ax_in1):
+        for a in in1:
+            if in1.count(a) == 1:
+                continue
+            a = a if a != '0' else '...'
+            raise ValueError(f"einsum string {fn} is invalid: input1 subscripts includes '{a}' multiple times")
+    if len(sax_out) != len(ax_out):
+        for a in out:
+            if out.count(a) == 1:
+                continue
+            a = a if a != '0' else '...'
+            raise ValueError(f"einsum string {fn} is invalid: output subscripts includes '{a}' multiple times")
+
+    # Invalid broadcasting
+    if '0' in sax_in0 or '0' in sax_in1 or '0' in sax_out:
+        if '0' not in sax_out:
+            raise ValueError(f'einsum string {fn} is invalid: output does not allow broadcasting, but inputs do')
+        if '0' not in sax_in0 and '0' not in sax_in1:
+            raise ValueError(f'einsum string {fn} is invalid: output allows broadcasting, but inputs do not')
+
+    # Output index out of nowhere
+    if remaining := sax_out - sax_in0 - sax_in1:
+        raise ValueError(f'einsum string {fn} is invalid: output subscripts {remaining} not found in inputs')
+
+    _common_in = sax_in0 & sax_in1
+
+    if '0' in sax_in0 and '0' in sax_in1:
+        # Simultaneous axes expansion in both inputs
+        n_boardcast0 = len(shape0) - len(sax_in0) + 1
+        n_boardcast1 = len(shape1) - len(sax_in1) + 1
+        assert n_boardcast0 == n_boardcast1, f'... expands to {n_boardcast0} and {n_boardcast1}-axis in input0 and input1.'
+        # Replace expansion indices with free indices
+        in0 = in0.replace('0', free_indices[:n_boardcast0])
+        in1 = in1.replace('0', free_indices[:n_boardcast1])
+        out = out.replace('0', free_indices[:n_boardcast0])
+        ax_in0, ax_in1, ax_out = list(in0), list(in1), list(out)
+        _common_in = set(ax_in0) & set(ax_in1)
+
+    else:
+        # Axes expansion in input0 or input1 only
+        if '0' in sax_in0:
+            if len(sax_in0) - 1 > len(shape0):
+                raise ValueError(f'Input0 requires at least {len(sax_in0) - 1} dimensions, but only {len(shape0)} given')
+            # Replace auto expansion indices with free indices
+            n_broadcast = len(shape0) - len(sax_in0) + 1
+            in0 = in0.replace('0', free_indices[:n_broadcast])
+            out = out.replace('0', free_indices[:n_broadcast])
+            ax_in0 = list(in0)
+            ax_out = list(out)
+        else:
+            if len(sax_in0) != len(shape0):
+                raise ValueError(f'Input0 requires {len(sax_in0)} dimensions, but {len(shape0)} is given')
+
+        if '0' in sax_in1:
+            if len(sax_in1) - 1 > len(shape1):
+                raise ValueError(f'Input1 requires at least {len(sax_in1) - 1} dimensions, but only {len(shape1)} given')
+            # Replace expansion indices with free indices
+            n_broadcast = len(shape1) - len(sax_in1) + 1
+            in1 = in1.replace('0', free_indices[:n_broadcast])
+            out = out.replace('0', free_indices[:n_broadcast])
+            ax_in1 = list(in1)
+            ax_out = list(out)
+        else:
+            if len(sax_in1) != len(shape1):
+                raise ValueError(f'Input1 requires {len(sax_in1)} dimensions, but {len(shape1)} is given')
+
+    # Input dimension mismatch
+    for a in _common_in:
+        ax_0 = ax_in0.index(a)
+        ax_1 = ax_in1.index(a)
+        if shape0[ax_0] != shape1[ax_1]:
+            raise ValueError(f"Input dimension size mismatches for common subscript '{a}': {shape0[ax_0]} and {shape1[ax_1]}")
+
+    out_shape = tuple(shape0[ax_in0.index(a)] if a in ax_in0 else shape1[ax_in1.index(a)] for a in ax_out)
+    return f'{in0},{in1}->{out}', out_shape
+
+
+def parse_einsum(fn: str, input_shape0: tuple[int, ...], input_shape1: tuple[int, ...]) -> EinsumRecipe:
+    """Parse einsum operation on two input arrays, return a recipe for execution
+
+    Parameters
+    ----------
+    fn : str
+        einsum string, e.g. 'ij,jk->ik'
+    input : np.ndarray
+        input0, the first input array
+    input1 : np.ndarray
+        input1, the second input array
+
+    Returns
+    -------
+    EinsumRecipe
+        einsum recipe; executed by _exec_einsum
+    """
+
+    fn, _ = _validate_einsum_expr(fn, input_shape0, input_shape1)
+
+    _in, _out = fn.split('->')
+    _in0, _in1 = _in.split(',')
+
+    in0, in1, out = list(_in0), list(_in1), list(_out)
+    s_in0, s_in1, s_out = set(in0), set(in1), set(out)
+    _common = s_in0 & s_in1
+    _contract = _common - s_out
+    _inplace = _common & s_out
+    contract = sorted(_contract, key=lambda x: in1.index(x))
+    inplace = sorted(_inplace, key=lambda x: in1.index(x))
+    invariant0 = sorted((s_out - _common) & s_in0, key=lambda x: in0.index(x))
+    invariant1 = sorted((s_out - _common) & s_in1, key=lambda x: in1.index(x))
+    direct_sum0 = s_in0 - s_out - _common
+    direct_sum1 = s_in1 - s_out - _common
+    direct_sum_axis = (
+        tuple(sorted(in0.index(x) for x in direct_sum0)),
+        tuple(sorted(in1.index(x) for x in direct_sum1)),
+    )
+
+    contract_idxs = tuple(map(in0.index, contract)), tuple(map(in1.index, contract))
+    inplace_idxs = tuple(map(in0.index, inplace)), tuple(map(in1.index, inplace))
+    invariant_idxs = tuple(map(in0.index, invariant0)), tuple(map(in1.index, invariant1))
+
+    inplace_shape = tuple(input_shape0[i] for i in inplace_idxs[0])
+    inplace_size = prod(inplace_shape)
+    contract_size = prod(input_shape0[i] for i in contract_idxs[0])
+    invariant_shape0 = tuple(input_shape0[i] for i in invariant_idxs[0])
+    invariant_shape1 = tuple(input_shape1[i] for i in invariant_idxs[1])
+    invariant_size0, invariant_size1 = prod(invariant_shape0), prod(invariant_shape1)
+
+    transpose_idx0 = inplace_idxs[0] + invariant_idxs[0] + contract_idxs[0]
+    transpose_idx1 = inplace_idxs[1] + invariant_idxs[1] + contract_idxs[1]
+
+    out_shape_pretranspose = inplace_shape + invariant_shape0 + invariant_shape1
+    _out_transpose_idx = np.argsort(tuple(map(out.index, inplace + invariant0 + invariant1)))
+    out_transpose_idx = tuple(int(i) for i in _out_transpose_idx)
+
+    return EinsumRecipe(
+        direct_sum_axis=direct_sum_axis,
+        in_transpose_idxs=(transpose_idx0, transpose_idx1),
+        out_interpert_shape=out_shape_pretranspose,
+        out_transpose_idxs=out_transpose_idx,
+        L0=invariant_size0,
+        L1=invariant_size1,
+        I=inplace_size,
+        C=contract_size,
+    )
+
+
+def _exec_einsum(recipe: EinsumRecipe, input0: np.ndarray, input1: np.ndarray) -> np.ndarray:
+    """Execute einsum operation on two input arrays
+
+    Parameters
+    ----------
+    recipe : EinsumRecipe
+        einsum recipe
+    input0 : np.ndarray
+        input0, the first input array
+    input1 : np.ndarray
+        input1, the second input array
+
+    Returns
+    -------
+    np.ndarray
+        output array
+    """
+    sum_axis0, sum_axis1 = recipe['direct_sum_axis']
+    if sum_axis0:
+        input0 = np.sum(input0, axis=sum_axis0)
+    if sum_axis1:
+        input1 = np.sum(input1, axis=sum_axis1)
+    input0 = input0.transpose(recipe['in_transpose_idxs'][0]).ravel()
+    input1 = input1.transpose(recipe['in_transpose_idxs'][1]).ravel()
+    out_dtype = object if input0.dtype == object or input1.dtype == object else np.float64
+    output = np.zeros(recipe['L0'] * recipe['L1'] * recipe['I'], dtype=out_dtype)
+
+    L0, L1, I, C = recipe['L0'], recipe['L1'], recipe['I'], recipe['C']
+
+    for l0 in range(L0):
+        for i in range(I):
+            A = input1[i * L1 * C : (i + 1) * L1 * C].reshape((L1, C))
+            B = input0[(i * L0 + l0) * C : (i * L0 + l0 + 1) * C]
+            output[(i * L0 + l0) * L1 : (i * L0 + l0 + 1) * L1] = A @ B
+
+    return output.reshape(recipe['out_interpert_shape']).transpose(recipe['out_transpose_idxs'])
+
+
+def _einsum(fn: str, input0, input1) -> np.ndarray:
+    """Execute einsum operation on two input arrays.
+
+    WARNING: Order of multiplication is reversed -- watchout if you are using non-commutative operators
+
+    Parameters
+    ----------
+    fn : str
+        einsum string, e.g. 'ij,jk->ik'
+    input : np.ndarray
+        input0, the first input array
+    input1 : np.ndarray
+        input1, the second input array
+
+    Returns
+    -------
+    np.ndarray
+        output array
+    """
+    recipe = parse_einsum(fn, input0.shape, input1.shape)
+    return _exec_einsum(recipe, input0, input1)
+
+
+@overload
+def einsum(fn: str, input0: 'FixedVariableArray', input1: 'FixedVariableArray') -> 'FixedVariableArray': ...
+
+
+@overload
+def einsum(fn: str, input0: 'FixedVariableArray', input1: NDArray[np.integer | np.floating]) -> 'FixedVariableArray': ...
+
+
+@overload
+def einsum(fn: str, input0: NDArray[np.integer | np.floating], input1: 'FixedVariableArray') -> 'FixedVariableArray': ...
+
+
+@overload
+def einsum(
+    fn: str, input0: NDArray[np.integer | np.floating], input1: NDArray[np.integer | np.floating]
+) -> NDArray[np.integer | np.floating]: ...
+
+
+def einsum(fn: str, input0, input1):
+    from ..fixed_variable_array import FixedVariableArray
+
+    fg0 = isinstance(input0, FixedVariableArray)
+    fg1 = isinstance(input1, FixedVariableArray)
+
+    r = _einsum(fn, input0, input1)
+
+    if fg0:
+        return FixedVariableArray(r, input0.solver_options)
+    elif fg1:
+        return FixedVariableArray(r, input1.solver_options)
+    else:
+        return r

da4ml/trace/ops/quantization.py

@@ -0,0 +1,74 @@
+from typing import TYPE_CHECKING, TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+from quantizers.fixed_point.fixed_point_ops_np import get_fixed_quantizer_np
+
+from ..fixed_variable_array import FixedVariable
+
+if TYPE_CHECKING:
+    from ..fixed_variable_array import FixedVariableArray
+
+T = TypeVar('T', 'FixedVariableArray', NDArray[np.floating], list[FixedVariable])
+
+
+def relu(x: T, i: NDArray[np.integer] | None = None, f: NDArray[np.integer] | None = None, round_mode: str = 'TRN') -> T:
+    from ..fixed_variable_array import FixedVariableArray
+
+    if isinstance(x, FixedVariableArray):
+        return x.relu(i=i, f=f, round_mode=round_mode)
+    elif isinstance(x, list):
+        return [xx.relu(i=ii, f=ff, round_mode=round_mode) for xx, ii, ff in zip(x, i, f)]  # type: ignore
+    else:
+        round_mode = round_mode.upper()
+        assert round_mode in ('TRN', 'RND')
+        x = np.maximum(x, 0)
+        if f is not None:
+            if round_mode == 'RND':
+                x += 2.0 ** (-f - 1)
+            sf = 2.0**f
+            x = np.floor(x * sf) / sf
+        if i is not None:
+            x = x % 2.0**i
+        return x
+
+
+def _quantize(
+    x: NDArray[np.floating],
+    k: NDArray[np.integer] | np.integer | int,
+    i: NDArray[np.integer] | np.integer | int,
+    f: NDArray[np.integer] | np.integer | int,
+    overflow_mode: str = 'WRAP',
+    round_mode: str = 'TRN',
+) -> NDArray[np.floating]:
+    q = get_fixed_quantizer_np(round_mode=round_mode, overflow_mode=overflow_mode)
+    return np.where(k + i + f <= 0, 0, q(x, k=k, i=i, f=f))  # type: ignore
+
+
+def quantize(
+    x: T,
+    k: NDArray[np.integer] | np.integer | int,
+    i: NDArray[np.integer] | np.integer | int,
+    f: NDArray[np.integer] | np.integer | int,
+    overflow_mode: str = 'WRAP',
+    round_mode: str = 'TRN',
+) -> T:
+    from ..fixed_variable_array import FixedVariableArray
+
+    if isinstance(x, (FixedVariableArray, FixedVariable)):
+        return x.quantize(k=k, i=i, f=f, overflow_mode=overflow_mode, round_mode=round_mode)
+    elif isinstance(x, list):
+        ret: list[FixedVariable] = []
+        for i in range(len(x)):
+            ret.append(
+                x[i].quantize(
+                    k=int(k[i] if isinstance(k, (list, np.ndarray)) else k),
+                    i=int(i[i] if isinstance(i, (list, np.ndarray)) else i),
+                    f=int(f[i] if isinstance(f, (list, np.ndarray)) else f),
+                    overflow_mode=overflow_mode,
+                    round_mode=round_mode,
+                )
+            )
+        return ret  # type: ignore
+    else:
+        return _quantize(x, k, i, f, overflow_mode, round_mode)

da4ml/trace/ops/reduce_utils.py

@@ -0,0 +1,105 @@
+import heapq
+import typing
+from collections.abc import Callable, Sequence
+from math import prod
+from typing import TypeVar
+
+import numpy as np
+from numpy.typing import NDArray
+
+if typing.TYPE_CHECKING:
+    from ..fixed_variable import FixedVariable
+    from ..fixed_variable_array import FixedVariableArray
+
+
+T = typing.TypeVar('T', 'FixedVariable', float, np.floating)
+TA = TypeVar('TA', 'FixedVariableArray', NDArray[np.integer | np.floating])
+
+
+class Packet:
+    def __init__(self, v):
+        self.value = v
+
+    def __gt__(self, other: 'Packet') -> bool:  # type: ignore
+        from ..fixed_variable_array import FixedVariable
+
+        a, b = self.value, other.value
+
+        if isinstance(a, FixedVariable):
+            if isinstance(b, FixedVariable):
+                if b.latency > a.latency:
+                    return False
+                if b.latency < a.latency:
+                    return True
+                if b._factor > 0 and a._factor < 0:
+                    return False
+                if b._factor < 0 and a._factor > 0:
+                    return True
+                return sum(a.kif[:2]) > sum(b.kif[:2])
+            return True
+
+        return False
+
+    def __lt__(self, other: 'Packet') -> bool:  # type: ignore
+        return not self.__gt__(other)
+
+
+def _reduce(operator: Callable[[T, T], T], arr: Sequence[T]) -> T:
+    from ..fixed_variable_array import FixedVariable
+
+    if isinstance(arr, np.ndarray):
+        arr = list(arr.ravel())
+    assert len(arr) > 0, 'Array must not be empty'
+    if len(arr) == 1:
+        return arr[0]
+    dtype = arr[0].__class__
+    if not issubclass(dtype, FixedVariable):
+        r = operator(arr[0], arr[1])
+        for i in range(2, len(arr)):
+            r = operator(r, arr[i])
+        return r
+
+    heap = [Packet(v) for v in arr]  # type: ignore
+    heapq.heapify(heap)
+    while len(heap) > 1:
+        v1 = heapq.heappop(heap).value
+        v2 = heapq.heappop(heap).value
+        v = operator(v1, v2)
+        heapq.heappush(heap, Packet(v))  # type: ignore
+    return heap[0].value
+
+
+def reduce(operator: Callable[[T, T], T], x: TA, axis: int | Sequence[int] | None = None, keepdims: bool = False) -> TA:
+    """
+    Reduce the array by summing over the specified axis.
+    """
+    from ..fixed_variable_array import FixedVariableArray
+
+    if isinstance(x, FixedVariableArray):
+        solver_config = x.solver_options
+        arr = x._vars
+    else:
+        solver_config = None
+        arr = x
+    all_axis = tuple(range(arr.ndim))
+    axis = axis if axis is not None else all_axis
+    axis = (axis,) if isinstance(axis, int) else tuple(axis)
+    axis = tuple(a if a >= 0 else a + arr.ndim for a in axis)
+
+    xpose_axis = sorted(all_axis, key=lambda a: (a in axis) * 1000 + a)
+    if keepdims:
+        target_shape = tuple(d if ax not in axis else 1 for ax, d in enumerate(arr.shape))
+    else:
+        target_shape = tuple(d for ax, d in enumerate(arr.shape) if ax not in axis)
+
+    dim_contract = prod(arr.shape[a] for a in axis)
+    arr = np.transpose(arr, xpose_axis)  # type: ignore
+    _arr = arr.reshape(-1, dim_contract)
+    _arr = np.array([_reduce(operator, _arr[i]) for i in range(_arr.shape[0])])
+    r = _arr.reshape(target_shape)  # type: ignore
+
+    if isinstance(x, FixedVariableArray):
+        r = FixedVariableArray(r, solver_config)
+        if r.shape == ():
+            return r._vars.item()  # type: ignore
+    return r if r.shape != () or keepdims else r.item()  # type: ignore

da4ml/trace/pipeline.py

@@ -0,0 +1,181 @@
+from math import ceil, floor
+
+from ..cmvm.types import CombLogic, Op, Pipeline
+from .fixed_variable import FixedVariable, HWConfig
+from .tracer import comb_trace
+
+
+def retime_pipeline(csol: Pipeline, 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 _get_new_idx(
+    idx: int,
+    locator: list[dict[int, int]],
+    opd: dict[int, list[Op]],
+    out_idxd: dict[int, list[int]],
+    ops: list[Op],
+    stage: int,
+    latency_cutoff: float,
+):
+    if idx < 0:
+        return idx
+    p0_stages = locator[idx].keys()
+    if stage not in p0_stages:
+        # Need to copy parent to later states
+        p0_stage = max(p0_stages)
+        p0_idx = locator[idx][p0_stage]
+        for j in range(p0_stage, stage):
+            op0 = ops[idx]
+            latency = float(latency_cutoff * (j + 1))
+            out_idxd.setdefault(j, []).append(locator[idx][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[idx][j + 1] = p0_idx
+    else:
+        p0_idx = locator[idx][stage]
+    return p0_idx
+
+
+def to_pipeline(comb: CombLogic, latency_cutoff: float, retiming=True, verbose=True) -> Pipeline:
+    """Split the record into multiple stages based on the latency of the operations.
+    Only useful for HDL generation.
+
+    Parameters
+    ----------
+    sol : CombLogic
+        The combinational logic to be pipelined into multiple stages.
+    latency_cutoff : float
+        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(comb.ops) > 0, 'No operations in the record'
+    for i, op in enumerate(comb.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 = comb.ops.copy()
+    lat = max(ops[i].latency for i in comb.out_idxs)
+    for i in comb.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_idx = _get_new_idx(op.id0, locator, opd, out_idxd, ops, stage, latency_cutoff)
+        p1_idx = _get_new_idx(op.id1, locator, opd, out_idxd, ops, stage, latency_cutoff)
+        if op.opcode in (6, -6):
+            k = op.data & 0xFFFFFFFF
+            _shift = (op.data >> 32) & 0xFFFFFFFF
+            k = _get_new_idx(k, locator, opd, out_idxd, ops, stage, latency_cutoff)
+            data = _shift << 32 | k
+        else:
+            data = op.data
+
+        if p1_idx == -1001:
+            # Output to external buffer
+            out_idxd.setdefault(stage, []).append(p0_idx)
+        else:
+            _Op = Op(p0_idx, p1_idx, op.opcode, 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())
+    n_in = comb.shape[0]
+    for i, stage in enumerate(opd.keys()):
+        _ops = opd[stage]
+        _out_idx = out_idxd[stage]
+        n_out = len(_out_idx)
+
+        if i == max_stage:
+            out_shifts = comb.out_shifts
+            out_negs = comb.out_negs
+        else:
+            out_shifts = [0] * len(_out_idx)
+            out_negs = [False] * len(_out_idx)
+
+        if comb.lookup_tables is not None:
+            _ops, lookup_tables = remap_table_idxs(comb, _ops)
+        else:
+            lookup_tables = None
+        _sol = CombLogic(
+            shape=(n_in, n_out),
+            inp_shifts=[0] * n_in,
+            out_idxs=_out_idx,
+            out_shifts=out_shifts,
+            out_negs=out_negs,
+            ops=_ops,
+            carry_size=comb.carry_size,
+            adder_size=comb.adder_size,
+            lookup_tables=lookup_tables,
+        )
+        sols.append(_sol)
+
+        n_in = n_out
+    csol = Pipeline(tuple(sols))
+
+    if retiming:
+        csol = retime_pipeline(csol, verbose=verbose)
+    return csol
+
+
+def remap_table_idxs(comb: CombLogic, _ops):
+    assert comb.lookup_tables is not None
+    table_idxs = sorted(list({op.data for op in _ops if op.opcode == 8}))
+    remap = {j: i for i, j in enumerate(table_idxs)}
+    _ops_remap = []
+    for op in _ops:
+        if op.opcode == 8:
+            op = Op(op.id0, op.id1, op.opcode, remap[op.data], op.qint, op.latency, op.cost)
+        _ops_remap.append(op)
+    _ops = _ops_remap
+    lookup_tables = tuple(comb.lookup_tables[i] for i in table_idxs)
+    return _ops, lookup_tables