🎉 add support for HGQ-proxy-model

🎉 add support for HGQ-proxy-model

more consistent type naming

revert example model version

Author: calad0i

hls4ml/backends/fpga/passes/hgq_proxy_model.py

@@ -0,0 +1,89 @@
+import numpy as np
+
+from hls4ml.backends import Backend
+from hls4ml.backends.template import FunctionCallTemplate
+from hls4ml.model.layers import Layer
+from hls4ml.model.optimizer import OptimizerPass
+from hls4ml.model.optimizer.passes.hgq_proxy_model import FixedPointQuantizer
+from hls4ml.model.types import Source
+
+
+def to_apfixed(k, b, i, RND, SAT):
+    u = 'u' if k == 0 else ''
+    return f'ap_{u}fixed<{b},{i},AP_{RND},AP_{SAT}>'
+
+
+def to_acfixed(k, b, i, RND, SAT):
+    k = 'false' if k == 0 else 'true'
+    return f'ac_fixed<{b},{i},{k},AC_{RND},AC_{SAT}>'
+
+
+def generate_mask_fn(
+    name: str, shape: tuple[int, ...], k: np.ndarray, b: np.ndarray, i: np.ndarray, RND: str, SAT: str, backend: str
+) -> str:
+    """Generate heterogenous quantization mask function, ONLY works for IOType=io_parallel"""
+    assert k.shape[0] == b.shape[0] == i.shape[0] == 1
+    assert backend.lower() in ('quartus', 'vivado', 'vitis'), f'Backend {backend} not tested'
+    Ks, Bs, Is = k[0], b[0], i[0]
+    Ks, Bs, Is = np.broadcast_to(Ks, shape), np.broadcast_to(Bs, shape), np.broadcast_to(Is, shape)
+    Ks, Bs, Is = Ks.ravel(), Bs.ravel(), Is.ravel()
+    masks = []
+    to_fixed = to_acfixed if backend.lower() == 'quartus' else to_apfixed
+    for idx, (k, b, i) in enumerate(zip(Ks, Bs, Is)):
+        if b == 0:
+            fn = f'out[{idx}] = 0;'
+        else:
+            fn = f'out[{idx}] = {to_fixed(k,b,i,RND,SAT)}(inp[{idx}]);'
+        masks.append(f'    {fn}')
+    body = "\n".join(masks)
+    mask_fn = f'''
+template<typename input_t, typename output_t>
+void {name}(input_t *inp, output_t *out) {{
+    #pragma HLS INLINE
+    #pragma HLS PIPELINE
+
+{body}
+}}
+'''
+    return mask_fn
+
+
+class ProcessFixedPointQuantizerLayer(OptimizerPass):
+    def match(self, node: Layer):
+        return isinstance(node, FixedPointQuantizer)
+
+    def transform(self, model, node: FixedPointQuantizer):
+        if node.fusible:
+            model.remove_node(node, rewire=True)
+            return True
+
+        if model.config.config['IOType'] != 'io_parallel':
+            raise NotImplementedError('Heterogenous quantization for activations is only supported with IOType=io_parallel')
+
+        backend = model.config.config['Backend']
+
+        name = node.name
+
+        assert node.mask_kbi is not None
+        k, b, i = node.mask_kbi
+        RND = node.RND
+        SAT = node.SAT
+        mask_fn: str = generate_mask_fn(name, node.get_input_variable().shape, k, b, i, RND, SAT, backend)
+
+        node.set_attr('mask_fn_codegen', Source(mask_fn))
+
+
+class ProcessFixedPointQuantizerCall(FunctionCallTemplate):
+    def __init__(self):
+        super().__init__(FixedPointQuantizer, include_header=[])
+        self.template = 'nnet::{name}<{input_t}, {output_t}>({input}, {output});'
+
+    def format(self, node):
+        params = self._default_function_params(node)
+
+        return self.template.format(**params)
+
+
+def register_hgq_proxy_model(backend: Backend):
+    backend.register_pass('process_fixed_point_quantizer_layer', ProcessFixedPointQuantizerLayer)
+    backend.register_template(ProcessFixedPointQuantizerCall)

hls4ml/backends/quartus/quartus_backend.py

@@ -1,5 +1,6 @@
 import os
 from contextlib import contextmanager
+from warnings import warn
 
 import numpy as np
 
