skip BatchNorm fusion when input/output is used multiple times (#481)

* skip bn fusion when input/ouput used mult. times

* format

Author: jmduarte

hls4ml/backends/fpga/passes/clone.py

@@ -55,13 +55,7 @@ def transform(self, model, node):
         if model.config.get_config_value('IOType') != 'io_stream':
             return False
 
-        output_map = {}
-        for output in node.outputs:
-            output_map[output] = []
-            for layer in model.get_layers():
-                for inp in layer.inputs:
-                    if output == inp:
-                        output_map[output].append(layer)
+        output_map = node.get_output_use_map()
 
         transformed = False
         for output in node.outputs:

hls4ml/model/layers.py

@@ -131,10 +131,28 @@ def get_input_variable(self, input_name=None):
         else:
             return self.model.get_layer_output_variable(self.inputs[0])
 
+    def get_output_use_map(self):
+        output_map = {}
+        for output in self.outputs:
+            output_map[output] = []
+            for layer in self.model.get_layers():
+                for inp in layer.inputs:
+                    if output == inp:
+                        output_map[output].append(layer)
+        return output_map
+
     def get_output_nodes(self, output_name=None):
-        if output_name is None:
-            output_name = self.outputs[0]
-        return [node for node in self.model.graph.values() if node.inputs[0] == output_name]
+        output_nodes = []
+        if output_name is not None:
+            outputs = [output_name]
+        else:
+            outputs = self.outputs
+        for output in outputs:
+            for layer in self.model.get_layers():
+                for inp in layer.inputs:
+                    if output == inp:
+                        output_nodes.append(layer)
+        return output_nodes
 
     def get_output_variable(self, output_name=None):
         if output_name is not None:

hls4ml/model/optimizer/passes/bn_fuse.py

@@ -12,6 +12,10 @@ def match(self, node):
     def transform(self, model, node):
         # Fuse weight and bias of Dense/Conv1D/Conv2D layer with BN values
         parent_node = node.get_input_node()
+        parent_map = parent_node.get_output_use_map()
+        node_map = node.get_output_use_map()
+        if len(parent_map[parent_node.name]) > 1 or len(node_map[node.name]) > 1:
+            return False
 
         parent_weight = parent_node.weights['weight']
         parent_bias = parent_node.weights['bias']

hls4ml/model/optimizer/passes/qkeras.py

@@ -226,6 +226,10 @@ def match(self, node):
     def transform(self, model, node):
         bn0 = node.get_input_node()
         bn1 = node
+        bn0_map = bn0.get_output_use_map()
+        bn1_map = bn1.get_output_use_map()
+        if len(bn0_map[bn0.name]) > 1 or len(bn1_map[bn1.name]) > 1:
+            return False
 
         s0 = bn0.weights['scale'].data
         b0 = bn0.weights['bias'].data