[microNPU] Adding rounding mode attribute to operators (apache#9514)

* [microNPU] Adding rounding mode attribute to operators Allows rounding mode to be specified for each supported operator. By default "TFL" is used, which matches that of the behavior of TFLite. Other rounding mode options include "NATURAL" which rounds to the nearest value and "TRUNCATE" which rounds towards zero.
mehrdadh · Dec 1, 2021 · 02b3290 · 02b3290
1 parent 75b3353
commit 02b3290
Show file tree

Hide file tree

Showing 29 changed files with 371 additions and 114 deletions.
diff --git a/python/tvm/relay/backend/contrib/ethosu/op/binary_elementwise.py b/python/tvm/relay/backend/contrib/ethosu/op/binary_elementwise.py
@@ -45,6 +45,7 @@ def _extract_ethosu_binary_elementwise_params(attrs, args):
     activation = attrs.activation
     clip_min = attrs.clip_min
     clip_max = attrs.clip_max
+    rounding_mode = attrs.rounding_mode
     ifm_layout = attrs.ifm_layout
     ifm2_layout = attrs.ifm2_layout
     ofm_layout = attrs.ofm_layout
@@ -67,6 +68,7 @@ def _extract_ethosu_binary_elementwise_params(attrs, args):
         activation,
         clip_min,
         clip_max,
+        rounding_mode,
         ifm_layout,
         ifm2_layout,
         ofm_layout,
@@ -111,6 +113,7 @@ def ethosu_binary_elementwise(
     activation: Optional[str] = "NONE",
     clip_min: Optional[int] = 0,
     clip_max: Optional[int] = 0,
+    rounding_mode: Optional[str] = "TFL",
     ifm_layout: Optional[str] = "NHWC",
     ifm2_layout: Optional[str] = "NHWC",
     ofm_layout: Optional[str] = "NHWC",
@@ -179,6 +182,11 @@ def ethosu_binary_elementwise(
         The minimum clipping value if activation = "CLIP".
     clip_max : int, optional
         The maximum clipping value if activation = "CLIP".
+    rounding_mode : str, optional
+        The rounding mode to apply to the Output Feature Map tensor.
+            "TFL" - Tensorflow Lite rounding scheme.
+            "TRUNCATE" - Truncate towards zero.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     ifm_layout : str, optional
         The layout of the Input Feature Map tensor. Can be "NHWC" or "NHCWB16".
     ifm2_layout : str, optional
@@ -208,6 +216,7 @@ def ethosu_binary_elementwise(
         activation,
         clip_min,
         clip_max,
+        rounding_mode,
         ifm_layout,
         ifm2_layout,
         ofm_layout,

diff --git a/python/tvm/relay/backend/contrib/ethosu/op/convolution.py b/python/tvm/relay/backend/contrib/ethosu/op/convolution.py
@@ -45,6 +45,7 @@ def _extract_ethosu_conv2d_params(attrs, args):
     activation = attrs.activation
     clip_min = attrs.clip_min
     clip_max = attrs.clip_max
+    rounding_mode = attrs.rounding_mode
     upscale = attrs.upscale
     ifm_layout = attrs.ifm_layout
     ofm_layout = attrs.ofm_layout
@@ -65,6 +66,7 @@ def _extract_ethosu_conv2d_params(attrs, args):
         activation,
         clip_min,
         clip_max,
+        rounding_mode,
         upscale,
         ifm_layout,
         ofm_layout,
@@ -108,6 +110,7 @@ def ethosu_conv2d(
     activation: str = "NONE",
     clip_min: int = 0,
     clip_max: int = 0,
+    rounding_mode: str = "TFL",
     upscale: str = "NONE",
     ifm_layout: str = "NHWC",
     ofm_layout: str = "NHWC",
@@ -164,6 +167,11 @@ def ethosu_conv2d(
         The minimum clipping value if activation = "CLIP"
     clip_max : int, optional,
         The maximum clipping value if activation = "CLIP"
+    rounding_mode : str, optional
+        The rounding mode to apply to the Output Feature Map tensor.
+            "TFL" - Tensorflow Lite rounding scheme.
+            "TRUNCATE" - Truncate towards zero.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     upscale : str, optional
         The 2x2 upscaling mode to apply to the Input Feature Map tensor.
             "NONE" - no upscaling.
@@ -198,6 +206,7 @@ def ethosu_conv2d(
         activation,
         clip_min,
         clip_max,
+        rounding_mode,
         upscale,
         ifm_layout,
         ofm_layout,

diff --git a/python/tvm/relay/backend/contrib/ethosu/op/depthwise.py b/python/tvm/relay/backend/contrib/ethosu/op/depthwise.py
@@ -46,6 +46,7 @@ def _extract_ethosu_depthwise_conv2d_params(attrs, args):
     activation = attrs.activation
     clip_min = attrs.clip_min
     clip_max = attrs.clip_max
+    rounding_mode = attrs.rounding_mode
     upscale = attrs.upscale
     ifm_layout = attrs.ifm_layout
     ofm_layout = attrs.ofm_layout
@@ -66,6 +67,7 @@ def _extract_ethosu_depthwise_conv2d_params(attrs, args):
         activation,
         clip_min,
         clip_max,
+        rounding_mode,
         upscale,
         ifm_layout,
         ofm_layout,
@@ -109,6 +111,7 @@ def ethosu_depthwise_conv2d(
     activation: str = "NONE",
     clip_min: int = 0,
     clip_max: int = 0,
+    rounding_mode: str = "TFL",
     upscale: str = "NONE",
     ifm_layout: str = "NHWC",
     ofm_layout: str = "NHWC",
@@ -166,6 +169,11 @@ def ethosu_depthwise_conv2d(
         The minimum clipping value if activation = "CLIP"
     clip_max : int, optional,
         The maximum clipping value if activation = "CLIP"
+    rounding_mode : str, optional
+        The rounding mode to apply to the Output Feature Map tensor.
+            "TFL" - Tensorflow Lite rounding scheme.
+            "TRUNCATE" - Truncate towards zero.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     upscale : str, optional
         The 2x2 upscaling mode to apply to the Input Feature Map tensor.
             "NONE" - no upscaling.
@@ -200,6 +208,7 @@ def ethosu_depthwise_conv2d(
         activation,
         clip_min,
         clip_max,
+        rounding_mode,
         upscale,
         ifm_layout,
         ofm_layout,

diff --git a/python/tvm/relay/backend/contrib/ethosu/op/pooling.py b/python/tvm/relay/backend/contrib/ethosu/op/pooling.py
@@ -44,6 +44,7 @@ def _extract_ethosu_pooling_params(attrs, args):
     activation = attrs.activation
     clip_min = attrs.clip_min
     clip_max = attrs.clip_max
+    rounding_mode = attrs.rounding_mode
     upscale = attrs.upscale
     ifm_layout = attrs.ifm_layout
     ofm_layout = attrs.ofm_layout
@@ -63,6 +64,7 @@ def _extract_ethosu_pooling_params(attrs, args):
         activation,
         clip_min,
         clip_max,
+        rounding_mode,
         upscale,
         ifm_layout,
         ofm_layout,
@@ -103,6 +105,7 @@ def ethosu_pooling(
     activation: str = "NONE",
     clip_min: int = 0,
     clip_max: int = 0,
+    rounding_mode: str = "TFL",
     upscale: str = "NONE",
     ifm_layout: str = "NHWC",
     ofm_layout: str = "NHWC",
@@ -146,6 +149,11 @@ def ethosu_pooling(
         The minimum clipping value if activation = "CLIP".
     clip_max : int, optional
         The maximum clipping value if activation = "CLIP".
+    rounding_mode : str, optional
+        The rounding mode to apply to the Output Feature Map tensor.
+            "TFL" - Tensorflow Lite rounding scheme.
+            "TRUNCATE" - Truncate towards zero.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     upscale: str, optional
         The 2x2 upscaling mode to apply to the Input Feature Map tensor.
             "NONE" - no upscaling.
@@ -176,6 +184,7 @@ def ethosu_pooling(
         activation,
         clip_min,
         clip_max,
+        rounding_mode,
         upscale,
         ifm_layout,
         ofm_layout,

diff --git a/python/tvm/relay/backend/contrib/ethosu/te/binary_elementwise.py b/python/tvm/relay/backend/contrib/ethosu/te/binary_elementwise.py
@@ -38,6 +38,7 @@ def binary_elementwise_compute(
     activation: str,
     clip_min: int,
     clip_max: int,
+    rounding_mode: str,
     ifm_layout: str,
     ifm2_layout: str,
     ofm_layout: str,
@@ -94,6 +95,11 @@ def binary_elementwise_compute(
         The minimum clipping value if activation = "CLIP".
     clip_max : int
         The maximum clipping value if activation = "CLIP".
+    rounding_mode : str
+        The rounding mode to apply to the Output Feature Map tensor.
+            "TFL" - Tensorflow Lite rounding scheme.
+            "TRUNCATE" - Truncate towards zero.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     ifm_layout : str, optional
         The layout of the Input Feature Map tensor. Can be "NHWC" or "NHCWB16".
     ifm2_layout : str, optional
@@ -136,6 +142,7 @@ def binary_elementwise_compute(
         "activation": activation,
         "clip_min": clip_min,
         "clip_max": clip_max,
+        "rounding_mode": rounding_mode,
     }
 
     operators = {

diff --git a/python/tvm/relay/backend/contrib/ethosu/te/convolution.py b/python/tvm/relay/backend/contrib/ethosu/te/convolution.py
@@ -38,6 +38,7 @@ def conv2d_compute(
     activation: str,
     clip_min: int,
     clip_max: int,
+    rounding_mode: str,
     upscale: str,
     ifm_layout: str,
     ofm_layout: str,
@@ -81,11 +82,17 @@ def conv2d_compute(
         The minimum clipping value if activation = "CLIP".
     clip_max : int
         The maximum clipping value if activation = "CLIP".
+    rounding_mode : str
+        The rounding mode to apply to the Output Feature Map tensor.
+            "TFL" - Tensorflow Lite rounding scheme.
+            "TRUNCATE" - Truncate towards zero.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     upscale : str
         The 2x2 upscaling mode to apply to the Input Feature Map tensor.
             "NONE" - no upscaling.
             "NEAREST" - upscale using nearest neighbour.
             "ZEROS" - upscale using zeros.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     ifm_layout : str
         The layout of the Input Feature Map tensor. Can be "NHWC" or "NHCWB16".
     ofm_layout : str
@@ -126,6 +133,7 @@ def conv2d_compute(
         "upscale": upscale,
         "clip_min": clip_min,
         "clip_max": clip_max,
+        "rounding_mode": rounding_mode,
         "stride_h": stride_h,
         "stride_w": stride_w,
         "dilation_h": dilation_h,

diff --git a/python/tvm/relay/backend/contrib/ethosu/te/depthwise.py b/python/tvm/relay/backend/contrib/ethosu/te/depthwise.py
@@ -38,6 +38,7 @@ def depthwise_conv2d_compute(
     activation: str,
     clip_min: int,
     clip_max: int,
+    rounding_mode: str,
     upscale: str,
     ifm_layout: str,
     ofm_layout: str,
@@ -81,6 +82,11 @@ def depthwise_conv2d_compute(
         The minimum clipping value if activation = "CLIP".
     clip_max : int
         The maximum clipping value if activation = "CLIP".
+    rounding_mode : str
+        The rounding mode to apply to the Output Feature Map tensor.
+            "TFL" - Tensorflow Lite rounding scheme.
+            "TRUNCATE" - Truncate towards zero.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     upscale : str
         The 2x2 upscaling mode to apply to the Input Feature Map tensor.
             "NONE" - no upscaling.
@@ -120,9 +126,10 @@ def depthwise_conv2d_compute(
         "op": "ethosu_depthwise_conv2d",
         "weight_zero_point": weight_zero_point,
         "activation": activation,
-        "upscale": upscale,
         "clip_min": clip_min,
         "clip_max": clip_max,
+        "rounding_mode": rounding_mode,
+        "upscale": upscale,
         "stride_h": stride_h,
         "stride_w": stride_w,
         "dilation_h": dilation_h,

diff --git a/python/tvm/relay/backend/contrib/ethosu/te/pooling.py b/python/tvm/relay/backend/contrib/ethosu/te/pooling.py
@@ -37,6 +37,7 @@ def pooling_compute(
     activation: str,
     clip_min: int,
     clip_max: int,
+    rounding_mode: str,
     upscale: str,
     ifm_layout: str,
     ofm_layout: str,
@@ -78,6 +79,11 @@ def pooling_compute(
         The minimum clipping value if activation = "CLIP".
     clip_max : int
         The maximum clipping value if activation = "CLIP".
+    rounding_mode : str
+        The rounding mode to apply to the Output Feature Map tensor.
+            "TFL" - Tensorflow Lite rounding scheme.
+            "TRUNCATE" - Truncate towards zero.
+            "NATURAL" - Round to nearest value, with x.5 rounded up towards +infinity.
     upscale : str
         The 2x2 upscaling mode to apply to the Input Feature Map tensor.
             "NONE" - no upscaling.
@@ -113,6 +119,7 @@ def pooling_compute(
         "activation": activation,
         "clip_min": clip_min,
         "clip_max": clip_max,
+        "rounding_mode": rounding_mode,
         "upscale": upscale,
     }
 

diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/binary_elementwise.py b/python/tvm/relay/backend/contrib/ethosu/tir/binary_elementwise.py
@@ -96,6 +96,7 @@ def get_binary_elementwise_params(
             operator_type=attrs["operator_type"],
             reversed_operands=reversed_operands,
             activation=serial_activation,
+            rounding_mode=attrs["rounding_mode"],
         ),
         output_pointer,
         replace_pointer,

diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/convolution.py b/python/tvm/relay/backend/contrib/ethosu/tir/convolution.py
@@ -99,6 +99,7 @@ def get_conv2d_params(stmt, producers, consumers):
             scale_bias=serial_scale_bias,
             padding=serial_padding,
             activation=serial_activation,
+            rounding_mode=attrs["rounding_mode"],
             upscale="NONE",
         ),
         output_pointer,

diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/depthwise.py b/python/tvm/relay/backend/contrib/ethosu/tir/depthwise.py
@@ -109,6 +109,7 @@ def get_depthwise_conv2d_params(
             scale_bias=serial_scale_bias,
             padding=serial_padding,
             activation=serial_activation,
+            rounding_mode=attrs["rounding_mode"],
             upscale="NONE",
         ),
         output_pointer,

diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/identity.py b/python/tvm/relay/backend/contrib/ethosu/tir/identity.py
@@ -148,6 +148,7 @@ def get_identity_params(
             padding=SerialPadding(0, 0, 0, 0),
             activation=serial_activation,
             upscale="NONE",
+            rounding_mode="TFL",
         ),
         output_pointer,
         replace_pointer,

diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/pooling.py b/python/tvm/relay/backend/contrib/ethosu/tir/pooling.py
@@ -83,6 +83,7 @@ def get_pooling_params(
             pool_shape=serial_kernel,
             padding=serial_padding,
             activation=serial_activation,
+            rounding_mode=attrs["rounding_mode"],
             upscale="NONE",
         ),
         output_pointer,

diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/spec.py b/python/tvm/relay/backend/contrib/ethosu/tir/spec.py
@@ -188,6 +188,7 @@ def __init__(
         scale_bias: SerialAddressRange,
         padding: SerialPadding,
         activation: SerialActivation,
+        rounding_mode: str,
         upscale: str,
     ):
         self.ifm = ifm
@@ -198,6 +199,7 @@ def __init__(
         self.scale_bias = scale_bias
         self.padding = padding
         self.activation = activation
+        self.rounding_mode = rounding_mode
         self.upscale = upscale
 
 
@@ -215,6 +217,7 @@ def __init__(
         scale_bias: SerialAddressRange,
         padding: SerialPadding,
         activation: SerialActivation,
+        rounding_mode: str,
         upscale: str,
     ):
         self.ifm = ifm
@@ -225,6 +228,7 @@ def __init__(
         self.scale_bias = scale_bias
         self.padding = padding
         self.activation = activation
+        self.rounding_mode = rounding_mode
         self.upscale = upscale
 
 
@@ -252,6 +256,7 @@ def __init__(
         pool_shape: SerialKernel,
         padding: SerialPadding,
         activation: SerialActivation,
+        rounding_mode: str,
         upscale: str,
     ):
         self.ifm = ifm
@@ -260,6 +265,7 @@ def __init__(
         self.pool_shape = pool_shape
         self.padding = padding
         self.activation = activation
+        self.rounding_mode = rounding_mode
         self.upscale = upscale
 
 
@@ -275,10 +281,12 @@ def __init__(
         operator_type: str,
         reversed_operands: bool,
         activation: SerialActivation,
+        rounding_mode: str,
     ):
         self.ifm = ifm
         self.ifm2 = ifm2
         self.ofm = ofm
         self.operator_type = operator_type
         self.reversed_operands = reversed_operands
         self.activation = activation
+        self.rounding_mode = rounding_mode