@@ -73,6 +74,7 @@ def _register_flows(self):
             'quartus:inplace_stream_flatten',
             'quartus:skip_softmax',
             'quartus:fix_softmax_table_size',
+            'quartus:process_fixed_point_quantizer_layer',
             'infer_precision_types',
         ]
         optimization_flow = register_flow('optimize', optimization_passes, requires=[init_flow], backend=self.name)
@@ -265,7 +267,17 @@ def init_conv1d(self, layer):
         n_in, n_out = self.get_layer_mult_size(layer)
         self.set_target_reuse_factor(layer)
         self.set_closest_reuse_factor(layer, n_in, n_out)
-        layer.set_attr('parallelization', layer.model.config.get_layer_config_value(layer, 'ParallelizationFactor', 1))
+
+        # Not overriding user parallelization factor, if already set and user has not specified a value
+        user_pf = layer.model.config.get_layer_config_value(layer, 'ParallelizationFactor', None)
+        layer_pf = layer.get_attr('parallelization_factor', None)
+        chosen_pf = user_pf or layer_pf or 1
+        if user_pf is not None and layer_pf is not None:
+            if user_pf != layer_pf:
+                warn(
+                    f'For layer {layer.name}, parallelization factor of {layer_pf} is defined in the proxy-model, but is overridden by the user to {user_pf}.'  # noqa: E501
+                )
+        layer.set_attr('parallelization', chosen_pf)
 
         # impl_filt_width determines the filter size post-Winograd transformation
         layer.set_attr('impl_filt_width', layer.get_attr('filt_width'))
@@ -295,7 +307,17 @@ def init_conv2d(self, layer):
         n_in, n_out = self.get_layer_mult_size(layer)
         self.set_target_reuse_factor(layer)
         self.set_closest_reuse_factor(layer, n_in, n_out)
-        layer.set_attr('parallelization', layer.model.config.get_layer_config_value(layer, 'ParallelizationFactor', 1))
+
+        # Not overriding user parallelization factor, if already set and user has not specified a value
+        user_pf = layer.model.config.get_layer_config_value(layer, 'ParallelizationFactor', None)
+        layer_pf = layer.get_attr('parallelization_factor', None)
+        chosen_pf = user_pf or layer_pf or 1
+        if user_pf is not None and layer_pf is not None:
+            if user_pf != layer_pf:
+                warn(
+                    f'For layer {layer.name}, parallelization factor of {layer_pf} is defined in the proxy-model, but is overridden by the user to {user_pf}.'  # noqa: E501
+                )
+        layer.set_attr('parallelization', chosen_pf)
 
         # impl_filt_width & impl_filt_height determine the filter size post-Winograd transformation
         layer.set_attr('impl_filt_height', layer.get_attr('filt_height'))

hls4ml/backends/vivado/vivado_backend.py

@@ -1,5 +1,6 @@
 import os
 import sys
+from warnings import warn
 
 import numpy as np
 
@@ -107,6 +108,7 @@ def _register_flows(self):
             'vivado:inplace_stream_flatten',
             'vivado:skip_softmax',
             'vivado:fix_softmax_table_size',
+            'vivado:process_fixed_point_quantizer_layer',
             'infer_precision_types',
         ]
         optimization_flow = register_flow('optimize', optimization_passes, requires=[init_flow], backend=self.name)
@@ -266,7 +268,17 @@ def init_conv1d(self, layer):
             layer.set_attr('strategy', 'latency')
 
         out_width = layer.get_output_variable().shape[0]
-        chosen_pf = layer.model.config.get_layer_config_value(layer, 'ParallelizationFactor', 1)
+
+        # Not overriding user parallelization factor, if already set and user has not specified a value
+        user_pf = layer.model.config.get_layer_config_value(layer, 'ParallelizationFactor', None)
+        layer_pf = layer.get_attr('parallelization_factor', None)
+        chosen_pf = user_pf or layer_pf or 1
+        if user_pf is not None and layer_pf is not None:
+            if user_pf != layer_pf:
+                warn(
+                    f'For layer {layer.name}, parallelization factor of {layer_pf} is defined in the proxy-model, but is overridden by the user to {user_pf}.'  # noqa: E501
+                )
+
         valid_pf = self.get_valid_conv_partition_splits(1, out_width)
         if chosen_pf not in valid_pf:
             closest_pf = self.get_closest_reuse_factor(valid_pf, chosen_pf)
