Include Shape Information Into DecomposeComplexOps Patterns

[zjgarvey]

Many pass pipelines rely on --torch-decompose-complex-ops, but do not include shape inference passes afterwards. For large models, the complex computations performed in shape inference passes take a significant amount of time to perform and are not as reliable as explicitly specifying the result types by hand. The only downside to adding these shapes by hand is the minor expense of code complexity, but with the development of appropriate helper functions such as Torch::unsqueezeTensor as a replacement for rewriter.create<AtenUnsqueezeOp>(loc, baseType,...) as seen here, the potential downside to code readability should rendered completely irrelevant.

As a design decision for the short term (before larger changes like fx decompositions are added), I'd like to come to a consensus with the community to avoid using

baseType = ValueType::getWithLeastStaticInformation(context);

in decompose complex ops conversions, and to request authors of recent conversions to go back and add shape information into their code.

Here is the very short list of decompose complex ops passes which currently rely on shape inference:

• "DecomposeAtenTriuOp" here (see also issue #573) @Mi-Jiazhi @qingyunqu
• "DecomposeAtenNanToNumOp" here @linuxlonelyeagle
• "DecomposeAtenGroupNormOp" here @JianzheXiao @qingyunqu @vivekkhandelwal1
• "DecomposeAtenNativeGroupNormOp" here @JianzheXiao @qingyunqu
• "DecomposeAtenLinspaceOp" here @qingyunqu
• "DecomposeAtenFakeQuantizePerTensorAffineOp" here @qingyunqu

If anyone in the community has strong feelings in opposition to this change, please feel free to discuss here.

[qingyunqu]

I think it's ok to apply shape infer in decomposition when it's a simple op. But we find that it's difficult to apply shape infer on aten.native_group_norm's decomposition, specifically when dynamic shape.

[Xinyu302]

I complete the case "DecomposeAtenTriuOp" in #3330. Please give me some suggestions, and then I will complete the other cases in the same way.

[zjgarvey]

I complete the case "DecomposeAtenTriuOp" in #3330. Please give me some suggestions, and then I will complete the other cases in the same way.

Thanks for taking this up @Xinyu302 ! I'll take a look tomorrow (in U.S. time).

[zjgarvey]

I think it's ok to apply shape infer in decomposition when it's a simple op. But we find that it's difficult to apply shape infer on aten.native_group_norm's decomposition, specifically when dynamic shape.

I see. If you want, I can take a look and see if there is a nice way to add the shape information in without clogging up the code too much.