[microNPU] Adding rounding mode attribute to operators (#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.

Author: lhutton1

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,

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,

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,

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,

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 = {

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,

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,

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,
     }
 

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,

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,