[mlir] Update LICM to support Graph Regions

[Mogball]

Summary

Changes the algorithm of LICM to support graph regions (no guarantee of topologically sorted order). Also fixes an issue where ops with recursive side effects and regions would not be hoisted if any nested ops used operands that were defined within the nested region.

[Mogball]

mlir/test/Transforms/loop-invariant-code-motion.mlir

171

@nicolasvasilache Can you verify whether this is acceptable behaviour?

[Mogball]

change to worklist algorithm

[mehdi_amini]

Adds default implementations of isDefinedOutsideOfLoop and moveOutOfLoop since 99% of all implementations of these functions were identical

Can you split this as a NFC commit?

[Mogball]

rebase (fix bazel build maybe)

[Mogball]

Yeah let me split the patch

[Mogball]

split out simplifications to LoopLikeOpInterface

[bondhugula]

mlir/lib/Interfaces/LoopLikeInterface.cpp

35

This would be a useful function and can directly be added in MLIRAnalysis library?

mlir/test/lib/Dialect/Test/TestOps.td

23

Sorted order here.

2761

Nit: ... with a graph region.

2780

Nit: extra blank line.

[Mogball]

review comments

mlir/lib/Interfaces/LoopLikeInterface.cpp

35

I will be submitting a patch after this one that moves moveLoopInvariantCode to TransformUtils. At the same time, I'll move this function to MLIRAnalysis (the same function exists elsewhere in the codebase and I will refactor those as well)

[rriddle]

mlir/lib/Interfaces/LoopLikeInterface.cpp

58

-> /

58

Sigh, I hate phab. Use /// comments here.

mlir/test/lib/Dialect/Test/TestOps.td

14

Do you know why this dialect is depending on Linalg?

2773–2775

[Mogball]

review comments

[Mogball]

mlir/test/lib/Dialect/Test/TestOps.td

14

There's a test op for LinalgConvolutionOpInterface

[rriddle]

mlir/lib/Interfaces/LoopLikeInterface.cpp

68

Probably more idiomatic to check the trait for IsTerminator.

90

nit: Add a newline before the comment.

mlir/test/Transforms/loop-invariant-code-motion.mlir

362

Can you add a test that checks that we don't hoist from graph regions into non-graph regions?

366

Do we have a test in a non-graph region where operations that can be hoisted feed into each other? Something like:

%v0 = arith.addi %c0, %c0 : i32
%v1 = arith.addi %c0, %v0 : i32

to make sure that dominance still holds after hoisting?

[Mogball]

mlir/test/Transforms/loop-invariant-code-motion.mlir

362

Is the behaviour that we want that ops don't cross between non-graph and graph regions (in either direction)?

Right now, LICM can move ops from graph regions to non-graph regions but only in topological order (dominance will be preserved). Ops that form cycles won't be moved because it doesn't detect them.

366

Oddly enough, I don't see one. I'll add one.

[rriddle]

mlir/test/Transforms/loop-invariant-code-motion.mlir

362

It's more about being conservatively correct. We can hoist from a graph into a non-graph(and vice-versa), but we, for example, have to ensure that the dominance of the hoisted operations is valid. We have to be careful about the contract changes between the two regions.

[Mogball]

mlir/test/Transforms/loop-invariant-code-motion.mlir

362

We will be okay to move ops from graph regions to non-graph regions; since LICM only considers one operation at a time, ops are moved only when all their operands are defined outside the graph region body. If the outer region is a non-graph region, the defining ops must precede the loop op.

I'll add some test cases.

[Mogball]

review comments -- adding some tests

flang/lib/Optimizer/Dialect/FIROps.cpp

p.printOptionalAttrDictWithKeyword((*this)->getAttrs(),
                                      {getFinalValueAttrNameStr()});
   p << ' ';
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/true);
 }
 
 mlir::Region &fir::IterWhileOp::getLoopBody() { return getRegion(); }
 
