This is an archive of the discontinued LLVM Phabricator instance.

[MLIR] Introduce utility for affine loop unswitching / hoisting if/else
ClosedPublic

Authored by bondhugula on Apr 10 2020, 6:16 AM.

Details

Summary

This revision introduces a utility to unswitch affine.for/parallel loops
by hoisting affine.if operations past surrounding affine.for/parallel.
The hoisting works for both perfect/imperfect nests and in the presence
of else blocks. The hoisting is currently to as outermost a level as
possible. Uses a test pass to test the utility.
Add convenience method Operation::getParentWithTrait<Trait>.

Depends on D77487.

Diff Detail

Event Timeline

bondhugula created this revision.Apr 10 2020, 6:16 AM

Format fixes.

Harbormaster completed remote builds in B52665: Diff 256560.
andydavis1 added inline comments.Apr 10 2020, 8:24 AM
mlir/lib/Dialect/Affine/Utils/Utils.cpp
68

This function is well documented but a little long. Could it be broken up a bit? For example, the canonicalization bit a the top of the function could be pulled out into a separate function.

This seems like it would be easy to generalize.
What facilities are concretely missing to make this transformation work on more general loops and conditionals?
Could it be implemented it in a more general fashion with LoopLikeInterface and other interfaces, if appropriate?

The same questions apply to full/partial tile separation, unrolling, L/S forwarding and maybe others.

This seems like it would be easy to generalize.
What facilities are concretely missing to make this transformation work on more general loops and conditionals?

Take a look at the checks on invariance / validity - you'll have to do all kinds of value provenance to do something equivalent on other loop ops - depending on whether you already know what hoisting you'd like to do. Conditionals/loops returning values also add additional complexity. With affine, all of this becomes clean and trivial. That said, the mechanical movement of blocks would be similar - but that's about a couple of dozen lines of code here.

Could it be implemented it in a more general fashion with LoopLikeInterface and other interfaces, if appropriate?

The same questions apply to full/partial tile separation, unrolling, L/S forwarding and maybe others.

Reg. unrolling - at the time op interfaces were introduced, I thought about whether something could be moved to the loop like op interface - it just was not at all worth it; the amount of sharing was tiny enough a fraction to make it pointless. More so with full/partial tile separation (the code moving the blocks is a small part embedded in). L/S fwd'ing relies on affine dependence analysis - the actual mechanics of it which is the last step are trivial. In all these cases, you'll find the mechanics at the end could be reused, but you'll have to weigh how heavy those are. One should still have op interfaces like LoadLikeOp, StoreLikeOp, IfLikeOp for future purposes to share mechanics.

rriddle added inline comments.Apr 10 2020, 11:10 AM
mlir/include/mlir/Dialect/Affine/Utils.h
25

Why not LogicalResult?

mlir/lib/Dialect/Affine/Utils/Utils.cpp
67

Please use ///

119

Cache the Identifier to avoid string comparison in attribute lookup.

160

I would have expected the if you want to remove dead else blocks you could specifically target the ops you know about, instead of running on the entire function.

