clone streams up to 3 times

Author: jmduarte

hls4ml/model/optimizer/passes/clone.py

@@ -26,19 +26,44 @@ def config_cpp(self):
 clone_function_template = 'nnet::clone_stream<{input_t}, {output_t}, {size}>({input}, {output1}, {output2});'
 clone_include_list = ['nnet_utils/nnet_stream.h']
 
+class CloneThree(Layer):
+    ''' Inserted after the layer whose output is used more than twice.'''
+
+    def initialize(self):
+        inp = self.get_input_variable()
+        self.add_output_variable(inp.shape, inp.dim_names, out_name=self.outputs[0], var_name='layer{index}_cpy1')
+        self.add_output_variable(inp.shape, inp.dim_names, out_name=self.outputs[1], var_name='layer{index}_cpy2')
+        self.add_output_variable(inp.shape, inp.dim_names, out_name=self.outputs[2], var_name='layer{index}_cpy3')
+
+    def function_cpp(self):
+        params = self._default_function_params()
+        params['size'] = self.get_attr('size')
+        params['output1'] = self.variables[self.outputs[0]].name
+        params['output2'] = self.variables[self.outputs[1]].name
+        params['output3'] = self.variables[self.outputs[2]].name
+        return [self._function_template.format(**params)]
+
+    def config_cpp(self):
+        return None
+
+clonethree_function_template = 'nnet::clone_stream<{input_t}, {output_t}, {size}>({input}, {output1}, {output2}, {output3});'
+clonethree_include_list = ['nnet_utils/nnet_stream.h']
+
 # Register the layer types to the layer map
 register_layer('Clone', Clone)
+register_layer('CloneThree', CloneThree)
 
 # Register the templates for config and function
 for backend in ['Vivado', 'VivadoAccelerator']:
     templates.get_backend(backend).register_templates('Clone', clone_function_template, None, clone_include_list)
+    templates.get_backend(backend).register_templates('CloneThree', clonethree_function_template, None, clonethree_include_list)
 
 
 class CloneOutput(OptimizerPass):
     ''' Clones streams that are used multiple times '''
     def match(self, node):
         # We may have already inserted the Clone layer
-        if node.__class__.__name__ == 'Clone':
+        if node.__class__.__name__ in ['Clone', 'CloneThree']:
             return False
 
         return True
@@ -59,8 +84,8 @@ def transform(self, model, node):
         transformed = False
         for output in node.outputs:
             if len(output_map[output]) > 1:
-                if len(output_map[output]) > 2:
-                    print('WARN: Cannot clone output {} of {} ({})'.format(output, node.__class__.__name__, node.name))
+                if len(output_map[output]) > 3:
+                    print('WARNING: Cloning output {} of {} ({}) more than 3 times not currently supported'.format(output, node.__class__.__name__, node.name))
                     return False
                 out_var = node.get_output_variable(output)
                 for i, layer in enumerate(output_map[output], 1):
@@ -69,7 +94,10 @@ def transform(self, model, node):
                     }
                     idx = layer.inputs.index(output)
                     layer.inputs[idx] = output + '_cpy' + str(i)
-                clone_layer = model.make_node('Clone', 'clone_' + node.name, attrs, [output], [output + '_cpy1', output + '_cpy2'])
+                if len(output_map[output]) == 3:
+                    clone_layer = model.make_node('CloneThree', 'clone_' + node.name, attrs, [output], [output + '_cpy1', output + '_cpy2', output + '_cpy3'])
+                else:
+                    clone_layer = model.make_node('Clone', 'clone_' + node.name, attrs, [output], [output + '_cpy1', output + '_cpy2'])
                 model.insert_node(clone_layer)
                 transformed = True
 

hls4ml/templates/vivado/nnet_utils/nnet_stream.h

@@ -37,6 +37,32 @@ void clone_stream(hls::stream<data_T> &data, hls::stream<res_T> &res1, hls::stre
     }
 }
 
+template<class data_T, class res_T, int N>
+void clone_stream(hls::stream<data_T> &data, hls::stream<res_T> &res1, hls::stream<res_T> &res2, hls::stream<res_T> &res3) {
+    CloneLoop: for (int i = 0; i < N / data_T::size; i++) {
+        #pragma HLS PIPELINE
+
+        data_T in_data = data.read();
+        res_T out_data1;
+        res_T out_data2;
+        res_T out_data3;
+        #pragma HLS DATA_PACK variable=out_data1
+        #pragma HLS DATA_PACK variable=out_data2
+        #pragma HLS DATA_PACK variable=out_data3
+
+        ClonePack: for (int j = 0; j < data_T::size; j++) {
+            #pragma HLS UNROLL
+            out_data1[j] = in_data[j];
+            out_data2[j] = in_data[j];
+            out_data3[j] = in_data[j];
+        }
+
+        res1.write(out_data1);
+        res2.write(out_data2);
+        res3.write(out_data3);
+    }
+}
+
 template<class data_T, class res_T, int N>
 void repack_stream(hls::stream<data_T> &data, hls::stream<res_T> &res) {
     if (data_T::size == res_T::size) {