[TensorExpr] Add support for Axis reordering in LoopNest

[nickgg]

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;
*/

[dr-ci]

💊 Build failures summary and remediations

As of commit 5ae30e4 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker.

See how this bot performed.

This comment has been revised 19 times.

[ZolotukhinM]

Looks good! I didn't follow all the logic but I trust your tests :) I left some questions inline - they are not to request changes but rather to double check my own understanding.

[ZolotukhinM]

Do you feel this methods will generally be helpful in other use-cases? If not, I'd rather make it a local helper method rather than pulling it into the For class itself.

[nickgg]

I did use it in the ReductionExpander as well - but they're both in loopnest.cpp so for now could be local (could be complicated if we pull the mutators out of loopnest). I think this is generally a useful helper though, I think we'll reuse it again.

[ZolotukhinM]

Nit: according to std convention, it's probably better to change nparams to return size_t.

[nickgg]

You're right. That would touch quite a few different places in the code though so rather not roll into this diff.

[ZolotukhinM]

Just curious: why do we clone here?

[nickgg]

root_stmt returns the internal pointer, which is modified by reorderAxis below. We need to clone to preserve the original behaviour.

[ZolotukhinM]

We don't really use the tensor here, only its loops: maybe take std::vector<For*> as the first parameter instead (on the callsite we'll get it by l.getLoopStmtsFor(t))?

[nickgg]

Not sure about this. It would save the extra call to getLoopStatementsFor but it's not longer really a loopNest thing. Prefer to refactor when we pull this and maybe rfactor/split out

[ZolotukhinM]

I was thinking about it mostly from the data we're accessing here: we pass the tensor, but only use part of it that's responsible for the corresponding loops. That made me think that we should pass only the loops since we're not using anything else from the tensor. I don't have a strong preference on it though since we don't have clear guidelines on it yet anyways :)

[nickgg]

addressed comments

[facebook-github-bot]

@nickgg is landing this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@nickgg merged this pull request in ee3d046.