Remove final reshapes even for Quartus

Author: jmitrevs

hls4ml/backends/fpga/passes/final_reshape.py

@@ -0,0 +1,21 @@
+import numpy as np
+
+from hls4ml.model.optimizer import OptimizerPass
+from hls4ml.model.layers import Reshape
+
+
+class RemoveFinalReshape(OptimizerPass):
+    ''' Remove reshape if final layer '''
+    def match(self, node):
+        # match if reshape is final node
+        return isinstance(node, Reshape) and not node.get_output_nodes()
+
+    def transform(self, model, node):
+        if model.config.get_config_value('IOType') == 'io_parallel':
+            print('WARNING: Final layer is a Reshape, which does not affect the output for io_parallel; removing it')
+            # remove, but don't rewire because it's the output layer
+            model.remove_node(node, rewire=False) 
+            return True
+        elif model.config.get_config_value('IOType') == 'io_stream':
+            print('WARNING: Final layer is a Reshape, which may incur a large resource cost for io_stream; consider removing it')
+        return False

hls4ml/backends/quartus/quartus_backend.py

@@ -60,6 +60,7 @@ 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',
         ]
         optimization_flow = register_flow('optimize', optimization_passes, requires=[init_flow], backend=self.name)

hls4ml/backends/vivado/passes/repack_stream.py

@@ -2,7 +2,6 @@
 
 from hls4ml.model.optimizer import OptimizerPass
 from hls4ml.model.layers import Layer, Merge, Reshape, Concatenate, register_layer
-from hls4ml.backends import get_backend
 from hls4ml.backends.template import FunctionCallTemplate, LayerConfigTemplate
 
 class Repack(Layer):
@@ -82,7 +81,6 @@ def register_repack_stream(backend):
     register_layer('Broadcast', Broadcast)
 
     # Register the optimization passes
-    backend.register_pass('remove_final_reshape', RemoveFinalReshape)
     backend.register_pass('reshape_stream', ReshapeStream)
     backend.register_pass('broadcast_stream', BroadcastStream)
 
@@ -162,19 +160,3 @@ def supported_broadcast(inp_shape, target_shape):
         node.inputs[idx] = brdcst_out
 
         return True
-
-class RemoveFinalReshape(OptimizerPass):
-    ''' Remove reshape if final layer '''
-    def match(self, node):
-        # match if reshape is final node
-        return isinstance(node, Reshape) and not node.get_output_nodes()
-
-    def transform(self, model, node):
-        if model.config.get_config_value('IOType') == 'io_parallel':
-            print('WARNING: Final layer is a Reshape, which does not affect the output for io_parallel; removing it')
-            # remove, but don't rewire because it's the output layer
-            model.remove_node(node, rewire=False) 
-            return True
-        elif model.config.get_config_value('IOType') == 'io_stream':
-            print('WARNING: Final layer is a Reshape, which may incur a large resource cost for io_stream; consider removing it')
-        return False

hls4ml/backends/vivado/vivado_backend.py

@@ -34,7 +34,6 @@ def _register_flows(self):
         init_flow = register_flow('init_layers', initializers, requires=['optimize'], backend=self.name)
 
         streaming_passes = [
-            'vivado:remove_final_reshape',
             'vivado:reshape_stream',
             'vivado:clone_output',
             'vivado:insert_zero_padding_before_conv1d',
@@ -51,6 +50,7 @@ def _register_flows(self):
         quantization_flow = register_flow('quantization', quantization_passes, requires=[init_flow], backend=self.name)
 
         optimization_passes = [
+            'vivado:remove_final_reshape',
             'vivado:optimize_pointwise_conv',
         ]
         optimization_flow = register_flow('optimize', optimization_passes, requires=[init_flow], backend=self.name)