This is an archive of the discontinued LLVM Phabricator instance.

[mlir][Linalg] Implement the tiling interface for softmax
ClosedPublic

Authored by qcolombet on Jun 21 2023, 6:32 AM.

Diff Detail

Event Timeline

qcolombet created this revision.Jun 21 2023, 6:32 AM
Herald added a project: Restricted Project. · View Herald TranscriptJun 21 2023, 6:32 AM
qcolombet requested review of this revision.Jun 21 2023, 6:32 AM
rengolin added inline comments.Jun 21 2023, 6:54 AM
mlir/test/Dialect/Linalg/tile-softmax.mlir
25

The semantics of this is to take the max of each tile, which is not always the same as the softmax of the original dimension.

If the tile is a whole head, this may be what you want. If not, you'll get different results.

I'm not sure how to restrict this in a meaningful way, or perhaps this is up to the compiler to "do the right thing".

At the very least, this should be documented somewhere, perhaps in the op description?

40

I'm assuming this also works with the tile-and-fuse pass. Could there be a test for that, too?

qcolombet added inline comments.Jun 21 2023, 8:23 AM
mlir/test/Dialect/Linalg/tile-softmax.mlir
25

Good point.
In this example I went with the easiest transformation to apply, but you're right this may not be correct.
Documenting this somewhere is sensible, but it may make more sense to reject invalid tiling.
At the same time maybe the compiler/ir author knows something we don't and we should let this go through...

@ftynse, @nicolasvasilache do you have a recommendation on what we should do here?

Technically getTiledImplementation can return a failure, but I don't know if that's doable/desirable to do that check here.

40

Let me give a try!

qcolombet added inline comments.Jun 22 2023, 9:48 AM
mlir/test/Dialect/Linalg/tile-softmax.mlir
25

Spoke with @nicolasvasilache offline and the conclusion was the compiler has to do the right thing.
I'll add some wording around that.

nicolasvasilache accepted this revision.Jun 29 2023, 3:17 AM
This revision is now accepted and ready to land.Jun 29 2023, 3:17 AM
rengolin accepted this revision.Jun 29 2023, 3:34 AM

Should be fine with just transform tests for now. Once we incorporate on existing infra we'll see how it plays with passes.

jpienaar added inline comments.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
53

This is the deprecated form https://mlir.llvm.org/deprecation/

2239

s/mlir:://

qcolombet added inline comments.Jun 30 2023, 12:23 AM
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
53

Thanks I didn't know about that.

qcolombet updated this revision to Diff 536274.Jun 30 2023, 8:45 AM
  • Add a tile and fuse test
  • Remove deprecated use of cast
  • Fix the way we generate the dimension (value vs. attribute)
qcolombet marked 2 inline comments as done.Jun 30 2023, 8:52 AM
qcolombet added a subscriber: springerm.
qcolombet added inline comments.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
55

FYI, the "fix" I mentioned in my previous update is here. We use an attribute instead of a value here and that allows the fusing to work properly. I'll highlight what was broken without that.

65

@springerm Do you know if such utility function already exists somewhere?

2239

I had to keep that otherwise clang grab the clone from Operation:: which doesn't take any argument.

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
531

This change is not technically required for this PR, but when I "screwed up" my tiling by returning a value instead of an attribute, this led to slice ops with dynamic dimensions and these were breaking ExtractSliceOp::rankReduceIfNeeded.

By changing this assert into a note, this makes the compiler more robust in case someone does the same mistake as I did. I.e., it will not fuse instead of crashing the compiler.

mlir/test/Dialect/Linalg/tile-softmax.mlir
40

@rengolin I added a test for tile and fuse and it exposed a weakness in the fuse algorithm, which I fixed here as well (it was asserting.)

akuegel added a subscriber: akuegel.Jul 3 2023, 3:30 AM
akuegel added inline comments.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
20

This introduces a cycle between LinalgUtils and Linalg. Is this intended?

qcolombet added inline comments.Jul 3 2023, 3:43 AM
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
20

Ah good catch, I forgot to remove that line.
I saw this problem and in the end I moved the utility functions directly in this file, so the include is actually not needed.

Let me fix that.

qcolombet added inline comments.Jul 3 2023, 3:52 AM
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
20

Fixed at cf244b55c3d2