This is an archive of the discontinued LLVM Phabricator instance.

[mlir][linalg] Add pattern to split reduction dimension in a linalg op
ClosedPublic

Authored by ThomasRaoux on Mar 18 2022, 3:38 PM.

Details

Summary

This transformation allow to break up a reduction dimension in a
parallel and a reduction dimension. This is followed by a separate
reduction op. This allows to generate tree reduction which is beneficial
on target allowing to take advantage parallelism.

Diff Detail

Event Timeline

ThomasRaoux created this revision.Mar 18 2022, 3:38 PM
Herald added a project: Restricted Project. · View Herald TranscriptMar 18 2022, 3:38 PM
ThomasRaoux requested review of this revision.Mar 18 2022, 3:38 PM
mravishankar requested changes to this revision.Mar 18 2022, 5:02 PM

Just did a skim for now. Will come back for a more detailed review.

mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp
59
76

Have seen this come up many places. Maybe worth making this an interface function.

This revision now requires changes to proceed.Mar 18 2022, 5:02 PM
nirvedhmeshram added inline comments.
mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
627

Is this the current controlSplitReductionFn? If so dont we need something better than returning 4?

ThomasRaoux added inline comments.Mar 19 2022, 12:10 PM
mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
627

this is just a dummy heuristic for the lit tests above :)
The transformation will eventually be controlled by an heuristic or a codegen strategy.

address review comments

ThomasRaoux marked an inline comment as done.Mar 22 2022, 1:16 PM
ThomasRaoux added inline comments.
mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp
76

Added a function to the interface. I can replace other usages of such pattern in a follow up patch if everybody is fine with the new interface.

Fine tune the control

Allen added a subscriber: Allen.Mar 22 2022, 9:38 PM

I addressed the style comments and changed a bit the control to allow deciding the shape of the intermediate tensor. Please take a look.

mravishankar requested changes to this revision.Mar 24 2022, 1:42 PM

Mostly looks fine. Some of the output handling code , im not able to fully understand. Maybe worth taking a look again?

mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp
26

Could we make this Optional<Attribute> or FailureOr<Attribute> . Returning nullptr can just lead to segfaults if not handled correctly everywhere.

58

Nit: General not, better practice to return rewriter.notifyMatchFailure() instead of failure(). Makes debug log easier to read.

61

Do we also want to check that the reduction is innermost, or does that not matter?

100

Nit: use llvm::seq<unsigned>(0, map.getNumResults())

129

Not sure I understand the comment. AFAICS, the ControlSplitReductionFn returned the place to add the new dimension. That was account for in the reshape above. Why is the same not done for the outputs. I would expect both to be handled the same way.

136

You already checked the num outputs to be 1. So this loop is dead?

147

Actually the code looks fine, the comment is misleading to me. What is exactly the "most inner dimension" in the comment marked above?

This revision now requires changes to proceed.Mar 24 2022, 1:42 PM
ThomasRaoux marked 2 inline comments as done.

Address review comments

ThomasRaoux marked an inline comment as done.Mar 24 2022, 2:04 PM
ThomasRaoux added inline comments.
mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp
61

It shouldn't matter, I have a test where it isn't

129

Fixed the comment

136

this is the number of results of the output map which is the rank of the destination

147

yes the comment was wrong :( thanks for catching this

mravishankar accepted this revision.Mar 24 2022, 3:05 PM
This revision is now accepted and ready to land.Mar 24 2022, 3:05 PM