[mlir][GPU] Don't look into neighboring functions for barrier elimination

[krzysz00]

If a func.func is nested in some other operation, the barrier eliminator's recursion into parents will examine the neighbors of each function. Therefore, don't recurse into the parent of an operation if that operation is IsolatedFromAbove, like a func.func is.

Furthermore, define functions as a region that executes only once, since, within the context of this pass (which runs on functions) it is true.

[llvmbot]

@llvm/pr-subscribers-mlir-gpu

Author: Krzysztof Drewniak (krzysz00)

Changes

If a func.func is nested in some other operation, the barrier eliminator's recursion into parents will examine the neighbors of each function. Therefore, don't recurse into the parent of an operation if that operation is IsolatedFromAbove, like a func.func is.

Furthermore, define functions as a region that executes only once, since, within the context of this pass (which runs on functions) it is true.

---

Full diff: https://github.com/llvm/llvm-project/pull/135293.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp (+11-6)

diff --git a/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp b/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp
index 2178555cb62f7..84c12c0ba05e5 100644
--- a/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp
@@ -68,7 +68,7 @@ static bool isSequentialLoopLike(Operation *op) { return isa<scf::ForOp>(op); }
 /// most once. Thus, if an operation in one of the nested regions of `op` is
 /// executed than so are all the other operations in this region.
 static bool hasSingleExecutionBody(Operation *op) {
-  return isa<scf::IfOp, memref::AllocaScopeOp>(op);
+  return isa<FunctionOpInterface, scf::IfOp, memref::AllocaScopeOp>(op);
 }
 
 /// Returns `true` if the operation is known to produce a pointer-like object
@@ -182,8 +182,10 @@ getEffectsBefore(Operation *op,
   if (isParallelRegionBoundary(op->getParentOp()))
     return true;
 
+  Operation *parent = op->getParentOp();
   // Otherwise, keep collecting above the parent operation.
-  if (!getEffectsBefore(op->getParentOp(), effects, stopAtBarrier))
+  if (!parent->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+      !getEffectsBefore(parent, effects, stopAtBarrier))
     return false;
 
   // If the op is loop-like, collect effects from the trailing operations until
@@ -200,7 +202,7 @@ getEffectsBefore(Operation *op,
   // the operation `op2` at iteration `i` is known to be executed before the
   // operation `op1` at iteration `i+1` and the side effects must be ordered
   // appropriately.
-  if (isSequentialLoopLike(op->getParentOp())) {
+  if (isSequentialLoopLike(parent)) {
     // Assuming loop terminators have no side effects.
     return getEffectsBeforeInBlock(op->getBlock()->getTerminator(), effects,
                                    /*stopAtBarrier=*/true);
@@ -268,12 +270,15 @@ getEffectsAfter(Operation *op,
   // Collect all effects after the op.
   getEffectsAfterInBlock(op, effects, stopAtBarrier);
 
+  Operation *parent = op->getParentOp();
   // Stop if reached the parallel region boundary.
-  if (isParallelRegionBoundary(op->getParentOp()))
+  if (isParallelRegionBoundary(parent))
     return true;
 
   // Otherwise, keep collecting below the parent operation.
-  if (!getEffectsAfter(op->getParentOp(), effects, stopAtBarrier))
+  // Don't look into, for example, neighboring functions
+  if (!parent->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+      !getEffectsAfter(parent, effects, stopAtBarrier))
     return false;
 
   // If the op is loop-like, collect effects from the leading operations until
@@ -290,7 +295,7 @@ getEffectsAfter(Operation *op,
   // the operation `op1` at iteration `i` is known to be executed after the
   // operation `op2` at iteration `i-1` and the side effects must be ordered
   // appropriately.
-  if (isSequentialLoopLike(op->getParentOp())) {
+  if (isSequentialLoopLike(parent)) {
     if (isa<BarrierOp>(op->getBlock()->front()))
       return true;
 

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Krzysztof Drewniak (krzysz00)

Changes

If a func.func is nested in some other operation, the barrier eliminator's recursion into parents will examine the neighbors of each function. Therefore, don't recurse into the parent of an operation if that operation is IsolatedFromAbove, like a func.func is.

Furthermore, define functions as a region that executes only once, since, within the context of this pass (which runs on functions) it is true.

---

Full diff: https://github.com/llvm/llvm-project/pull/135293.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp (+11-6)

diff --git a/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp b/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp
index 2178555cb62f7..84c12c0ba05e5 100644
--- a/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/EliminateBarriers.cpp
@@ -68,7 +68,7 @@ static bool isSequentialLoopLike(Operation *op) { return isa<scf::ForOp>(op); }
 /// most once. Thus, if an operation in one of the nested regions of `op` is
 /// executed than so are all the other operations in this region.
 static bool hasSingleExecutionBody(Operation *op) {
-  return isa<scf::IfOp, memref::AllocaScopeOp>(op);
+  return isa<FunctionOpInterface, scf::IfOp, memref::AllocaScopeOp>(op);
 }
 
 /// Returns `true` if the operation is known to produce a pointer-like object
@@ -182,8 +182,10 @@ getEffectsBefore(Operation *op,
   if (isParallelRegionBoundary(op->getParentOp()))
     return true;
 
+  Operation *parent = op->getParentOp();
   // Otherwise, keep collecting above the parent operation.
-  if (!getEffectsBefore(op->getParentOp(), effects, stopAtBarrier))
+  if (!parent->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+      !getEffectsBefore(parent, effects, stopAtBarrier))
     return false;
 
   // If the op is loop-like, collect effects from the trailing operations until
@@ -200,7 +202,7 @@ getEffectsBefore(Operation *op,
   // the operation `op2` at iteration `i` is known to be executed before the
   // operation `op1` at iteration `i+1` and the side effects must be ordered
   // appropriately.
-  if (isSequentialLoopLike(op->getParentOp())) {
+  if (isSequentialLoopLike(parent)) {
     // Assuming loop terminators have no side effects.
     return getEffectsBeforeInBlock(op->getBlock()->getTerminator(), effects,
                                    /*stopAtBarrier=*/true);
@@ -268,12 +270,15 @@ getEffectsAfter(Operation *op,
   // Collect all effects after the op.
   getEffectsAfterInBlock(op, effects, stopAtBarrier);
 
+  Operation *parent = op->getParentOp();
   // Stop if reached the parallel region boundary.
-  if (isParallelRegionBoundary(op->getParentOp()))
+  if (isParallelRegionBoundary(parent))
     return true;
 
   // Otherwise, keep collecting below the parent operation.
-  if (!getEffectsAfter(op->getParentOp(), effects, stopAtBarrier))
+  // Don't look into, for example, neighboring functions
+  if (!parent->hasTrait<OpTrait::IsIsolatedFromAbove>() &&
+      !getEffectsAfter(parent, effects, stopAtBarrier))
     return false;
 
   // If the op is loop-like, collect effects from the leading operations until
@@ -290,7 +295,7 @@ getEffectsAfter(Operation *op,
   // the operation `op1` at iteration `i` is known to be executed after the
   // operation `op2` at iteration `i-1` and the side effects must be ordered
   // appropriately.
-  if (isSequentialLoopLike(op->getParentOp())) {
+  if (isSequentialLoopLike(parent)) {
     if (isa<BarrierOp>(op->getBlock()->front()))
       return true;