Pass that removes reshapes post LowerTE

[asparkhi]

Introduces a Pass for removing intermediate reshapes post
LowerTE() in AOT compiler. This commit adds pass specific
tests and updates usmp generated workspace pools due to
reduction in number of allocations post reshape removals.

Note: this pass at present does not support first reshape
appearing in the graph. If seen as a useful case, it can be
added in the future.

cc: @manupa-arm @grant-arm

[manupak]

Thanks for this work!
Also worth highlighting in some docs of the pass the value of it and why it should be enabled.

cc : @areusch @altanh @lhutton1

[manupak]

Thanks @ashutosh-arm. looking good!

I think we need unit tests for the pass as well.
(E.g. https://github.com/apache/tvm/blob/main/tests/python/relay/test_pass_partition_graph.py)

[lhutton1]

Great work @ashutosh-arm! Just some small things I picked up on..

[areusch]

don't have tons of context here, but left a couple suggestions

[areusch]

i'm probably missing some context here, but what about just returning the args to reshape()?

[asparkhi]

Graph contains let nodes in between the call_lowered(). I've included the following piece as part of the Rewrite_() as well.

    /*
    %1 = call_lowered(@tvmgen_default_non_reshape_function, %input, ...);
    let %x: = on_device(%1, ...);
    %2 = (%x,);
    %3 = call_lowered(@tvmgen_default_fused_reshape, %2, ...,
    "relay_attrs"=__dict__="relay.reshape_only"=1, ...);
    */

[manupak]

LGTM!

@lhutton1 @areusch maybe have a look ?

[lhutton1]

LGTM!

[lhutton1]

Thanks @ashutosh-arm @manupa-arm @areusch!