@@ -278,6 +290,7 @@ def init_conv1d(self, layer):
         else:
             closest_pf = chosen_pf
         layer.set_attr('n_partitions', out_width // closest_pf)
+        layer.set_attr('parallelization_factor', closest_pf)
 
         layer.set_attr('implementation', layer.model.config.get_conv_implementation(layer).lower())
 
@@ -332,7 +345,17 @@ def init_conv2d(self, layer):
 
         out_height = layer.get_output_variable().shape[0]
         out_width = layer.get_output_variable().shape[1]
-        chosen_pf = layer.model.config.get_layer_config_value(layer, 'ParallelizationFactor', 1)
+
+        # Not overriding user parallelization factor, if already set and user has not specified a value
+        user_pf = layer.model.config.get_layer_config_value(layer, 'ParallelizationFactor', None)
+        layer_pf = layer.get_attr('parallelization_factor', None)
+        chosen_pf = user_pf or layer_pf or 1
+        if user_pf is not None and layer_pf is not None:
+            if user_pf != layer_pf:
+                warn(
+                    f'For layer {layer.name}, parallelization factor of {layer_pf} is defined in the proxy-model, but is overridden by the user to {user_pf}.'  # noqa: E501
+                )
+
         valid_pf = self.get_valid_conv_partition_splits(out_height, out_width)
         if chosen_pf not in valid_pf:
             closest_pf = self.get_closest_reuse_factor(valid_pf, chosen_pf)
@@ -344,6 +367,7 @@ def init_conv2d(self, layer):
         else:
             closest_pf = chosen_pf
         layer.set_attr('n_partitions', out_height * out_width // closest_pf)
+        layer.set_attr('parallelization_factor', closest_pf)
 
         layer.set_attr('implementation', layer.model.config.get_conv_implementation(layer).lower())
 

hls4ml/converters/keras/hgq_proxy_model.py

@@ -0,0 +1,21 @@
+from hls4ml.converters.keras_to_hls import keras_handler, parse_default_keras_layer
+
+
+@keras_handler('FixedPointQuantizer')
+def fixedpoint_quantizer_handler(keras_layer, input_names, input_shapes, data_reader):
+    config = parse_default_keras_layer(keras_layer, input_names)
+
+    name = config['name']
+    fusible = keras_layer['config']['fusible']
+    config['RND'] = keras_layer['config']['RND']
+    config['SAT'] = keras_layer['config']['SAT']
+    config['fusible'] = fusible
+    if not fusible:
+        k = data_reader.get_weights_data(name, 'keep_negative')
+        b = data_reader.get_weights_data(name, 'bits')
+        i = data_reader.get_weights_data(name, 'integers')
+        config['mask_kbi'] = k, b, i
+    config['overrides'] = keras_layer['config']['overrides']
+
+    layer = config
+    return layer, input_shapes[0]

hls4ml/model/optimizer/__init__.py

@@ -44,6 +44,7 @@
         'qkeras_factorize_alpha',
         'extract_ternary_threshold',
         'fuse_consecutive_batch_normalization',
+        'enforce_proxy_model_embedded_config',
     ],
 )  # TODO Maybe not all QKeras optmizers belong here?
 

hls4ml/model/optimizer/passes/hgq_proxy_model.py

