Page MenuHomePhabricator

[mlir][Linalg] Rethink fusion ot linalg ops with reshape ops.
ClosedPublic

Authored by mravishankar on Oct 7 2020, 12:52 PM.

Details

Summary

The current fusion on tensors fuses reshape ops with generic ops by
linearizing the indexing maps of the fused tensor in the generic
op. This has some limitations

  • It only works for static shapes
  • The resulting indexing map has a linearization that would be potentially prevent fusion later on (for ex. tile + fuse).

Instead, try to fuse the reshape consumer (producer) with generic op
producer (consumer) by expanding the dimensionality of the generic op
when the reshape is expanding (folding).
This approach conflicts with the linearization approach. The use of
the linearization approach is still kept as the default, but the
approach added in this change will be made the default after further
experimentation.

Diff Detail

Event Timeline

mravishankar created this revision.Oct 7 2020, 12:52 PM
mravishankar requested review of this revision.Oct 7 2020, 12:53 PM

First pass...

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
319–320

nit: explicitly return if(!reshapeOp.getSrcType().hasStaticShape()) return false; to simplify condition below.

379

I wonder the case when reshape is a collapsing dims, how do we avoid linearization ?

392

Very nice! and the assumption here pushing reshapes up in the chain the top most reshape will be folded away (e.g IREE or in general reshape of arguments)

441

nit: s/folded/collapsed

nicolasvasilache added a comment.EditedOct 8 2020, 12:57 AM

I am unclear why this is implemented as part of fusion ?
My thinking was that reshapes that expand are pulled upwards, reshapes that contract are pushed downwards.
Applying this reshape folding onto producer and consumer ops enables tiling and fusion in higher-D.

Inside the tiles, depending on the tile sizes and the contiguity aspects, we may or may not be able to just "reshape-without-copy" to convert a matmul in 4-D form to a regular 2-D matmul (think fusing the output of the last conv layer with the following linear layer).

Can we revisit this to just introduce 2 patterns for each case and keep fusion separate from this?

mravishankar marked an inline comment as done.Oct 8 2020, 10:28 AM

I am unclear why this is implemented as part of fusion ?
My thinking was that reshapes that expand are pulled upwards, reshapes that contract are pushed downwards.
Applying this reshape folding onto producer and consumer ops enables tiling and fusion in higher-D.

Inside the tiles, depending on the tile sizes and the contiguity aspects, we may or may not be able to just "reshape-without-copy" to convert a matmul in 4-D form to a regular 2-D matmul (think fusing the output of the last conv layer with the following linear layer).

Can we revisit this to just introduce 2 patterns for each case and keep fusion separate from this?

We can do this as folding, but the "fusion" added here conflicts with the fusion that uses "linearization" by folding a collapsing reshape with its producer by linearizing the collapsed dimensions in the indexing map of the fused tensor in the producer. This is opposite to the pattern here where collapsing reshapes get fused with its consumer. (In the previous two statements you can replace collapsing with expanding and producer with consumer and the statement still holds). So I think it makes sense to drop the linearization approach. Just trying to stage this here. Here is what I plan

  1. Land this change in its current state
  2. The reshape + (generic/indexed-generic) op fusion can be replace with patterns. Both for the expansion approach here and the linearization approach. The "expansion' pattern can be moved into a folding pattern with the linearization patterns dropped from the fusion pass here. I just need to verify that this works in IREE (which is probably the main user of these fusions). So far it seems to be the case, but the pain will be lesser if the switch is staged a bit

WDYT?

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
379

They conflict. So eventually want to deprecate linearization. No one is really using it, so maybe I will refactor these into separate patterns (they kind of already are). Will wait to get resolution on Nicolas' comment.

392

Thats right. The "collapsing" reshapes get pushed down the function to the returns, and the "expanding" reshape get pushed all the way to the function arguments.

asaadaldien accepted this revision.Oct 8 2020, 1:30 PM

L(very)GTM, thanks Mahesh!. I will leave the pattern splitting/refactoring discussion between you and Nicolas.

This revision is now accepted and ready to land.Oct 8 2020, 1:30 PM

SGTM, I am fine with evolving after this CL.

Just added a small comment.
I am tight on time this week and will not have time to review more deeply.
I'll take a deeper pass when this gets rewritten as independent patterns.

Thanks for pushing this!

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp
547

There is a clone in the interface LinalgStructuredOpsInterface.td.
Can we (improve and) reuse that ?

mravishankar marked an inline comment as done.

Rebase

Just added a small comment.
I am tight on time this week and will not have time to review more deeply.
I'll take a deeper pass when this gets rewritten as independent patterns.

Thanks for pushing this!

I added a couple of patches on top of this that does the refactoring to use patterns (particularly : https://reviews.llvm.org/D89201). PTAL. I would like to squash and submit them in one go.

nicolasvasilache accepted this revision.Oct 14 2020, 7:59 AM

Please go ahead and squash merge once child comments are addressed.
Thanks for these improvements!