-bool fir::IterWhileOp::isDefinedOutsideOfLoop(mlir::Value value) {
-  return !getRegion().isAncestor(value.getParentRegion());
-}
-
-mlir::LogicalResult
-fir::IterWhileOp::moveOutOfLoop(llvm::ArrayRef<mlir::Operation *> ops) {
-  for (auto *op : ops)
-    op->moveBefore(*this);
-  return success();
-}
-
 mlir::BlockArgument fir::IterWhileOp::iterArgToBlockArg(mlir::Value iterArg) {
   for (auto i : llvm::enumerate(getInitArgs()))
     if (iterArg == i.value())
       return getRegion().front().getArgument(i.index() + 1);
   return {};
 }
 
 void fir::IterWhileOp::resultToSourceOps(
p.printOptionalAttrDictWithKeyword((*this)->getAttrs(),
                                      {"unordered", "finalValue"});
   p << ' ';
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
                 printBlockTerminators);
 }
 
 mlir::Region &fir::DoLoopOp::getLoopBody() { return getRegion(); }
 
-bool fir::DoLoopOp::isDefinedOutsideOfLoop(mlir::Value value) {
-  return !getRegion().isAncestor(value.getParentRegion());
-}
-
-mlir::LogicalResult
-fir::DoLoopOp::moveOutOfLoop(llvm::ArrayRef<mlir::Operation *> ops) {
-  for (auto op : ops)
-    op->moveBefore(*this);
-  return success();
-}
-
 /// Translate a value passed as an iter_arg to the corresponding block
 /// argument in the body of the loop.
 mlir::BlockArgument fir::DoLoopOp::iterArgToBlockArg(mlir::Value iterArg) {
   for (auto i : llvm::enumerate(getInitArgs()))
     if (iterArg == i.value())
       return getRegion().front().getArgument(i.index() + 1);
   return {};
 }

mlir/include/mlir/Interfaces/LoopLikeInterface.h

 #include "mlir/Interfaces/LoopLikeInterface.h.inc"
 
 //===----------------------------------------------------------------------===//
 // LoopLike Utilities
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 /// Move loop invariant code out of a `looplike` operation.
-LogicalResult moveLoopInvariantCode(LoopLikeOpInterface looplike);
+void moveLoopInvariantCode(LoopLikeOpInterface looplike);
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_LOOPLIKEINTERFACE_H_

mlir/include/mlir/Interfaces/LoopLikeInterface.td

def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
   let methods = [
     InterfaceMethod<[{
         Returns true if the given value is defined outside of the loop.
         A sensible implementation could be to check whether the value's defining
         operation lies outside of the loops body region. If the loop uses
         explicit capture of dependencies, an implementation could check whether
         the value corresponds to a captured dependency.
       }],
-      "bool", "isDefinedOutsideOfLoop", (ins "::mlir::Value ":$value)
+      "bool", "isDefinedOutsideOfLoop", (ins "::mlir::Value ":$value), [{}], [{
+        return value.getParentRegion()->isProperAncestor(&$_op.getLoopBody());
+      }]
     >,
     InterfaceMethod<[{
         Returns the region that makes up the body of the loop and should be
         inspected for loop-invariant operations.
       }],
       "::mlir::Region &", "getLoopBody"
     >,
     InterfaceMethod<[{
-        Moves the given vector of operations out of the loop. The vector is
-        sorted topologically.
+        Moves the given loop-invariant operation out of the loop.
       }],
-      "::mlir::LogicalResult", "moveOutOfLoop", (ins "::mlir::ArrayRef<::mlir::Operation *>":$ops)
+      "void", "moveOutOfLoop",
+      (ins "::mlir::Operation *":$op), [{}], [{
+        op->moveBefore($_op);
+      }]
     >,
     InterfaceMethod<[{
         If there is a single induction variable return it, otherwise return
         llvm::None.
       }],
       /*retTy=*/"::mlir::Optional<::mlir::Value>",
       /*methodName=*/"getSingleInductionVar",
       /*args=*/(ins),

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

for (unsigned i = 0, e = lbMap.getNumInputs(); i < e; i++) {
     if (getOperand(i) != getOperand(numOperands + i))
       return false;
   }
   return true;
 }
 
 Region &AffineForOp::getLoopBody() { return region(); }
 
