Da custom layer (#1429)

* use metaclass for handler reg

* feat: general da4ml fallback

* bump da4ml version for LUT-Layer support

* fix syntax err

* api and docstr

Author: calad0i

docs/advanced/extension.rst

@@ -81,7 +81,7 @@ In this case, there a single output with the same shape as the input.
 
             from hls4ml.converters.keras_v3._base import register, KerasV3LayerHandler
 
-            @register
+            # Handlers are registered by metaclass
             class KReverseHandler(KerasV3LayerHandler):
                 '''Keras v3 layer handler for KReverse'''
 

hls4ml/backends/vivado/passes/distributed_arithmetic.py

@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from hls4ml.model.layers import Conv1D, Conv2D, Dense, EinsumDense, Layer
+from hls4ml.model.layers import Conv1D, Conv2D, DACombinational, Dense, EinsumDense, Layer
 from hls4ml.model.optimizer import OptimizerPass
 from hls4ml.model.optimizer.passes.bit_exact import get_input_layers, get_output_layers, im2col, pad_arrs, stride_arrs
 from hls4ml.model.optimizer.passes.hgq_proxy_model import FixedPointQuantizer
@@ -408,3 +408,52 @@ def transform(self, model: 'ModelGraph', node: Layer):
         del node.attributes['weight_data']
         del node.attributes['weight']
         del node.attributes['weight_t']
+
+
+class DACombinationalTemplate(OptimizerPass):
+    def match(self, node):
+        return isinstance(node, DACombinational)
+
+    def transform(self, model: 'ModelGraph', node: DACombinational):
+        from da4ml.codegen.hls import hls_logic_and_bridge_gen
+        from da4ml.trace import FixedVariableArrayInput, comb_trace
+
+        io_type = model.config.get_config_value('IOType')
+        if io_type != 'io_parallel':
+            original_type = node.attributes['original_type']
+            raise ValueError(f'DACombinational layer (from {original_type}) only supports io_parallel.')
+
+        if not node.attributes.get('bit_exact_transformed', False):
+            inp_p: FixedPrecisionType = node.get_input_variable().type.precision
+            B, I, s = inp_p.width, inp_p.integer, inp_p.signed
+            i, f = I - s, B - I
+            comb = node.attributes['da_comb_logic']
+            inp = FixedVariableArrayInput(comb.shape[0]).quantize(s, i, f)
+            out = comb(inp)
+            comb = comb_trace(inp, out)
+            node.attributes['da_comb_logic'] = comb
+
+        comb = node.attributes['da_comb_logic']
+
+        backend = model.config.get_config_value('Backend').lower()
+        if backend in ('vitis', 'vivado'):
+            flavor = 'vitis'
+        elif backend == 'oneapi':
+            flavor = 'oneapi'
+        else:
+            raise ValueError(f'Unsupported backend {backend} for DACombinational layer.')
+
+        fn_name = f'da_comblogic_{node.index}'
+        comb_logic, _ = hls_logic_and_bridge_gen(
+            comb, fn_name, flavor=flavor, pragmas=['#pragma HLS INLINE'], print_latency=True
+        )
+        namespace = model.config.get_writer_config().get('Namespace', None) or 'nnet'
+
+        inp_t: str = node.get_input_variable().type.name
+        out_t: str = node.get_output_variable().type.name
+        inp_name: str = node.get_input_variable().name
+        out_name: str = node.get_output_variable().name
+
+        fn_cpp = f'{namespace}::{fn_name}<{inp_t}, {out_t}>({inp_name}, {out_name});'
+        node.attributes['da_codegen'] = Source(comb_logic)
+        node.attributes['function_cpp'] = fn_cpp

hls4ml/backends/vivado/vivado_backend.py

@@ -187,6 +187,7 @@ def _register_flows(self):
             'vivado:set_pipeline_style',
             'vivado:d_a_latency_dense_template',
             'vivado:d_a_latency_conv_template',
+            'vivado:d_a_combinational_template',
         ]
         vivado_types_flow = register_flow('specific_types', vivado_types, requires=[init_flow], backend=self.name)
 

hls4ml/converters/__init__.py

