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

[bondhugula]

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.

[bondhugula]

Format fixes.

[andydavis1]

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.

[nicolasvasilache]

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.

[bondhugula]

In D77870#1974419, @nicolasvasilache wrote:

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]

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

[dcaballe]

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]

In D77870#1975018, @dcaballe wrote:

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]

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]

Address review comments.

[bondhugula]

In D77870#1975507, @bondhugula wrote:

In D77870#1975018, @dcaballe wrote:

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.

[bondhugula]

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

[jdoerfert]

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

3

Nit: ^

[dcaballe]

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

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

28

Better capture arg0 and arg1 as variables

[bondhugula]

Address review comments.

[bondhugula]

Thanks for the review.

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

3

Thanks.

[bondhugula]

@andydavis1 - did you have any more comments here?

[rriddle]

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

140

Please don't call static functions as members.

[bondhugula]

Address review comments.

[bondhugula]

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

141–142

@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]

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

29

This builder looks unused.

97

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

141–142

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.

[bondhugula]

Address another round of review comments from @rriddle.

[bondhugula]

Thanks for the review.

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

141–142

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.

[jpienaar]

This seems to be failing to build https://buildkite.com/mlir/mlir-core/builds/4201#90281b55-af17-4d42-837d-e3cc8f5eed12

[mehdi_amini]

In D77870#1984865, @jpienaar wrote:

This seems to be failing to build https://buildkite.com/mlir/mlir-core/builds/4201#90281b55-af17-4d42-837d-e3cc8f5eed12

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

[bondhugula]

In D77870#1985749, @mehdi_amini wrote:

In D77870#1984865, @jpienaar wrote:

This seems to be failing to build https://buildkite.com/mlir/mlir-core/builds/4201#90281b55-af17-4d42-837d-e3cc8f5eed12

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.

[zero9178]

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&)'

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

let extraClassDeclaration = [{
     IntegerSet getIntegerSet();
     void setIntegerSet(IntegerSet newSet);
 
     /// Sets the integer set with its operands. The size of 'operands' must not
     /// exceed the current number of operands for this instance, as the operands
     /// list of AffineIf is not resizable.
     void setConditional(IntegerSet set, ValueRange operands);
 
+    /// Returns true if an else block exists.
+    bool hasElse() { return !elseRegion().empty(); }
+
     Block *getThenBlock() {
       assert(!thenRegion().empty() && "Unexpected empty 'then' region.");
       return &thenRegion().front();
     }
 
     Block *getElseBlock() {
-      assert(!elseRegion().empty() && "Empty 'else' region.");
+      assert(hasElse() && "Empty 'else' region.");
       return &elseRegion().front();
     }
 
     OpBuilder getThenBodyBuilder() {
       assert(!thenRegion().empty() && "Unexpected empty 'then' region.");
       Block &body = thenRegion().front();
       return OpBuilder(&body, std::prev(body.end()));
     }
     OpBuilder getElseBodyBuilder() {
-      assert(!elseRegion().empty() && "Unexpected empty 'else' region.");
+      assert(hasElse() && "No 'else' block");
       Block &body = elseRegion().front();
       return OpBuilder(&body, std::prev(body.end()));
     }
   }];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
let extraClassDeclaration = [{
     AffineValueMap getUpperBoundsValueMap();
     AffineValueMap getRangesValueMap();
 
     /// Get ranges as constants, may fail in dynamic case.
     Optional<SmallVector<int64_t, 8>> getConstantRanges();
 
     Block *getBody();
     OpBuilder getBodyBuilder();
+    MutableArrayRef<BlockArgument> getIVs() {
+      return getBody()->getArguments();
+    }
     void setSteps(ArrayRef<int64_t> newSteps);
 
     static StringRef getLowerBoundsMapAttrName() { return "lowerBoundsMap"; }
     static StringRef getUpperBoundsMapAttrName() { return "upperBoundsMap"; }
     static StringRef getStepsAttrName() { return "steps"; }
   }];
 }
 

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

