da4ml 0.3.0.post1__py3-none-any.whl → 0.3.2__py3-none-any.whl

da4ml/converter/hgq2/replica.py

@@ -7,14 +7,21 @@ import hgq
 import keras
 import numpy as np
 from hgq.layers import (
+    QAdd,
     QBatchNormalization,
     QBatchNormDense,
     QConv1D,
     QConv2D,
     QConv3D,
     QDense,
+    QDot,
+    QEinsum,
     QEinsumDense,
     QEinsumDenseBatchnorm,
+    QMaximum,
+    QMeanPow2,
+    QMinimum,
+    QSubtract,
     QSum,
 )
 from hgq.layers.core.base import MultipleQuantizers, Quantizer
@@ -23,10 +30,19 @@ from keras.layers import ReLU
 from keras.src.layers.pooling.base_global_pooling import BaseGlobalPooling
 from keras.src.layers.pooling.base_pooling import BasePooling
 from keras.src.ops.numpy import (
+    Abs,
+    Absolute,
     Add,
     Concatenate,
     Divide,
+    Dot,
+    Einsum,
     GetItem,
+    Matmul,
+    Max,
+    Maximum,
+    Min,
+    Minimum,
     Moveaxis,
     Multiply,
     Ravel,
@@ -49,13 +65,13 @@ def mirror_quantizer(q: Quantizer, v: FixedVariableArray) -> FixedVariableArray:
     return quantize(v, k, i, f, overflow_mode=overflow_mode, round_mode=round_mode)
 
 
-_registry: dict[type, 'type[MirrorOperationBase]'] = {}
+_registry: dict[type, 'type[ReplayOperationBase]'] = {}
 
 
-class MirrorOperationMeta(type):
+class ReplayOperationMeta(type):
     def __new__(mcs, name: str, bases: tuple[type, ...], namespace: dict[str, typing.Any]):
         cls = super().__new__(mcs, name, bases, namespace)
-        if name == 'MirrorOperationBase':
+        if name == 'ReplayOperationBase':
             return cls
 
         handles: type | tuple[type, ...] = namespace['handles']
@@ -67,7 +83,7 @@ class MirrorOperationMeta(type):
         return cls
 
 
-class MirrorOperationBase(metaclass=MirrorOperationMeta):
+class ReplayOperationBase(metaclass=ReplayOperationMeta):
     handles: tuple[type, ...] = ()
 
     def __init__(self, layer: 'keras.Operation'):
@@ -124,7 +140,7 @@ class MirrorOperationBase(metaclass=MirrorOperationMeta):
         return outputs
 
 
-class MirrorQuantizer(MirrorOperationBase):
+class ReplayQuantizer(ReplayOperationBase):
     handles = (Quantizer,)
 
     def __init__(self, op: 'Quantizer'):
@@ -135,8 +151,8 @@ class MirrorQuantizer(MirrorOperationBase):
         return mirror_quantizer(self.op, inputs)
 
 
-class MirrorQDense(MirrorOperationBase):
-    handles = (QDense, QEinsumDense, QEinsumDenseBatchnorm, QBatchNormDense, QBatchNormalization, keras.layers.EinsumDense)
+class ReplayQDense(ReplayOperationBase):
+    handles = (QDense, QEinsumDense, QEinsumDenseBatchnorm, QBatchNormDense, keras.layers.EinsumDense)
 
     def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
         op = self.op
@@ -152,14 +168,6 @@ class MirrorQDense(MirrorOperationBase):
             qkernel = op.kernel
             qbias = op.bias
             eq = op.equation
-        elif isinstance(op, QBatchNormalization):
-            qkernel, qbias = op.qscaler_and_qoffset
-            dim = inputs._vars.ndim
-            axis = op.axis
-            assert axis != 0, 'Cannot normalizing on batch axis'
-            axis -= 1
-            idx = ''.join(chr(ord('a') + i) for i in range(dim))
-            eq = f'...{idx},{idx[axis]}->...{idx}'
         else:
             raise TypeError(f'Unsupported layer type: {type(op)}')
 
@@ -168,7 +176,28 @@ class MirrorQDense(MirrorOperationBase):
         return (einsum(eq, inputs[None], qkernel) + qbias)[0]
 
 
-class MirrorQConv(MirrorOperationBase):
+class ReplayQDot(ReplayOperationBase):
+    handles = (QDot, keras.layers.Dot)
+
+    def call(self, inputs: tuple[FixedVariableArray, FixedVariableArray]) -> FixedVariableArray:
+        layer: QDot | keras.layers.Dot = self.op
+        assert not layer.normalize, 'normalize is not supported in mirror operation'
+
+        axes = layer.axes
+        return np.dot(inputs[0][None], inputs[1][None], axes=axes)[0]  # type: ignore
+
+
+class ReplayQBatchNormalization(ReplayOperationBase):
+    handles = (QBatchNormalization,)
+
+    def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
+        layer: QBatchNormalization = self.op
+        scale, bias = map(np.array, layer.qscaler_and_qoffset)
+        shape = layer._shape
+        return inputs * scale.reshape(shape) + bias.reshape(shape)
+
+
+class ReplayQConv(ReplayOperationBase):
     handles = (QConv1D, QConv2D, QConv3D)
 
     def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
@@ -190,14 +219,14 @@ class MirrorQConv(MirrorOperationBase):
         return outputs
 
 
-class MirrorReLU(MirrorOperationBase):
+class ReplayReLU(ReplayOperationBase):
     handles = (ReLU,)
 
     def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
         return relu(inputs)
 
 
-class MirrorReshape(MirrorOperationBase):
+class ReplayReshape(ReplayOperationBase):
     handles = (keras.layers.Reshape, keras.layers.Flatten, Reshape, Ravel)
 
     def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
@@ -211,8 +240,8 @@ class MirrorReshape(MirrorOperationBase):
             raise TypeError(f'Unsupported layer type: {type(self.op)}')
 
 
-class MirrorMerge(MirrorOperationBase):
-    handles = (keras.layers.Add, keras.layers.Concatenate, hgq.layers.QAdd)
+class ReplayMerge(ReplayOperationBase):
+    handles = (keras.layers.Add, keras.layers.Concatenate, QAdd)
 
     def call(self, inputs: tuple[FixedVariableArray, FixedVariableArray]) -> FixedVariableArray:
         op: keras.Operation = self.op
@@ -226,7 +255,7 @@ class MirrorMerge(MirrorOperationBase):
             raise TypeError(f'Unsupported layer type: {type(op)}')
 
 
-class MirrorPool(MirrorOperationBase):
+class ReplayPool(ReplayOperationBase):
     handles = (
         hgq.layers.QAvgPool1D,
         hgq.layers.QAvgPool2D,
@@ -295,7 +324,7 @@ class MirrorPool(MirrorOperationBase):
         return out  # type: ignore
 
 
-class MirrorRepeatVector(MirrorOperationBase):
+class ReplayRepeatVector(ReplayOperationBase):
     handles = (keras.layers.RepeatVector,)
 
     def call(self, inputs: FixedVariableArray) -> FixedVariableArray:
@@ -306,7 +335,7 @@ class MirrorRepeatVector(MirrorOperationBase):
         return np.repeat(inputs[None], layer.n, axis=0)[0]  # type: ignore
 
 
-class MirrorGetItem(MirrorOperationBase):
+class ReplayGetItem(ReplayOperationBase):
     handles = (GetItem,)
 
     def call(self, x: FixedVariableArray, key):
@@ -315,15 +344,21 @@ class MirrorGetItem(MirrorOperationBase):
         return x[None][key][0]
 
 
-class MirrorSum(MirrorOperationBase):
-    handles = (Sum,)
+class ReplayReduction(ReplayOperationBase):
+    handles = (Sum, Max, Min)
 
     def call(self, x: FixedVariableArray, axis=None, keepdims=False):
-        return np.sum(x[None], axis=axis, keepdims=keepdims)[0]  # type: ignore
+        if isinstance(self.op, Sum):
+            op = np.sum
+        elif isinstance(self.op, Max):
+            op = np.amax
+        elif isinstance(self.op, Min):
+            op = np.amin
+        return op(x[None], axis=axis, keepdims=keepdims)[0]  # type: ignore
 
 
-class MirrorQSum(MirrorOperationBase):
-    handles = (QSum,)
+class ReplayQReduction(ReplayOperationBase):
+    handles = (QSum, QMeanPow2)
 
     def call(self, x: FixedVariableArray):
         layer: QSum = self.op
@@ -331,11 +366,14 @@ class MirrorQSum(MirrorOperationBase):
         return np.sum(x[None], axis=axes, keepdims=keepdims)[0] * scale  # type: ignore
 
 
-class MirrorArithmetic(MirrorOperationBase):
-    handles = (Add, Subtract, Multiply, TrueDivide, Divide)
+class ReplayArithmetic(ReplayOperationBase):
+    handles = (Add, Subtract, Multiply, TrueDivide, Divide, QSubtract, QMaximum, QMinimum, Maximum, Minimum)
 
     def call(self, x1: FixedVariableArray, x2: FixedVariableArray):
-        match self.op.__class__.__name__:
+        name = self.op.__class__.__name__
+        if name.startswith('Q'):
+            name = name[1:]
+        match name:
             case 'Add':
                 return x1 + x2
             case 'Subtract':
@@ -344,11 +382,15 @@ class MirrorArithmetic(MirrorOperationBase):
                 return x1 * x2
             case 'TrueDivide' | 'Divide':
                 return x1 / x2
+            case 'Maximum':
+                return np.maximum(x1, x2)  # type: ignore
+            case 'Minimum':
+                return np.minimum(x1, x2)  # type: ignore
             case _:
                 raise TypeError(f'Unsupported arithmetic operation: {type(self.op)}')
 
 
-class MirrorConcatenate(MirrorOperationBase):
+class ReplayConcatenate(ReplayOperationBase):
     handles = (Concatenate,)
 
     def call(self, xs: Sequence[FixedVariableArray]):
@@ -358,7 +400,7 @@ class MirrorConcatenate(MirrorOperationBase):
         return np.concatenate([x[None] for x in xs], axis=axis)[0]  # type: ignore
 
 
-class MirrorRepeat(MirrorOperationBase):
+class ReplayRepeat(ReplayOperationBase):
     handles = (Repeat,)
 
     def call(self, x: FixedVariableArray):
@@ -367,7 +409,7 @@ class MirrorRepeat(MirrorOperationBase):
         return np.repeat(x[None], repeats, axis=axis)[0]  # type: ignore
 
 
-class MirrorTranspose(MirrorOperationBase):
+class ReplayTranspose(ReplayOperationBase):
     handles = (Transpose,)
 
     def call(self, x: FixedVariableArray):
@@ -375,9 +417,57 @@ class MirrorTranspose(MirrorOperationBase):
         return np.transpose(x, axes)  # type: ignore
 
 
-class MirrorMoveaxis(MirrorOperationBase):
+class ReplayMoveaxis(ReplayOperationBase):
     handles = (Moveaxis,)
 
     def call(self, x: FixedVariableArray):
         source, destination = self.op.source, self.op.destination
         return np.moveaxis(x[None], source, destination)[0]  # type: ignore
+
+
+noop_layers = []
+for k, v in keras.layers.__dict__.items():
+    name = k.lower()
+    if 'dropout' in name or 'random' in name or 'noise' in name:
+        noop_layers.append(v)
+
+
+class ReplayNoOp(ReplayOperationBase):
+    handles = tuple(noop_layers)
+
+    def call(self, x: FixedVariableArray, training=False) -> FixedVariableArray:
+        assert not training, 'Training mode is not supported in mirror operation'
+        return x
+
+
+class ReplayQEinsum(ReplayOperationBase):
+    handles = (QEinsum,)
+
+    def call(self, inputs: tuple[FixedVariableArray, ...]) -> FixedVariableArray:
+        layer: QEinsum = self.op
+        eq = layer.equation
+        return einsum(eq, *inputs)
+
+
+class ReplayEinsum(ReplayOperationBase):
+    handles = (Einsum,)
+
+    def call(self, *operands: FixedVariableArray) -> FixedVariableArray:
+        layer: Einsum = self.op
+        eq = layer.subscripts
+        operands = [operand[None] for operand in operands]  # type: ignore
+        return einsum(eq, *operands)[0]
+
+
+class ReplayMatmul(ReplayOperationBase):
+    handles = (Matmul, Dot)
+
+    def call(self, x1: FixedVariableArray, x2: FixedVariableArray) -> FixedVariableArray:
+        return x1 @ x2
+
+
+class ReplayAbs(ReplayOperationBase):
+    handles = (Absolute, Abs)
+
+    def call(self, x: FixedVariableArray) -> FixedVariableArray:
+        return np.abs(x)  # type: ignore

da4ml/trace/fixed_variable.py

@@ -1,11 +1,15 @@
+import random
+from collections.abc import Generator
 from decimal import Decimal
 from math import ceil, floor, log2
 from typing import NamedTuple
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from ..cmvm.core import cost_add
 from ..cmvm.types import QInterval
 
+rd = random.SystemRandom()
+
 
 class HWConfig(NamedTuple):
     adder_size: int
@@ -26,6 +30,23 @@ def _const_f(const: float | Decimal):
     return _high
 
 
+def to_csd_powers(x: float) -> Generator[float, None, None]:
+    if x == 0:
+        return
+    f = _const_f(abs(x))
+    x = x * 2**f
+    s = 2**-f
+    N = ceil(log2(abs(x) * 1.5 + 1e-19))
+    for n in range(N - 1, -1, -1):
+        _2pn = 2**n
+        thres = _2pn / 1.5
+        bit = int(x > thres) - int(x < -thres)
+        v = _2pn * bit
+        x -= v
+        if v != 0:
+            yield v * s
+
+
 class FixedVariable:
     def __init__(
         self,
@@ -43,13 +64,14 @@ class FixedVariable:
     ) -> None:
         assert low <= high, f'low {low} must be less than high {high}'
 
-        if low == high and opr != 'new':
+        if low != high and opr == 'const':
+            raise ValueError('Constant variable must have low == high')
+
+        if low == high:
             opr = 'const'
-            _factor = _factor
             _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
@@ -58,7 +80,7 @@ class FixedVariable:
         opr = opr
         self.opr = opr
         self._data = _data
-        self.id = _id or uuid4()
+        self.id = _id or UUID(int=rd.getrandbits(128), version=4)
         self.hwconf = hwconf
 
         if opr == 'cadd':
@@ -81,7 +103,7 @@ class FixedVariable:
     def get_cost_and_latency(self):
         if self.opr == 'const':
             return 0.0, 0.0
-        if self.opr in ('vadd', 'cadd', 'min', 'max'):
+        if self.opr in ('vadd', 'cadd', 'min', 'max', 'vmul'):
             adder_size = self.hwconf.adder_size
             carry_size = self.hwconf.carry_size
             latency_cutoff = self.hwconf.latency_cutoff
@@ -92,13 +114,25 @@ class FixedVariable:
                 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:
+            elif self.opr == 'cadd':
                 assert len(self._from) == 1
                 assert self._data is not None, 'cadd must have data'
                 _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
+            elif self.opr == 'vmul':
+                assert len(self._from) == 2
+                v0, v1 = self._from
+                b0, b1 = sum(v0.kif), sum(v1.kif)
+                int0, int1 = v0.qint, v1.qint
+                dlat0, _cost0 = cost_add(int0, int0, 0, False, adder_size, carry_size)
+                dlat1, _cost1 = cost_add(int1, int1, 0, False, adder_size, carry_size)
+                dlat = max(dlat0 * b1, dlat1 * b0)
+                _cost = min(_cost0 * b1, _cost1 * b0)
+                base_latency = max(v0.latency, v1.latency)
+            else:
+                raise NotImplementedError(f'Operation {self.opr} is unknown')
 
             _latency = dlat + base_latency
             if latency_cutoff > 0 and ceil(_latency / latency_cutoff) > ceil(base_latency / latency_cutoff):
@@ -107,6 +141,7 @@ class FixedVariable:
                     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
@@ -154,6 +189,7 @@ class FixedVariable:
         return f'({self._factor}) FixedVariable({self.low}, {self.high}, {self.step})'
 
     def __neg__(self):
+        opr = self.opr if self.low != self.high else 'const'
         return FixedVariable(
             -self.high,
             -self.low,
@@ -162,7 +198,7 @@ class FixedVariable:
             _factor=-self._factor,
             latency=self.latency,
             cost=self.cost,
-            opr=self.opr,
+            opr=opr,
             _id=self.id,
             _data=self._data,
             hwconf=self.hwconf,
@@ -195,7 +231,7 @@ class FixedVariable:
             hwconf=self.hwconf,
         )
 
-    def _const_add(self, other: float | Decimal | None):
+    def _const_add(self, other: float | Decimal | None) -> 'FixedVariable':
         if other is None:
             return self
         if not isinstance(other, (int, float, Decimal)):
@@ -229,29 +265,66 @@ class FixedVariable:
     def __sub__(self, other: 'FixedVariable|int|float|Decimal'):
         return self + (-other)
 
-    def __mul__(
-        self,
-        other: 'float|Decimal',
-    ):
+    def __mul__(self, other: 'FixedVariable|int|float|Decimal') -> 'FixedVariable':
         if other == 0:
             return FixedVariable(0, 0, 1, hwconf=self.hwconf, opr='const')
 
-        assert log2(abs(other)) % 1 == 0, 'Only support pow2 multiplication'
+        if isinstance(other, FixedVariable):
+            return self._var_mul(other)
+
+        if log2(abs(other)) % 1 == 0:
+            return self._pow2_mul(other)
+
+        variables = [self._pow2_mul(v) for v in to_csd_powers(float(other))]
+        while len(variables) > 1:
+            v = variables.pop() + variables.pop()
+            variables.append(v)
+        return variables[0]
+
+    def _var_mul(self, other: 'FixedVariable') -> 'FixedVariable':
+        if other is not self:
+            a, b, c, d = self.high * other.low, self.low * other.high, self.high * other.high, self.low * other.low
+            low = min(a, b, c, d)
+            high = max(a, b, c, d)
+        else:
+            a, b = self.low * other.low, self.high * other.high
+            if self.low < 0 and self.high > 0:
+                low = min(a, b, 0)
+                high = max(a, b, 0)
+            else:
+                low = min(a, b)
+                high = max(a, b)
 
+        step = self.step * other.step
+        _factor = self._factor * other._factor
+        opr = 'vmul'
+        return FixedVariable(
+            low,
+            high,
+            step,
+            _from=(self, other),
+            _factor=_factor,
+            opr=opr,
+        )
+
+    def _pow2_mul(
+        self,
+        other: float | Decimal,
+    ):
         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
-
+        opr = self.opr
         return FixedVariable(
             low,
             high,
             step,
             _from=self._from,
             _factor=_factor,
-            opr=self.opr,
+            opr=opr,
             latency=self.latency,
             cost=self.cost,
             _id=self.id,
@@ -268,6 +341,21 @@ class FixedVariable:
     def __rmul__(self, other: 'float|Decimal|int|FixedVariable'):
         return self * other
 
+    def __pow__(self, other):
+        _power = int(other)
+        assert _power == other, 'Power must be an integer'
+        assert _power >= 0, 'Power must be non-negative'
+        if _power == 0:
+            return FixedVariable(1, 1, 1, hwconf=self.hwconf, opr='const')
+        if _power == 1:
+            return self
+
+        pow0 = _power // 2
+        ret = (self**pow0) * (self ** (_power - pow0))
+        if other % 2 == 0:
+            ret.low = max(ret.low, 0)
+        return ret
+
     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')
@@ -317,18 +405,21 @@ class FixedVariable:
         round_mode: str = 'TRN',
     ):
         overflow_mode, round_mode = overflow_mode.upper(), round_mode.upper()
-        assert overflow_mode in ('WRAP', 'SAT', 'SAT_SM')
+        assert overflow_mode in ('WRAP', 'SAT', 'SAT_SYM')
         assert round_mode in ('TRN', 'RND')
 
+        if k + i + f <= 0:
+            return FixedVariable(0, 0, 1, hwconf=self.hwconf, opr='const')
         _k, _i, _f = self.kif
 
         if k >= _k and i >= _i and f >= _f:
-            return self
+            if overflow_mode != 'SAT_SYM' or i > _i:
+                return self
 
         if f < _f and round_mode == 'RND':
             return (self + 2.0 ** (-f - 1)).quantize(k, i, f, overflow_mode, 'TRN')
 
-        if overflow_mode in ('SAT', 'SAT_SM'):
+        if overflow_mode in ('SAT', 'SAT_SYM'):
             step = Decimal(2) ** -f
             _high = Decimal(2) ** i
             high = _high - step
@@ -458,21 +549,43 @@ class FixedVariableInput(FixedVariable):
         self._from: tuple[FixedVariable, ...] = ()
         self.opr = 'new'
         self._data = None
-        self.id = uuid4()
+        self.id = UUID(int=rd.getrandbits(128), version=4)
         self.hwconf = hwconf
 
         self.latency = latency if latency is not None else 0.0
         self.cost = 0.0
 
     def __add__(self, other):
+        if other == 0:
+            return self
         raise ValueError('Cannot operate on unquantized input variable')
 
     def __sub__(self, other):
+        if other == 0:
+            return self
         raise ValueError('Cannot operate on unquantized input variable')
 
     def __neg__(self):
         raise ValueError('Cannot negate unquantized input variable')
 
+    def __mul__(self, other):
+        if other == 1:
+            return self
+        raise ValueError('Cannot multiply unquantized input variable')
+
+    def __rmul__(self, other):
+        if other == 1:
+            return self
+        raise ValueError('Cannot multiply unquantized input variable')
+
+    def __radd__(self, other):
+        if other == 0:
+            return self
+        raise ValueError('Cannot add unquantized input variable')
+
+    def __rsub__(self, other):
+        raise ValueError('Cannot subtract unquantized input variable')
+
     def relu(self, *args, **kwargs):
         raise ValueError('Cannot apply relu on unquantized input variable')