diff --git a/python/tvm/relay/backend/contrib/ethosu/legalize.py b/python/tvm/relay/backend/contrib/ethosu/legalize.py index 242d3c2d0cc5..2806ef8a4699 100644 --- a/python/tvm/relay/backend/contrib/ethosu/legalize.py +++ b/python/tvm/relay/backend/contrib/ethosu/legalize.py @@ -1216,6 +1216,7 @@ def callback( ifm_zero_point=int(params.ifm.q_params.zero_point), ofm_scale=float(params.ofm.q_params.scale_f32), ofm_zero_point=int(params.ofm.q_params.zero_point), + rounding_mode="NATURAL", ) diff --git a/python/tvm/relay/backend/contrib/ethosu/op/identity.py b/python/tvm/relay/backend/contrib/ethosu/op/identity.py index f070144ba294..d91de971dbf2 100644 --- a/python/tvm/relay/backend/contrib/ethosu/op/identity.py +++ b/python/tvm/relay/backend/contrib/ethosu/op/identity.py @@ -36,8 +36,16 @@ def create_ethosu_identity_compute(attrs, args, out_type): ofm_scale = attrs.ofm_scale ofm_zero_point = attrs.ofm_zero_point activation = attrs.activation + rounding_mode = attrs.rounding_mode op = identity_compute( - ifm, lut, ifm_scale, ifm_zero_point, ofm_scale, ofm_zero_point, activation + ifm, + lut, + ifm_scale, + ifm_zero_point, + ofm_scale, + ofm_zero_point, + activation, + rounding_mode, ) return [op] @@ -61,6 +69,7 @@ def ethosu_identity( ofm_scale: float = 1, ofm_zero_point: int = 0, activation: str = "NONE", + rounding_mode: str = "TFL", ) -> tvm.relay.Call: """The Identity operator that runs on the NPU. @@ -87,6 +96,11 @@ def ethosu_identity( "TANH" - tanh activation function. "SIGMOID" - sigmoid activation function. "LUT" - use a look-up table to perform the activation function. + 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. Returns ------- @@ -94,5 +108,5 @@ def ethosu_identity( A call to the ethosu_identity op. """ return _make.ethosu_identity( - ifm, lut, ifm_scale, ifm_zero_point, ofm_scale, ofm_zero_point, activation + ifm, lut, ifm_scale, ifm_zero_point, ofm_scale, ofm_zero_point, activation, rounding_mode ) diff --git a/python/tvm/relay/backend/contrib/ethosu/te/identity.py b/python/tvm/relay/backend/contrib/ethosu/te/identity.py index d2ffcee085ac..7f9bcebf70d4 100644 --- a/python/tvm/relay/backend/contrib/ethosu/te/identity.py +++ b/python/tvm/relay/backend/contrib/ethosu/te/identity.py @@ -31,6 +31,7 @@ def identity_compute( ofm_scale: float, ofm_zero_point: int, activation: str, + rounding_mode: str, ) -> te.Tensor: """A compute operator for the NPU identity operator. @@ -54,6 +55,11 @@ def identity_compute( "TANH" - tanh activation function. "SIGMOID" - sigmoid activation function. "LUT" - use a look-up table to perform the activation function. + 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. Returns ------- @@ -61,7 +67,7 @@ def identity_compute( The Output Feature Map tensor. """ dmaed_ifm = read_compute(ifm, ifm_zero_point, ifm_scale) - id_attrs = {"op": "ethosu_identity", "activation": activation} + id_attrs = {"op": "ethosu_identity", "activation": activation, "rounding_mode": rounding_mode} has_lut = activation in ("TANH", "LUT", "SIGMOID") diff --git a/python/tvm/relay/backend/contrib/ethosu/tir/identity.py b/python/tvm/relay/backend/contrib/ethosu/tir/identity.py index 43ae52b3bae7..9610c8dd3cdc 100644 --- a/python/tvm/relay/backend/contrib/ethosu/tir/identity.py +++ b/python/tvm/relay/backend/contrib/ethosu/tir/identity.py @@ -166,7 +166,7 @@ def get_identity_params( padding=SerialPadding(0, 0, 0, 0), activation=serial_activation, upscale="NONE", - rounding_mode="TFL", + rounding_mode=attrs["rounding_mode"], block_config=SerialBlockConfig(0, 0, 0), ), output_pointer, diff --git a/src/relay/op/contrib/ethosu/identity.cc b/src/relay/op/contrib/ethosu/identity.cc index 9b00978d43d8..f808e8c21902 100644 --- a/src/relay/op/contrib/ethosu/identity.cc +++ b/src/relay/op/contrib/ethosu/identity.cc @@ -63,13 +63,14 @@ bool EthosuIdentityRel(const Array& types, int num_inputs, const Attrs& at } Expr MakeEthosuIdentity(Expr ifm, Expr lut, double ifm_scale, int ifm_zero_point, double ofm_scale, - int ofm_zero_point, String activation) { + int ofm_zero_point, String activation, String rounding_mode) { auto attrs = make_object(); attrs->ifm_scale = ifm_scale; attrs->ifm_zero_point = ifm_zero_point; attrs->ofm_scale = ofm_scale; attrs->ofm_zero_point = ofm_zero_point; attrs->activation = std::move(activation); + attrs->rounding_mode = std::move(rounding_mode); static const Op& op = Op::Get("contrib.ethosu.identity"); return Call(op, {ifm, lut}, Attrs(attrs), {}); } diff --git a/src/relay/op/contrib/ethosu/op_attrs.h b/src/relay/op/contrib/ethosu/op_attrs.h index 74e7fe856e89..868d9d6ad42f 100644 --- a/src/relay/op/contrib/ethosu/op_attrs.h +++ b/src/relay/op/contrib/ethosu/op_attrs.h @@ -319,6 +319,7 @@ struct EthosuIdentityAttrs : public tvm::AttrsNode { double ofm_scale; int ofm_zero_point; String activation; + String rounding_mode; TVM_DECLARE_ATTRS(EthosuIdentityAttrs, "relay.attrs.EthosuIdentityAttrs") { TVM_ATTR_FIELD(ifm_scale).describe("The quantization scale for the Input Feature Map tensor."); @@ -335,6 +336,13 @@ struct EthosuIdentityAttrs : public tvm::AttrsNode { "'SIGMOID' - sigmoid activation function. " "'LUT' - use a look-up table to perform the activation function.") .set_default("NONE"); + TVM_ATTR_FIELD(rounding_mode) + .describe( + "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.") + .set_default("TFL"); } }; diff --git a/tests/python/contrib/test_ethosu/cascader/test_ethosu_identity_matcher.py b/tests/python/contrib/test_ethosu/cascader/test_ethosu_identity_matcher.py index 4bdccfced33f..11d76ab2b8dd 100644 --- a/tests/python/contrib/test_ethosu/cascader/test_ethosu_identity_matcher.py +++ b/tests/python/contrib/test_ethosu/cascader/test_ethosu_identity_matcher.py @@ -39,6 +39,7 @@ def test_ethosu_identity_matcher(): ofm_scale=1, ofm_zero_point=0, activation="NONE", + rounding_mode="TFL", ) length = len(ifm.shape) diff --git a/tests/python/contrib/test_ethosu/test_codegen.py b/tests/python/contrib/test_ethosu/test_codegen.py index fde2e284347e..66809a775f48 100644 --- a/tests/python/contrib/test_ethosu/test_codegen.py +++ b/tests/python/contrib/test_ethosu/test_codegen.py @@ -1233,6 +1233,7 @@ def split_func(x): [ [(1, 8, 8, 3), 1.0, 0, 1.0, 0], [(1, 20, 30, 3), 1.345, 34, 0.32, -23], + [(1, 1, 4, 8), 0.0078125, 0, 0.00997, -30], ], ) def test_ethosu_requantize(accel_type, ifm_shape, ifm_scale, ifm_zp, ofm_scale, ofm_zp):