Add support for early exit walk methods.

This is done by providing a walk callback that returns a WalkResult. This result is either `advance` or `interrupt`. `advance` means that the walk should continue, whereas `interrupt` signals that the walk should stop immediately. An example is shown below:

auto result = op->walk([](Operation *op) {
  if (some_invariant)
    return WalkResult::interrupt();
  return WalkResult::advance();
});

if (result.wasInterrupted())
  ...;

PiperOrigin-RevId: 266436700

1 files changed, 16 insertions, 19 deletions

diff --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp
index aaefd98d1bd..660b77e0f96 100644
--- a/mlir/lib/Analysis/Utils.cpp
+++ b/mlir/lib/Analysis/Utils.cpp
@@ -905,11 +905,10 @@ static Optional<int64_t> getMemoryFootprintBytes(Block &block,
   SmallDenseMap<Value *, std::unique_ptr<MemRefRegion>, 4> regions;
 
   // Walk this 'affine.for' operation to gather all memory regions.
-  bool error = false;
-  block.walk(start, end, [&](Operation *opInst) {
+  auto result = block.walk(start, end, [&](Operation *opInst) -> WalkResult {
     if (!isa<AffineLoadOp>(opInst) && !isa<AffineStoreOp>(opInst)) {
       // Neither load nor a store op.
-      return;
+      return WalkResult::advance();
     }
 
     // Compute the memref region symbolic in any IVs enclosing this block.
@@ -917,23 +916,20 @@ static Optional<int64_t> getMemoryFootprintBytes(Block &block,
     if (failed(
             region->compute(opInst,
                             /*loopDepth=*/getNestingDepth(*block.begin())))) {
-      opInst->emitError("Error obtaining memory region\n");
-      error = true;
-      return;
+      return opInst->emitError("Error obtaining memory region\n");
     }
+
     auto it = regions.find(region->memref);
     if (it == regions.end()) {
       regions[region->memref] = std::move(region);
     } else if (failed(it->second->unionBoundingBox(*region))) {
-      opInst->emitWarning(
+      return opInst->emitWarning(
           "getMemoryFootprintBytes: unable to perform a union on a memory "
           "region");
-      error = true;
-      return;
     }
+    return WalkResult::advance();
   });
-
-  if (error)
+  if (result.wasInterrupted())
     return None;
 
   int64_t totalSizeInBytes = 0;
@@ -969,17 +965,18 @@ void mlir::getSequentialLoops(
 bool mlir::isLoopParallel(AffineForOp forOp) {
   // Collect all load and store ops in loop nest rooted at 'forOp'.
   SmallVector<Operation *, 8> loadAndStoreOpInsts;
-  bool hasSideEffectingOps = false;
-  forOp.getOperation()->walk([&](Operation *opInst) {
+  auto walkResult = forOp.walk([&](Operation *opInst) {
     if (isa<AffineLoadOp>(opInst) || isa<AffineStoreOp>(opInst))
-      return loadAndStoreOpInsts.push_back(opInst);
-    if (!isa<AffineForOp>(opInst) && !isa<AffineTerminatorOp>(opInst) &&
-        !isa<AffineIfOp>(opInst) && !opInst->hasNoSideEffect()) {
-      hasSideEffectingOps = true;
-    }
+      loadAndStoreOpInsts.push_back(opInst);
+    else if (!isa<AffineForOp>(opInst) && !isa<AffineTerminatorOp>(opInst) &&
+             !isa<AffineIfOp>(opInst) && !opInst->hasNoSideEffect())
+      return WalkResult::interrupt();
+
+    return WalkResult::advance();
   });
+
   // Stop early if the loop has unknown ops with side effects.
-  if (hasSideEffectingOps)
+  if (walkResult.wasInterrupted())
     return false;
 
   // Dep check depth would be number of enclosing loops + 1.