move loop ordering after fusion (#126254)
Summary:
Restart the work from PR https://github.com/pytorch/pytorch/pull/100331 in this new PR since it's hard to rebase. It would be expected that some code is copy/pasted from the previous PR and main idea is the same.
Previously we see relatively large compilation time increase due to too many loop orders being considered. This PR tries to continue the work by doing pruning and only considering loop orders that we know for sure are relevant (i.e. do it on demand).
Some manually created cases that loop ordering matters are added as unit tests. The PR can make sure inductor does not miss fusion opportunities for them.
This PR should solve the not-able to fusion problem in https://github.com/pytorch/pytorch/issues/130015
Right now there is still significant increase of compilation time. I'll disable the feature by default. Later on after the compilation time issue is resolved, I'll enable it by default.
X-link: https://github.com/pytorch/pytorch/pull/126254
Approved by: https://github.com/jansel
Reviewed By: ZainRizvi
Differential Revision: D62008970
Pulled By: shunting314
fbshipit-source-id: ce4c7c7003b93a2faccd2c65d78eeee0300b6bff