-bool AffineForOp::isDefinedOutsideOfLoop(Value value) {
-  return !region().isAncestor(value.getParentRegion());
-}
-
 Optional<Value> AffineForOp::getSingleInductionVar() {
   return getInductionVar();
 }
 
 Optional<OpFoldResult> AffineForOp::getSingleLowerBound() {
   if (!hasConstantLowerBound())
     return llvm::None;
   OpBuilder b(getContext());
   return OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()));
 }
 
 Optional<OpFoldResult> AffineForOp::getSingleStep() {
   OpBuilder b(getContext());
   return OpFoldResult(b.getI64IntegerAttr(getStep()));
 }
 
-LogicalResult AffineForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
-  for (auto *op : ops)
-    op->moveBefore(*this);
-  return success();
-}
-
 /// Returns true if the provided value is the induction variable of a
 /// AffineForOp.
 bool mlir::isForInductionVar(Value val) {
   return getForInductionVarOwner(val) != AffineForOp();
 }
 
 /// Returns the loop parent of an induction variable. If the provided value is
 /// not an induction variable, then return nullptr.
for (unsigned i = 0, e = steps.size(); i < e; ++i)
     body->addArgument(IndexType::get(builder.getContext()), result.location);
   bodyRegion->push_back(body);
   if (resultTypes.empty())
     ensureTerminator(*bodyRegion, builder, result.location);
 }
 
 Region &AffineParallelOp::getLoopBody() { return region(); }
 
