feat: add sparse CNN support (sparsepixels)

hls4ml/backends/vivado/vivado_backend.py

@@ -163,6 +163,8 @@ def _register_flows(self):
         quantization_flow = register_flow('quantization', quantization_passes, requires=[init_flow], backend=self.name)
 
         optimization_passes = [
+            'vivado:sparse_graph_optimizer',
+            'vivado:sparse_fix_input_precision',
             'vivado:remove_final_reshape',
             'vivado:optimize_pointwise_conv',
             'vivado:inplace_parallel_reshape',

hls4ml/converters/keras_v3/__init__.py

@@ -8,6 +8,7 @@
     pquant,  # noqa: F401
     qkeras,  # noqa: F401
     recurrent,  # noqa: F401
+    sparsepixels,  # noqa: F401
 )
 from ._base import registry as layer_handlers
 

hls4ml/converters/keras_v3/sparsepixels.py

@@ -0,0 +1,215 @@
+import math
+import typing
+from collections.abc import Sequence
+from typing import Any
+
+import numpy as np
+
+from ._base import KerasV3LayerHandler
+
+if typing.TYPE_CHECKING:
+    import keras
+    from keras import KerasTensor
+
+_sparse_context: dict[str, Any] = {}
+
+
+def _extract_sparse_iq_config(conv_layer, in_tensor_name: str, n_sparse: int, n_chan: int) -> dict[str, Any]:
+    """Extract input quantizer config from QConv2D, adapted for sparse tensor shape."""
+    from keras import ops
+
+    internal_q = conv_layer._iq.quantizer
+    kif_k, kif_i, kif_f = internal_q.kif
+    kif_k = np.ravel(ops.convert_to_numpy(kif_k)).astype(np.int16)
+    kif_i = np.ravel(ops.convert_to_numpy(kif_i)).astype(np.int16)
+    kif_f = np.ravel(ops.convert_to_numpy(kif_f)).astype(np.int16)
+
+    # HGQ quantizers may be per-element (H*W*C); reduce to per-channel
+    # Take max of each component independently to get the envelope type
+    if kif_k.size > n_chan:
+        kif_k = np.max(kif_k.reshape(-1, n_chan), axis=0)
+        kif_i = np.max(kif_i.reshape(-1, n_chan), axis=0)
+        kif_f = np.max(kif_f.reshape(-1, n_chan), axis=0)
+
+    # Reconstruct KBI from KIF: B = k + i + f, I_bits = k + i
+    k = kif_k
+    B = kif_k + kif_i + kif_f
+    I_bits = kif_k + kif_i
+
+    if k.size > 1:
+        k = np.tile(k, n_sparse).reshape(1, -1)
+        B = np.tile(B, n_sparse).reshape(1, -1)
+        I_bits = np.tile(I_bits, n_sparse).reshape(1, -1)
+
+    overflow_mode: str = internal_q.overflow_mode
+    round_mode: str = internal_q.round_mode
+    if round_mode.startswith('S_'):
+        round_mode = round_mode[2:]
+
+    return {
+        'name': conv_layer._iq.name,
+        'class_name': 'FixedPointQuantizer',
+        'mask_kbi': (k, B, I_bits),
+        'SAT': overflow_mode,
+        'RND': round_mode,
+        'fusible': None,
+        'input_keras_tensor_names': [in_tensor_name],
+        'output_keras_tensor_names': [f'{in_tensor_name}_q'],
+        'overrides': {},
+    }
+
+
+class InputReduceHandler(KerasV3LayerHandler):
+    handles = ('sparsepixels.layers.InputReduce',)
+
+    def handle(
+        self,
+        layer: 'keras.Layer',
+        in_tensors: Sequence['KerasTensor'],
+        out_tensors: Sequence['KerasTensor'],
+    ):
+        in_shape: tuple[int, ...] = in_tensors[0].shape[1:]  # type: ignore
+        in_height, in_width, n_chan = in_shape
+
+        n_sparse = layer.n_max_pixels
+        threshold = float(layer.threshold) if layer.threshold is not None else 0.0
+
+        # Clear any stale state from a previous conversion in the same Python process
+        _sparse_context.clear()
+        _sparse_context['n_sparse'] = n_sparse
+
+        # Hash stores 1-based H and W coordinates separately (see nnet_sparsepixels.h::sparse_input_reduce).
+        # Spatial dims only shrink through the network (pooling), so input H/W bound the required bits.
+        max_dim = max(int(in_height), int(in_width))
+        hash_bits = max(1, math.ceil(math.log2(max_dim + 1)))
+
+        return {
+            'class_name': 'SparseInputReduce',
+            'in_height': int(in_height),
+            'in_width': int(in_width),
+            'n_chan': int(n_chan),
+            'n_sparse': n_sparse,
+            'threshold': threshold,
+            'hash_bits': hash_bits,
+        }
+
+
+class QConv2DSparseHandler(KerasV3LayerHandler):
+    handles = ('sparsepixels.layers.QConv2DSparse',)
+
+    def handle(
+        self,
+        layer: 'keras.Layer',
+        in_tensors: Sequence['KerasTensor'],
+        out_tensors: Sequence['KerasTensor'],
+    ):
+        import keras
+        from keras import ops
+
+        conv = layer.conv
+        n_chan = int(conv.kernel.shape[2])
+        n_filt = int(conv.filters)
+        kernel_size = int(conv.kernel_size[0])
+        n_sparse = _sparse_context.get('n_sparse', 0)
+
+        if hasattr(conv, 'qkernel'):
+            weight_data = ops.convert_to_numpy(conv.qkernel)
+        else:
+            weight_data = ops.convert_to_numpy(conv.kernel)
+
+        bias_data = None
+        if layer._use_bias and hasattr(layer, 'sparse_bias'):
+            if hasattr(layer, '_bq'):
+                bias_data = ops.convert_to_numpy(layer._bq(layer.sparse_bias))
+            else:
+                bias_data = ops.convert_to_numpy(layer.sparse_bias)
+
+        name = layer.name
+        in_tensor_names = [t.name for t in in_tensors]
+        out_tensor_names = [t.name for t in out_tensors]
+
+        iq_conf = None
+        has_iq = hasattr(conv, '_iq') and hasattr(conv, '_enable_iq') and conv._enable_iq
+        if has_iq:
+            iq_conf = _extract_sparse_iq_config(conv, in_tensors[0].name, n_sparse, n_chan)
+            in_tensor_names = [f'{in_tensors[0].name}_q']
+
+        config: dict[str, Any] = {
+            'class_name': 'SparseConv2D',
+            'name': name,
+            'n_sparse': n_sparse,
+            'n_chan': n_chan,
+            'n_filt': n_filt,
+            'kernel_size': kernel_size,
+            'weight_data': weight_data,
+            'bias_data': bias_data,
+            'input_keras_tensor_names': in_tensor_names,
+            'output_keras_tensor_names': out_tensor_names,
+        }
+
+        activation = layer._activation
+        results: list[dict[str, Any]] = []
+        if iq_conf is not None:
+            results.append(iq_conf)
+
+        if activation not in (None, keras.activations.linear):
+            act_name = activation.__name__
+            intermediate = f'{out_tensors[0].name}_sparse_act'
+
+            config['output_keras_tensor_names'] = [intermediate]
+
+            act_config: dict[str, Any] = {
+                'class_name': 'SparseActivation',
+                'name': f'{name}_{act_name}',
+                'activation': act_name,
+                'n_sparse': n_sparse,
+                'n_chan': n_filt,
+                'input_keras_tensor_names': [intermediate],
+                'output_keras_tensor_names': out_tensor_names,
+            }
+            results.extend([config, act_config])
+            return tuple(results)
+
+        results.append(config)
+        return tuple(results)
+
+
+def _sparse_pooling_config(
+    in_tensors: Sequence['KerasTensor'], out_tensors: Sequence['KerasTensor'], pool_size: int, pool_op: str
+) -> dict[str, Any]:
+    """Shared config for the average/max sparse pooling handlers (differ only in pool_op)."""
+    feat_shape: tuple[int, ...] = in_tensors[0].shape[1:]  # type: ignore
+    n_chan = int(feat_shape[-1])
+    n_sparse = _sparse_context.get('n_sparse', 0)
+
+    return {
+        'class_name': 'SparsePooling2D',
+        'n_sparse': n_sparse,
+        'n_chan': n_chan,
+        'pool_size': pool_size,
+        'pool_op': pool_op,
+    }
+
+
+class AveragePooling2DSparseHandler(KerasV3LayerHandler):
+    handles = ('sparsepixels.layers.AveragePooling2DSparse',)
+
+    def handle(
+        self,
+        layer: 'keras.Layer',
+        in_tensors: Sequence['KerasTensor'],
+        out_tensors: Sequence['KerasTensor'],
+    ):
+        return _sparse_pooling_config(in_tensors, out_tensors, int(layer.avg_pool.pool_size[0]), 'avg')
+
+
+class MaxPooling2DSparseHandler(KerasV3LayerHandler):
+    handles = ('sparsepixels.layers.MaxPooling2DSparse',)
+
+    def handle(
+        self,
+        layer: 'keras.Layer',
+        in_tensors: Sequence['KerasTensor'],
+        out_tensors: Sequence['KerasTensor'],
+    ):
+        return _sparse_pooling_config(in_tensors, out_tensors, int(layer.max_pool.pool_size[0]), 'max')

