Add BatchNorm, Mul, Add, Sub, Mean, and more ops

Author: mbartling

utensor_cgen/backend/utensor/_graph_lower/_op_lower.py

@@ -41,6 +41,8 @@ def __init__(self, config):
   class OptypeRenameManager(object):
     NAME_MAP = {
       'Add': 'AddOperator',
+      'Mul': 'MulOperator',
+      'Sub': 'SubOperator',
       'Conv2D': 'ConvOperator',
       'MatMul': 'MatrixMultOperator'
     }

utensor_cgen/backend/utensor/code_generator/rearch/_operators/_impls.py

@@ -436,6 +436,39 @@ def get_eval_snippet(self, op_var_name, op_info, tensor_var_map):
       tensor_var_map=tensor_var_map,
     )
 
+@OperatorFactory.register
+class _ConvOperator(_CommonParams):
+  op_type = "ConvOperator"
+
+  @classmethod
+  @must_return_type(Hashable)
+  def get_constructor_parameters(cls, op_info):
+
+    strides = [
+        1,
+        op_info.op_attr['StrideW'],
+        op_info.op_attr['StrideH'],
+        1,
+      ]
+    padding = cls._PADDING_MAP[op_info.op_attr['Padding']]
+    strides_str = ','.join(map(str, strides))
+    return ("{{ {} }}".format(strides_str), padding)
+
+  def get_declare_snippet(self, op_var_name, tensor_var_map):
+    return DeclareOpSnippet(
+      op=self,
+      templ_dtypes=[self.out_dtypes[0]],
+      op_var_name=op_var_name,
+    )
+
+  def get_eval_snippet(self, op_var_name, op_info, tensor_var_map):
+    return ConvOpEvalSnippet(
+      op_info=op_info,
+      templ_dtypes=[self.out_dtypes[0]],
+      op_name=op_var_name,
+      tensor_var_map=tensor_var_map,
+    )
+
 
 @OperatorFactory.register
 class _QuantizedFullyConnectedOperator(_CommonParams):
@@ -464,3 +497,123 @@ def get_eval_snippet(self, op_var_name, op_info, tensor_var_map):
       op_name=op_var_name,
       tensor_var_map=tensor_var_map,
     )
