[Trunc] add v1 and v2 versions of the op separately

Author: maltanar

src/qonnx/custom_op/general/__init__.py

@@ -35,7 +35,7 @@
 from qonnx.custom_op.general.maxpoolnhwc import MaxPoolNHWC
 from qonnx.custom_op.general.multithreshold import MultiThreshold
 from qonnx.custom_op.general.quantavgpool2d import QuantAvgPool2d
-from qonnx.custom_op.general.trunc import Trunc
+from qonnx.custom_op.general.trunc import Trunc_v1, Trunc_v2
 from qonnx.custom_op.general.xnorpopcount import XnorPopcountMatMul
 
 custom_op = dict()
@@ -49,7 +49,7 @@
 custom_op["Im2Col"] = Im2Col
 custom_op["IntQuant"] = IntQuant
 custom_op["Quant"] = IntQuant
-custom_op["Trunc"] = Trunc
+custom_op["Trunc"] = Trunc_v1
 custom_op["BipolarQuant"] = BipolarQuant
 custom_op["FloatQuant"] = FloatQuant
 
@@ -62,6 +62,8 @@
 custom_op["Im2Col_v1"] = Im2Col
 custom_op["IntQuant_v1"] = IntQuant
 custom_op["Quant_v1"] = IntQuant
-custom_op["Trunc_v1"] = Trunc
+custom_op["Trunc_v1"] = Trunc_v1
 custom_op["BipolarQuant_v1"] = BipolarQuant
 custom_op["FloatQuant_v1"] = FloatQuant
+
+custom_op["Trunc_v2"] = Trunc_v2

src/qonnx/custom_op/general/trunc.py

@@ -32,9 +32,10 @@
 from qonnx.core.datatype import DataType
 from qonnx.custom_op.base import CustomOp
 from qonnx.custom_op.general.quant import max_int, min_int, resolve_rounding_mode
+from qonnx.util.basic import get_preferred_qonnx_opset
 
 
-def trunc(inp_tensor, scale, zeropt, input_bit_width, narrow, signed, output_scale, output_bit_width, rounding_mode):
+def trunc_v2(inp_tensor, scale, zeropt, input_bit_width, narrow, signed, output_scale, output_bit_width, rounding_mode):
     # Port of TruncIntQuant class from Brevitas: https://bit.ly/3wzIpTR
 
     # Scaling
@@ -65,18 +66,23 @@ def trunc(inp_tensor, scale, zeropt, input_bit_width, narrow, signed, output_sca
     return y
 
 
-class Trunc(CustomOp):
-    """Generic truncation operation for QONNX. Takes four inputs:  
-    - input tensor to truncate  
-    - the scale  
-    - the zero-point  
-    - the truncation scale  
+class Trunc_v2(CustomOp):
+    """Generic truncation operation for QONNX. Takes four inputs:
+    - input tensor to truncate
+    - the scale
+    - the zero-point
+    - the truncation scale
     - the truncation bit-width
 
     The output is a tensor of the same shape as the input tensor, with truncated
     values.
     """
 
+    def __init__(self, onnx_node, onnx_opset_version=get_preferred_qonnx_opset()):
+        super().__init__(onnx_node, onnx_opset_version)
+        # override any specified opset version, this instance is v2
+        self.onnx_opset_version = 2
+
     def get_nodeattr_types(self):
         return {
             # The rounding mode, which is used for the trunc function
@@ -107,11 +113,79 @@ def execute_node(self, context, graph):
         narrow = self.get_nodeattr("narrow")
         signed = self.get_nodeattr("signed")
         # calculate output
-        ret = trunc(
+        ret = trunc_v2(
             inp_tensor, scale, zeropt, input_bit_width, narrow, signed, output_scale, output_bit_width, rounding_mode
         )
         # set context according to output name
         context[node.output[0]] = ret
 
     def verify_node(self):
         pass
+
+
+def trunc_v1(inp_tensor, scale, zeropt, input_bit_width, output_bit_width, rounding_mode):
+    # Port of TruncIntQuant class from Brevitas: https://bit.ly/3wzIpTR
+
+    # Scaling
+    y = inp_tensor / scale
+    y = y + zeropt
+    # Rounding
+    y = np.round(y)
+    # Truncate
+    trunc_bit_width = input_bit_width - output_bit_width
+    trunc_scale = 2.0**trunc_bit_width
+    y = y / trunc_scale
+
+    # To int
+    rounding_fx = resolve_rounding_mode(rounding_mode)
+    y = rounding_fx(y)
+
+    # Rescale
+    y = y - zeropt
+    y = y * scale
+
+    return y
+
+
+class Trunc_v1(CustomOp):
+    """Generic truncation operation for QONNX. Takes four inputs:
+    - input tensor to truncate
+    - the scale
+    - the zero-point
+    - the truncation bit-width
+
+    The output is a tensor of the same shape as the input tensor, with truncated
+    values.
+    """
+
+    def get_nodeattr_types(self):
+        return {
+            # The rounding mode, which is used for the trunc function
+            "rounding_mode": ("s", True, "FLOOR"),
+        }
+
+    def make_shape_compatible_op(self, model):
+        node = self.onnx_node
+        return helper.make_node("Identity", [node.input[0]], [node.output[0]])
+
+    def infer_node_datatype(self, model):
+        node = self.onnx_node
+        model.set_tensor_datatype(node.output[0], DataType["FLOAT32"])
+
+    def execute_node(self, context, graph):
+        node = self.onnx_node
+        # save inputs
+        inp_tensor = context[node.input[0]]
+        scale = context[node.input[1]]
+        zeropt = context[node.input[2]]
+        input_bit_width = context[node.input[3]]
+        output_bit_width = context[node.input[4]]
+        # save attributes
+        rounding_mode = self.get_nodeattr("rounding_mode")
+        # calculate output
+        ret = trunc_v1(inp_tensor, scale, zeropt, input_bit_width, output_bit_width, rounding_mode)
+        # set context according to output name
+        context[node.output[0]] = ret
+
+    def verify_node(self):
+        pass