+//===- Utils.h - Affine dialect utilities -----------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file declares a set of utilities for the affine dialect ops.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_AFFINE_UTILS_H
+#define MLIR_DIALECT_AFFINE_UTILS_H
+
+namespace mlir {
+
+class AffineIfOp;
+
+/// Hoists out affine.if/else to as high as possible, i.e., past all invariant
+/// affine.fors/parallel's. Returns true if any hoisting happened or the `ifOp`
+/// was folded away; in the latter case `folded` is set to true. This
+/// hoisting could lead to significant code expansion in some cases.
+bool hoistAffineIfOp(AffineIfOp ifOp, bool *folded = nullptr);
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_AFFINE_UTILS_H

mlir/include/mlir/IR/OpDefinition.h

public:
   /// or nullptr if this is a top-level operation.
   Operation *getParentOp() { return getOperation()->getParentOp(); }
 
   /// Return the closest surrounding parent operation that is of type 'OpTy'.
   template <typename OpTy> OpTy getParentOfType() {
     return getOperation()->getParentOfType<OpTy>();
   }
 
+  /// Returns the closest surrounding parent operation with trait `Trait`.
+  template <template <typename T> class Trait>
+  Operation *getParentWithTrait() {
+    return getOperation()->getParentWithTrait<Trait>();
+  }
+
   /// Return the context this operation belongs to.
   MLIRContext *getContext() { return getOperation()->getContext(); }
 
   /// Print the operation to the given stream.
   void print(raw_ostream &os, OpPrintingFlags flags = llvm::None) {
     state->print(os, flags);
   }
   void print(raw_ostream &os, AsmState &asmState,

mlir/include/mlir/IR/Operation.h

public:
   template <typename OpTy> OpTy getParentOfType() {
     auto *op = this;
     while ((op = op->getParentOp()))
       if (auto parentOp = dyn_cast<OpTy>(op))
         return parentOp;
     return OpTy();
   }
 
+  /// Returns the closest surrounding parent operation with trait `Trait`.
+  template <template <typename T> class Trait>
+  Operation *getParentWithTrait() {
+    auto *op = this;
+    while ((op = op->getParentOp()))
+      if (op->hasTrait<Trait>())
+        return op;
+    return nullptr;
+  }
+
   /// Return true if this operation is a proper ancestor of the `other`
   /// operation.
   bool isProperAncestor(Operation *other);
 
   /// Return true if this operation is an ancestor of the `other` operation. An
   /// operation is considered as its own ancestor, use `isProperAncestor` to
   /// avoid this.
   bool isAncestor(Operation *other) {

mlir/lib/Dialect/Affine/CMakeLists.txt

target_link_libraries(MLIRAffine
   MLIREDSC
   MLIRIR
   MLIRLoopLikeInterface
   MLIRSideEffects
   MLIRStandardOps
   )
 
 add_subdirectory(Transforms)
+add_subdirectory(Utils)

mlir/lib/Dialect/Affine/Utils/CMakeLists.txt

+add_mlir_dialect_library(MLIRAffineUtils
+  Utils.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
+
+  )
+target_link_libraries(MLIRAffineUtils
+  PUBLIC
+  MLIRAffine
+  )

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

+//===- Utils.cpp ---- Utilities for affine dialect transformation ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements miscellaneous transformation utilities for the Affine
+// dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/IntegerSet.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+/// Promotes the `then` or the `else` block of an AffineIfOp into the latter's
+/// containing block, and discards the rest of the op.
+static void promoteIfBlock(AffineIfOp ifOp, bool elseBlock) {
+  OpBuilder b(ifOp);
+
+  if (elseBlock)
+    assert(ifOp.hasElse() && "else block expected");
+
+  Block *destBlock = ifOp.getOperation()->getBlock();
+  Block *srcBlock = elseBlock ? ifOp.getElseBlock() : ifOp.getThenBlock();
+  destBlock->getOperations().splice(
+      Block::iterator(ifOp), srcBlock->getOperations(), srcBlock->begin(),
+      std::prev(srcBlock->end()));
+  ifOp.erase();
+}
+
+/// Returns the outermost affine.for/parallel op that the `ifOp` is invariant
+/// on. The `ifOp` could be hoisted and placed right before such an operation.
+/// This method assumes that the ifOp has been canonicalized (to be correct and
+/// effective).
+static Operation *getOutermostInvariantForOp(AffineIfOp ifOp) {
+  // Walk up the parents past all for op that this conditional is invariant on.
+  auto ifOperands = ifOp.getOperands();
+  auto res = ifOp.getOperation();
+  while (!isa<FuncOp>(res->getParentOp())) {
+    auto *parentOp = res->getParentOp();
+    if (auto forOp = dyn_cast<AffineForOp>(parentOp)) {
+      if (llvm::is_contained(ifOperands, forOp.getInductionVar()))
+        break;
+    } else if (auto parallelOp = dyn_cast<AffineParallelOp>(parentOp)) {
+      for (auto iv : parallelOp.getIVs())
+        if (llvm::is_contained(ifOperands, iv))
+          break;
+    } else if (!isa<AffineIfOp>(parentOp)) {
+      // Won't walk up past anything other than affine.for/if ops.
+      break;
+    }
+    // You can always hoist up past any affine.if ops.
+    res = parentOp;
+  }
+  return res;
+}
+
+/// A helper for the mechanics of hoistAffineIfOp. Hoists `ifOp just over
+/// `hoistOverOp`. Returns true if any hoisting happened.
+static AffineIfOp hoistAffineIfOp(AffineIfOp ifOp, Operation *hoistOverOp) {
+  // No hoisting to do.
+  if (hoistOverOp == ifOp)
+    return ifOp;
+
+  // Create the hoisted 'if' first. Then, clone the op we are hoisting over for
+  // the else block. Then drop the else block of the original 'if' in the 'then'
+  // branch while promoting its then block, and analogously drop the 'then'
+  // block of the original 'if' from the 'else' branch while promoting its else
+  // block.
+  BlockAndValueMapping operandMap;
+  OpBuilder b(hoistOverOp);
+  auto hoistedIfOp = b.create<AffineIfOp>(ifOp.getLoc(), ifOp.getIntegerSet(),
+                                          ifOp.getOperands(),
+                                          /*elseBlock=*/true);
+
+  // Create a clone of hoistOverOp to use for the else branch of the hoisted
+  // conditional. The else block may get optimized away if empty.
+  Operation *hoistOverOpClone = nullptr;
+  // We use this unique name to identify/find  `ifOp`'s clone in the else
+  // version.
+  Identifier idForIfOp = b.getIdentifier("__mlir_if_hoisting");
+  operandMap.clear();
+  b.setInsertionPointAfter(hoistOverOp);
+  // We'll set an attribute to identify this op in a clone of this sub-tree.
+  ifOp.setAttr(b.getIdentifier(idForIfOp), b.getBoolAttr(true));
+  hoistOverOpClone = b.clone(*hoistOverOp, operandMap);
+
+  // Promote the 'then' block of the original affine.if in the then version.
+  promoteIfBlock(ifOp, /*elseBlock=*/false);
+
+  // Move the then version to the hoisted if op's 'then' block.
+  auto *thenBlock = hoistedIfOp.getThenBlock();
+  thenBlock->getOperations().splice(thenBlock->begin(),
+                                    hoistOverOp->getBlock()->getOperations(),
+                                    Block::iterator(hoistOverOp));
+
+  // Find the clone of the original affine.if op in the else version.
+  AffineIfOp ifCloneInElse;
+  hoistOverOpClone->walk([&](AffineIfOp ifClone) {
+    if (!ifClone.getAttr(idForIfOp))
+      return WalkResult::advance();
+    ifCloneInElse = ifClone;
+    return WalkResult::interrupt();
+  });
+  assert(ifCloneInElse && "if op clone should exist");
+  // For the else block, promote the else block of the original 'if' if it had
+  // one; otherwise, the op itself is to be erased.
+  if (!ifCloneInElse.hasElse())
+    ifCloneInElse.erase();
+  else
+    promoteIfBlock(ifCloneInElse, /*elseBlock=*/true);
+
+  // Move the else version into the else block of the hoisted if op.
+  auto *elseBlock = hoistedIfOp.getElseBlock();
+  elseBlock->getOperations().splice(
+      elseBlock->begin(), hoistOverOpClone->getBlock()->getOperations(),
+      Block::iterator(hoistOverOpClone));
+
+  return hoistedIfOp;
+}
+
+// Returns true if the hoisting happened or if the op was folded away.
+bool mlir::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
+  // Apply canonicalization patterns and folding - this is necessary for the
+  // hoisting check to be correct (operands should be composed), and to be more
+  // effective (no unused operands). Since the pattern rewriter's folding is
+  // entangled with application of patterns, we may fold/end up erasing the op,
+  // in which case we return with `folded` being set.
+  OwningRewritePatternList patterns;
+  ifOp.getCanonicalizationPatterns(patterns, ifOp.getContext());
+  bool erased;
+  applyOpPatternsAndFold(ifOp, patterns, &erased);
+  if (erased) {
+    if (folded)
+      *folded = true;
+    return true;
+  }
+  if (folded)
+    *folded = false;
+
+  // The folding above should have ensured this, but the affine.if's
+  // canonicalization is missing composition of affine.applys into it.
+  assert(llvm::all_of(ifOp.getOperands(),
+                      [](Value v) {
+                        return isTopLevelValue(v) || isForInductionVar(v);
+                      }) &&
+         "operands not composed");
+
+  // We are going hoist as high as possible.
+  // TODO: this could be customized in the future.
+  auto *hoistOverOp = getOutermostInvariantForOp(ifOp);
+
+  AffineIfOp hoistedIfOp = ::hoistAffineIfOp(ifOp, hoistOverOp);
+
+  // Nothing to hoist over.
+  if (hoistedIfOp == ifOp)
+    return false;
+
+  // Canonicalize to remove dead else blocks (happens whenever an 'if' moves up
+  // a sequence of affine.fors that are all perfectly nested).
+  applyPatternsAndFoldGreedily(
+      hoistedIfOp.getParentWithTrait<OpTrait::IsIsolatedFromAbove>(),
+      std::move(patterns));
+
+  return true;
+}

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

+// RUN: mlir-opt %s -split-input-file -test-affine-loop-unswitch | FileCheck %s
+
+// CHECK-DAG: #[[SET:.*]] = affine_set<(d0) : (d0 - 2 >= 0)>
+
+// CHECK-LABEL: func @if_else_imperfect
+func @if_else_imperfect(%A : memref<100xi32>, %B : memref<100xi32>, %v : i32) {
+// CHECK: %[[A:.*]]: memref<100xi32>, %[[B:.*]]: memref
+  affine.for %i = 0 to 100 {
+    affine.load %A[%i] : memref<100xi32>
+    affine.for %j = 0 to 100 {
+      affine.load %A[%j] : memref<100xi32>
+      affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) {
+        affine.load %B[%j] : memref<100xi32>
+      }
+      call @external() : () -> ()
+    }
+    affine.load %A[%i] : memref<100xi32>
+  }
+  return
+}
+func @external()
+
+// CHECK:       affine.for %[[I:.*]] = 0 to 100 {
+// CHECK-NEXT:    affine.load %[[A]][%[[I]]]
+// CHECK-NEXT:    affine.if #[[SET]](%[[I]]) {
+// CHECK-NEXT:      affine.for %[[J:.*]] = 0 to 100 {
+// CHECK-NEXT:        affine.load %[[A]][%[[J]]]
+// CHECK-NEXT:        affine.load %[[B]][%[[J]]]
+// CHECK-NEXT:        call
+// CHECK-NEXT:      }
+// CHECK-NEXT:    } else {
+// CHECK-NEXT:      affine.for %[[JJ:.*]] = 0 to 100 {
+// CHECK-NEXT:        affine.load %[[A]][%[[JJ]]]
+// CHECK-NEXT:        call
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:    affine.load %[[A]][%[[I]]]
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return
+
+// -----
+
+func @foo()
+func @bar()
+func @abc()
+func @xyz()
+
+// CHECK-LABEL: func @if_then_perfect
+func @if_then_perfect(%A : memref<100xi32>, %v : i32) {
+  affine.for %i = 0 to 100 {
+    affine.for %j = 0 to 100 {
+      affine.for %k = 0 to 100 {
+        affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%i) {
+          affine.load %A[%i] : memref<100xi32>
+        }
+      }
+    }
+  }
+  return
+}
+// CHECK:      affine.for
+// CHECK-NEXT:   affine.if
+// CHECK-NEXT:     affine.for
+// CHECK-NEXT:       affine.for
+// CHECK-NOT:    else
+
+
+// CHECK-LABEL: func @if_else_perfect
+func @if_else_perfect(%A : memref<100xi32>, %v : i32) {
+  affine.for %i = 0 to 99 {
+    affine.for %j = 0 to 100 {
+      affine.for %k = 0 to 100 {
+        call @foo() : () -> ()
+        affine.if affine_set<(d0, d1) : (d0 - 2 >= 0, -d1 + 80 >= 0)>(%i, %j) {
+          affine.load %A[%i] : memref<100xi32>
+          call @abc() : () -> ()
+        } else {
+          affine.load %A[%i + 1] : memref<100xi32>
+          call @xyz() : () -> ()
+        }
+        call @bar() : () -> ()
+      }
+    }
+  }
+  return
+}
+// CHECK:      affine.for
+// CHECK-NEXT:   affine.for
+// CHECK-NEXT:     affine.if
+// CHECK-NEXT:       affine.for
+// CHECK-NEXT:         call @foo
+// CHECK-NEXT:         affine.load %{{.*}}[%{{.*}}]
+// CHECK-NEXT:         call @abc
+// CHECK-NEXT:         call @bar
+// CHECK-NEXT:       }
+// CHECK-NEXT:     else
+// CHECK-NEXT:       affine.for
+// CHECK-NEXT:         call @foo
+// CHECK-NEXT:         affine.load %{{.*}}[%{{.*}} + 1]
+// CHECK-NEXT:         call @xyz
+// CHECK-NEXT:         call @bar
+// CHECK-NEXT:       }
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+
+// CHECK-LABEL: func @if_then_imperfect
+func @if_then_imperfect(%A : memref<100xi32>, %N : index) {
+  affine.for %i = 0 to 100 {
+    affine.load %A[0] : memref<100xi32>
+    affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%N) {
+      affine.load %A[%i] : memref<100xi32>
+    }
+  }
+  return
+}
+// CHECK:       affine.if
+// CHECK-NEXT:    affine.for
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:    }
+// CHECK-NEXT:  } else {
+// CHECK-NEXT:    affine.for
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// CHECK-NEXT:  return
+
+// Check if unused operands are dropped: hence, hoisting is possible.
+// CHECK-LABEL: func @hoist_after_canonicalize
+func @hoist_after_canonicalize() {
+  affine.for %i = 0 to 100 {
+    affine.for %j = 0 to 100 {
+      affine.if affine_set<(d0) : (d0 - 2 >= 0)>(%j)  {
+      affine.if affine_set<(d0, d1) : (d0 - 1 >= 0, -d0 + 99 >= 0)>(%i, %j)  {
+        // The call to external is to avoid DCE on affine.if.
+        call @foo() : () -> ()
+      }
+      }
+    }
+  }
+  return
+}
+// CHECK:      affine.for
+// CHECK-NEXT:   affine.if
+// CHECK-NEXT:     affine.for
+// CHECK-NEXT:       affine.if
+// CHECK-NEXT:         call
+// CHECK-NEXT:       }
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+// CHECK-NEXT: return
+
+// CHECK-LABEL: func @handle_dead_if
+func @handle_dead_if(%N : index) {
+  affine.for %i = 0 to 100 {
+      affine.if affine_set<(d0) : (d0 - 1 >= 0, -d0 + 99 >= 0)>(%N)  {
+      }
+  }
+  return
+}
+// CHECK-NEXT: affine.for
+// CHECK-NEXT: }
+// CHECK-NEXT: return
+
+// -----
+
+#flb1 = affine_map<(d0) -> (d0 * 3)>
+#fub1 = affine_map<(d0) -> (d0 * 3 + 3)>
+#flb0 = affine_map<(d0) -> (d0 * 16)>
+#fub0 = affine_map<(d0) -> (d0 * 16 + 16)>
+#pub1 = affine_map<(d0)[s0] -> (s0, d0 * 3 + 3)>
+#pub0 = affine_map<(d0)[s0] -> (s0, d0 * 16 + 16)>
+#lb1 = affine_map<(d0) -> (d0 * 480)>
+#ub1 = affine_map<(d0)[s0] -> (s0, d0 * 480 + 480)>
+#lb0 = affine_map<(d0) -> (d0 * 110)>
+#ub0 = affine_map<(d0)[s0] -> (d0 * 110 + 110, s0 floordiv 3)>
+
+#set0 = affine_set<(d0, d1)[s0, s1] : (d0 * -16 + s0 - 16 >= 0, d1 * -3 + s1 - 3 >= 0)>
+
+// CHECK-LABEL: func @perfect_if_else
+func @perfect_if_else(%arg0 : memref<?x?xf64>, %arg1 : memref<?x?xf64>, %arg4 : index,
+            %arg5 : index, %arg6 : index, %sym : index) {
+  affine.for %arg7 = #lb0(%arg5) to min #ub0(%arg5)[%sym] {
+    affine.parallel (%i0, %j0) = (0, 0) to (symbol(%sym), 100) step (10, 10) {
+      affine.for %arg8 = #lb1(%arg4) to min #ub1(%arg4)[%sym] {
+        affine.if #set0(%arg6, %arg7)[%sym, %sym] {
+          affine.for %arg9 = #flb0(%arg6) to #fub0(%arg6) {
+            affine.for %arg10 = #flb1(%arg7) to #fub1(%arg7) {
+              affine.load %arg0[0, 0] : memref<?x?xf64>
+            }
+          }
+        } else {
+          affine.for %arg9 = #lb0(%arg6) to min #pub0(%arg6)[%sym] {
+            affine.for %arg10 = #lb1(%arg7) to min #pub1(%arg7)[%sym] {
+              affine.load %arg0[0, 0] : memref<?x?xf64>
+            }
+          }
+        }
+      }
+    }
+  }
+  return
+}
+
+// CHECK:       affine.for
+// CHECK-NEXT:    affine.if
+// CHECK-NEXT:      affine.parallel
+// CHECK-NEXT:        affine.for
+// CHECK-NEXT:          affine.for
+// CHECK-NEXT:            affine.for
+// CHECK-NEXT:              affine.load
+// CHECK-NEXT:            }
+// CHECK-NEXT:          }
+// CHECK-NEXT:        }
+// CHECK-NEXT:      }
+// CHECK-NEXT:    } else {
+// CHECK-NEXT:      affine.parallel
+// CHECK-NEXT:        affine.for
+// CHECK-NEXT:          affine.for
+// CHECK-NEXT:            affine.for
+// CHECK-NEXT:              affine.load
+// CHECK-NEXT:            }
+// CHECK-NEXT:          }
+// CHECK-NEXT:        }
+// CHECK-NEXT:      }
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }

