Enable for qnn operations for const folding transformation

[apeskov]

Current the sequence cons -> qnn.quantize is not treated like a constant subgraph. Suggestion is to allow FoldConstant pass to replace this pattern with single int8 constant tensor.

Reason: Some BYOC runtimes may has a limitation to have a weight like a constant tensor. Pointed FoldConstant pass limitation may breaks BYOC runtimes applicability.

[manupak]

Thanks @apeskov !

It is an interesting change.

However, looking at the test case, I am struggling to understand why this can not be achieved via running Legalize before ConstantFold pass. Would you be able to share some motivation?

[manupak]

How is it different from running Legalize followed by FoldConstant ?

[masahi]

@apeskov Please see this PR #9135. I understand why you want to do this, namely, constant fold quantize(weight_fp32) in a QNN graph. Returning float32 weights from the PyTorch frontend and relying on Relay constant folding to recover quantized weights was my design mistake. Now you can directly obtain quantized weights from the frontend (we do quantize at numpy level).

@manupa-arm Running lowering before const fold is not acceptable when we want to keep the rest of QNN graphs (BYOC), while selectively lower constant subgraphs and evaluate them.

[manupak]

@masahi What is stopping us from the running the legalization on the IRModule with just the external function ? i.e. in the relay.ext.<codegen>

[masahi]

I'm thinking of use cases in BYOC where we want to pattern match against QNN ops, in which case we don't want to run QNN legalization. Not sure if this answers your question @manupa-arm

In my prev job, I directly took QNN subgraphs and send them to external codegen. I believe ethos-N does something similar. We had to develop constant folding on the external codegen side.

[manupak]

Hmmm, I was not suggesting to run the legalization before the partitioning.

We could identify the patterns with QNN ops and then we partition the external function.
I believe the requirement to do the legalization + constant folding comes post partitioning.

So there are two places we could do this :

1. post PartitionGraph pass in the partition_for_* function
2. relay.ext. which gets the partitioned function passed in via relay.build process.

I believe, this particular requirement is to mutate only the external function (thus, we could do it 2) and not 1) ) and not the 'main'. Therefore, why cant we achieve the same effect running the two passes -- legalization + constant folding -- in 2) ?

[masahi]

Hmm interesting, I never thought about doing constant folding on partitioned functions. My use cases have always been doing constant folding on main, before partitioning. For example, that was the case in PyTorch frontend before #9135 which always produced something like qnn.quantize(const_weight_fp32). The other case is QNN produced by FakeQuantizationToInteger pass, which also generates many qnn.quantize with constant weights.

In 2), if we run legalization on partitioned functions, wouldn't that decompose all QNN ops? I couldn't easily extract qparams anymore, for example. I needed to retain QNN ops all the way until I translated them to the external IR, so running legalization had never been my option. I did wish that we could selectively lower const-foldable QNN subgraphs only. Maybe I'm missing something.

[manupak]

In 2), if we run legalization on partitioned functions, wouldn't that decompose all QNN ops?

I think the constant folding pass is supposed to work in the IRModule (with the external function). Therefore, everything in the IRModule will be affected. However, we could create IRModules with what is in-scope for the transformation.

I needed to retain QNN ops all the way until I translated them to the external IR, so running legalization had never been my option. I did wish that we could selectively lower const-foldable QNN subgraphs only. Maybe I'm missing something.

It is about further granularity one would to do further partitioning. Today, I think we need to do further partitioning to achieve this. However, whether we want to annotations to block constant folding seems like an interesting but an orthogonal conversation to this one.

In the scope of changes in this PR, I feel it does the same thing (destroys QNN info in the process of constant folding). However, we could control what we want to pass into the Constant Folding Pass.

[masahi]

@apeskov please update or close this PR

[masahi]

@apeskov Please update or close this PR

[areusch]

closing due to inactivity, feel free to reopen