This is an archive of the discontinued LLVM Phabricator instance.

[mlir][Linalg] Allow collapsing subset of the reassociations when fusing by collapsing.
ClosedPublic

Authored by mravishankar on Apr 5 2022, 1:03 PM.

Details

Summary

This change generalizes the fusion of tensor.expand_shape ->
linalg.generic op by collapsing to handle cases where only a subset
of the reassociations specified in the tensor.expand_shape are valid
to be collapsed.
The method that does the collapsing is refactored to allow it to be a
generic utility when required.

Diff Detail

Event Timeline

mravishankar created this revision.Apr 5 2022, 1:03 PM
Herald added a project: Restricted Project. · View Herald TranscriptApr 5 2022, 1:03 PM
mravishankar requested review of this revision.Apr 5 2022, 1:03 PM
okkwon added a subscriber: okkwon.Apr 5 2022, 7:41 PM

I left a few comments mostly related to naming and comments. I will review the rest later.

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
1158

The indexing map is already checked outside, we can simply return a non-optional vector.
To guard a case in future for dev, we can use assert on the projected permutation condition.

Also, the result is a vector of dimensions on the iteration space, so I am not clear of using ReassociationIndices.

1219

It would be great if you can add fuse-able and un-fuse-able examples here.

1220

The function now returns a vector, so the function name seems to need a change.

1246

This can never happen knowing its implementation because the condition is already bailed out earlier.

1250

Please add why.

1253

We are mixing of the terms of domain and iteration. Sticking to one would make the code more readable.

1270

Please add TODO.

gysit added inline comments.Apr 6 2022, 5:52 AM
mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
1170

nit: repetition

1269

nit:associative

1303

origNumLoops since we have a lot of dims already?

1311

nit:Repetition

1368

...OpMapping[0] = ...
...OpMapping[1] = ...
...OpMapping[2] = ...
etc.?

1371

I wonder if a struct could help to clarify things a bit by naming the members collapsedOpDim and foldingIdx or similar?

mlir/test/Dialect/Linalg/fuse-with-reshape-by-collapsing.mlir
426

nit: probably should be [2, 3]{{\]}} with an additional closing bracket? (below as well?)

okkwon edited reviewers, added: okkwon; removed: okwank.Apr 6 2022, 8:55 AM
okkwon added a comment.Apr 6 2022, 3:49 PM

Adding more comments.

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
1267

nit: dimension -> dimensions

1313

Is there a case that the number of loops is different from the rank of the iteration space? It would be great to give an example as comment.

1315

Do we hit this case? I thought the case is already bailed out when producing foldedIterationDims.

1395

It is irrelevant to this PR, but can we make it const CollapsingInfo &?
Same for getCollapsedOpIndexingMap() below.

1538

nit: collapse

1647

Check control function first before calling isFusableWith...().

mravishankar marked 13 inline comments as done.

Address comments.

okkwon added a comment.Apr 8 2022, 3:59 PM

Adding more comments.

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
1233

%2 -> %3 in the examples.

Also, need an expand_shape after the generic op for this generic op.

1292

Can we just move the whole vector like emplace_back(std::move(foldedIterationSpaceDims))?

1310

nit: contains -> contain, thats -> that's

1358

nit: [0, 1, 2], [3, 4] and remove the last tick at the end.

1414

use const AffineMap &.

1437

use const AffineMap &.

mravishankar marked 6 inline comments as done.

Address comments.

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
1158

Dropped the Optional and made it assert that map is projected permutation. I am using Reassociation here cause just like expand/collapse shape op. the listed dimensions are "folded".

1219

Those are typically documented in lit tests? There are quite a few examples there.. Added descrition.

1270

Actually not sure about TODO. Its not a gaurantee of the linalg op that this is allowed. So TODO would be misleading.

1313

I wrote this method with a view that this could be moved into a separate utility method by itself. So added some extra checks here independent of whats checked earlier.

1315

It is possible if this method is made a utility and then passed in dims for collapsing iteration space has an error.

1371

I think pair does the job here. It is documented what the pair represents. I found the naming and description of such a struct hard to write (let alone convey meaning effectively)

1414

For such objects use of const I think is not recommended. (https://mlir.llvm.org/docs/Rationale/UsageOfConst/)

1437

See above.

1647

I first check if it is fusable at all structurally, then call the control function...

gysit accepted this revision.Apr 12 2022, 12:57 AM
gysit added inline comments.
mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
1371

Alright!

mlir/test/Dialect/Linalg/fuse-with-reshape-by-collapsing.mlir
501–502

nit: I think the dims and the constants they use are unused?

526

I think 2 and 3 are not collapsed since they are not in order? If yes it may be helpful to add a short comment here?

This revision is now accepted and ready to land.Apr 12 2022, 12:57 AM
mravishankar marked 2 inline comments as done.

Address comments, and make reduction dimension folding predicated on them being contiguous instead of monotonic.

mlir/test/Dialect/Linalg/fuse-with-reshape-by-collapsing.mlir
526

Actually, took a look again at this case. The reduction dimensions folded have to be contiguous in terms of reduction dims. So fixed that in code and modified test accordingly.

gysit added inline comments.Apr 12 2022, 11:24 AM
mlir/test/Dialect/Linalg/fuse-with-reshape-by-collapsing.mlir
526

Ah right that makes sense! That way the reduction order should be unchanged...