Page MenuHomePhabricator

[mlir][Linalg] Enable vectorization of explicit broadcasts
ClosedPublic

Authored by dcaballe on Oct 11 2021, 10:41 AM.

Details

Summary

This patch teaches isProjectedPermutation and inverseAndBroadcastProjectedPermutation
utilities to deal with maps representing an explicit broadcast, e.g., (d0, d1) -> (d0, 0).
This extension is needed to enable vectorization of such explicit broadcast in Linalg.

Diff Detail

Event Timeline

dcaballe created this revision.Oct 11 2021, 10:41 AM
dcaballe requested review of this revision.Oct 11 2021, 10:41 AM
ThomasRaoux added inline comments.Oct 11 2021, 12:21 PM
mlir/lib/IR/AffineMap.cpp
498–526

Does that mean (d0, d1) -> (0, 0) would be a projected permuation? That doesn't sound right.
I think you need a different function to check for something like prjected permutation with broadcast like a generalization of isPermutationOfMinorIdentityWithBroadcasting?

Could you please verify if all tests still pass? For example, we might need to modify applyPermutationMap in AffineMap.h.

template <typename T>
SmallVector<T> applyPermutationMap(AffineMap map, llvm::ArrayRef<T> source) {
  assert(map.isProjectedPermutation());
  assert(map.getNumInputs() == source.size());
  SmallVector<T> result;
  result.reserve(map.getNumResults());
  for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
    /* fix
    if (auto expr = map.getResult(i).dyn_cast<AffineConstantExpr>()) {
      result.push_back(0);
      continue;
    }
    */
    unsigned dim = map.getDimPosition(i);
    result.push_back(source[dim]);
  }
  return result;
}

I have something similar locally but I am not clear this is actually needed for now.
Would you mind holding off a bit until I clear up my stack and avoid conflicts?

dcaballe updated this revision to Diff 378855.Oct 11 2021, 7:07 PM

Addressed Alex's feedback for now.

Could you please verify if all tests still pass? For example, we might need to modify applyPermutationMap in AffineMap.h.

Good catch! Yes, I had run all the lit tests and nothing failed. I fixed this, thanks!

I have something similar locally but I am not clear this is actually needed for now.
Would you mind holding off a bit until I clear up my stack and avoid conflicts?

Sure, I'll wait for you. After looking at your code, I'm addressing the review comments since they will also apply to your code.

mlir/lib/IR/AffineMap.cpp
498–526

Good point! Note that this utility, not sure if intentionally, is already returning true for (d0, d1) -> (). (d0, d1) -> (0, 0) would be the equivalent preserving the dims. I think it boils down to what the definition of "projected permutation" is. I asked @nicolasvasilache and it seems to be ok to have zeros in a projected permutation. We could add logic to skip (d0, d1) -> () and (d0, d1) -> (0, 0) but they look like a valid corner case to me. WDYT?

ThomasRaoux added inline comments.Oct 11 2021, 7:20 PM
mlir/lib/IR/AffineMap.cpp
498–526

In my mind projected permutation meant it has a permutation of the input dimensions and some dimensions may be missing (projection), so (d0, d1) -> () sounds like a valid projected permutation but having zeros instead of dimensions seems a bit odd. @nicolasvasilache is the referenced so if he thinks it is correct to allow extra zeros it is fine with me but this seems less intuitive to me.
It would be good to look at how this helper is used right now and check if any places assume that it means that all the results are dimensions.

pifon2a accepted this revision.Oct 12 2021, 12:14 AM
This revision is now accepted and ready to land.Oct 12 2021, 12:14 AM
nicolasvasilache accepted this revision.Oct 12 2021, 6:49 AM

My stack is now flushed, feel free to land this @dcaballe and thanks for your patience!

Well technically, orthogonal projection of a vector v=(v0, .. vn) on the canonical axis e_k is the vector v=(v0, .. 0_k .. vn).
So far we have liberally jumped the gun a bit by also dropping those dimensions so that other invariants agree.

This LGTM; I would just make the behavior optional with a flag to avoid surprises (for now).

dcaballe updated this revision to Diff 379170.Oct 12 2021, 1:40 PM

Adding a flag and enable it only for vectorization.
This should be ready to go.

This revision was landed with ongoing or failed builds.Oct 12 2021, 2:10 PM
This revision was automatically updated to reflect the committed changes.