Fix inplace variables (#714)

* extract inplace variables from qonnx branch

* fix up reshape batch dimension handling

* add repack stream for quartus

* fix typo in HLS pragma

* pre-commint cleanup, and add missing check for pragma not None

* Minor cosmetic fixes

---------

Co-authored-by: Vladimir Loncar <[email protected]>

Author: jmitrevs

hls4ml/backends/fpga/fpga_types.py

@@ -0,0 +1,24 @@
+from hls4ml.model.layers import Reshape
+from hls4ml.model.optimizer import OptimizerPass
+from hls4ml.model.types import InplaceTensorVariable
+
+
+class InplaceParallelReshape(OptimizerPass):
+    """
+    Replaces the output variable of Reshape layer with an inplace variable when using io_parallel.
+
+    This is done because in io_parallel tensors are stored as flat arrays, requiring no reshaping.
+    """
+
+    def match(self, node):
+        return isinstance(node, Reshape)
+
+    def transform(self, model, node):
+        if model.config.get_config_value('IOType') != 'io_parallel':
+            return False
+
+        outvar = node.get_output_variable()
+        invar = node.get_input_variable()
+        newoutvar = InplaceTensorVariable(outvar, invar)
+        node.set_attr(node.outputs[0], newoutvar)
+        return False

hls4ml/backends/fpga/passes/inplace_parallel_reshape.py

@@ -0,0 +1,25 @@
+from hls4ml.model.layers import Reshape
+from hls4ml.model.optimizer import OptimizerPass
+from hls4ml.model.types import InplaceTensorVariable
+
+
+class InplaceStreamFlatten(OptimizerPass):
+    """
+    Replaces the output variable of Reshape (flatten) layer with an inplace variable when using io_stream.
+
+    This optimizer avoids the expensive repacking of the stream when Reshape layer flattens the tensor to 1d.
+    """
+
+    def match(self, node):
+        # Reshape acts as a Flatten layer when the result has 1 dimension
+        return isinstance(node, Reshape) and len(node.get_output_variable().shape) == 1
+
+    def transform(self, model, node):
+        if model.config.get_config_value('IOType') != 'io_stream':
+            return False
+
+        outvar = node.get_output_variable()
+        invar = node.get_input_variable()
+        newoutvar = InplaceTensorVariable(outvar, invar)
+        node.set_attr(node.outputs[0], newoutvar)
+        return False

hls4ml/backends/fpga/passes/inplace_stream_flatten.py

@@ -0,0 +1,64 @@
+import numpy as np
+
+from hls4ml.backends.template import FunctionCallTemplate
+from hls4ml.model.layers import Layer, Reshape, register_layer
+from hls4ml.model.optimizer import OptimizerPass
+
+
+class Repack(Layer):
+    '''Inserted between layers with different packing factors.'''
+
+    def initialize(self):
+        shape = self.attributes['target_shape']
+        if shape[0] is None:
+            shape = shape[1:]
+        dims = [f'N_SIZE_{i}_{self.index}' for i in range(1, len(shape) + 1)]
+
+        self.add_output_variable(shape, dims)
+
+
+repack_function_template = 'nnet::repack_stream<{input_t}, {output_t}, {size}>({input}, {output});'
+repack_include_list = ['nnet_utils/nnet_stream.h']
+
+
+class RepackFunctionTemplate(FunctionCallTemplate):
+    def __init__(self):
+        super().__init__(Repack, include_header=repack_include_list)
+        self.template = repack_function_template
+
+    def format(self, node):
+        params = self._default_function_params(node)
+        params['size'] = np.prod(node.get_output_variable().shape)
+
+        return self.template.format(**params)
+
+
+def register_repack_stream(backend):
+    # Register the layer types to the layer map
+    register_layer('Repack', Repack)
+
+    # Register the optimization passes
+    backend.register_pass('reshape_stream', ReshapeStream)
+
+    # Register template passes
+    backend.register_template(RepackFunctionTemplate)
+
+
+class ReshapeStream(OptimizerPass):
+    '''Repacks stream for Reshape layer'''
+
+    def match(self, node):
+        # do not run optimizer pass for a flatten layer (1 output dimension)
+        return isinstance(node, Reshape) and len(node.get_output_variable().shape) > 1
+
+    def transform(self, model, node):
+        if model.config.get_config_value('IOType') != 'io_stream':
+            return False
+
+        attrs = {'target_shape': node.get_attr('target_shape')}
+
+        # Insert new Repack node instead of Reshape
+        repack_layer = model.make_node(Repack, 'repack_' + node.name, attrs, node.inputs.copy())
+        model.replace_node(node, repack_layer)
+
+        return True

hls4ml/backends/fpga/passes/repack_stream.py

@@ -1,34 +1,47 @@
-
+from hls4ml.backends.fpga.fpga_types import (
+    ACTypeConverter,
+    HLSTypeConverter,
+    QuartusArrayVariableConverter,
+    QuartusInplaceArrayVariableConverter,
+    QuartusInplaceStreamVariableConverter,
+    QuartusStreamVariableConverter,
+    QuartusStructMemberVariableConverter,
+    StaticWeightVariableConverter,
+)
 from hls4ml.model.optimizer import GlobalOptimizerPass
-from hls4ml.model.types import InplaceVariable
-from hls4ml.backends.fpga.fpga_types import ACTypeConverter, QuartusArrayVariableConverter, HLSTypeConverter, QuartusInplaceVariableConverter, QuartusStreamVariableConverter, QuartusStructMemberVariableConverter, StaticWeightVariableConverter
+from hls4ml.model.types import InplaceTensorVariable
+
 
 class TransformTypes(GlobalOptimizerPass):
     def __init__(self):
         self.type_converter = HLSTypeConverter(precision_converter=ACTypeConverter())
         self.array_var_converter = QuartusArrayVariableConverter(type_converter=self.type_converter)
+        self.inplace_array_var_converter = QuartusInplaceArrayVariableConverter(type_converter=self.type_converter)
         self.struct_var_converter = QuartusStructMemberVariableConverter(type_converter=self.type_converter)
         self.stream_var_converter = QuartusStreamVariableConverter(type_converter=self.type_converter)
+        self.inplace_stream_var_converter = QuartusInplaceStreamVariableConverter(type_converter=self.type_converter)
         self.weight_var_converter = StaticWeightVariableConverter(type_converter=self.type_converter)
-        self.inplace_var_converter = QuartusInplaceVariableConverter(type_converter=self.type_converter)
 
     def transform(self, model, node):
         io_type = node.model.config.get_config_value('IOType')
 
         for out_name, var in node.variables.items():
-            if isinstance(var, InplaceVariable):
-                new_var = self.inplace_var_converter.convert(var, io_type)
             if io_type == 'io_stream':
-                new_var = self.stream_var_converter.convert(var)
+                if isinstance(var, InplaceTensorVariable):
+                    new_var = self.inplace_stream_var_converter.convert(var)
+                else:
+                    new_var = self.stream_var_converter.convert(var)
             elif io_type == 'io_parallel':
-                if node.name in node.model.inputs:
+                if out_name in node.model.inputs:
                     new_var = self.struct_var_converter.convert(var, pragma='hls_register', struct_name='inputs')
-                elif node.name in node.model.outputs:
+                elif out_name in node.model.outputs:
                     new_var = self.struct_var_converter.convert(var, pragma='hls_register', struct_name='outputs')
+                elif isinstance(var, InplaceTensorVariable):
+                    new_var = self.inplace_array_var_converter.convert(var, pragma='')
                 else:
                     new_var = self.array_var_converter.convert(var, pragma='hls_register')
             else:
-                raise Exception('Unknown IOType {} in {} ({})'.format(io_type, node.name, node.class_name))
+                raise Exception(f'Unknown IOType {io_type} in {node.name} ({node.class_name})')
 
             node.set_attr(out_name, new_var)
 

hls4ml/backends/quartus/passes/transform_types.py

@@ -45,7 +45,7 @@ def _register_flows(self):
         initializers = self._get_layer_initializers()
         init_flow = register_flow('init_layers', initializers, requires=['optimize'], backend=self.name)
 
-        streaming_passes = ['quartus:clone_output']
+        streaming_passes = ['quartus:reshape_stream', 'quartus:clone_output']
         streaming_flow = register_flow('streaming', streaming_passes, requires=[init_flow], backend=self.name)
 
         quartus_types = [
@@ -63,7 +63,13 @@ def _register_flows(self):
         ]
         quantization_flow = register_flow('quantization', quantization_passes, requires=[init_flow], backend=self.name)
 
-        optimization_passes = ['quartus:remove_final_reshape', 'quartus:optimize_pointwise_conv', 'quartus:skip_softmax']
+        optimization_passes = [
+            'quartus:remove_final_reshape',
+            'quartus:optimize_pointwise_conv',
+            'quartus:inplace_parallel_reshape',
+            'quartus:inplace_stream_flatten',
+            'quartus:skip_softmax',
+        ]
         optimization_flow = register_flow('optimize', optimization_passes, requires=[init_flow], backend=self.name)
 
         templates = self._get_layer_templates()

hls4ml/backends/quartus/quartus_backend.py

@@ -1,4 +1,3 @@
-
 from hls4ml.model.optimizer.optimizer import OptimizerPass
 
 
@@ -9,7 +8,7 @@ def __init__(self, name, layer_class, attribute_name):
         if not isinstance(self.layer_class, (list, tuple, set)):
             self.layer_class = [self.layer_class]
         self.attribute_name = attribute_name
-    
+
     def match(self, node):
         for layer_cls in self.layer_class:
             if node.class_name == layer_cls.__name__:
@@ -20,13 +19,14 @@ def transform(self, model, node):
         formatted_template = self.format(node)
         node.set_attr(self.attribute_name, formatted_template)
         return False
-    
+
     def format(self, node):
         raise NotImplementedError
 
     def get_name(self):
         return self.name
-    
+
+
 class LayerConfigTemplate(Template):
     def __init__(self, layer_class):
         if isinstance(layer_class, (list, tuple, set)):
@@ -35,7 +35,7 @@ def __init__(self, layer_class):
             name = layer_class.__name__.lower()
         name += '_config_template'
         super().__init__(name, layer_class, 'config_cpp')
-    
+
     def _default_config_params(self, layer):
         params = {}
         params.update(layer.attributes)
@@ -44,6 +44,7 @@ def _default_config_params(self, layer):
 
         return params
 
+
 class FunctionCallTemplate(Template):
     def __init__(self, layer_class, include_header=None):
         if isinstance(layer_class, (list, tuple, set)):
@@ -52,12 +53,15 @@ def __init__(self, layer_class, include_header=None):
             name = layer_class.__name__.lower()
         name += '_function_template'
         super().__init__(name, layer_class, 'function_cpp')
-        self.include_header = include_header
-    
+        if include_header is None:
+            self.include_header = ()
+        else:
+            self.include_header = include_header
+
     def _default_function_params(self, layer):
         params = {}
         params.update(layer.attributes)
-        params['config'] = 'config{}'.format(layer.index)
+        params['config'] = f'config{layer.index}'
         params['input_t'] = layer.get_input_variable().type.name
         params['output_t'] = layer.get_output_variable().type.name
         params['input'] = layer.get_input_variable().name