-bool AffineParallelOp::isDefinedOutsideOfLoop(Value value) {
-  return !region().isAncestor(value.getParentRegion());
-}
-
-LogicalResult AffineParallelOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
-  for (Operation *op : ops)
-    op->moveBefore(*this);
-  return success();
-}
-
 unsigned AffineParallelOp::getNumDims() { return steps().size(); }
 
 AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() {
   return getOperands().take_front(lowerBoundsMap().getNumInputs());
 }
 
 AffineParallelOp::operand_range AffineParallelOp::getUpperBoundsOperands() {
   return getOperands().drop_front(lowerBoundsMap().getNumInputs());

mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp

assert(((read.extractSliceOp && write.insertSliceOp) ||
          "expected matching extract_slice / insert_slice");
   LLVM_DEBUG(DBGS() << "In forOp:\n"
                     << *forOp.getOperation()
                     << "\nHoist: " << *read.transferReadOp.getOperation()
                     << "\nHoist: " << *write.transferWriteOp.getOperation()
                     << "\nInvolving: " << tensorBBArg << "\n");
 
   // If a read slice is present, hoist it.
-  if (read.extractSliceOp && failed(forOp.moveOutOfLoop({read.extractSliceOp})))
-    llvm_unreachable("Unexpected failure moving extract_slice out of loop");
+  if (read.extractSliceOp)
+    forOp.moveOutOfLoop(read.extractSliceOp);
 
   // Hoist the transfer_read op.
-  if (failed(forOp.moveOutOfLoop({read.transferReadOp})))
-    llvm_unreachable("Unexpected failure moving transfer read out of loop");
+  forOp.moveOutOfLoop(read.transferReadOp);
 
   // TODO: don't hardcode /*numIvs=*/1.
   assert(tensorBBArg.getArgNumber() >= /*numIvs=*/1);
   unsigned initArgNumber = tensorBBArg.getArgNumber() - /*numIvs=*/1;
 
   // Update the source tensor.
   if (read.extractSliceOp)
     read.extractSliceOp.sourceMutable().assign(

 }
 
 void mlir::linalg::hoistRedundantVectorTransfers(FuncOp func) {
   bool changed = true;
   while (changed) {
     changed = false;
     // First move loop invariant ops outside of their loop. This needs to be
     // done before as we cannot move ops without interputing the function walk.
-    func.walk([&](LoopLikeOpInterface loopLike) {
-      if (failed(moveLoopInvariantCode(loopLike)))
-        llvm_unreachable(
-            "Unexpected failure to move invariant code out of loop");
-    });
+    func.walk(
+        [&](LoopLikeOpInterface loopLike) { moveLoopInvariantCode(loopLike); });
 
     func.walk([&](vector::TransferReadOp transferRead) {
       if (!transferRead.getShapedType().isa<MemRefType>())
         return WalkResult::advance();
 
       LLVM_DEBUG(DBGS() << "Candidate for hoisting: "
                         << *transferRead.getOperation() << "\n");
       auto loop = dyn_cast<scf::ForOp>(transferRead->getParentOp());
func.walk([&](vector::TransferReadOp transferRead) {
         } else {
           // Unknown use, we cannot prove that it doesn't alias with the
           // transferRead/transferWrite operations.
           return WalkResult::advance();
         }
       }
 
       // Hoist read before.
-      if (failed(loop.moveOutOfLoop({transferRead})))
-        llvm_unreachable(
-            "Unexpected failure to move transfer read out of loop");
+      loop.moveOutOfLoop(transferRead);
 
       // Hoist write after.
       transferWrite->moveAfter(loop);
 
       // Rewrite `loop` with new yields by cloning and erase the original loop.
       OpBuilder b(transferRead);
       auto newForOp = cloneWithNewYields(b, loop, transferRead.vector(),
                                          transferWrite.vector());

mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp

void runOnOperation() override {
     MLIRContext *context = funcOp.getContext();
     RewritePatternSet patterns =
         linalg::getLinalgTilingCanonicalizationPatterns(context);
     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
       return signalPassFailure();
 
     if (options.licm) {
-      if (funcOp
-              ->walk([&](LoopLikeOpInterface loopLike) {
-                if (failed(moveLoopInvariantCode(loopLike)))
-                  return WalkResult::interrupt();
-                return WalkResult::advance();
-              })
-              .wasInterrupted())
-        return signalPassFailure();
+      funcOp->walk([&](LoopLikeOpInterface loopLike) {
+        moveLoopInvariantCode(loopLike);
+      });
     }
 
     // Gathers all innermost loops through a post order pruned walk.
     funcOp.walk([](Operation *op) {
       if (auto forOp = dyn_cast<AffineForOp>(op))
         (void)promoteIfSingleIteration(forOp);
       else if (auto forOp = dyn_cast<scf::ForOp>(op))
         (void)promoteIfSingleIteration(forOp);

mlir/lib/Dialect/SCF/SCF.cpp

ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
   if (parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
   return success();
 }
 
 Region &ForOp::getLoopBody() { return getRegion(); }
 
-bool ForOp::isDefinedOutsideOfLoop(Value value) {
-  return !getRegion().isAncestor(value.getParentRegion());
-}
-
-LogicalResult ForOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
-  for (auto *op : ops)
-    op->moveBefore(*this);
-  return success();
-}
-
 ForOp mlir::scf::getForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   if (!ivArg)
     return ForOp();
   assert(ivArg.getOwner() && "unlinked block argument");
   auto *containingOp = ivArg.getOwner()->getParentOp();
   return dyn_cast_or_null<ForOp>(containingOp);
 }
void ParallelOp::print(OpAsmPrinter &p) {
   p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
   p.printOptionalAttrDict(
       (*this)->getAttrs(),
       /*elidedAttrs=*/ParallelOp::getOperandSegmentSizeAttr());
 }
 
 Region &ParallelOp::getLoopBody() { return getRegion(); }
 
-bool ParallelOp::isDefinedOutsideOfLoop(Value value) {
-  return !getRegion().isAncestor(value.getParentRegion());
-}
-
-LogicalResult ParallelOp::moveOutOfLoop(ArrayRef<Operation *> ops) {
-  for (auto *op : ops)
-    op->moveBefore(*this);
-  return success();
-}
-
 ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
   auto ivArg = val.dyn_cast<BlockArgument>();
   if (!ivArg)
     return ParallelOp();
   assert(ivArg.getOwner() && "unlinked block argument");
   auto *containingOp = ivArg.getOwner()->getParentOp();
   return dyn_cast<ParallelOp>(containingOp);
 }

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

 
 //===----------------------------------------------------------------------===//
 // TOSA control flow support.
 //===----------------------------------------------------------------------===//
 
 /// Returns the while loop body.
 Region &tosa::WhileOp::getLoopBody() { return body(); }
 
-bool tosa::WhileOp::isDefinedOutsideOfLoop(Value value) {
-  return !body().isAncestor(value.getParentRegion());
-}
-
-LogicalResult WhileOp::moveOutOfLoop(ArrayRef<mlir::Operation *> ops) {
-  if (ops.empty())
-    return success();
-
-  Operation *tosaWhileOp = this->getOperation();
-  for (auto *op : ops)
-    op->moveBefore(tosaWhileOp);
-
-  return success();
-}
-
 //===----------------------------------------------------------------------===//
 // Tosa dialect initialization.
 //===----------------------------------------------------------------------===//
 
 void TosaDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"

mlir/lib/Interfaces/LoopLikeInterface.cpp

 //===- LoopLikeInterface.cpp - Loop-like operations in MLIR ---------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
+#include <queue>
 
 using namespace mlir;
 
 #define DEBUG_TYPE "loop-like"
 
 //===----------------------------------------------------------------------===//
 // LoopLike Interfaces
 //===----------------------------------------------------------------------===//
 
 /// Include the definitions of the loop-like interfaces.
 #include "mlir/Interfaces/LoopLikeInterface.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // LoopLike Utilities
 //===----------------------------------------------------------------------===//
 
-// Checks whether the given op can be hoisted by checking that
-// - the op and any of its contained operations do not depend on SSA values
-//   defined inside of the loop (by means of calling definedOutside).
-// - the op has no side-effects. If sideEffecting is Never, sideeffects of this
-//   op and its nested ops are ignored.
-static bool canBeHoisted(Operation *op,
-                         function_ref<bool(Value)> definedOutside) {
-  // Check that dependencies are defined outside of loop.
-  if (!llvm::all_of(op->getOperands(), definedOutside))
-    return false;
-  // Check whether this op is side-effect free. If we already know that there
-  // can be no side-effects because the surrounding op has claimed so, we can
-  // (and have to) skip this step.
+/// Returns true if the given operation is side-effect free as are all of its
+/// nested operations.
+///
+/// TODO: There is a duplicate function in ControlFlowSink. Move
+/// `moveLoopInvariantCode` to TransformUtils and then factor out this function.
+static bool isSideEffectFree(Operation *op) {
   if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+    // If the op has side-effects, it cannot be moved.
     if (!memInterface.hasNoEffect())
       return false;
-    // If the operation doesn't have side effects and it doesn't recursively
-    // have side effects, it can always be hoisted.
+    // If the op does not have recursive side effects, then it can be moved.
     if (!op->hasTrait<OpTrait::HasRecursiveSideEffects>())
       return true;
-
-    // Otherwise, if the operation doesn't provide the memory effect interface
-    // and it doesn't have recursive side effects we treat it conservatively as
-    // side-effecting.
   } else if (!op->hasTrait<OpTrait::HasRecursiveSideEffects>()) {
+    // Otherwise, if the op does not implement the memory effect interface and
+    // it does not have recursive side effects, then it cannot be known that the
+    // op is moveable.
     return false;
   }
 
-  // Recurse into the regions for this op and check whether the contained ops
-  // can be hoisted.
-  for (auto &region : op->getRegions()) {
-    for (auto &block : region) {
-      for (auto &innerOp : block)
-        if (!canBeHoisted(&innerOp, definedOutside))
-          return false;
-    }
-  }
+  // Recurse into the regions and ensure that all nested ops can also be moved.
+  for (Region &region : op->getRegions())
+    for (Operation &op : region.getOps())
+      if (!isSideEffectFree(&op))
   return true;
 }
 
-LogicalResult mlir::moveLoopInvariantCode(LoopLikeOpInterface looplike) {
-  auto &loopBody = looplike.getLoopBody();
+// Checks whether the given op can be hoisted by checking that
+// - the op and none of its contained operations depend on values inside of the
+//   loop (by means of calling definedOutside).
+// - the op has no side-effects.
+static bool canBeHoisted(Operation *op,
+                         function_ref<bool(Value)> definedOutside) {
+  if (!isSideEffectFree(op))
+    return false;
 
-  // We use two collections here as we need to preserve the order for insertion
-  // and this is easiest.
-  SmallPtrSet<Operation *, 8> willBeMovedSet;
-  SmallVector<Operation *, 8> opsToMove;
-
-  // Helper to check whether an operation is loop invariant wrt. SSA properties.
-  auto isDefinedOutsideOfBody = [&](Value value) {
-    auto *definingOp = value.getDefiningOp();
-    return (definingOp && !!willBeMovedSet.count(definingOp)) ||
-           looplike.isDefinedOutsideOfLoop(value);
-  };
+  // Do not move terminators.
+  if (op->getBlock()->getTerminator() == op)
+    return false;
 
-  // Do not use walk here, as we do not want to go into nested regions and hoist
-  // operations from there. These regions might have semantics unknown to this
-  // rewriting. If the nested regions are loops, they will have been processed.
-  for (auto &block : loopBody) {
-    for (auto &op : block.without_terminator()) {
-      if (canBeHoisted(&op, isDefinedOutsideOfBody)) {
-        opsToMove.push_back(&op);
-        willBeMovedSet.insert(&op);
-      }
+  // Walk the nested operations and check that all used values are either
+  // defined outside of the loop or in a nested region, but not at the level of
+  // the loop body.
+  auto walkFn = [&](Operation *child) {
+    for (Value operand : child->getOperands()) {
+      // Ignore values defined in a nested region.
+      if (op->isAncestor(operand.getParentRegion()->getParentOp()))
+        continue;
+      if (!definedOutside(operand))
+        return WalkResult::interrupt();
     }
-  }
+    return WalkResult::advance();
+  };
+  return !op->walk(walkFn).wasInterrupted();
 
-  // For all instructions that we found to be invariant, move outside of the
-  // loop.
-  LogicalResult result = looplike.moveOutOfLoop(opsToMove);
-  LLVM_DEBUG(looplike.print(llvm::dbgs() << "\n\nModified loop:\n"));
-  return result;
+void mlir::moveLoopInvariantCode(LoopLikeOpInterface looplike) {
+  Region *loopBody = &looplike.getLoopBody();
+
+  std::queue<Operation *> worklist;
+  // Add top-level operations in the loop body to the worklist.
+  for (Operation &op : loopBody->getOps())
+    worklist.push(&op);
+
+  auto definedOutside = [&](Value value) {
+    return looplike.isDefinedOutsideOfLoop(value);
+  };
+
+  while (!worklist.empty()) {
+    Operation *op = worklist.front();
+    worklist.pop();
+    // Skip ops that have already been moved. Check if the op can be hoisted.
+    if (op->getParentRegion() != loopBody || !canBeHoisted(op, definedOutside))
+      continue;
+
+    looplike.moveOutOfLoop(op);
+    // Since the op has been moved, we need to check its users within the
+    // top-level of the loop body.
+    for (Operation *user : op->getUsers())
+      if (user->getParentRegion() == loopBody)
+        worklist.push(user);
+  }
 }

mlir/lib/Transforms/LoopInvariantCodeMotion.cpp

 } // namespace
 
 void LoopInvariantCodeMotion::runOnOperation() {
   // Walk through all loops in a function in innermost-loop-first order. This
   // way, we first LICM from the inner loop, and place the ops in
   // the outer loop, which in turn can be further LICM'ed.
   getOperation()->walk([&](LoopLikeOpInterface loopLike) {
     LLVM_DEBUG(loopLike.print(llvm::dbgs() << "\nOriginal loop:\n"));
-    if (failed(moveLoopInvariantCode(loopLike)))
-      signalPassFailure();
+    moveLoopInvariantCode(loopLike);
   });
 }
 
 std::unique_ptr<Pass> mlir::createLoopInvariantCodeMotionPass() {
   return std::make_unique<LoopInvariantCodeMotion>();
 }

mlir/test/Transforms/loop-invariant-code-motion.mlir

 // -----
 
 func @invariant_affine_nested_if() {
   %m = memref.alloc() : memref<10xf32>
   %cf8 = arith.constant 8.0 : f32
   affine.for %arg0 = 0 to 10 {
     affine.for %arg1 = 0 to 10 {
       affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
-          %cf9 = arith.addf %cf8, %cf8 : f32
-          affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
-            %cf10 = arith.addf %cf9, %cf9 : f32
-          }
       }
     }
   }
 
   // CHECK: memref.alloc
   // CHECK-NEXT: arith.constant
   // CHECK-NEXT: affine.for
+  // CHECK-NEXT: }
   // CHECK-NEXT: affine.for
   // CHECK-NEXT: affine.if
   // CHECK-NEXT: arith.addf
   // CHECK-NEXT: affine.if
   // CHECK-NEXT: arith.addf
   // CHECK-NEXT: }
   // CHECK-NEXT: }
-  // CHECK-NEXT: }
 
 
   return
 }
 
 // -----
 
 func @invariant_affine_nested_if_else() {
scf.for %i = %lb to %ub step %step {
     %val = call @make_val() : () -> (index)
     %val2 = call @make_val() : () -> (index)
     %r = scf.if %true -> (index) {
       scf.yield %val: index
     } else {
       scf.yield %val2: index
     }
   }
-  return 
+  return
 }
 
+// -----
+
+// Test invariant nested loop is hoisted.
+// CHECK-LABEL: func @test_invariant_nested_loop
+func @test_invariant_nested_loop() {
+  // CHECK: %[[C:.*]] = arith.constant
+  %0 = arith.constant 5 : i32
+  // CHECK: %[[V0:.*]] = arith.addi %[[C]], %[[C]]
+  // CHECK-NEXT: %[[V1:.*]] = arith.addi %[[V0]], %[[C]]
+  // CHECK-NEXT: test.graph_loop
+  // CHECK-NEXT: ^bb0(%[[ARG0:.*]]: i32)
+  // CHECK-NEXT: %[[V2:.*]] = arith.subi %[[ARG0]], %[[ARG0]]
+  // CHECK-NEXT: test.region_yield %[[V2]]
+  // CHECK: test.graph_loop
+  // CHECK-NEXT: test.region_yield %[[V1]]
+  test.graph_loop {
+    %1 = arith.addi %0, %0 : i32
+    %2 = arith.addi %1, %0 : i32
+    test.graph_loop {
+    ^bb0(%arg0: i32):
+      %3 = arith.subi %arg0, %arg0 : i32
+      test.region_yield %3 : i32
+    } : () -> ()
+    test.region_yield %2 : i32
+  } : () -> ()
+  return
+}
+
+
+// -----
+
+// Test ops in a graph region are hoisted.
+// CHECK-LABEL: func @test_invariants_in_graph_region
+func @test_invariants_in_graph_region() {
+  // CHECK: test.single_no_terminator_op
+  test.single_no_terminator_op : {
+    // CHECK-NEXT: %[[C:.*]] = arith.constant
+    // CHECK-NEXT: %[[V1:.*]] = arith.addi %[[C]], %[[C]]
+    // CHECK-NEXT: %[[V0:.*]] = arith.addi %[[C]], %[[V1]]
+    test.graph_loop {
+      %v0 = arith.addi %c0, %v1 : i32
+      %v1 = arith.addi %c0, %c0 : i32
+      %c0 = arith.constant 5 : i32
+      test.region_yield %v0 : i32
+    } : () -> ()
+  }
+  return
+}

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

void runOnOperation() override {
     func.walk([&](Operation *fakeRead) {
       if (fakeRead->getName().getStringRef() != "fake_read")
         return;
       auto *fakeCompute = fakeRead->getResult(0).use_begin()->getOwner();
       auto *fakeWrite = fakeCompute->getResult(0).use_begin()->getOwner();
       auto loop = fakeRead->getParentOfType<scf::ForOp>();
 
       OpBuilder b(loop);
-      (void)loop.moveOutOfLoop({fakeRead});
+      loop.moveOutOfLoop(fakeRead);
       fakeWrite->moveAfter(loop);
       auto newLoop = cloneWithNewYields(b, loop, fakeRead->getResult(0),
                                         fakeCompute->getResult(0));
       fakeCompute->getResult(0).replaceAllUsesWith(
           newLoop.getResults().take_back()[0]);
       toErase.push_back(loop);
     });
     for (auto loop : llvm::reverse(toErase))

mlir/test/lib/Dialect/Test/TestOps.td

 //===-- TestOps.td - Test dialect operation definitions ----*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TEST_OPS
 #define TEST_OPS
 
 include "TestDialect.td"
 include "mlir/Dialect/DLTI/DLTIBase.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/PatternBase.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
+include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
 include "TestInterfaces.td"
 
 
 // Include the attribute definitions.

 
 //===----------------------------------------------------------------------===//
 // Test Ops with effects
 //===----------------------------------------------------------------------===//
 
 def TestResource : Resource<"TestResource">;
 
 def TestEffectsOpA : TEST_Op<"op_with_effects_a"> {
-    let arguments = (ins
-        Arg<Variadic<AnyMemRef>, "", [MemRead]>,
-        Arg<FlatSymbolRefAttr, "", [MemRead]>:$first,
-        Arg<SymbolRefAttr, "", [MemWrite]>:$second,
-        Arg<OptionalAttr<SymbolRefAttr>, "", [MemRead]>:$optional_symbol
-        );
 
-    let results = (outs Res<AnyMemRef, "", [MemAlloc<TestResource>]>);
 }
 
 def TestEffectsOpB : TEST_Op<"op_with_effects_b",
     [MemoryEffects<[MemWrite<TestResource>]>]>;
 
+//===----------------------------------------------------------------------===//
+// Test Loop Op
+//===----------------------------------------------------------------------===//
+
+// Test loop op with graph regions.
+def TestGraphLoopOp : TEST_Op<"graph_loop",
+                         [LoopLikeOpInterface, NoSideEffect,
+                          RecursiveSideEffects, SingleBlock,
+                          RegionKindInterface, HasOnlyGraphRegion]> {
+  let arguments = (ins Variadic<AnyType>:$args);
+  let results = (outs Variadic<AnyType>:$rets);
+  let regions = (region SizedRegion<1>:$body);
+
+  let assemblyFormat = [{ $args $body attr-dict `:`
+                          functional-type(operands, results) }];
+
+  let extraClassDeclaration = [{
+    mlir::Region &getLoopBody() { return getBody(); }
+  }];
+}
+
+
 #endif // TEST_OPS