From 6269e4e52c9cd5005deec834a7c4263d941270b1 Mon Sep 17 00:00:00 2001
From: "Javier M. Duarte" <jduarte@ucsd.edu>
Date: Fri, 23 Apr 2021 22:03:48 -0700
Subject: [PATCH] generalize for Conv2DBatchnorm and Pynq

---
 .../optimizer/passes/conv_single_output.py    | 51 ++++++++++++-------
 1 file changed, 33 insertions(+), 18 deletions(-)

diff --git a/hls4ml/model/optimizer/passes/conv_single_output.py b/hls4ml/model/optimizer/passes/conv_single_output.py
index 94abbd97b6..ec3e6c67ea 100644
--- a/hls4ml/model/optimizer/passes/conv_single_output.py
+++ b/hls4ml/model/optimizer/passes/conv_single_output.py
@@ -2,7 +2,7 @@
 import re
 
 from hls4ml.model.optimizer import OptimizerPass
-from hls4ml.model.hls_model import Conv1D, Conv2D, register_layer
+from hls4ml.model.hls_model import Conv1D, Conv2D, Conv2DBatchnorm, register_layer
 from hls4ml.templates import templates
 
 class SingleOutputConv1D(Conv1D):
@@ -17,6 +17,12 @@ class SingleOutputConv2D(Conv2D):
     # Nothing to do, will pick up function and config from class name
     pass
 
+class SingleOutputConv2DBatchnorm(Conv2DBatchnorm):
+    ''' Optimized Conv2DBatchnorm implementation for kernel_size = input_size resulting in single output pixel. '''
+
+    # Nothing to do, will pick up function and config from class name
+    pass
+
 single_out_conv1d_function_template = 'nnet::single_output_conv_1d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
 single_out_conv2d_function_template = 'nnet::single_output_conv_2d_{data_format}<{input_t}, {output_t}, {config}>({input}, {output}, {w}, {b});'
 
@@ -26,32 +32,41 @@ class SingleOutputConv2D(Conv2D):
 # Register the layer types to the layer map
 register_layer('SingleOutputConv1D', SingleOutputConv1D)
 register_layer('SingleOutputConv2D', SingleOutputConv2D)
+register_layer('SingleOutputConv2DBatchnorm', SingleOutputConv2DBatchnorm)
+
+# Register the templates for config and function                                                                                
+for backend in ['Vivado', 'Pynq']:
+    templates.get_backend(backend).register_templates(
+        'SingleOutputConv1D',
+        single_out_conv1d_function_template,
+        templates.get_backend(backend).get_config_template('Conv1D'),
+        single_out_conv1d_include_list
+    )
+
+    templates.get_backend(backend).register_templates(
+        'SingleOutputConv2D',
+        single_out_conv2d_function_template,
+        templates.get_backend(backend).get_config_template('Conv2D'),
+        single_out_conv2d_include_list
+    )
 
-# Register the templates for config and function
-templates.get_backend('Vivado').register_templates(
-    'SingleOutputConv1D',
-    single_out_conv1d_function_template,
-    templates.get_backend('Vivado').get_config_template('Conv1D'),
-    single_out_conv1d_include_list
-)
-
-templates.get_backend('Vivado').register_templates(
-    'SingleOutputConv2D',
-    single_out_conv2d_function_template,
-    templates.get_backend('Vivado').get_config_template('Conv2D'),
-    single_out_conv2d_include_list
-)
+    templates.get_backend(backend).register_templates(
+        'SingleOutputConv2DBatchnorm',
+        single_out_conv2d_function_template,
+        templates.get_backend(backend).get_config_template('Conv2DBatchnorm'),
+        single_out_conv2d_include_list
+    )
 
 class OptimizeSingleOutConv(OptimizerPass):
     def match(self, node):
-        return node.__class__.__name__ in ['Conv1D', 'Conv2D'] and \
+        return node.__class__.__name__ in ['Conv1D', 'Conv2D', 'Conv2DBatchnorm'] and \
             node.get_attr('filt_height', 1) == node.get_attr('in_height', 1) and \
             node.get_attr('filt_width') == node.get_attr('in_width') and \
             node.get_attr('out_height', 1) == 1 and node.get_attr('out_width') == 1
 
     def transform(self, model, node):
-        dim = node.__class__.__name__[-2:] # '1D' or '2D'
-        pw_node = model.make_node('SingleOutputConv' + dim, node.name, node.attributes.copy(), node.inputs.copy())
+        class_name = 'SingleOutput' + node.__class__.__name__
+        pw_node = model.make_node(class_name, node.name, node.attributes.copy(), node.inputs.copy())
         model.replace_node(node, pw_node)
         
         return True