[microNPU] Move optimization passes to be a module pass and ensure they are running

[lhutton1]

Moves LayoutOptimizer and LUTOptimizer passes to be a module pass, rather than a function pass. This is because it was found that these passes were not running in the NPU compilation flow. In addition, a test for both LayoutOptimizer and LUTOptimizer has been added to check that the passes are running in the compilation pipeline of the NPU.

cc @ekalda @manupa-arm

[ekalda]

Thanks @lhutton1 for spotting and fixing this, it is a very subtle problem! As discussed offline, the fact that these passes silently ceased running as a part of a pipeline (they certainly did use to run when I was implementing the LUT pass) is quite worrying since it means there has been an underlying change to how a relay.transform.function_pass works (or maybe even a bug since the docstring of the implementation implies that it should work on the modules), so there could be other optimization passes that are silently not running now. I think @lhutton1 mentioned though that he didn't grep out many use cases of that function_pass outside of tutorials, but I think in that case we should properly retire it and update the tutorials. Re-tagging @manupa-arm for opinion...

[ekalda]

(Forgot to say that dealing with function_pass is a stuff for other PRs, this one looks good to me!)

[lhutton1]

Thanks @ekalda for taking a look, yes it seems quite strange to me. Just wanted to note that function_pass is used outside of tutorials, I was foolishly searching for relay.transform.function_pass rather than just function_pass :)

I did a little bit more digging and it seems like the function_pass doesn't get run on a function when there is a "Compiler" attribute associated with it. I wonder if this is intended behavior and if/when this changed?

[manupak]

kCompiler (a.k.a. "Compiler") attribute generally blocks most passes from visiting the function body. Are we saying there was a time where the pass worked even when the kCompiler is set ?

[ekalda]

I saw the LUT pass working as a part of a pipeline when I was developing it (I first developed it by running an end to end codegen test with a breakpoint before the pass and checked that the Relay had changed as expected after stepping over it). Out of interest, why does the kCompiler attribute block the passes visiting the function body?

[manupak]

It is used to annotate relay functions meant to be compiled with BYOC. Therefore, the compilation passes of core compiler ignores it. IIRC, that was the reason.

I believe this affects function pass because (and not module pass) the pass it self does not choose which functions to mutate as opposed to a module pass.

[lhutton1]

Thanks @manupa-arm, in that case it sounds like migrating to a module pass is the correct thing to do here. It's still a bit of a mystery as to why @ekalda and I believe this used to work, although its probably worth not dwelling on it for now

[manupak]

The implementation LGTM!

I think we need to be a specific with the title and/or commit message to say what this PR enables. I believe that should be :

1. Moves LUTOptimizer to be a module pass
2. Ensure LUTOptimizer Pass is run

rather than generically adding the feature to ensure all optimization passes are run.

[lhutton1]

Agreed, that makes more sense :)

[lhutton1]

friendly ping for approval/comments

[manupak]

LGTM!

[manupak]

Thanks @lhutton1 @ekalda !