[microNPU] Add support for TFLite FULLY_CONNECTED by dchauhan-arm · Pull Request #10345 · apache/tvm

dchauhan-arm · 2022-02-22T12:53:59Z

This is primarily a legalization to an NPU Conv2d operator. The
legalization target is Conv2d with 1 1 I O (HWIO)

This is primarily a legalization to an NPU Conv2d operator. The legalization target is Conv2d with 1 1 I O (HWIO)

Test TVM runtime against TFLite for codegen and operator legalization.

dchauhan-arm · 2022-02-22T12:54:50Z

cc @ekalda @manupa-arm @lhutton1 @NicolaLancellotti

lhutton1

Thanks @dchauhan-arm, this will be a very useful addition. Although this is still WIP I just wanted to offer a couple of suggestions which may help get this in!

lhutton1 · 2022-02-22T19:14:06Z

python/tvm/relay/op/contrib/ethosu.py

+    )
+    bias_add = is_op("nn.bias_add")(dense, is_constant())
+    req = is_op("qnn.requantize")(
+        dense | bias_add, is_constant(), is_constant(), is_constant(), is_constant()


Currently I the legalization will fall over if there is not a bias present. We should make bias optional in FullyConnectedParams, see QnnTransposeConv2dParams for an idea

ack! (and thanks for the pointer on transpose conv2d)

lhutton1 · 2022-02-24T09:29:15Z

tests/python/contrib/test_ethosu/test_codegen.py

+):
+    @tf.function
+    def fully_connected():
+        return tf.keras.layers.Dense(


I'm not too familiar with the Keras API, but I'm not sure this will work. One thing we could do instead is use tf.matmul which gets legalized to fully connected in TFLite under the conditions we will use it for. e.g. something like this would be a starting point:

@tf.function def dense_layer(x): w = tf.constant( np.random.uniform(size=[units, units]), dtype=tf.float32, ) return tf.matmul(x, w) _compare_tvm_with_tflite(dense_layer, [(1, units)], accel_type)

Happy to keep the Keras implementation if we get it working though, just wanted to offer an alternative :)

this is a very welcome change, I'l try and make this work!

lhutton1 · 2022-02-24T09:39:22Z

python/tvm/relay/op/contrib/ethosu.py

+            requantize_op.args[RequantArgs.IFM_ZERO_POINT.value],
+        )
+        self.ifm = TensorParams(
+            qnn_dense.args[QDenseArgs.ifm.value],


ifm should be capitals i.e. QDenseArgs.IFM.value

Fix linting

ekalda

Thanks @dchauhan-arm, broadly looks good, I left some suggestions for improvements :)

ekalda · 2022-02-24T11:40:34Z

python/tvm/relay/backend/contrib/ethosu/legalize.py

+        # IFM reshapes
+        ifm = post.args[0]
+        if len(params.ifm.shape) != 4 or not params.ifm.shape[1] == params.ifm.shape[2] == 1:
+            ifm = relay.reshape(ifm, (-1, 1, 1, params.ifm.shape[-1]))


Suggested change

ifm = relay.reshape(ifm, (-1, 1, 1, params.ifm.shape[-1]))

ifm = relay.reshape(ifm, (1, 1, 1, params.ifm.shape[-1]))

should be safer since the NPU doesn't support IFMs with a batch size anything other than 1 and this kind of fully connected wouldn't be offloaded anyway

ekalda · 2022-02-24T11:42:40Z

python/tvm/relay/backend/contrib/ethosu/legalize.py

+    def callback(self, pre, post, node_map):
+        params = ethosu_patterns.FullyConnectedParams(post.op.body)
+        params.ifm.tensor = post.args[0]
+        activation_map = {"clip": "CLIP"}


nit: we don't expect that dict to expand, so we can just do if activation == "clip": etc

ekalda · 2022-02-24T11:45:31Z

python/tvm/relay/backend/contrib/ethosu/legalize.py

+        if len(params.ofm.shape) != 4 or not params.ofm.shape[1] == params.ofm.shape[2] == 1:
+            ethosu_fc = relay.reshape(ethosu_fc, params.ofm.shape)


I suspect there isn't a test case that exercises this case since on line 1700 this pass runs after the no op legalizer, so the last reshape won't have a following identity op and will fall over in TE

ekalda · 2022-02-24T11:46:11Z

python/tvm/relay/backend/contrib/ethosu/legalize.py