@@ -0,0 +1,128 @@
+import re
+from warnings import warn
+
+from hls4ml.backends.fpga.fpga_types import NamedType
+from hls4ml.model.layers import Layer, register_layer
+from hls4ml.model.optimizer import OptimizerPass, register_pass
+from hls4ml.model.types import FixedPrecisionType, WeightVariable
+
+re_purge_prefix = re.compile(r'(?<!\w)(?:ap_|ac_)', re.IGNORECASE)
+re_parse_fixed = re.compile(r'\s*(u?)fixed<([^>]+)>\s*', re.IGNORECASE)
+
+
+class FixedPointQuantizer(Layer):
+    def initialize(self):
+        inp = self.get_input_variable()
+        shape = inp.shape
+        dims = inp.dim_names
+        self.add_output_variable(shape, dims)
+        self.set_attr('n_in', self.get_input_variable().size())
+        self.overrides = self.attributes['overrides']
+        self.fusible = self.attributes['fusible']
+        self.SAT, self.RND = self.attributes['SAT'], self.attributes['RND']
+        self.mask_kbi = self.attributes.get('mask_kbi', None)
+
+
+def to_hls4ml_fixed(fixed: str):
+    matched = re_parse_fixed.match(re_purge_prefix.sub('', fixed))
+    assert matched is not None, f'Cannot parse {fixed}'
+    signed = matched.group(1) != 'u'
+    b, i, *args = matched.group(2).split(',')
+    b, i = int(b), int(i)
+    args = [arg.upper() for arg in args]
+    new_type = FixedPrecisionType(b, i, signed, *args)
+    # For some reason, __class__ is overwritten in hls4ml
+    return new_type
+
+
+def userconf_ifdef(key: str, layer_name: str, model):
+    hls_config: dict = model.config.config['HLSConfig']
+    layer_confs: dict = hls_config.get('LayerName', None)
+    if not layer_confs:
+        return False
+    layer_conf = layer_confs.get(layer_name, None)
+    if not layer_conf:
+        return False
+    # return key in layer_conf # Ideal case. Not for now.
+    if key.endswith('_t') and key != 'table_t':
+        # table_t cannot be defined in Precision, for some reason.
+        # On the other hand, result_t, weight_t, bias_t, accum_t cannot be decleared explicitly outside Precision, for now.
+        # However, still assume that they can be defined explicitly outside Precision.
+        precision_conf = layer_conf.get('Precision', None)
+        if not precision_conf:
+            return key in layer_conf
+        return key[:-2] in precision_conf or key in layer_conf
+
+    if key == 'parallelization_factor':
+        # Irregular config key name.
+        return 'ParallelizationFactor' in layer_conf
+
+    return key in layer_conf
+
+
+class EnforceProxyModelEmbeddedConfig(OptimizerPass):
+    def match(self, node: Layer):
+        if not isinstance(node, FixedPointQuantizer):
+            return False
+        if not node.overrides:
+            return False
+        return True
+
+    def transform(self, model, node: FixedPointQuantizer):
+        if 'layers' not in node.overrides:
+            return False
+
+        graph_changed = False
+        layers = node.overrides['layers']
+        for name, conf in layers.items():
+            conf: dict[str, str]
+            name: str
+            if name not in model.graph:
+                # Some layer may be removed by other passes. (e.g. Final flatten layer)
+                continue
+            target_node: Layer = model.graph[name]
+            for k, v in conf.items():
+                if userconf_ifdef(k, name, model):
+                    warn(
+                        f'Config key {k} is defined in hls_config for layer {name} by user. Proxy model config is ignored.',
+                        stacklevel=1,
+                    )
+                    continue
+
+                if k.endswith('_t'):
+                    var_type = target_node.get_attr(k)  # type: ignore
+                    if var_type is None:
+                        continue
+                    var_type: NamedType
+                    precision = to_hls4ml_fixed(v)
+                    var_type.precision = precision
+                    if k == 'result_t':
+                        type_name = f'{name}_t'
+                    else:
+                        type_name = f'{name}_{k}'
+                    var_type.name = type_name
+                    # Need to overwrite kernel/bias writing precision also, or written weights will likely be wrong.
+                    if k[:-2] in target_node.attributes.keys():
+                        weight_var: WeightVariable = target_node.attributes[k[:-2]]
+                        # weight_var should be a StaticWeightVariable, which is again, defined with meta programming
+                        # Type hinting using StaticWeightVariableDefinition which is the base class.
+                        weight_var.update_precision(precision)
+                    # Well, it turned out that there is yet ANOTHER copy saved in config.
+                    model.config.layer_name_precision[f'{name}_{k[:-2]}'] = v
+                elif k in target_node.attributes.attributes:
+                    target_node.set_attr(k, v)
+                elif k == 'parallelization_factor':
+                    target_node.set_attr(k, int(v))
+
+            if linear_node := model.graph.get(f'{name}_linear'):
+                # Proxy model does not assume any extra linear layer.
+                # Purge them on sight
+                model.remove_node(linear_node)
+                graph_changed = True
+
+        return graph_changed
+
+
+def register_hgq_proxy_model():
+    register_layer('FixedPointQuantizer', FixedPointQuantizer)
+    register_pass('enforce_proxy_model_embedded_config', EnforceProxyModelEmbeddedConfig)