[TensorExpr] Add support for Axis reordering in LoopNest (#36540)
Summary:
Adds a capability for reordering axes in the LoopNest. This was fairly straightforward except when handling Reduction initializers which required more changes, UPDATE: actually the complicated bit was preserving the ordering of statements in the loopnest which should not be reordered.
Usage looks something like this:
```
Tensor* tensor = Compute(
"f", {{2, "x"}, {3, "y"}}, [](const VarHandle& x, const VarHandle& y) {
return ExprHandle(1.0f) + cast<float>(x) * x + cast<float>(y) * y;
});
LoopNest l({tensor});
/* LoopNest looks like:
for x in ...
for y in ...
f[x,y] = 1 + x * x + y * y;
*/
auto loops = l.getLoopStmtsFor(tensor);
l.reorderAxis(tensor, loops[0], loops[1])
/* LoopNest looks like:
for y in ...
for x in ...
f[x,y] = 1 + x * x + y * y;
*/
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/36540
Differential Revision: D21068143
Pulled By: nickgg
fbshipit-source-id: f02c29004376df4f5a9bedff366c075772726618