+
+@OperatorFactory.register
+class _BatchNormOperator(_CommonParams):
+  op_type = "BatchNormOperator"
+
+  @classmethod
+  @must_return_type(Hashable)
+  def get_constructor_parameters(cls, op_info):
+    strides = [
+        1,
+        op_info.op_attr['StrideW'],
+        op_info.op_attr['StrideH'],
+        1,
+      ]
+    padding = cls._PADDING_MAP[op_info.op_attr['Padding']]
+    strides_str = ','.join(map(str, strides))
+    return ("{{ {} }}".format(strides_str), padding)
+
+  def get_declare_snippet(self, op_var_name, tensor_var_map):
+    return DeclareOpSnippet(
+      op=self,
+      templ_dtypes=[self.out_dtypes[0]],
+      op_var_name=op_var_name,
+    )
+
+  def get_eval_snippet(self, op_var_name, op_info, tensor_var_map):
+    return BatchNormSnippet(
+      op_info=op_info,
+      templ_dtypes=[self.out_dtypes[0]],
+      op_name=op_var_name,
+      tensor_var_map=tensor_var_map,
+    )
+
+@OperatorFactory.register
+class _MeanOperator(_CommonParams):
+  op_type = "MeanOperator"
+
+  @classmethod
+  @must_return_type(Hashable)
+  def get_constructor_parameters(cls, op_info):
+    keep_dims = str(op_info.op_attr["keep_dims"])
+    return (" {} ".format(keep_dims), )
+
+  def get_declare_snippet(self, op_var_name, tensor_var_map):
+    return DeclareOpSnippet(
+      op=self,
+      templ_dtypes=[self.out_dtypes[0]],
+      op_var_name=op_var_name,
+    )
+
+  def get_eval_snippet(self, op_var_name, op_info, tensor_var_map):
+    return BatchNormSnippet(
+      op_info=op_info,
+      templ_dtypes=[self.out_dtypes[0]],
+      op_name=op_var_name,
+      tensor_var_map=tensor_var_map,
+    )
+
+@OperatorFactory.register
+class _SoftmaxOperator(_CommonParams):
+  op_type = "SoftmaxOperator"
+
+  @classmethod
+  @must_return_type(Hashable)
+  def get_constructor_parameters(cls, op_info):
+    Beta = op_info.op_attr["Beta"]
+    return (" %f " % Beta,)
+
+  def get_declare_snippet(self, op_var_name, tensor_var_map):
+    return DeclareOpSnippet(
+      op=self,
+      templ_dtypes=[self.out_dtypes[0]],
+      op_var_name=op_var_name,
+    )
+
+  def get_eval_snippet(self, op_var_name, op_info, tensor_var_map):
+    return BatchNormSnippet(
+      op_info=op_info,
+      templ_dtypes=[self.out_dtypes[0]],
+      op_name=op_var_name,
+      tensor_var_map=tensor_var_map,
+    )
+
+@OperatorFactory.register
+class _MulOperator(_Operator):
+  op_type = 'MulOperator'
+
+  def get_declare_snippet(self, op_var_name, tensor_var_map):
+    return DeclareOpSnippet(
+      op=self,
+      templ_dtypes=[self.in_dtypes[0]],
+      op_var_name=op_var_name,
+    )
+
+  def get_eval_snippet(self, op_var_name, op_info, tensor_var_map):
+    return MulOpEvalSnippet(
+      op_info=op_info,
+      templ_dtypes=[self.in_dtypes[0]],
+      op_name=op_var_name,
+      tensor_var_map=tensor_var_map,
+    )
+
+@OperatorFactory.register
+class _SubOperator(_Operator):
+  op_type = 'SubOperator'
+
+  def get_declare_snippet(self, op_var_name, tensor_var_map):
+    return DeclareOpSnippet(
+      op=self,
+      templ_dtypes=[self.in_dtypes[0]],
+      op_var_name=op_var_name,
+    )
+
+  def get_eval_snippet(self, op_var_name, op_info, tensor_var_map):
+    return SubOpEvalSnippet(
+      op_info=op_info,
+      templ_dtypes=[self.in_dtypes[0]],
+      op_name=op_var_name,
+      tensor_var_map=tensor_var_map,
+    )

utensor_cgen/backend/utensor/snippets/rearch/_snippets.py

@@ -27,6 +27,12 @@
   "MinPoolEvalSnippet",
   "MaxPoolEvalSnippet",
   "QuantizedFullyConnectedSnippet",
+  "BatchNormSnippet",
+  "MulOpEvalSnippet",
+  "SubOpEvalSnippet",
+  "ConvOpEvalSnippet",
+  "MeanOpEvalSnippet",
+  "SoftmaxOpEvalSnippet",
   "SimpleContainer",
 ]
 
@@ -143,6 +149,9 @@ class DepthwiseSeperateConvOpEvalSnippet(OpEvalSnippet):
   __inputs__ = ["in", "depthwise_filter", "pointwise_filter"]
   __outputs__ = ["out"]
 
+class ConvOpEvalSnippet(OpEvalSnippet):
+  __inputs__ = ["in", "filter"]
+  __outputs__ = ["out"]
 
 class QuantDepthwiseSeperateConvOpEvalSnippet(OpEvalSnippet):
   __inputs__ = ["in", "filter", "bias"]
@@ -218,6 +227,23 @@ class QuantizedFullyConnectedSnippet(OpEvalSnippet):
   __inputs__ = ["input", "filter", "bias"]
   __outputs__ = ["output"]
 
+class BatchNormSnippet(OpEvalSnippet):
+  __inputs__ = ["x", "mean", "variance", "offset", "scale"]
+  __outputs__ = ["output"]
+
+class MulOpEvalSnippet(OpEvalSnippet):
+  __inputs__ = ['a', 'b']
+  __outputs__ = ['c']
+class SubOpEvalSnippet(OpEvalSnippet):
+  __inputs__ = ['a', 'b']
+  __outputs__ = ['c']
+class MeanOpEvalSnippet(OpEvalSnippet):
+  __inputs__ = ['input', 'axis']
+  __outputs__ = ['output']
+class SoftmaxOpEvalSnippet(OpEvalSnippet):
+  __inputs__ = ['input']
+  __outputs__ = ['output']
+
 
 class SimpleContainer(SnippetBase):
   __headers__ = set(['"uTensor.h"', "<vector>"])