mlir/test/lib/Dialect/Affine/CMakeLists.txt

 add_llvm_library(MLIRAffineTransformsTestPasses
   TestAffineDataCopy.cpp
+  TestAffineLoopUnswitching.cpp
   TestLoopPermutation.cpp
   TestParallelismDetection.cpp
   TestVectorizationUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
   )
 
 target_link_libraries(MLIRAffineTransformsTestPasses PRIVATE
   MLIRIR
   MLIRPass
   MLIRAffineTransforms
   MLIRSupport
   )

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

+//===- TestAffineLoopUnswitching.cpp - Test affine if/else hoisting -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to hoist affine if/else structures.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/Utils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
+
+#define PASS_NAME "test-affine-loop-unswitch"
+
+using namespace mlir;
+
+namespace {
+
+/// This pass applies the permutation on the first maximal perfect nest.
+struct TestAffineLoopUnswitching
+    : public PassWrapper<TestAffineLoopUnswitching, FunctionPass> {
+  TestAffineLoopUnswitching() = default;
+  TestAffineLoopUnswitching(const TestAffineLoopUnswitching &pass) {}
+
+  void runOnFunction() override;
+
+  /// The maximum number of iterations to run this for.
+  constexpr static unsigned kMaxIterations = 5;
+};
+
+} // end anonymous namespace
+
+void TestAffineLoopUnswitching::runOnFunction() {
+  // Each hoisting invalidates a lot of IR around. Just stop the walk after the
+  // first if/else hoisting, and repeat until no more hoisting can be done, or
+  // the maximum number of iterations have been run.
+  auto func = getFunction();
+  bool changed;
+  unsigned i = 0;
+  do {
+    changed = false;
+    func.walk([&](AffineIfOp op) {
+      if (hoistAffineIfOp(op)) {
+        changed = true;
+        return WalkResult::interrupt();
+      }
+      return WalkResult::advance();
+    });
+  } while (changed && ++i < kMaxIterations);
+}
+
+namespace mlir {
+void registerTestAffineLoopUnswitchingPass() {
+  PassRegistration<TestAffineLoopUnswitching>(
+      PASS_NAME, "Tests affine loop unswitching / if/else hoisting");
+}
+
+} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerPassManagerTestPass();
 void registerPatternsTestPass();
 void registerPrintOpAvailabilityPass();
 void registerSideEffectTestPasses();
 void registerSimpleParametricTilingPass();
 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAllReduceLoweringPass();
+void registerTestAffineLoopUnswitchingPass();
 void registerTestLinalgMatmulToVectorPass();
 void registerTestLoopPermutationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestDominancePass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
void registerTestPasses() {
   registerPassManagerTestPass();
   registerPatternsTestPass();
   registerPrintOpAvailabilityPass();
   registerSideEffectTestPasses();
   registerSimpleParametricTilingPass();
   registerSymbolTestPasses();
   registerTestAffineDataCopyPass();
   registerTestAllReduceLoweringPass();
+  registerTestAffineLoopUnswitchingPass();
   registerTestLinalgMatmulToVectorPass();
   registerTestLoopPermutationPass();
   registerTestCallGraphPass();
   registerTestConstantFold();
 #if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
   registerTestDominancePass();