Use HGQ's trained softmax lookup tables in the generated HLS

hls4ml/backends/fpga/passes/hgq_proxy_model.py

@@ -34,7 +34,8 @@ def generate_mask_fn(
     masks = []
     to_fixed = to_acfixed if backend.lower() in ['oneapi', 'quartus'] else to_apfixed
     for idx, (k, b, i) in enumerate(zip(Ks, Bs, Is)):
-        if b == 0:
+        if b <= 0:
+            # b <= 0 means nothing is representable: the channel is a constant zero
             fn = f'out[{idx}] = 0;'
         else:
             fn = f'out[{idx}] = {to_fixed(k, b, i, RND, SAT)}(inp[{idx}]);'

hls4ml/backends/vivado/passes/core_templates.py

@@ -250,6 +250,9 @@ def format(self, node):
 param_activ_function_template = (
     'nnet::{activation}<{input_t}, {param_t.name}, {output_t}, {config}>({input}, {param}, {output});'
 )
+softmax_lut_function_template = (
+    'nnet::{activation}<{input_t}, {output_t}, {config}>({input}, {output}, {exp_table}, {inv_table});'
+)
 
 activ_include_list = ['nnet_utils/nnet_activation.h', 'nnet_utils/nnet_activation_stream.h']
 
@@ -349,6 +352,14 @@ def format(self, node):
         params['activation'] = 'softmax' if not use_multidim else 'softmax_multidim'
         params['config'] = '{}_config{}'.format(node.get_attr('activation'), node.index)
 
+        # HGQ2 softmax carries its trained tables as weights; every other softmax lets the
+        # HLS code build them at runtime.
+        if 'exp_table' in node.weights and 'inv_table' in node.weights:
+            params['activation'] += '_lut'
+            params['exp_table'] = node.get_weights('exp_table').name
+            params['inv_table'] = node.get_weights('inv_table').name
+            return softmax_lut_function_template.format(**params)
+
         return self.template.format(**params)
 
 

hls4ml/backends/vivado/passes/hgq_softmax_tables.py

@@ -0,0 +1,68 @@
+from warnings import warn
+
+from hls4ml.model.layers import Layer, Softmax
+from hls4ml.model.optimizer import OptimizerPass
+from hls4ml.model.types import FixedPrecisionType
+
+
+class MaterializeSoftmaxTables(OptimizerPass):
+    """Turn HGQ's trained softmax exp/reciprocal lookup tables into weight arrays.
+
+    Without this the generated C++ rebuilds both tables at runtime from std::exp and 1/x
+    (nnet::init_exp_table / nnet::init_invert_table), which does not reproduce HGQ's own
+    quantized values. QSoftmaxHandler stashes the builders as _exp_table_fn / _inv_table_fn;
+    both take the (k, i, f) of the type the C++ addresses the table with, which is known only
+    here: for the latency implementation it is the softmax input precision, decided by
+    bit_exact. This pass therefore runs after bit_exact and before transform_types.
+    """
+
+    def match(self, node: Layer):
+        # Keyed on the stash rather than on the layer type: only the HGQ frontend leaves it,
+        # and it must always be removed again, being a callable no saved model can hold.
+        return isinstance(node, Softmax) and '_exp_table_fn' in node.attributes
+
+    def transform(self, model, node: Layer):
+        exp_table_fn = node.attributes.pop('_exp_table_fn')
+        inv_table_fn = node.attributes.pop('_inv_table_fn')
+
+        impl = node.get_attr('implementation')
+        if impl not in ('latency', 'stable'):
+            return False  # argmax and legacy use no lookup table
+
+        # The exp table is indexed by the normalized input (x_max - x) for the stable
+        # implementation and by the layer input itself for the latency one. Reading the
+        # element precision off the variable is only unambiguous because this runs before
+        # transform_types wraps io_stream variables.
+        if impl == 'stable':
+            exp_inp_t: FixedPrecisionType = node.attributes['inp_norm_t'].precision
+        else:
+            exp_inp_t = node.get_input_variable().type.precision
+        inv_inp_t: FixedPrecisionType = node.attributes['inv_inp_t'].precision
+
+        for name, addr_t, fn in (('exp_table', exp_inp_t, exp_table_fn), ('inv_table', inv_inp_t, inv_table_fn)):
+            # softmax_idx_from_real_val slices the top ceillog2(table_size) bits of the address
+            # word. Sizing the table as 2**width makes that the whole word, so no low bits are
+            # dropped and the slice cannot run off the end of a narrower word.
+            size = 1 << int(addr_t.width)
+            if int(node.get_attr(f'{name}_size') or 0) != size:
+                warn(
+                    f'{node.name}: {name}_size {node.get_attr(f"{name}_size")} does not match the {addr_t.width}-bit '
+                    f'address type {addr_t}; overriding to {size} to keep table indexing bit-exact.',
+                    stacklevel=1,
+                )
+            node.set_attr(f'{name}_size', size)
+            k = int(bool(addr_t.signed))
+            data = fn((k, addr_t.integer - k, int(addr_t.width) - addr_t.integer))
+
+            # Capture the type bit_exact derived first: storing a WeightVariable under `name`
+            # makes AttributeDict overwrite `{name}_t` with the variable's own type.
+            named_t = node.attributes[f'{name}_t']
+            node.set_attr(f'{name}_data', data)
+            node.add_weights_variable(
+                name=name, var_name=name + '{index}', precision=named_t.precision, type_name=named_t.name
+            )
+            node.get_weights(name).type = named_t
+            node.attributes[f'{name}_t'] = named_t
+            model.config.layer_name_precision[f'{node.name}_{name}'] = str(named_t.precision)
+
+        return False

hls4ml/backends/vivado/vivado_backend.py

@@ -171,6 +171,7 @@ def _register_flows(self):
             'vivado:inplace_stream_flatten',
             'vivado:skip_softmax',
             'vivado:fix_softmax_table_size',
+            'vivado:materialize_softmax_tables',
             'infer_precision_types',
             'vivado:distributed_arithmetic_codegen',
             'vivado:distributed_arithmetic_einsum_codegen',

hls4ml/converters/keras_v3/hgq2/_base.py

@@ -1,6 +1,7 @@
 from collections.abc import Sequence
 from math import prod
 from typing import TYPE_CHECKING, Any
+from warnings import warn
 
 import numpy as np
 
@@ -39,6 +40,17 @@ def extract_fixed_quantizer_config(q, tensor: 'KerasTensor', is_input: bool) ->
         B = np.ravel(B).astype(np.int16)
         I = np.ravel(I).astype(np.int16)  # noqa: E741
 
+    # A channel can be trained down to a negative total width. Clamp to 0, not 1: B == 0 is the
+    # constant-zero encoding generate_mask_fn understands, while B == 1 would bring the channel
+    # back as a live 1-bit one. I is legitimately negative and already zeroed for these above.
+    if np.any(B < 0):
+        warn(
+            f'Quantizer {q.name} has {int(np.sum(B < 0))} channel(s) with total bitwidth < 0; '
+            'treating them as constant zero.',
+            stacklevel=2,
+        )
+    B = np.maximum(B, 0)
+
     overflow_mode: str = internal_q.overflow_mode
     round_mode: str = internal_q.round_mode
     if round_mode.startswith('S_'):

hls4ml/converters/keras_v3/hgq2/softmax.py

@@ -1,10 +1,12 @@
 import typing
 from collections.abc import Sequence
+from functools import partial
 from math import prod
 
 from hls4ml.model.types import FixedPrecisionType, RoundingMode, SaturationMode
 
 from ._base import QLayerHandler
+from .unary_lut import extract_lut_table
 
 if typing.TYPE_CHECKING:
     import hgq
@@ -118,6 +120,16 @@ def handle(
                 'parallelization_factor': parallelization_factor,
                 'class_name': class_name,
                 '_bit_exact': True,
+                # Not materialized here: for the latency implementation the domain is the
+                # softmax input precision, only final after bit_exact. materialize_softmax_tables
+                # calls these with the (k, i, f) of the address type, and pops them again -- they
+                # are not JSON-serializable, so they must not reach a saved model.
+                '_exp_table_fn': partial(
+                    extract_lut_table, layer.exp_table.activation, layer.exp_table.oq if layer.exp_table.enable_oq else None
+                ),
+                '_inv_table_fn': partial(
+                    extract_lut_table, layer.inv_table.activation, layer.inv_table.oq if layer.inv_table.enable_oq else None
+                ),
             }
         )
 