+
+@ir.transform.module_pass(opt_level=1)
+class LegalizeFullyConnected:
+    """This is the pass that wraps the AddRewriter"""


Suggested change

"""This is the pass that wraps the AddRewriter"""

"""This is the pass that wraps the FullyConnectedRewriter"""

python/tvm/relay/backend/contrib/ethosu/legalize.py

ekalda · 2022-02-24T11:48:12Z

python/tvm/relay/op/contrib/ethosu.py

+    """
+
+    composite_name = "ethosu.fully_connected"
+    activation_map = {"clip": "CLIP"}


Same nit about the clip dict as before :)

ekalda · 2022-02-24T12:08:28Z

tests/python/contrib/test_ethosu/test_legalize.py

+        if activation_function == "RELU":
+            assert str(op.attrs.activation) == "CLIP"
+
+    dense_pattern_table = [


nit: it would be better to keep the naming consistent, so maybe rename this to fc_pattern_table or fully_connected_pattern_table

ekalda · 2022-02-24T12:08:58Z

tests/python/contrib/test_ethosu/test_legalize.py

+
+        # check IFM
+        ifm = op.args[0].checked_type
+        assert list([1, 3, units, 1]) == list([1, 3, units, 1])


This assert doesn't check anything... Some things to potentially check:

That we have ended up with a ethosu_conv2d op (taking into account that there might be reshape ops before and after the conv2d)

That the IFM is in a shape of (1, 1, 1, c)

That the weights are in a shape (o, 1, 1, c) with o being the output channels of the weights

That the kernel and dilation are (1, 1)

Address comments, update codegen test, fix linting.

Address more comments, ensure qnn.dense is lowered to NPU, fix linting

Fix linting, update legalization test and codegen test for completeness.

lhutton1

Thanks for the update @dchauhan-arm, looking much better! Mostly just a few stylistic suggestions below

lhutton1 · 2022-03-04T11:48:42Z

python/tvm/relay/op/contrib/ethosu.py

 import tvm  # type: ignore
 from tvm import relay
-from tvm.relay.expr import Constant, Call  # type: ignore
+from tvm.relay.expr import Constant, Call


Nit: I don't think we need to change this

python/tvm/relay/op/contrib/ethosu.py

lhutton1 · 2022-03-04T12:01:19Z

python/tvm/relay/op/contrib/ethosu.py

+        if not np.all(np.array(self.ifm.shape[:-1]) == 1):
+            # As we reshape the ifm from
+            # [n0, n1, ... , n_m] to [n0 * n1 * ... * n_{m-1}, n_m]
+            # all except the last dims need to be 1.
+            return False


I don't think we need this due to reasoning in the above comment and since we already check that the batch size == 1 with check_batch_size above and we know that the ifm must be 2D

lhutton1 · 2022-03-04T12:04:02Z

python/tvm/relay/op/contrib/ethosu.py

+        return True  # optional_bias_add = (
+
+    #     is_op("nn.bias_add")(dense, is_constant()) | dense
+    # )


Nit: remove comments :)

lhutton1 · 2022-03-04T12:04:55Z

python/tvm/relay/op/contrib/ethosu.py

+    dense = is_op("qnn.dense")(
+        wildcard(), is_constant(), is_constant(), is_constant(), is_constant(), is_constant()
+    )
+    optional_bias_add = is_op("nn.bias_add")(dense, is_constant()) | dense


I think this should just be optional_bias_add = is_op("nn.bias_add")(dense, is_constant())

lhutton1 · 2022-03-04T12:15:24Z

tests/python/contrib/test_ethosu/test_legalize.py

+
+        # check OFM
+        ofm = op.checked_type
+        assert [ofm.shape[2], ofm.shape[3]] == [1, ofm_channels]


Same as above, lets alter this to check the whole ofm shape assert ofm.shape == ...

This would need to be assert list(ofm.shape) == [1, 1, 1, ofm_channels]

lhutton1 · 2022-03-04T12:15:51Z

tests/python/contrib/test_ethosu/test_legalize.py

+        assert [ofm.shape[2], ofm.shape[3]] == [1, ofm_channels]
+        # assert list(ofm.shape) == list(expected_ofm_shape)
+        assert str(ofm.dtype) == dtype
+        # assert ofm.shape[3] == ofm_channels


Nit: remove some comments

lhutton1 · 2022-03-04T12:17:01Z

tests/python/contrib/test_ethosu/test_legalize.py

+        assert weights_ohwi.shape[0] == ofm_channels
+        assert weights_ohwi.shape[1] == 1
+        assert weights_ohwi.shape[2] == 1
+        assert weights_ohwi.shape[3] == ifm_shape[1]


Nit: we could do this with one assert e.g. assert list(weights_ohwi) == [ofm_channels, 1, 1, ifm_shape[1]

lhutton1 · 2022-03-04T12:18:31Z

tests/python/contrib/test_ethosu/test_legalize.py

+        #     (1, 1),
+        #     (1, 1),
+        #     (1, 1),
+        # )


Nit: comments

lhutton1 · 2022-03-04T13:05:38Z

tests/python/contrib/test_ethosu/test_legalize.py

+        #     (1, 1),
+        #     (1, 1),
+        # )
+        assert list(op.attrs.padding) == [0, 0, 0, 0]


Nit: might also be worth checking the op name is an NPU convolution here as well

lhutton1 · 2022-03-04T13:31:42Z

Looks like CI only failed because of a flaky GPU test :)

Address comments, fix linting. Certain legalization test assertions were updated. Co-authored-by: Rishabh Jain <rishabh.jain2@arm.com>

Fix assertion in legalization test.

lhutton1 · 2022-03-09T11:34:19Z

tests/python/contrib/test_ethosu/test_legalize.py

+
+        # check IFM
+        ifm = op.args[0].checked_type
+        assert [ifm.shape[2], ifm.shape[3]] == list(ifm_shape)


To get this change to work we would need assert list(ifm.shape) == [1, 1] + list(ifm_shape), since ifm.shape is not a list

lhutton1 · 2022-03-09T11:34:42Z

tests/python/contrib/test_ethosu/test_legalize.py

+
+        # check OFM
+        ofm = op.checked_type
+        assert [ofm.shape[2], ofm.shape[3]] == [1, ofm_channels]


This would need to be assert list(ofm.shape) == [1, 1, 1, ofm_channels]

lhutton1 · 2022-03-09T11:36:33Z

tests/python/contrib/test_ethosu/test_legalize.py

+        # check weights
+        weights_ohwi = op.args[1].data.asnumpy()
+        assert str(weights_ohwi.dtype) == dtype
+        assert list(weights_ohwi) == [ofm_channels, 1, 1, ifm_shape[1]]


This ones just missing .shape i.e. assert list(weights_ohwi.shape) == [ofm_channels, 1, 1, ifm_shape[1]]

ack! good spot on the list(ifm.shape)

Address comments, fixing assertion on ifm and ofm shape.

lhutton1

Thanks @dchauhan-arm, LGTM!.. Providing CI stays green

manupak

LGTM!

Thanks @dchauhan-arm @lhutton1 @ekalda @jainris !

* [microNPU] Add support for TFLite FULLY_CONNECTED This is primarily a legalization to an NPU Conv2d operator. The legalization target is Conv2d with 1 1 I O (HWIO) * [microNPU] Add support for TFLite FULLY_CONNECTED Test TVM runtime against TFLite for codegen and operator legalization. * [microNPU] Add support for TFLite FULLY_CONNECTED Fix linting * [microNPU] Add support for TFLite FULLY_CONNECTED Address comments, update codegen test, fix linting. * [microNPU] Add support for TFLite FULLY_CONNECTED Address more comments, ensure qnn.dense is lowered to NPU, fix linting * [microNPU] Add support for TFLite FULLY_CONNECTED Fix linting, update legalization test and codegen test for completeness. * [microNPU] Add support for TFLite FULLY_CONNECTED Address comments, fix linting. Certain legalization test assertions were updated. Co-authored-by: Rishabh Jain <rishabh.jain2@arm.com> * [microNPU] Add support for TFLite FULLY_CONNECTED Fix assertion in legalization test. * [microNPU] Add support for TFLite FULLY_CONNECTED Address comments, fixing assertion on ifm and ofm shape. Co-authored-by: Rishabh Jain <rishabh.jain2@arm.com>

dchauhan-arm added 2 commits February 22, 2022 11:29

[microNPU] Add support for TFLite FULLY_CONNECTED

8c0ea73

This is primarily a legalization to an NPU Conv2d operator. The legalization target is Conv2d with 1 1 I O (HWIO)

[microNPU] Add support for TFLite FULLY_CONNECTED

dab5c5e

Test TVM runtime against TFLite for codegen and operator legalization.

lhutton1 mentioned this pull request Feb 22, 2022

[Tracking Issue] Continued Arm(R) Ethos(TM)-U NPU operator support #9961

Closed

11 tasks

lhutton1 reviewed Feb 24, 2022

View reviewed changes

dchauhan-arm and others added 2 commits February 24, 2022 09:50

Merge branch 'apache:main' into pers-FCsupport

b79fec2

[microNPU] Add support for TFLite FULLY_CONNECTED

ef1d576

Fix linting

ekalda reviewed Feb 24, 2022

View reviewed changes

dchauhan-arm and others added 6 commits February 28, 2022 09:33

Merge branch 'apache:main' into pers-FCsupport

d80302a

Merge branch 'apache:main' into pers-FCsupport

7b177bd

[microNPU] Add support for TFLite FULLY_CONNECTED

a4741b9

Address comments, update codegen test, fix linting.

Merge branch 'apache:main' into pers-FCsupport

83d9ee1

Merge branch 'apache:main' into pers-FCsupport

97cd5ce

Merge branch 'apache:main' into pers-FCsupport

3bc81a5

dchauhan-arm marked this pull request as ready for review March 4, 2022 10:42

dchauhan-arm requested review from MarisaKirisame, ZihengJiang, anijain2305, comaniac, icemelon, jroesch, junrushao, jwfromm, mbrookhart, slyubomirsky, tqchen, vinx13, wweic, yzhliu and zhiics as code owners March 4, 2022 10:42

dchauhan-arm requested review from areusch and merrymercy as code owners March 4, 2022 10:42

dchauhan-arm added 2 commits March 4, 2022 10:43

[microNPU] Add support for TFLite FULLY_CONNECTED

ae6827c

Address more comments, ensure qnn.dense is lowered to NPU, fix linting

[microNPU] Add support for TFLite FULLY_CONNECTED

18bd546

Fix linting, update legalization test and codegen test for completeness.

lhutton1 reviewed Mar 4, 2022

View reviewed changes

dchauhan-arm and others added 5 commits March 4, 2022 16:46

Merge branch 'apache:main' into pers-FCsupport

efbe30e

Merge branch 'apache:main' into pers-FCsupport

fcca2d2

[microNPU] Add support for TFLite FULLY_CONNECTED

06aa4d9

Address comments, fix linting. Certain legalization test assertions were updated. Co-authored-by: Rishabh Jain <rishabh.jain2@arm.com>

Merge branch 'apache:main' into pers-FCsupport

04fd825

[microNPU] Add support for TFLite FULLY_CONNECTED

82dc516

Fix assertion in legalization test.

lhutton1 reviewed Mar 9, 2022

View reviewed changes

[microNPU] Add support for TFLite FULLY_CONNECTED

f924021

Address comments, fixing assertion on ifm and ofm shape.

lhutton1 approved these changes Mar 9, 2022

View reviewed changes

manupak approved these changes Mar 10, 2022

View reviewed changes

manupak merged commit a1fb44d into apache:main Mar 10, 2022

dchauhan-arm deleted the pers-FCsupport branch March 10, 2022 10:00

ekalda mentioned this pull request Mar 21, 2022

[Tracking Issue] Arm(R) Ethos(TM)-U NPU operator support #9410

Closed

15 tasks

driazati mentioned this pull request Jul 14, 2022

TVM v0.9.0.rc0 Release Candidate Notes #12102

Closed

	ifm = relay.reshape(ifm, (-1, 1, 1, params.ifm.shape[-1]))
	ifm = relay.reshape(ifm, (1, 1, 1, params.ifm.shape[-1]))

		if len(params.ofm.shape) != 4 or not params.ofm.shape[1] == params.ofm.shape[2] == 1:
		ethosu_fc = relay.reshape(ethosu_fc, params.ofm.shape)

	"""This is the pass that wraps the AddRewriter"""
	"""This is the pass that wraps the FullyConnectedRewriter"""

Conversation

dchauhan-arm commented Feb 22, 2022

Uh oh!

dchauhan-arm commented Feb 22, 2022

Uh oh!

lhutton1 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dchauhan-arm Feb 24, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lhutton1 Feb 24, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ekalda left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lhutton1 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lhutton1 commented Mar 4, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lhutton1 left a comment

Choose a reason for hiding this comment

Uh oh!

manupak left a comment

dchauhan-arm Feb 24, 2022 •

edited

Loading

lhutton1 Feb 24, 2022 •

edited

Loading