fixup! fixup! Enable LICM for ops with read side effects in scf.for wrapped by a guard

Author: ardaunal

mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

@@ -40,7 +40,7 @@ def SCF_Dialect : Dialect {
     and then lowered to some final target like LLVM or SPIR-V.
   }];
 
-  let dependentDialects = ["arith::ArithDialect"];
+  let dependentDialects = ["arith::ArithDialect", "ub::UBDialect"];
 }
 
 // Base class for SCF dialect ops.
@@ -138,6 +138,7 @@ def ForOp : SCF_Op<"for",
        ["getInitsMutable", "getLoopResults", "getRegionIterArgs",
         "getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
         "getLoopUpperBounds", "getYieldedValuesMutable",
+        "moveOutOfLoopWithGuard",
         "promoteIfSingleIteration", "replaceWithAdditionalYields",
         "wrapInTripCountCheck", "unwrapTripCountCheck",
         "yieldTiledValuesAndReplace"]>,

mlir/include/mlir/Interfaces/LoopLikeInterface.td

@@ -80,23 +80,15 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
       /*defaultImplementation=*/"op->moveBefore($_op);"
     >,
     InterfaceMethod<[{
-        Wraps the loop into a trip-count check.
+        Moves the given loop-invariant operation out of the loop with a
+        trip-count guard.
       }],
