This is an archive of the discontinued LLVM Phabricator instance.

[mlir][Linalg] Add a pattern to decompose `linalg.generic` ops.
ClosedPublic

Authored by mravishankar on Jul 13 2022, 4:26 PM.

Details

Summary

This patch adds a pattern to decompose a linalg.generic operations
that

  • has only parallel iterator types
  • has more than 2 statements (including the yield)

into multiple linalg.generic operation such that each operation has
a single statement and a yield.
The pattern added here just splits the matching linalg.generic into
two linalg.generics, one containing the first statement, and the
other containing the remaining. The same pattern can be applied
repeatedly on the second op to ultimately fully decompose the generic
op.

Diff Detail

Event Timeline

mravishankar created this revision.Jul 13 2022, 4:26 PM

Fix bug related to output indexing map handling.

Add lit tests and finish corner cases.

Fix lit test.

mravishankar published this revision for review.Jul 15 2022, 1:42 AM

A few small comments. I need to look at the lit test on a bigger screen and will follow-up. Nothing obvious, but looking to see if it can be simplified.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
704

Thanks - I've open-coded this a few times.

mlir/include/mlir/IR/Types.h
97 ↗(On Diff #444915)

I'm trying to spot the diff here. Is it just formatting fallout?

132 ↗(On Diff #444915)

I'm not sure these helpers are in the "this was a great idea" bucket. I usually just say t.isa<IntegerType>. A few extra characters but less core coupling.

Minimally: this new function is now not consistent with the comment.

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
139

River would say: Remove trivial braces (here and below)

(if River is watching: see, I can be taught)

152

I'm trying to understand why this restriction needs to be this strict. Isn't the restriction actually that the peeled operation cannot use the out operand (and let us guarantee that we have not violated ordering semantics)?

I'm fine with the stricter form as a starting point but making sure I understand (and if we think it can be loosened, to note that in a comment).

164

Nit: "operation has less than 2 statements"

179

I would have expected that a BlockAndValueMapper would be needed to remap the operands? I'm a bit surprised this works without leaving the uses pointing at the original block (which will be deleted and should assert).

(edit: I see you being clever below. disregard)

241

Stale comment?

mravishankar marked 3 inline comments as done.

Address comments and remove all changes to Types.h/Types.cpp

Actually drop the changes to Types.h

mravishankar marked 3 inline comments as done.Jul 15 2022, 10:46 AM

Addressed the comments, but they were not in the place they were initially intended to be, so hopefully I read the comments correctly.

mlir/include/mlir/IR/Types.h
97 ↗(On Diff #444915)

Yeah I think so. Ill revert these changes.

132 ↗(On Diff #444915)

Ah thats right. Ill clean up this and all related changes.

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
152

Left a comment saying it might be relaxed. It is definitely being conservative. Mostly the outs values are only read when the op has a reduction iterator type. If you have a case where this is done, worth looking into. IIUC we should be able to "just drop the condition" and still not violate any dependency.

164

its less than 3? It is meant to abort on only "yield" or "op + yield"

mravishankar marked 2 inline comments as done.

Rebase

stellaraccident accepted this revision.Jul 15 2022, 3:46 PM
stellaraccident added inline comments.
mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp
164

Oh right. Had my original impl in mind.

mlir/test/Dialect/Linalg/decompose-ops.mlir
2

I wish this test were able to write workout being so extensive. But this is also a tricky/detailed pattern and I can't see an immediate simplification to how you have it.

This revision is now accepted and ready to land.Jul 15 2022, 3:46 PM

Rebase to retrigger pre-merge check

This revision was landed with ongoing or failed builds.Jul 15 2022, 4:01 PM
This revision was automatically updated to reflect the committed changes.