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.
Details
Diff Detail
Event Timeline
Just did a skim for now. Will come back for a more detailed review.
mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp | ||
---|---|---|
58 | Could probably replace uses of this with this method https://github.com/llvm/llvm-project/blob/251d062e4e2737d88ba53e65969b210f2830c4df/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L1006 | |
75 | Have seen this come up many places. Maybe worth making this an interface function. |
mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp | ||
---|---|---|
627 | Is this the current controlSplitReductionFn? If so dont we need something better than returning 4? |
mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp | ||
---|---|---|
627 | this is just a dummy heuristic for the lit tests above :) |
mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp | ||
---|---|---|
75 | 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. |
I addressed the style comments and changed a bit the control to allow deciding the shape of the intermediate tensor. Please take a look.
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 | ||
---|---|---|
27 | Could we make this Optional<Attribute> or FailureOr<Attribute> . Returning nullptr can just lead to segfaults if not handled correctly everywhere. | |
59 | Nit: General not, better practice to return rewriter.notifyMatchFailure() instead of failure(). Makes debug log easier to read. | |
62 | Do we also want to check that the reduction is innermost, or does that not matter? | |
101 | Nit: use llvm::seq<unsigned>(0, map.getNumResults()) | |
130 | 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. | |
137 | You already checked the num outputs to be 1. So this loop is dead? | |
148 | Actually the code looks fine, the comment is misleading to me. What is exactly the "most inner dimension" in the comment marked above? |
Could we make this Optional<Attribute> or FailureOr<Attribute> . Returning nullptr can just lead to segfaults if not handled correctly everywhere.