Page MenuHomePhabricator

[mlir][linalg] Add vectorization for element-wise linalg ops
ClosedPublic

Authored by ThomasRaoux on Dec 2 2020, 9:12 PM.

Details

Summary

Add support for vectorization for linalg.generic representing element-wise ops. Those are converted to transfer_read + vector ops + transfer_write.
Also re-organize the vectorization tests to be together.

Implementation derived from @burmako's work.

Diff Detail

Event Timeline

ThomasRaoux created this revision.Dec 2 2020, 9:12 PM
ThomasRaoux requested review of this revision.Dec 2 2020, 9:12 PM
nicolasvasilache accepted this revision.Dec 3 2020, 3:46 AM
nicolasvasilache added a subscriber: silvas.

Very cool, thanks for pushing this @ThomasRaoux !

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
92

I thought we had some LLVM interface already for elementwise like (ElementsMappable something from @silvas' work?)

106
if (llvm::any_of(genericOp.getOutputIndexingMaps(), lambda)) return false;``` ?
110

same

135

Ah cool that the insertion point has a loc now, can you just modify the ScopedContext to just take a builder and simplify uses in this file?

140

I'd move up 2 lines to keep C++ together and meta-programmed IR together.

170

early return and drop the else to reduce indentation entropy?

207

The type-only logic here seems like it could be refactored into VectorType?

235
return llvm::TypeSwitch<Operation*, Operation*>(...)
.Case(...)
.Default(...);
This revision is now accepted and ready to land.Dec 3 2020, 3:46 AM
aartbik added inline comments.Dec 3 2020, 9:01 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
124

note that load and transfer-read are slightly different terms (see doc),
so perhaps call this transferLoadVector?

134

same, transferStoreVector?

150

"to reasds" and "to writes" at a quick glance reads like bad grammar, but then you realize it is used as a noun

so how about "to transfer reads" and "to transfer writes" to make that more clear right away

aartbik added inline comments.Dec 3 2020, 10:20 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
124

Ooop, I meant transferRead and transferWrite obviously, not Load/Store ;-)

ThomasRaoux marked 9 inline comments as done.

Address review comments.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
92

Well, I want to restrict it to operations supported below so I'm not sure there is a better solution?

106

There is no getOutputIndexingMaps in linalg op. I can get a list of all indexingmaps but I want to distinguish the input and outputs so I couldn't find a better way to do that. Any ideas?

207

Good point, changed it to just skip if none of the types are vector, it makes logic simpler and should be more generic.

silvas added inline comments.Dec 3 2020, 12:18 PM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
92

Just left a comment below. This should fit nicely with ElementwiseMappable.

234

You can do this for all ops that have ElementwiseMappable trait using a single piece of code like we do for elementwise on tensors ->linalg.generic: https://github.com/llvm/llvm-project/blob/f24c642178a5a31aa1c44585537c08fdfc3fdfd5/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp#L53
You just copy all the attributes and convert all the types to vector types.

The invariants of ElementwiseMappable guarantee that this can always be done: https://github.com/llvm/llvm-project/blob/f24c642178a5a31aa1c44585537c08fdfc3fdfd5/mlir/include/mlir/IR/OpDefinition.h#L1243

mravishankar added inline comments.Dec 3 2020, 2:06 PM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
107

In general this could be a permutation. Is that not supported?

111

This could also be a permuted projection. Is that a todo?

Use ElementwiseMappable.

ThomasRaoux marked an inline comment as done.Dec 3 2020, 2:26 PM
ThomasRaoux added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
107

Yes I think permutation could be supported, I think we can just propagate the affineMap into the transfer read/write. Do you think this needs to go into this patch?

111

I was trying to make it work with permuted projection but it isn't trivial as the type of the memref may not be compatible with the type of the vector we want out of the transfer_read.
For instance if we have a map like affine_map<(d0, d1) -> (d0)> with 1D memref I would have to load the vector as 1D and then broadcast it along the first dimension which isn't supported by the broadcast.
Maybe I'm missing something, I can add a TODO if you want.

234

Thanks Sean! I used ElementwiseMappable and it does clean up significantly.

Adding todo to relax indexing map restrictions.

pravnar added a subscriber: pravnar.Dec 3 2020, 2:41 PM
pravnar added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
206

Typo? Probably should be "broadcastIfNeeded" :)

Fix typo

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
206

Oops, yes. Fixed it. Thanks.

burmako accepted this revision.Dec 3 2020, 2:48 PM

Thanks a lot for your work! This will be really useful :)

I'd also like to share the credit that you kindly gave to me with fedelebron@ and agrue@.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
100

Can you talk more about the context behind the definition of isElementwise, e.g. in the comment above?

I'm asking because in one of the internal tests we have affine_map<(d0, d1) -> (d0)> as one of the supported indexing maps, and, if I understand correctly, that doesn't seem to work with isElementwise.

This comment doesn't need to block this patch (thanks a lot for working on it!). I'm mostly curious about what use cases you're targeting, and what use cases we'll need to add support for later on.

206

Typo: broadcastIfNeeded.

mlir/test/Dialect/Linalg/vectorization.mlir
187

What do you think about spelling out the types for vector.transfer_read/transfer_write, since these types can affect semantics? (This is something that I regret not doing in the original code).

Add comment and spell out vector types in test.

ThomasRaoux marked 2 inline comments as done.Dec 3 2020, 3:08 PM

Thanks a lot for your work! This will be really useful :)

I'd also like to share the credit that you kindly gave to me with fedelebron@ and agrue@.

Thanks for the review.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
100

I added a comment explaining why I added this restriction. I probably need to talk to you to understand how you handle affine_map<(d0, d1) -> (d0)>. I limited it to minor identity so that I can convert the input vector to a larger vector using just a vector.broadcast. With such an affine map I believe I need to do broadcast follow by transpose and it can get more complicated which is why I limited the scope in this case.
Most likely this will need to be relaxed in the future. I'd be happy to learn about the use cases you have that would require that.

mlir/test/Dialect/Linalg/vectorization.mlir
187

Good point, added it.