mlir/test/lib/Dialect/Affine/TestIfHoist.cpp
28 ↗(On Diff #256560)

Remove the semicolon

Thanks for adding this optimization!
One general comment: this is called loop unswitching in LLVM: https://llvm.org/docs/Passes.html#loop-unswitch-unswitch-loops
It would be great if we could use that terminology for consistency.

mlir/include/mlir/Dialect/Affine/Utils.h
22

Wondering what happens if the if is loop invariant to only some of the IVs in an affine.parallel. Do we plan to split the affine.parallel?

mlir/lib/Dialect/Affine/Utils/Utils.cpp
61

could you please add a couple of tests for the last two scenarios (line 56 and 60)?

102

Comments refer to hoistOverOp but variable is hoistPoint? Please, align terminology.

mlir/test/Dialect/Affine/if-hoist.mlir
12 ↗(On Diff #256560)

Could we use different loads so that we can check that the guarded load is placed in the right place?

57 ↗(On Diff #256560)

Maybe interesting to add a CHECK-NOT for else?

70 ↗(On Diff #256560)

Could you please add another test that covers if-else + extra code in the innermost loop and check for the right order of the resulting instructions? Similar to the first test

183 ↗(On Diff #256560)

could you please add some tests with parallel?

bondhugula marked 5 inline comments as done.Apr 10 2020, 5:18 PM

Thanks for adding this optimization!
One general comment: this is called loop unswitching in LLVM: https://llvm.org/docs/Passes.html#loop-unswitch-unswitch-loops
It would be great if we could use that terminology for consistency.

Sure, thanks.

mlir/include/mlir/Dialect/Affine/Utils.h
22

Good question. At the moment, we don't intend to split the affine.parallel. If needed, this could be made configurable.

25

It would always return success() then - this utility can't fail. (I didn't think 'no hoisting' should be called a failure, but if we do, we could use LogicalResult.)

mlir/lib/Dialect/Affine/Utils/Utils.cpp
67

This is an implementation comment to help understand the return value. The longer doc comment is on the declaration in the header - didn't want to duplicate it.

68

I think this is a good idea. But instead of the top part which is short, we could break the bottom half which has the mechanics.

bondhugula marked 17 inline comments as done.Apr 10 2020, 6:14 PM
bondhugula added inline comments.
mlir/lib/Dialect/Affine/Utils/Utils.cpp
61

Thanks - this wasn't covered earlier.

160

I actually forgot to update this call to use the 'op local' apply patterns, which I used at the entry. But looking back, that won't yield the desired simplification because in some cases, we want some dead loops in the else version to go away (even if the 'else' isn't all dead). So, it's something between 'op local' or whole function level that we need here. Example: the following will yield a dead loop in the else branch after hoisting unless you run the rewriter on all ops inside the hoisted if (and on only those to be precise). We should support such a method in the rewrite driver.

for 
  x;
  for 
    if 
      y
mlir/test/Dialect/Affine/if-hoist.mlir
57 ↗(On Diff #256560)

Thanks!

bondhugula marked 3 inline comments as done.

Address review comments.

Thanks for adding this optimization!
One general comment: this is called loop unswitching in LLVM: https://llvm.org/docs/Passes.html#loop-unswitch-unswitch-loops
It would be great if we could use that terminology for consistency.

Done. The test pass is now called loop unswitching (-test-affine-loop-unswitch) although the utility is still called hoistAffineIf since it's really meant for on an 'if' op.

Rename test pass - affine loop unswitching (-test-affine-loop-unswitch)

jdoerfert added inline comments.
mlir/test/lib/Dialect/Affine/TestAffineLoopUnswitching.cpp
2

Nit: ^

dcaballe accepted this revision.Apr 13 2020, 11:48 AM

Thanks for addressing the comments. Just a minor comment. Otherwise, LGTM. Please, wait for the other reviewers.

mlir/test/Dialect/Affine/loop-unswitch.mlir
27

Better capture arg0 and arg1 as variables

bondhugula marked 3 inline comments as done.

Address review comments.

Thanks for the review.

mlir/test/lib/Dialect/Affine/TestAffineLoopUnswitching.cpp
2

Thanks.

@andydavis1 - did you have any more comments here?

rriddle added inline comments.Apr 14 2020, 11:40 PM
mlir/include/mlir/Dialect/Affine/Utils.h
25

I would say that success/failure takes on different meanings depending on the context. Here I would read success as the if was successfully hoisted. This is different from the failure mode that was decided on for passes. e.g., the folding methods returns success if the operation was folded, and failure otherwise.

mlir/include/mlir/IR/Operation.h
132

nit: auto -> Operation

mlir/lib/Dialect/Affine/Utils/Utils.cpp
141

Please don't call static functions as members.

bondhugula marked 5 inline comments as done.

Address review comments.

Thanks for the review @rriddle

mlir/include/mlir/Dialect/Affine/Utils.h
25

Sure - changed to LogicalResult then.

mlir/lib/Dialect/Affine/Utils/Utils.cpp
142–143

@rriddle Here is where we need the erased parameter. I can't think of an alternative. You need to know whether the op was erased if this canonicalization is called. The user of this utility may not call the canonicalization nor could we expect them. Even if they do, they'd themselves have to potentially deal with the aftermath of the op getting erased.

rriddle accepted this revision.Apr 15 2020, 10:24 AM
rriddle added inline comments.
mlir/lib/Dialect/Affine/Utils/Utils.cpp
29

This builder looks unused.

97

b.getIdentifier should be unnecessary here, you already have an identifier.

142–143

Sorry, I don't think I was clear in the message in the other revision. I think we should remove the bool "converged" result that the methods are currently returning. I agree that knowing if the ops was erased or not can be important.

mlir/test/lib/Dialect/Affine/TestAffineLoopUnswitching.cpp
49

nit: You could write this as:

auto walkFn = ([&](AffineIfOp op) -> WalkResult {
  return hoistAffineIfOp(op);
};
if (func.walk(walkFn).wasInterrupted())
  break;

You may even be able to use hoistAffineIfOp directly as walkFn, but I can't remember.

64

nit: drop the newline here.

This revision is now accepted and ready to land.Apr 15 2020, 10:24 AM
andydavis1 accepted this revision.Apr 15 2020, 10:39 AM
bondhugula marked 7 inline comments as done.

Address another round of review comments from @rriddle.

Thanks for the review.

mlir/lib/Dialect/Affine/Utils/Utils.cpp
142–143

I see; sure, but it appears as if it would definitely be needed some day.

mlir/test/lib/Dialect/Affine/TestAffineLoopUnswitching.cpp
49

Thanks - I didn't know about wasInterrupted(). Note that the walk continues here if the hoist fails, and is interrupted when the hoist succeeds. So I can't implicitly convert as you show above, but have used the pattern.

bondhugula retitled this revision from [MLIR] Introduce utility to hoist affine if/else conditions to [MLIR] Introduce utility for affine loop unswitching / hoisting if/else.Apr 15 2020, 11:59 AM
bondhugula edited the summary of this revision. (Show Details)
bondhugula edited the summary of this revision. (Show Details)Apr 15 2020, 12:01 PM
This revision was automatically updated to reflect the committed changes.

Fixed in a07e5b857425a8d411dbf1cdfca5ba5d6521549d ; feel free to revert immediately on build failure in the future, this can hide other bugs otherwise.

Fixed in a07e5b857425a8d411dbf1cdfca5ba5d6521549d ; feel free to revert immediately on build failure in the future, this can hide other bugs otherwise.

Thank you for fixing this - I had recently switched to LLD and missed this. I should be paying more attention to the build failure notifications here - sorry.

This patch also breaks building with BUILD_SHARED_LIBS=ON on both Linux and windows:

FAILED: lib/libMLIRAffineUtils.so.11git 
: && /usr/bin/c++ -fPIC -fPIC -fvisibility-inlines-hidden -Werror=date-time -Wall -Wextra -Wno-unused-parameter -Wwrite-strings -Wcast-qual -Wno-missing-field-initializers -pedantic -Wno-long-long -Wimplicit-fallthrough -Wno-maybe-uninitialized -Wno-noexcept-type -Wdelete-non-virtual-dtor -Wno-comment -fdiagnostics-color -g  -Wl,-z,defs -Wl,-z,nodelete -shared -Wl,-soname,libMLIRAffineUtils.so.11git -o lib/libMLIRAffineUtils.so.11git tools/mlir/lib/Dialect/Affine/Utils/CMakeFiles/MLIRAffineUtils.dir/Utils.cpp.o  -Wl,-rpath,"\$ORIGIN/../lib" lib/libMLIRAffine.so.11git lib/libMLIRLoopLikeInterface.so.11git lib/libMLIRStandardOps.so.11git lib/libMLIREDSC.so.11git lib/libMLIRSideEffects.so.11git lib/libMLIRCallInterfaces.so.11git lib/libMLIRControlFlowInterfaces.so.11git lib/libMLIRIR.so.11git lib/libMLIRSupport.so.11git -lpthread lib/libLLVMSupport.so.11git && :
tools/mlir/lib/Dialect/Affine/Utils/CMakeFiles/MLIRAffineUtils.dir/Utils.cpp.o: In function `mlir::hoistAffineIfOp(mlir::AffineIfOp, bool*)':
/mnt/c/Users/Markus/Documents/CLion/llvm-project/mlir/lib/Dialect/Affine/Utils/Utils.cpp:142: undefined reference to `mlir::applyOpPatternsAndFold(mlir::Operation*, mlir::OwningRewritePatternList const&, bool*)'
/mnt/c/Users/Markus/Documents/CLion/llvm-project/mlir/lib/Dialect/Affine/Utils/Utils.cpp:170: undefined reference to `mlir::applyPatternsAndFoldGreedily(mlir::Operation*, mlir::OwningRewritePatternList const&)'