hls4ml/converters/keras_v3/hgq2/unary_lut.py

@@ -1,5 +1,5 @@
 import typing
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 
 import numpy as np
 from quantizers import get_fixed_quantizer_np
@@ -10,11 +10,76 @@
 
 if typing.TYPE_CHECKING:
     import hgq
+    from hgq.quantizer import Quantizer
+    from hgq.quantizer.internal import FixedPointQuantizerBase
     from keras import KerasTensor
 
 from decimal import Decimal
 
 
+def fixed_grid_by_bit_pattern(k: int, i: int, f: int) -> np.ndarray:
+    """All values a fixed-point type can hold, indexed by their two's-complement bit pattern.
+
+    This is the addressing convention of ``nnet::get_index_unary_lut`` and
+    ``nnet::softmax_real_val_from_idx`` in the generated C++.
+    """
+    K, I, F = Decimal(int(k)), Decimal(int(i)), Decimal(int(f))  # noqa: E741
+    _eps = Decimal(2) ** -F
+    _min = -K * Decimal(2) ** I
+    _max = Decimal(2) ** I - _eps
+    N = (_max - _min) / _eps + 1
+    assert float(N).is_integer(), 'Invalid quantizer range'
+    N = int(N)
+    assert N <= 1e6, 'Too large quantizer range'
+    assert np.log2(N).is_integer(), f'Invalid quantizer range: N must be power of 2, got {N}'
+
+    grid = np.linspace(float(_min), float(_max), N, dtype=np.float32)
+    if k:
+        # idx by binary repr, move the positive part to the front
+        grid = np.concatenate([grid[N // 2 :], grid[: N // 2]])
+    return grid
+
+
+def kif_of(q: 'FixedPointQuantizerBase') -> tuple[int, int, int]:
+    """(k, i, f) of a homogeneous quantizer, with entries representing nothing masked out."""
+    from keras import ops
+
+    k, i, f = q.kif
+    mask = k + i + f > 0
+    i, f = np.where(mask, i, -32), np.where(mask, f, -32)  # type: ignore
+    return int(ops.max(k)), int(ops.max(i)), int(ops.max(f))  # type: ignore
+
+
+def extract_lut_table(activation: Callable, oq: 'Quantizer|None', kif: tuple[int, int, int]) -> np.ndarray:
+    """Tabulate activation over the fixed-point type kif describes, quantized by oq.
+
+    kif is the type the generated C++ addresses the table with, passed in rather than read off
+    a layer: QSoftmax.exp_table has no input quantizer when stable=False, and for the latency
+    softmax the domain is the softmax input precision, only final after the bit_exact pass.
+    """
+    from hgq.quantizer.internal import FixedPointQuantizerBase
+    from keras import ops
+
+    grid = fixed_grid_by_bit_pattern(*kif)
+    table = activation(grid)
+
+    if oq is not None:
+        internal_q = oq.quantizer
+        if not isinstance(internal_q, FixedPointQuantizerBase):
+            raise NotImplementedError('FloatPointQuantizer is not supported yet')
+
+        # Not oq(table): the Quantizer layer broadcasts against the shape it was built for,
+        # which is not the rank-1 grid here. Homogeneous, so this is the same operation.
+        round_mode = internal_q.round_mode
+        if round_mode.startswith('S_'):
+            round_mode = round_mode[2:]
+        fixed_q = get_fixed_quantizer_np(round_mode, internal_q.overflow_mode)
+        k, i, f = (ops.convert_to_numpy(x).ravel().item() for x in internal_q.kif)
+        table = fixed_q(table, k, i, f)  # type: ignore
+
+    return np.asarray(ops.convert_to_numpy(table))
+
+
 class QUnaryLUTHandler(QLayerHandler, KerasV3LayerHandler):
     handles = ('hgq.layers.activation.QUnaryFunctionLUT',)
 
@@ -24,61 +89,28 @@ def handle(
         in_tensors: Sequence['KerasTensor'],
         out_tensors: Sequence['KerasTensor'],
     ):
-        from hgq.quantizer.internal import FixedPointQuantizerBase, FloatPointQuantizer
+        from hgq.quantizer.internal import FixedPointQuantizerBase
         from keras import ops
 
         if not layer.enable_iq and not layer.enable_oq:
             raise ValueError('Currently only support input_quantizer enabled UnaryFunctionLUT layer')
         assert not layer._allow_heterogeneous_table, 'Heterogeneous table is not supported in QUnaryFunctionLUT layer'
 
         iq = layer.iq.quantizer
-        if isinstance(iq, FixedPointQuantizerBase):
-            k, i, f = iq.kif
-            mask = k + i + f > 0
-            i, f = np.where(mask, i, -32), np.where(mask, f, -32)  # type: ignore
-            k, i, f = (Decimal(int(ops.max(x))) for x in (k, i, f))  # type: ignore
-            _min = -k * 2**i
-            _eps = 2**-f
-            _max = 2**i - _eps
-            N = (_max - _min) / _eps + 1
-            assert float(N).is_integer(), 'Invalid quantizer range'
-            N = int(N)
-            assert N <= 1e6, 'Too large quantizer range'
-            assert np.log2(N).is_integer(), f'Invalid quantizer range: N must be power of 2, got {N}'
-
-            all_inputs = np.linspace(float(_min), float(_max), N, dtype=np.float32)
-
-            config = {}
-            config.update(self.default_config)
-            table = layer.activation(all_inputs)
-            if layer.enable_oq:
-                table = layer.oq(table[None, ...])[0]
-            table = ops.convert_to_numpy(table)
-            if k:
-                # idx by binary repr, move the positive part to the front
-                table_pos, table_neg = table[N // 2 :], table[: N // 2]
-                table = np.concatenate([table_pos, table_neg])
-        else:
+        if not isinstance(iq, FixedPointQuantizerBase):
             raise NotImplementedError('FloatPointQuantizer is not supported yet')
 
+        table = extract_lut_table(layer.activation, layer.oq if layer.enable_oq else None, kif_of(iq))
+
         oq = layer.oq.quantizer
-        if isinstance(oq, FixedPointQuantizerBase):
-            round_mode = oq.round_mode
-            if round_mode.startswith('S_'):
-                round_mode = round_mode[2:]
-            overflow_mode = oq.overflow_mode
-            fixed_q = get_fixed_quantizer_np(round_mode, overflow_mode)
-            k, i, f = (ops.convert_to_numpy(x).ravel().item() for x in oq.kif)
-            table = fixed_q(table, k, i, f)  # type: ignore
-
-            k, b, I = bool(k), k + i + f, k + i  # noqa: E741
-            table_t = FixedPrecisionType(b, I, k)
-        else:
-            assert isinstance(oq, FloatPointQuantizer)
+        if not isinstance(oq, FixedPointQuantizerBase):
             raise NotImplementedError('FloatPointQuantizer is not supported yet')
+        k, i, f = (ops.convert_to_numpy(x).ravel().item() for x in oq.kif)
+        k, b, I = bool(k), k + i + f, k + i  # noqa: E741
+        table_t = FixedPrecisionType(b, I, k)
 
-        table = ops.convert_to_numpy(table)
-
+        config = {}
+        config.update(self.default_config)
         config.update(
             {
                 'class_name': 'UnaryLUT',