bit exact extension

Author: calad0i

hls4ml/converters/__init__.py

@@ -165,6 +165,7 @@ def convert_from_keras_model(
     output_data_tb=None,
     backend='Vivado',
     hls_config=None,
+    bit_exact=None,
     **kwargs,
 ):
     """Convert Keras model to hls4ml model based on the provided configuration.
@@ -214,6 +215,7 @@ def convert_from_keras_model(
 
     model_config = hls_config.get('Model', None)
     config['HLSConfig']['Model'] = _check_model_config(model_config)
+    config['HLSConfig']['Model']['BitExact'] = bit_exact
 
     _check_hls_config(config, hls_config)
     if 'KerasModel' in config:

hls4ml/converters/keras/qkeras.py

@@ -174,6 +174,7 @@ def get_activation_quantizer(keras_layer, input_names, activation_name='activati
         layer[activation_name] = activation_config['class_name'].replace('quantized_', '')
 
     layer[f'{activation_name}_quantizer'] = activation_config
+    layer['trusted'] = True
 
     return layer
 

hls4ml/model/optimizer/passes/bit_exact.py

@@ -133,6 +133,30 @@ def _(layer: Reshape):
 @_request_kif.register
 def _(layer: Activation):
     fn_name = layer.attributes.get('activation')
+
+    if layer.attributes.get('trusted', False):
+        result_t = layer.get_output_variable().type.precision
+        if fn_name in ('linear', 'relu'):
+            output_shape = get_output_shape(layer)
+            k, w, f = result_t.signed, result_t.width, result_t.fractional
+            i = w - k - f
+            k = np.full(output_shape, k, dtype=np.int8)
+            i = np.full(output_shape, i, dtype=np.int8)
+            f = np.full(output_shape, f, dtype=np.int8)
+            if result_t.rounding_mode == RoundingMode.RND:
+                f += 1
+            elif result_t.rounding_mode != RoundingMode.TRN:
+                f = np.full(output_shape, 126, dtype=np.int8)
+            if result_t.saturation_mode != SaturationMode.WRAP:
+                k = np.ones(output_shape, dtype=np.int8)
+                i = np.full(output_shape, 126, dtype=np.int8)
+            if fn_name == 'linear':
+                return ((k, i, f),)
+            else:
+                k = np.ones(output_shape, dtype=np.int8)
+                i = np.full(output_shape, 126, dtype=np.int8)
+                return ((k, i, f),)
+
     if fn_name == 'linear':
         return (requested_kif(layer),)
     if fn_name == 'relu':
@@ -531,6 +555,16 @@ def _(layer: Concatenate):
 @_produce_kif.register
 def _(layer: Activation):
     fn_name = layer.attributes['activation'].lower()
+    if layer.attributes.get('trusted', False):
+        output_shape = get_output_shape(layer)
+        result_t = layer.get_output_variable().type.precision
+        k, w, f = result_t.signed, result_t.width, result_t.fractional
+        i = w - k - f
+        k = np.full(output_shape, k, dtype=np.int8)
+        i = np.full(output_shape, i, dtype=np.int8)
+        f = np.full(output_shape, f, dtype=np.int8)
+        return k, i, f
+
     k, i, f = get_input_kifs(layer)[0]
 
     match fn_name:
@@ -603,6 +637,10 @@ def requested_by_non_saturating_quantizer(layer: Layer) -> bool:
 
 
 def default_register_precision(layer: Layer):
+    if layer.attributes.get('trusted', False):
+        # Trusted layers have their precision already set
+        return
+
     _pk, _pi, _pf = produce_kif(layer)  # Maximum possible k,i,f output from this layer
     _rk, _ri, _rf = requested_kif(layer)  # Maximum possible k,i,f may be utilized by the next layer
     _oi, _of = np.minimum(_pi, _ri), np.minimum(_pf, _rf)
@@ -791,7 +829,11 @@ def has_fixed_quantizer(self, model: 'ModelGraph'):
         return True
 
     def _match(self, model: 'ModelGraph'):
-        return self.has_fixed_quantizer(model)
+        enabled = model.config.config['HLSConfig']['Model'].get('BitExact', None)
+        if enabled is None:
+            # Enable by default if any FixedPointQuantizer is present
+            enabled = self.has_fixed_quantizer(model)
+        return enabled
 
     def transform(self, model: 'ModelGraph'):
         if not self._match(model):

hls4ml/model/optimizer/passes/hgq_proxy_model.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 from hls4ml.model.attributes import Attribute, TypeAttribute, WeightAttribute
-from hls4ml.model.layers import Layer, Reshape, register_layer
+from hls4ml.model.layers import Activation, Layer, Reshape, register_layer
 from hls4ml.model.optimizer import OptimizerPass, register_pass
 from hls4ml.model.types import FixedPrecisionType, UnspecifiedPrecisionType
 
@@ -79,11 +79,13 @@ def userconf_ifdef(key: str, layer_name: str, model):
 
 class FuseFixedPointQuantizer(OptimizerPass):
     def match(self, node: Layer):
-        if not isinstance(node, FixedPointQuantizer):
-            return False
-        if any(np.unique(x).size > 1 for x in node.mask_kbi):
-            return False
-        return True
+        if isinstance(node, FixedPointQuantizer):
+            return all(np.unique(x).size == 1 for x in node.mask_kbi)
+
+        if isinstance(node, Activation):
+            return node.get_attr('activation') == 'linear' and node.get_attr('trusted', False)
+
+        return False
 
     def propagate(self, node: Layer, precision: FixedPrecisionType):
         from hls4ml.model.optimizer.passes.bit_exact import get_input_layers, get_output_layers
@@ -115,13 +117,16 @@ def propagate(self, node: Layer, precision: FixedPrecisionType):
     def transform(self, model: 'ModelGraph', node: FixedPointQuantizer):
         from hls4ml.model.optimizer.passes.bit_exact import get_input_layers, get_output_layers
 
-        # Rounding and saturation for FixedPointQuantizer are applied in generated code, thus not reflected in result_t.
-        if node.RND == 'TRN' and node.SAT == 'WRAP':
-            precision: FixedPrecisionType = copy(node.get_output_variable().type.precision)
+        if isinstance(node, FixedPointQuantizer):
+            # Rounding and saturation for FixedPointQuantizer are applied in generated code, thus not reflected in result_t.
+            if node.RND == 'TRN' and node.SAT == 'WRAP':
+                precision: FixedPrecisionType = copy(node.get_output_variable().type.precision)
+            else:
+                k, b, i = node.mask_kbi
+                k, b, i = bool(k.ravel()[0]), max(int(b.ravel()[0]), 1), int(i.ravel()[0])
+                precision = FixedPrecisionType(b, i, k, node.RND, node.SAT)
         else:
-            k, b, i = node.mask_kbi
-            k, b, i = bool(k.ravel()[0]), max(int(b.ravel()[0]), 1), int(i.ravel()[0])
-            precision = FixedPrecisionType(b, i, k, node.RND, node.SAT)
+            precision = copy(node.get_output_variable().type.precision)
 
         inp_layer = get_input_layers(node)[0]
         can_fuse = len(get_output_layers(inp_layer)) == 1