@@ -167,6 +167,8 @@ def convert_from_keras_model(
     backend='Vivado',
     hls_config=None,
     bit_exact=None,
+    allow_da_fallback=True,
+    allow_v2_fallback=True,
     **kwargs,
 ):
     """Convert Keras model to hls4ml model based on the provided configuration.
@@ -195,12 +197,17 @@ def convert_from_keras_model(
         io_type (str, optional): Type of implementation used. One of
             'io_parallel' or 'io_stream'. Defaults to 'io_parallel'.
         hls_config (dict, optional): The HLS config.
-        kwargs** (dict, optional): Additional parameters that will be used to create the config of the specified backend
         bit_exact (bool, optional): If True, enable model-wise precision propagation
         with **only fixed-point data types**. If None, enable if there is at least one
         FixedPointQuantizer layer in the model (only resulting from converting HGQ1/2
         models for now). By default, None.
+        allow_da_fallback: Whether to allow fallback to DA combinational logic generation
+            for unsupported layers. Only affects keras v3 models. Defaults to True.
+        allow_v2_fallback: Whether to allow fallback to keras v2 layer handlers
+            for unsupported layers. Only affects keras v3 models. Defaults to True. If both this and
+            `allow_da_fallback` are True, DA fallback is attempted first.
 
+        kwargs** (dict, optional): Additional parameters that will be used to create the config of the specified backend
     Raises:
         Exception: If precision and reuse factor are not present in 'hls_config'.
 
@@ -227,7 +234,7 @@ def convert_from_keras_model(
         import keras
 
         if keras.__version__ >= '3.0':
-            return keras_v3_to_hls(config)
+            return keras_v3_to_hls(config, allow_da_fallback, allow_v2_fallback)
 
     return keras_v2_to_hls(config)
 

hls4ml/converters/keras_v3/_base.py

@@ -33,7 +33,16 @@ def maybe_add_attrs(config: dict[str, Any] | DefaultConfig, obj: Any, *attrs: st
             config[attr] = getattr(obj, attr)
 
 
-class KerasV3LayerHandler:
+class KerasV3LayerHandlerMeta(type):
+    def __new__(cls, name, bases, attrs):
+        new_class = super().__new__(cls, name, bases, attrs)
+        if 'handles' in attrs:
+            for handle in attrs['handles']:
+                registry[handle] = new_class()
+        return new_class
+
+
+class KerasV3LayerHandler(metaclass=KerasV3LayerHandlerMeta):
     """Base class for keras v3 layer handlers. Subclass this class to create a handler for a specific layer type."""
 
     handles = ()
@@ -160,26 +169,3 @@ def load_weight(self, layer: 'keras.Layer', key: str):
         import keras
 
         return keras.ops.convert_to_numpy(getattr(layer, key))
-
-
-def register(cls: type):
-    """Decorator to register a handler for a specific layer class. Suggested to decorate the `KerasV3LayerHandler` class.
-
-    Args:
-        cls: the class to register the handler for.
-
-    Examples:
-        ```python
-        @keras_dispatcher.register
-        class MyLayerHandler(KerasV3LayerHandler):
-            handles = ('my_package.src.submodule.MyLayer', 'MyLayer2')
-
-            def handle(self, layer, inp_tensors, out_tensors):
-                # handler code
-        ```
-    """
-
-    fn = cls()
-    for k in fn.handles:
-        registry[k] = fn
-    return cls

hls4ml/converters/keras_v3/conv.py

@@ -3,7 +3,7 @@
 from math import ceil
 from typing import Any
 
-from ._base import KerasV3LayerHandler, register
+from ._base import KerasV3LayerHandler
 
 if typing.TYPE_CHECKING:
     import keras
@@ -75,7 +75,6 @@ def gen_conv_config(
     return config
 
 
-@register
 class ConvHandler(KerasV3LayerHandler):
     handles = (
         'keras.src.layers.convolutional.conv1d.Conv1D',

hls4ml/converters/keras_v3/core.py

@@ -5,14 +5,13 @@
 
 import numpy as np
 
-from ._base import KerasV3LayerHandler, register
+from ._base import KerasV3LayerHandler
 
 if typing.TYPE_CHECKING:
     import keras
     from keras import KerasTensor
 
 
-@register
 class DenseHandler(KerasV3LayerHandler):
     handles = ('keras.src.layers.core.dense.Dense',)
 
@@ -36,7 +35,6 @@ def handle(
         return config
 
 
-@register
 class InputHandler(KerasV3LayerHandler):
     handles = ('keras.src.layers.core.input_layer.InputLayer',)
 
@@ -50,7 +48,6 @@ def handle(
         return config
 
 
-@register
 class ActivationHandler(KerasV3LayerHandler):
     handles = ('keras.src.layers.activations.activation.Activation',)
 
@@ -89,7 +86,6 @@ def handle(
         return (config,)
 
 
-@register
 class ReLUHandler(KerasV3LayerHandler):
     handles = (
         'keras.src.layers.activations.leaky_relu.LeakyReLU',
@@ -136,7 +132,6 @@ def handle(
         return (config,)
 
 
-@register
 class SoftmaxHandler(KerasV3LayerHandler):
     handles = ('keras.src.layers.activations.softmax.Softmax',)
 
@@ -169,7 +164,6 @@ def handle(
         return (config,)
 
 
-@register
 class EluHandler(KerasV3LayerHandler):
     handles = ('keras.src.layers.activations.elu.ELU',)
 
@@ -190,7 +184,6 @@ def handle(
         return (config,)
 
 
-@register
 class ReshapeHandler(KerasV3LayerHandler):
     handles = ('keras.src.layers.reshaping.reshape.Reshape', 'keras.src.layers.reshaping.flatten.Flatten')
 
@@ -206,7 +199,6 @@ def handle(
         }
 
 
-@register
 class PermuteHandler(KerasV3LayerHandler):
     handles = ('keras.src.layers.reshaping.permute.Permute',)
 
@@ -220,7 +212,6 @@ def handle(
         return config
 
 
-@register
 class NoOp(KerasV3LayerHandler):
     handles = (
         'keras.src.layers.preprocessing.image_preprocessing.random_brightness.RandomBrightness',

hls4ml/converters/keras_v3/einsum_dense.py

@@ -1,7 +1,7 @@
 import typing
 from collections.abc import Sequence
 
-from ._base import KerasV3LayerHandler, register
+from ._base import KerasV3LayerHandler
 
 if typing.TYPE_CHECKING:
     import keras
@@ -41,7 +41,6 @@ def strip_batch_dim(equation: str, einsum_dense: bool = True):
     return f'{inp0},{inp1}->{out}'
 
 
-@register
 class EinsumDenseHandler(KerasV3LayerHandler):
     handles = ('keras.src.layers.core.einsum_dense.EinsumDense',)
 

hls4ml/converters/keras_v3/hgq2/_base.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from hls4ml.converters.keras_v3._base import KerasV3LayerHandler, register
+from hls4ml.converters.keras_v3._base import KerasV3LayerHandler
 from hls4ml.converters.keras_v3.conv import ConvHandler
 from hls4ml.converters.keras_v3.core import ActivationHandler, DenseHandler
 from hls4ml.converters.keras_v3.einsum_dense import EinsumDenseHandler
@@ -67,7 +67,6 @@ def override_io_tensor_confs(confs: tuple[dict[str, Any], ...], overrides: dict[
         conf['output_keras_tensor_names'] = [overrides.get(name, name) for name in out_tensor_names]
 
 
-@register
 class QLayerHandler(KerasV3LayerHandler):
     def __call__(
         self,
@@ -121,15 +120,13 @@ def default_class_name(self, layer: 'Layer') -> str:
         return class_name
 
 
-@register
 class QEinsumDenseHandler(QLayerHandler, EinsumDenseHandler):
     handles = (
         'hgq.layers.core.einsum_dense.QEinsumDense',
         'hgq.layers.einsum_dense_batchnorm.QEinsumDenseBatchnorm',
     )
 
 
-@register
 class QStandaloneQuantizerHandler(KerasV3LayerHandler):
     handles = ('hgq.quantizer.quantizer.Quantizer',)
 
@@ -144,7 +141,6 @@ def handle(
         return conf
 
 
-@register
 class QConvHandler(QLayerHandler, ConvHandler):
     handles = (
         'hgq.layers.conv.QConv1D',
@@ -170,7 +166,6 @@ def handle(
         return conf
 
 
-@register
 class QDenseHandler(QLayerHandler, DenseHandler):
     handles = ('hgq.layers.core.dense.QDense', 'hgq.layers.core.dense.QBatchNormDense')
 
@@ -189,12 +184,10 @@ def handle(
         return conf
 
 
-@register
 class QActivationHandler(QLayerHandler, ActivationHandler):
     handles = ('hgq.layers.activation.QActivation',)
 
 
-@register
 class QBatchNormalizationHandler(QLayerHandler):
     handles = ('hgq.layers.batch_normalization.QBatchNormalization',)
 
@@ -220,7 +213,6 @@ def handle(
         }
 
 
-@register
 class QMergeHandler(QLayerHandler, MergeHandler):
     handles = (
         'hgq.layers.ops.merge.QAdd',

hls4ml/converters/keras_v3/hgq2/einsum.py

@@ -2,14 +2,13 @@
 from collections.abc import Sequence
 
 from ..einsum_dense import strip_batch_dim
-from ._base import QLayerHandler, register
+from ._base import QLayerHandler
 
 if typing.TYPE_CHECKING:
     import hgq
     from keras import KerasTensor
 
 
-@register
 class QEinsumHandler(QLayerHandler):
     handles = ('hgq.layers.ops.einsum.QEinsum',)