[mlir] fix memory effects in GPU barrier elimination (#117432)

Existing implementation may trigger infinite cycles when collecting
effects above or below the current block after wrapping around a
loop-like construct. Limit this case to only looking at the immediate
block (loop body). This is correct because wrap around is intended to
consider effects of different iterations of the same loop and shouldn't
be existing the loop block.

Reported-by: Fabian Mora <[email protected]>
Co-authored-by: Fabian Mora <[email protected]>

Author: ftynse

mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp

@@ -132,6 +132,29 @@ collectEffects(Operation *op,
   return false;
 }
 
+/// Get all effects before the given operation caused by other operations in the
+/// same block. That is, this will not consider operations beyond the block.
+static bool
+getEffectsBeforeInBlock(Operation *op,
+                        SmallVectorImpl<MemoryEffects::EffectInstance> &effects,
+                        bool stopAtBarrier) {
+  if (op == &op->getBlock()->front())
+    return true;
+
+  for (Operation *it = op->getPrevNode(); it != nullptr;
+       it = it->getPrevNode()) {
+    if (isa<BarrierOp>(it)) {
+      if (stopAtBarrier)
+        return true;
+      continue;
+    }
+
+    if (!collectEffects(it, effects))
+      return false;
+  }
+  return true;
+}
+
 /// Collects memory effects from operations that may be executed before `op` in
 /// a trivial structured control flow, e.g., without branches. Stops at the
 /// parallel region boundary or at the barrier operation if `stopAtBarrier` is
@@ -153,19 +176,7 @@ getEffectsBefore(Operation *op,
   }
 
   // Collect all effects before the op.
-  if (op != &op->getBlock()->front()) {
-    for (Operation *it = op->getPrevNode(); it != nullptr;
-         it = it->getPrevNode()) {
-      if (isa<BarrierOp>(it)) {
-        if (stopAtBarrier)
-          return true;
-        else
-          continue;
-      }
-      if (!collectEffects(it, effects))
-        return false;
-    }
-  }
+  getEffectsBeforeInBlock(op, effects, stopAtBarrier);
 
   // Stop if reached the parallel region boundary.
   if (isParallelRegionBoundary(op->getParentOp()))
@@ -191,8 +202,8 @@ getEffectsBefore(Operation *op,
   // appropriately.
   if (isSequentialLoopLike(op->getParentOp())) {
     // Assuming loop terminators have no side effects.
-    return getEffectsBefore(op->getBlock()->getTerminator(), effects,
-                            /*stopAtBarrier=*/true);
+    return getEffectsBeforeInBlock(op->getBlock()->getTerminator(), effects,
+                                   /*stopAtBarrier=*/true);
   }
 
   // If the parent operation is not guaranteed to execute its (single-block)
@@ -212,6 +223,28 @@ getEffectsBefore(Operation *op,
   return !conservative;
 }
 
+/// Get all effects after the given operation caused by other operations in the
+/// same block. That is, this will not consider operations beyond the block.
+static bool
+getEffectsAfterInBlock(Operation *op,
+                       SmallVectorImpl<MemoryEffects::EffectInstance> &effects,
+                       bool stopAtBarrier) {
+  if (op == &op->getBlock()->back())
+    return true;
+
+  for (Operation *it = op->getNextNode(); it != nullptr;
+       it = it->getNextNode()) {
+    if (isa<BarrierOp>(it)) {
+      if (stopAtBarrier)
+        return true;
+      continue;
+    }
+    if (!collectEffects(it, effects))
+      return false;
+  }
+  return true;
+}
+
 /// Collects memory effects from operations that may be executed after `op` in
 /// a trivial structured control flow, e.g., without branches. Stops at the
 /// parallel region boundary or at the barrier operation if `stopAtBarrier` is
@@ -233,17 +266,7 @@ getEffectsAfter(Operation *op,
   }
 
   // Collect all effects after the op.
-  if (op != &op->getBlock()->back())
-    for (Operation *it = op->getNextNode(); it != nullptr;
-         it = it->getNextNode()) {
-      if (isa<BarrierOp>(it)) {
-        if (stopAtBarrier)
-          return true;
-        continue;
-      }
-      if (!collectEffects(it, effects))
-        return false;
-    }
+  getEffectsAfterInBlock(op, effects, stopAtBarrier);
 
   // Stop if reached the parallel region boundary.
   if (isParallelRegionBoundary(op->getParentOp()))
@@ -272,8 +295,8 @@ getEffectsAfter(Operation *op,
       return true;
 
     bool exact = collectEffects(&op->getBlock()->front(), effects);
-    return getEffectsAfter(&op->getBlock()->front(), effects,
-                           /*stopAtBarrier=*/true) &&
+    return getEffectsAfterInBlock(&op->getBlock()->front(), effects,
+                                  /*stopAtBarrier=*/true) &&
            exact;
   }
 

mlir/test/Dialect/GPU/barrier-elimination.mlir

@@ -182,3 +182,20 @@ attributes {__parallel_region_boundary_for_test} {
   %4 = memref.load %C[] : memref<f32>
   return %0, %1, %2, %3, %4 : f32, f32, f32, f32, f32
 }
+
+// CHECK-LABEL: @nested_loop_barrier_only
+func.func @nested_loop_barrier_only() attributes {__parallel_region_boundary_for_test} {
+  %c0 = arith.constant 0 : index
+  %c42 = arith.constant 42 : index
+  %c1 = arith.constant 1 : index
+  // Note: the barrier can be removed and as consequence the loops get folded
+  // by the greedy rewriter.
+  // CHECK-NOT: scf.for
+  // CHECK-NOT: gpu.barrier
+  scf.for %j = %c0 to %c42 step %c1 {
+    scf.for %i = %c0 to %c42 step %c1 {
+      gpu.barrier
+    }
+  }
+  return
+}