[Unity][Transform] Some Improvements on pass DecomposeOps

[SiriusNEO]

Prior to this PR, we first have a pass named SimplifyNormInference (#14221), which is used to decompose batch_norm into simple operators for optimization during inference. And there are some follow-up changes:

• [Unity][Transform] Introduce data-dependent operation of reshape and its constant folding #14282 This PR introduces a policy to decompose tensor_to_shape op.
• [Pass][Fix] Merging DecomposeOps and SimplifyNorm mlc-ai/relax#170 In mlc, we make SimplifyNormInference support op simplification during training (Because batch_norm behaves differently in inference and training).
• And in the future maybe there will be more operators needed to be decomposed.

So this PR polishes the changes in #14282 mainly in the following aspects:

• Use one mutator (instead of two) to visit and decompose ops.
• Separate inference and training by registering two passes (Adviced by TQ).
• Shorten the name (DecomposeCompositeOps -> DecomposeOps)
• Func Pass => Module Pass with func_name to specify the function we want to apply the pass (Or let it be None if we want to simplify all functions). This is because sometimes we only want to decompose operators in a specified function.

Now the code is clear and easy to read. Now we use if-else if-else if ... pattern to recognize the op we want to decompose. In the future maybe we can introduce a map or register the decomposition policy in the op attribute, just like what we do in the LegalizeOps pass.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[SiriusNEO]

cc @tqchen @sunggg @Hzfengsy

[vinx13]

We don't need to decompose LayerNorm since it can be implemented more efficiently as a single op

[sunggg]

Is there any complication for fusion if we have a composite ops in general? (not specifically about LayerNorm). i.e., Is there a case where we cannot fuse the composite op but we can when we decompose?

[SiriusNEO]

Ah here we simplify LayerNorm for now just because it is easy to get its gradient. If it is more efficient as a single op, maybe in the future we can remove this from the pass after we finish the gradient function of the whole LayerNorm

[SiriusNEO]

Or maybe we can add some arguments for the pass to let the user specify what operators he want to simplify?

[Hzfengsy]

As decomposing LayerNorm is used for training and may influence the inference perf, let's only decompose it in training mode

[vinx13]

Is there any complication for fusion if we have a composite ops in general?

Yes, we need to be aware of the implementation when setting op type (kOpaque or kInjective). This will cause different fusion results.

[sunggg]

LGTM! Thank you for the improvement.
Just a few comments.

[Hzfengsy]

As decomposing LayerNorm is used for training and may influence the inference perf, let's only decompose it in training mode

[tqchen]

@sunggg @vinx13 please take another look