-      /*retTy=*/"FailureOr<std::pair<::mlir::Operation *, ::mlir::Region *>>",
-      /*methodName=*/"wrapInTripCountCheck",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/"return ::mlir::failure();"
-    >,
-    InterfaceMethod<[{
-        Unwraps the trip-count check.
-      }],
-      /*retTy=*/"::llvm::LogicalResult",
-      /*methodName=*/"unwrapTripCountCheck",
-      /*args=*/(ins),
+      /*retTy=*/"void",
+      /*methodName=*/"moveOutOfLoopWithGuard",
+      /*args=*/(ins "::mlir::Operation *":$op),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return ::mlir::failure();
+        return;
       }]
     >,
     InterfaceMethod<[{

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

@@ -433,9 +433,9 @@ bool wouldOpBeTriviallyDead(Operation *op);
 /// conditions are satisfied.
 bool isMemoryEffectFree(Operation *op);
 
-/// Returns true if the given operation implements `MemoryEffectOpInterface` and
-/// has only read effects.
-bool hasOnlyReadEffect(Operation *op);
+/// Returns true if the given operation is free of memory effects or has only
+/// read effect.
+bool isMemoryEffectFreeOrOnlyRead(Operation *op);
 
 /// Returns the side effects of an operation. If the operation has
 /// RecursiveMemoryEffects, include all side effects of child operations.

mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h

@@ -47,20 +47,19 @@ class Value;
 ///   }
 /// }
 /// ```
-///
-/// Users must supply five callbacks.
+/// Users must supply four callbacks.
 ///
 /// - `isDefinedOutsideRegion` returns true if the given value is invariant with
 ///   respect to the given region. A common implementation might be:
 ///   `value.getParentRegion()->isProperAncestor(region)`.
 /// - `shouldMoveOutOfRegion` returns true if the provided operation can be
 ///   moved of the given region, e.g. if it is side-effect free or has only read
 ///   side effects.
-/// - `wrapInGuard` wraps the given operation in a trip-count check guard.
-/// - `moveOutOfRegion` moves the operation out of the given region. A common
-///   implementation might be: `op->moveBefore(region->getParentOp())`.
-/// - `unwrapGuard` unwraps the trip-count check if there is no op guarded by
-///   this check.
+/// - `moveOutOfRegionWithoutGuard` moves the operation out of the given region
+///   without a guard. A common implementation might be:
+///   `op->moveBefore(region->getParentOp())`.
+/// - `moveOutOfRegionWithGuard` moves the operation out of the given region
+///   with a guard.
 ///
 /// An operation is moved if all of its operands satisfy
 /// `isDefinedOutsideRegion` and it satisfies `shouldMoveOutOfRegion`.
@@ -70,9 +69,8 @@ size_t moveLoopInvariantCode(
     ArrayRef<Region *> regions,
     function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
     function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
-    function_ref<FailureOr<std::pair<Operation *, Region *>>()> wrapInGuard,
-    function_ref<void(Operation *, Region *)> moveOutOfRegion,
-    function_ref<LogicalResult()> unwrapGuard);
+    function_ref<void(Operation *)> moveOutOfRegionWithoutGuard,
+    function_ref<void(Operation *)> moveOutOfRegionWithGuard);
 
 /// Move side-effect free loop invariant code out of a loop-like op using
 /// methods provided by the interface.

mlir/lib/Dialect/SCF/IR/SCF.cpp

@@ -15,6 +15,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/DeviceMappingInterface.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Matchers.h"
@@ -395,58 +396,38 @@ std::optional<SmallVector<OpFoldResult>> ForOp::getLoopUpperBounds() {
 
 std::optional<ResultRange> ForOp::getLoopResults() { return getResults(); }
 
-FailureOr<std::pair<Operation *, Region *>> ForOp::wrapInTripCountCheck() {
-
+/// Moves the op out of the loop with a guard that checks if the loop has at
+/// least one iteration.
+void ForOp::moveOutOfLoopWithGuard(Operation *op) {
   IRRewriter rewriter(this->getContext());
   OpBuilder::InsertionGuard insertGuard(rewriter);
-  rewriter.setInsertionPointAfter(this->getOperation());
-
-  auto loc = this->getLoc();
-  auto cmpIOp = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ugt,
-                                               this->getUpperBound(),
-                                               this->getLowerBound());
-  scf::YieldOp yieldInThen;
+  rewriter.setInsertionPoint(this->getOperation());
+  Location loc = this->getLoc();
+  arith::CmpIOp cmpIOp = rewriter.create<arith::CmpIOp>(
+      loc, arith::CmpIPredicate::ult, this->getLowerBound(),
+      this->getUpperBound());
   // Create the trip-count check.
-  auto ifOp = rewriter.create<scf::IfOp>(
+  scf::YieldOp thenYield;
+  scf::IfOp ifOp = rewriter.create<scf::IfOp>(
       loc, cmpIOp,
       [&](OpBuilder &builder, Location loc) {
-        yieldInThen = builder.create<scf::YieldOp>(loc, this->getResults());
+        thenYield = builder.create<scf::YieldOp>(loc, op->getResults());
       },
       [&](OpBuilder &builder, Location loc) {
-        builder.create<scf::YieldOp>(loc, this->getInitArgs());
+        SmallVector<Value> poisonResults;
+        poisonResults.reserve(op->getResults().size());
+        for (Type type : op->getResults().getTypes()) {
+          ub::PoisonOp poisonOp =
+              rewriter.create<ub::PoisonOp>(loc, type, nullptr);
+          poisonResults.push_back(poisonOp);
+        }
+        builder.create<scf::YieldOp>(loc, poisonResults);
       });
-
-  for (auto [forOpResult, ifOpResult] :
-       llvm::zip(this->getResults(), ifOp.getResults()))
-    rewriter.replaceAllUsesExcept(forOpResult, ifOpResult, yieldInThen);
-  // Move the scf.for into the then block.
-  rewriter.moveOpBefore(this->getOperation(), yieldInThen);
-  return std::make_pair(ifOp.getOperation(), &this->getRegion());
-}
-
-LogicalResult ForOp::unwrapTripCountCheck() {
-  auto ifOp = (*this)->getParentRegion()->getParentOp();
-  if (!isa<scf::IfOp>(ifOp))
-    return failure();
-
-  IRRewriter rewriter(ifOp->getContext());
-  OpBuilder::InsertionGuard insertGuard(rewriter);
-  rewriter.setInsertionPoint(ifOp);
-
-  auto cmpOp = ifOp->getOperand(0).getDefiningOp();
-  if (!isa<arith::CmpIOp>(cmpOp))
-    return failure();
-
-  auto wrappedForOp = this->getOperation();
-  rewriter.moveOpBefore(wrappedForOp, ifOp);
-
-  for (auto [forOpResult, ifOpResult] :
-       llvm::zip(wrappedForOp->getResults(), ifOp->getResults()))
-    rewriter.replaceAllUsesWith(ifOpResult, forOpResult);
-
-  rewriter.eraseOp(ifOp);
-  rewriter.eraseOp(cmpOp);
-  return success();
+  for (auto [opResult, ifOpResult] :
+       llvm::zip(op->getResults(), ifOp->getResults()))
+    rewriter.replaceAllUsesExcept(opResult, ifOpResult, thenYield);
+  // Move the op into the then block.
+  rewriter.moveOpBefore(op, thenYield);
 }
 
 /// Promotes the loop body of a forOp to its containing block if the forOp

mlir/lib/Interfaces/SideEffectInterfaces.cpp

@@ -306,26 +306,6 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
   return wouldOpBeTriviallyDeadImpl(op);
 }
 
-bool mlir::hasOnlyReadEffect(Operation *op) {
-  if (auto memEffects = dyn_cast<MemoryEffectOpInterface>(op)) {
-    if (!op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
-      return memEffects.onlyHasEffect<MemoryEffects::Read>();
-  } else if (!op->hasTrait<OpTrait::HasRecursiveMemoryEffects>()) {
-    // 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 and ensure that all nested ops are memory effect
-  // free.
-  for (Region &region : op->getRegions())
-    for (Operation &op : region.getOps())
-      if (!hasOnlyReadEffect(&op))
-        return false;
-  return true;
-}
-
 bool mlir::isMemoryEffectFree(Operation *op) {
   if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
     if (!memInterface.hasNoEffect())
@@ -383,6 +363,16 @@ mlir::getEffectsRecursively(Operation *rootOp) {
   return effects;
 }
 
+bool mlir::isMemoryEffectFreeOrOnlyRead(Operation *op) {
+  std::optional<SmallVector<MemoryEffects::EffectInstance>> effects =
+      getEffectsRecursively(op);
+  if (!effects)
+    return false;
+  return std::all_of(effects->begin(), effects->end(), [](auto &effect) {
+    return isa<MemoryEffects::Read>(effect.getEffect());
+  });
+}
+
 bool mlir::isSpeculatable(Operation *op) {
   auto conditionallySpeculatable = dyn_cast<ConditionallySpeculatable>(op);
   if (!conditionallySpeculatable)

mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp

@@ -56,117 +56,67 @@ static bool canBeHoisted(Operation *op,
       op, [&](OpOperand &operand) { return definedOutside(operand.get()); });
 }
 
-static bool dependsOnGuarded(Operation *op,
-                             function_ref<bool(OpOperand &)> condition) {
-  auto walkFn = [&](Operation *child) {
-    for (OpOperand &operand : child->getOpOperands()) {
-      if (!condition(operand))
-        return WalkResult::interrupt();
-    }
-    return WalkResult::advance();
-  };
-  return op->walk(walkFn).wasInterrupted();
-}
-
-static bool dependsOnGuarded(Operation *op,
-                             function_ref<bool(Value)> definedOutsideGuard) {
-  return dependsOnGuarded(op, [&](OpOperand &operand) {
-    return definedOutsideGuard(operand.get());
-  });
-}
-
-static bool loopSideEffectFreeOrHasOnlyReadEffect(Operation *loop) {
-  for (Region &region : loop->getRegions()) {
-    for (Block &block : region.getBlocks()) {
-      for (Operation &op : block.getOperations()) {
-        if (!isMemoryEffectFree(&op) && !hasOnlyReadEffect(&op))
-          return false;
-      }
-    }
-  }
-  return true;
-}
-
 size_t mlir::moveLoopInvariantCode(
     ArrayRef<Region *> regions,
     function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
     function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
-    function_ref<FailureOr<std::pair<Operation *, Region *>>()> wrapInGuard,
-    function_ref<void(Operation *, Region *)> moveOutOfRegion,
-    function_ref<LogicalResult()> unwrapGuard) {
+    function_ref<void(Operation *)> moveOutOfRegionWithoutGuard,
+    function_ref<void(Operation *)> moveOutOfRegionWithGuard) {
   size_t numMoved = 0;
 
   for (Region *region : regions) {
     LLVM_DEBUG(llvm::dbgs() << "Original loop:\n"
                             << *region->getParentOp() << "\n");
 
-    auto loopSideEffectFreeOrHasOnlyReadSideEffect =
-        loopSideEffectFreeOrHasOnlyReadEffect(region->getParentOp());
-
-    size_t numMovedWithoutGuard = 0;
-
-    FailureOr<std::pair<Operation *, Region *>> ifOpAndRegion = wrapInGuard();
-    Region *loopRegion = region;
-    auto isLoopWrapped = false;
-    if (succeeded(ifOpAndRegion)) {
-      loopRegion = ifOpAndRegion->second;
-      isLoopWrapped = true;
-    }
+    bool anyOpHoistedWithGuard = false;
+    bool loopSideEffectFreeOrHasOnlyReadSideEffect =
+        isMemoryEffectFreeOrOnlyRead(region->getParentOp());
 
     std::queue<Operation *> worklist;
     // Add top-level operations in the loop body to the worklist.
-    for (Operation &op : loopRegion->getOps())
+    for (Operation &op : region->getOps())
       worklist.push(&op);
 
     auto definedOutside = [&](Value value) {
-      return isDefinedOutsideRegion(value, loopRegion);
-    };
-
-    auto definedOutsideGuard = [&](Value value) {
-      return isDefinedOutsideRegion(value, loopRegion->getParentRegion());
+      return isDefinedOutsideRegion(value, region);
     };
 
     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() != loopRegion)
+      if (op->getParentRegion() != region)
         continue;
 
       LLVM_DEBUG(llvm::dbgs() << "Checking op: " << *op << "\n");
 
-      if (!shouldMoveOutOfRegion(op, loopRegion) ||
+      if (!shouldMoveOutOfRegion(op, region) ||
           !canBeHoisted(op, definedOutside))
         continue;
       // Can only hoist pure ops (side-effect free) when there is an op with
       // write and/or unknown side effects in the loop.
       if (!loopSideEffectFreeOrHasOnlyReadSideEffect && !isMemoryEffectFree(op))
         continue;
 
-      LLVM_DEBUG(llvm::dbgs() << "Moving loop-invariant op: " << *op << "\n");
-
-      auto moveWithoutGuard = isMemoryEffectFree(op) &&
-                              !dependsOnGuarded(op, definedOutsideGuard) &&
-                              isLoopWrapped;
-      numMovedWithoutGuard += moveWithoutGuard;
-
-      moveOutOfRegion(op, moveWithoutGuard ? loopRegion->getParentRegion()
-                                           : loopRegion);
+      bool moveWithoutGuard = !anyOpHoistedWithGuard && isMemoryEffectFree(op);
+      if (moveWithoutGuard) {
+        LLVM_DEBUG(llvm::dbgs() << "Moving loop-invariant op: " << *op
+                                << " without guard\n");
+        moveOutOfRegionWithoutGuard(op);
+      } else {
+        LLVM_DEBUG(llvm::dbgs()
+                   << "Moving loop-invariant op: " << *op << " with guard\n");
+        moveOutOfRegionWithGuard(op);
+        anyOpHoistedWithGuard = true;
+      }
       ++numMoved;
 
       // 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() == loopRegion)
+        if (user->getParentRegion() == region)
           worklist.push(user);
     }
-
-    // Unwrap the loop if it was wrapped but no ops were moved in the guard.
-    if (isLoopWrapped && numMovedWithoutGuard == numMoved) {
-      auto tripCountCheckUnwrapped = unwrapGuard();
-      if (failed(tripCountCheckUnwrapped))
-        llvm_unreachable("Should not fail unwrapping trip-count check");
-    }
   }
 
   return numMoved;
@@ -179,14 +129,10 @@ size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
         return !region->isAncestor(value.getParentRegion());
       },
       [&](Operation *op, Region *) {
-        return isSpeculatable(op) &&
-               (isMemoryEffectFree(op) || hasOnlyReadEffect(op));
-      },
-      [&]() { return loopLike.wrapInTripCountCheck(); },
-      [&](Operation *op, Region *region) {
-        op->moveBefore(region->getParentOp());
+        return isSpeculatable(op) && isMemoryEffectFreeOrOnlyRead(op);
       },
-      [&]() { return loopLike.unwrapTripCountCheck(); });
+      [&](Operation *op) { loopLike.moveOutOfLoop(op); },
+      [&](Operation *op) { loopLike.moveOutOfLoopWithGuard(op); });
 }
 
 namespace {