This is an archive of the discontinued LLVM Phabricator instance.

[mlir][memref] Add transformation to do loop multi-buffering
ClosedPublic

Authored by ThomasRaoux on Feb 9 2022, 9:01 PM.

Details

Summary

This transformation is useful to break dependency between consecutive loop iterations by increasing the size of a temporary buffer.
This is usually combined with heavy software pipelining.
Started a discourse thread in case we need this needs further discussions: https://discourse.llvm.org/t/loop-double-buffering-multi-buffering/59979

Diff Detail

Event Timeline

ThomasRaoux created this revision.Feb 9 2022, 9:01 PM
ThomasRaoux requested review of this revision.Feb 9 2022, 9:01 PM
ThomasRaoux edited the summary of this revision. (Show Details)Feb 9 2022, 9:15 PM

clang-format

bondhugula requested changes to this revision.Feb 10 2022, 10:35 PM
bondhugula added a subscriber: bondhugula.
bondhugula added inline comments.
mlir/lib/Dialect/MemRef/Transforms/MultiBuffering.cpp
68–70

This should in general have nothing to do with an scf::for op -- you aren't transforming the loop in any way. Just use LoopLikeOpInterface.

99–104

Avoid generating lower-level arithmetic ops. Instead, just generate a modulo using affine.apply.

This revision now requires changes to proceed.Feb 10 2022, 10:35 PM

Use affine.apply for arithmetic

ThomasRaoux added inline comments.Feb 11 2022, 12:31 AM
mlir/lib/Dialect/MemRef/Transforms/MultiBuffering.cpp
68–70

I need to access the single induction variable as well as the lower bound and step. This is not possible with LoopLikeInterface as it cannot be applied to any loops.

99–104

Good point, I didn't realize this was a valid affine map. Changed it.

bondhugula requested changes to this revision.Feb 14 2022, 9:11 AM
bondhugula added inline comments.
mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h
77

Please document the return value.

77

This is a utility and should be returning LogicalResult instead of bool I think.

mlir/lib/Dialect/MemRef/Transforms/MultiBuffering.cpp
66

Please document the return value.

67

You don't need a triple doc comment here. Just the more detailed version on the declaration will suffice.

68–70

This should be pretty straightforward: you can just have the op interface method returning the IV (say getInductionVar()) return None for multi-dim loops (like for scf.parallel) and have this method fail in such a case just like you already do at line 87. It's going to be as simple and strictly more general. For the other two things you mention (lower bound and step), that's straightforward as well but the interface method will have to materialize some trivial IR (for eg. when it's a constant as in the case of affine.for):

// For IV.
Optional<Value> getSingleInductionVar();

  1. lower bound

Optional<Value> getOrCreateLowerBound(OpBuilder &b);
Just return the lower bound SSA value (or materialize the step in case of affine.for)

  1. step

Optional<Value> getOrCreateStep(OpBuilder &b);

With the above, it would be strictly more general than restricting it to scf::ForOp and as simple. And in your case of interest here, you'll get output IR identical to what you have in the test cases below.

mlir/test/Dialect/MemRef/multibuffering.mlir
68–69

You don't need the %{{.*}} = part.

This revision now requires changes to proceed.Feb 14 2022, 9:11 AM
bondhugula added inline comments.Feb 14 2022, 9:21 AM
mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h
77

The utilities are almost always named imperatively instead of a gerund. multiBuffering -> multiBuffer?

Address review feedback

ThomasRaoux marked 6 inline comments as done.Feb 15 2022, 11:55 PM

I switched to using loop interface and added a test for affine.for op. Please take another look.

mlir/lib/Dialect/MemRef/Transforms/MultiBuffering.cpp
68–70

I added the interface and moved the code to use it. This means it now works on affine.for op. I added a test for it.

The comments that I had have been addressed. Dismissing change request.

mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp
23 ↗(On Diff #409160)

Define using dyn_cast to avoid a double check in assert mode.

bondhugula added inline comments.Feb 20 2022, 7:45 PM
mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp
29 ↗(On Diff #409160)

You are also collecting and erasing ops here.

36–53 ↗(On Diff #409160)

Invert the conditional?

if (!....) {
   push_back
   continue;
}

...
70 ↗(On Diff #409160)

Unchecked use of dyn_cast result. This should just be a single cast?

130–131 ↗(On Diff #409160)

strides: append -> assign (since this is empty at this stage).

Address review comments

ThomasRaoux marked 3 inline comments as done.Feb 23 2022, 4:01 PM
ThomasRaoux added inline comments.
mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp
70 ↗(On Diff #409160)

OpFoldResult/PointerUnion don't have cast they only have dyn_cast, I see it used like that in the rest of the code

mravishankar accepted this revision.Feb 23 2022, 4:02 PM

LGTM!

mlir/lib/Dialect/MemRef/Transforms/MultiBuffer.cpp
89 ↗(On Diff #409160)

Nit: if this user doesnt dominate...

This revision is now accepted and ready to land.Feb 23 2022, 4:02 PM
This revision was landed with ongoing or failed builds.Feb 24 2022, 9:43 AM
This revision was automatically updated to reflect the committed changes.