hls4ml/model/layers.py

@@ -1891,6 +1891,93 @@ def initialize(self):
         self.add_output_variable(shape)
 
 
+class SparseInputReduce(Layer):
+    _expected_attributes = [
+        Attribute('in_height'),
+        Attribute('in_width'),
+        Attribute('n_chan'),
+        Attribute('n_sparse'),
+        Attribute('threshold', value_type=float),
+        Attribute('hash_bits', value_type=int, default=10),
+    ]
+
+    def initialize(self):
+        shape = [self.attributes['n_sparse'] * self.attributes['n_chan']]
+        self.add_output_variable(shape)
+
+
+class SparseConv2D(Layer):
+    _expected_attributes = [
+        Attribute('n_sparse'),
+        Attribute('n_chan'),
+        Attribute('n_filt'),
+        Attribute('kernel_size'),
+        WeightAttribute('weight'),
+        WeightAttribute('bias'),
+        TypeAttribute('weight'),
+        TypeAttribute('bias'),
+        TypeAttribute('accum'),
+    ]
+
+    def initialize(self):
+        shape = [self.attributes['n_sparse'] * self.attributes['n_filt']]
+        self.add_output_variable(shape)
+        self.add_weights(quantizer=self.get_attr('weight_quantizer'))
+        self.add_bias(quantizer=self.get_attr('bias_quantizer'))
+
+    def add_bias(self, quantizer=None):
+        data = self.get_attr('bias_data', None)
+        precision = None
+        type_name = None
+        if data is None:
+            data = np.zeros(self.attributes['n_filt'])
+            precision = IntegerPrecisionType(width=1, signed=False)
+            type_name = 'bias{index}_t'
+            quantizer = None
+        self.add_weights_variable(
+            name='bias', var_name='b{index}', type_name=type_name, precision=precision, data=data, quantizer=quantizer
+        )
+
+
+class SparseActivation(Layer):
+    _expected_attributes = [
+        Attribute('n_sparse'),
+        Attribute('n_chan'),
+        Attribute('activation', value_type=str),
+    ]
+
+    def initialize(self):
+        shape = [self.attributes['n_sparse'] * self.attributes['n_chan']]
+        self.add_output_variable(shape)
+
+
+class SparsePooling2D(Layer):
+    _expected_attributes = [
+        Attribute('n_sparse'),
+        Attribute('n_chan'),
+        Attribute('pool_size'),
+        Attribute('pool_op', value_type=str, default='avg'),  # 'avg' or 'max'
+        TypeAttribute('accum'),
+    ]
+
+    def initialize(self):
+        shape = [self.attributes['n_sparse'] * self.attributes['n_chan']]
+        self.add_output_variable(shape)
+
+
+class SparseFlatten(Layer):
+    _expected_attributes = [
+        Attribute('n_sparse'),
+        Attribute('n_chan'),
+        Attribute('out_height'),
+        Attribute('out_width'),
+    ]
+
+    def initialize(self):
+        shape = [self.attributes['out_height'] * self.attributes['out_width'] * self.attributes['n_chan']]
+        self.add_output_variable(shape)
+
+
 layer_map = {
     'Input': Input,
     'InputLayer': Input,
@@ -1972,6 +2059,12 @@ def initialize(self):
     # TensorFlow-specific layers:
     'BiasAdd': BiasAdd,
     'DACombinational': DACombinational,
+    # Sparsepixels layers:
+    'SparseInputReduce': SparseInputReduce,
+    'SparseConv2D': SparseConv2D,
+    'SparseActivation': SparseActivation,
+    'SparsePooling2D': SparsePooling2D,
+    'SparseFlatten': SparseFlatten,
 }
 
 

hls4ml/model/optimizer/__init__.py

@@ -66,6 +66,7 @@
     [
         'channels_last_converter',
         'separable_to_depthwise_and_conv',
+        'convert_sparse_flatten',
         'remove_transpose_before_flatten',
         'remove_nop_transpose',
         'remove_single_channel_transpose',