This is an archive of the discontinued LLVM Phabricator instance.

[mlir][vector] Add lowering of Transfer_read with leading dim broadcast and permutation map
ClosedPublic

Authored by ThomasRaoux on Mar 20 2021, 9:41 AM.

Details

Summary

Convert transfer_read ops with permutation maps into simpler transfer_read with minority map + vector.braodcast and vector.transpose

Diff Detail

Event Timeline

ThomasRaoux created this revision.Mar 20 2021, 9:41 AM
ThomasRaoux requested review of this revision.Mar 20 2021, 9:41 AM
Herald added a project: Restricted Project. · View Herald TranscriptMar 20 2021, 9:41 AM
mlir/lib/Dialect/Vector/VectorTransforms.cpp
2849

I am not sure I understand this case, valid minor identity maps are (d0, d1, d2), (d1, d2) and (d2) ?

2850

s/permutate/permute/g (both comment and code please).

2852

Can you expose this new functionality in AffineMap.h?

2865

s/broadcasted/broadcast/g

ThomasRaoux added inline comments.Mar 24 2021, 9:15 AM
mlir/lib/Dialect/Vector/VectorTransforms.cpp
2849

Agree, but the idea is to get to a minor identity with broadcasting, meaning some dimensions can be replaced by 0 (this match the definition in isMinorIdentityWithBroadcasting). So (d1, 0) is (d1, d2) with broadcast on dimension d2. This means we can do a transfer read with no permute to <Nx1> and broadcast the lower dimension.

If the map is (d0, d1, d2) -> (d1) it would require first transforming it to (d0, d1, d2) -> (d1, 0) which needs more than a transpose. That's why I currently don't support it.

Address review comments

ThomasRaoux marked 3 inline comments as done.Mar 24 2021, 9:55 AM
sgrechanik added inline comments.Mar 24 2021, 8:29 PM
mlir/lib/Dialect/Vector/VectorTransforms.cpp
2876

typo

mlir/lib/IR/AffineMap.cpp
153

I think that a bitvector (SmallBitVector?) could be better here, but I don't insist.

158

The "or" counterpart of this comment is missing.

Address review comments

ThomasRaoux marked 2 inline comments as done.Mar 24 2021, 9:02 PM
ThomasRaoux added inline comments.
mlir/lib/IR/AffineMap.cpp
153

Makes sense, changed it.

mlir/include/mlir/IR/AffineMap.h
116

Just copy pasting my other comment here:

I think spelling it out a little more would be useful.
First I'd drop the N-D -> ND and focus on N-D -> (N-1)-D and/or N-D -> (N-2)-D as I think the asymmetry helps documentation.

Then I'd write something like:

* (d0, d1, d2) -> (d0, 0) => cannot be mapped to a minor identity by permutation + broadcast
* (d0, d1, d2) -> (0, d1) => maps to minor identity (d1, 0 = d2) with perm = [1, 0] and broadcast d2
...
a few more examples to build intuition
I would put that same documentation very loudly both in AffineMap.h and here
123

Would it be better API-wise to also pass the SmallBitVector &broadcastDims to fill that too?
Atm the broadcast is still implicit and subject to magic.

I'd also make it clear that it would be first permute by applying permuteDims then set the true dimensions of broadcastDims to 0.

mlir/lib/Dialect/Vector/VectorTransforms.cpp
2849

Right, so I think spelling it out a little more would be useful.
First I'd drop the N-D -> ND and focus on N-D -> (N-1)-D and/or N-D -> (N-2)-D as I think the asymmetry helps documentation.

Then I'd write something like:

* (d0, d1, d2) -> (d0, 0) => cannot be mapped to a minor identity by permutation + broadcast
* (d0, d1, d2) -> (0, d1) => maps to minor identity (d1, 0 = d2) with perm = [1, 0] and broadcast d2
...
a few more examples to build intuition

I would put that same documentation very loudly both in AffineMap.h and here

2853

Would also be nice to mention there is a dual transformation that takes vector.transfer + vector.transpose and turns it into linalg.transpose + vector.transfer to do stuff in memory.
This pattern and that transformations are related, it is target-specific what combination should be used.

2866

I think if isPermutationOfMinorIdentityWithBroadcasting also filled a broadcast bitvector, it would be a more natural read to create the map based on that.
It is not immediately clear to me what the swappy in-place behavior with modulos does.

2880

Can you use https://sourcegraph.com/github.com/llvm/llvm-project@d75a611/-/blob/mlir/include/mlir/Dialect/Linalg/Utils/Utils.h#L127:6 ?

void applyPermutationToVector(SmallVector<T, N> &inVec,
                              ArrayRef<unsigned> permutation)
mlir/test/Dialect/Vector/vector-transfer-lowering.mlir
230

can you add a few whitespaces to isolate the different parts of the test ?

mlir/lib/Dialect/Vector/VectorTransforms.cpp
2904

This is an interesting case where an internal broadcast is left.
Will be interesting to see how this composes with vector -> scf lowering (not for this revision).

2937

can this be written as something resembling newMask = rewriter.getBoolArrayAttr(op.masked().getAsValueRange<bool>().take_back(reducedShapeRank)); ?

mlir/lib/IR/AffineMap.cpp
145

same comment re adding a noisy, spelly comment.

ThomasRaoux marked 2 inline comments as done.

Address more review comments.

mlir/include/mlir/IR/AffineMap.h
116

Makes sense, I added more examples, hopefully this makes it easier to understand.

123

I tried to do it but I feel like it adds unnecessary complexity to this function and in my current case I am not able to use it. The broadcast dimensions can be accessed after permutation using the existing getMinorIdentityMap

mlir/lib/Dialect/Vector/VectorTransforms.cpp
2866

I'm not sure using the broadcast bitvector is easier. Here we cannot move all the broadcast dim to the front as we still want a minor identity map after the broadcasts. I added some comments explaining better what this does. Hopefully it is clearer.

2880

I would need to reverse the permutation first and I don't see any easy way to do it? Otherwise I could addd a different helper applyReversePermutationToVector? Is it okay to depend on linalg utils here?

2904

Makes sense. Note that it already compose well with the existing TransferReadToVectorLoadLowering pattern which load to Nx1xM then broadcast the 1.

2937

Yes, finally found the magic formula.

Simplify the logic based on comments

@nicolasvasilache I removed the confusing part we were discussing (the "rotating" the permutation map). This simplifies the code and should make it more clear what it does. Please take another look.

mlir/lib/Dialect/Vector/VectorTransforms.cpp
2866

This part was confusing and unnecessary complicated indeed. I removed it. It is better handled in isPermutationOfMinorIdentityWithBroadcasting by making sure we start the map with 0s in case numResults > numDims

nicolasvasilache accepted this revision.Mar 26 2021, 9:25 AM
This revision is now accepted and ready to land.Mar 26 2021, 9:25 AM
sgrechanik accepted this revision.Mar 28 2021, 11:55 AM
sgrechanik added inline comments.
mlir/include/mlir/IR/AffineMap.h
127

Could you write out the permutation for this last example too?

mlir/lib/Dialect/Vector/VectorTransforms.cpp
2885

nit: probably it would be better to rename it "permutation" because it's a single permutation.

ThomasRaoux marked an inline comment as done.

Address review comments

ThomasRaoux marked an inline comment as done.Mar 29 2021, 8:12 AM
This revision was landed with ongoing or failed builds.Mar 29 2021, 8:41 AM
This revision was automatically updated to reflect the committed changes.