[llvm-branch-commits] [mlir] [mlir][OpenMP] Don't allow loop bounds/step from inside the task (PR #190827)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-flang-openmp

Author: Tom Eccles (tblah)

<details>
<summary>Changes</summary>

The omp.taskloop.context region represents what goes inside the outlined task function. The loop bounds must be passed to the OpenMP runtime call for taskloop and so this cannot be supported in general.

In a follow up patch I will re-allow pure operations because sinking constants inside of the tasklooop context will be useful for something else I am prototyping.

Assisted-by: codex

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


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+3) 
- (modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+22) 
- (modified) mlir/test/Dialect/OpenMP/invalid.mlir (+17) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 76294cb86574f..40ec8904a136f 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -1016,6 +1016,9 @@ def TaskloopContextOp : OpenMP_Op<"taskloop.context", traits = [
     }
     ```
 
+    The wrapped `omp.loop_nest` lower bounds, upper bounds and steps must be
+    defined outside of the `omp.taskloop.context` region.
+
     For definitions of "undeferred task", "included task", "final task" and
     "mergeable task", please check OpenMP Specification.
 
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index dd2846f2cdb42..df74a316c7b66 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -3429,6 +3429,28 @@ LogicalResult TaskloopContextOp::verifyRegions() {
            << "expected exactly 1 TaskloopWrapperOp directly nested in "
               "the region, but "
            << count << " were found";
+  TaskloopWrapperOp loopWrapperOp = getLoopOp();
+
+  auto loopNestOp = dyn_cast<LoopNestOp>(loopWrapperOp.getWrappedLoop());
+  // This will fail the verifier for TaskloopWrapperOp and print an error
+  // message there.
+  if (!loopNestOp)
+    return failure();
+
+  auto isDefinedInTaskloopContext = [&](Value value) {
+    // A region is considered an ancestor of itself
+    return region.isAncestor(value.getParentRegion());
+  };
+  auto hasTaskloopLocalBound = [&](OperandRange range) {
+    return llvm::any_of(range, isDefinedInTaskloopContext);
+  };
+
+  if (hasTaskloopLocalBound(loopNestOp.getLoopLowerBounds()) ||
+      hasTaskloopLocalBound(loopNestOp.getLoopUpperBounds()) ||
+      hasTaskloopLocalBound(loopNestOp.getLoopSteps())) {
+    return emitOpError() << "expects loop bounds and steps to be defined "
+                            "outside of the taskloop.context region";
+  }
 
   return success();
 }
diff --git a/mlir/test/Dialect/OpenMP/invalid.mlir b/mlir/test/Dialect/OpenMP/invalid.mlir
index c7f2416b4c293..98087cc43588f 100644
--- a/mlir/test/Dialect/OpenMP/invalid.mlir
+++ b/mlir/test/Dialect/OpenMP/invalid.mlir
@@ -3035,6 +3035,23 @@ func.func @omp_taskloop_invalid_composite(%lb: index, %ub: index, %step: index)
   return
 }
 
+// -----
+func.func @omp_taskloop_local_loop_bounds() {
+  // expected-error @below {{'omp.taskloop.context' op expects loop bounds and steps to be defined outside of the taskloop.context region}}
+  omp.taskloop.context {
+    %lb = arith.constant 1 : index
+    %ub = arith.constant 10 : index
+    %step = arith.constant 1 : index
+    omp.taskloop.wrapper {
+      omp.loop_nest (%i) : index = (%lb) to (%ub) step (%step) {
+        omp.yield
+      }
+    }
+    omp.terminator
+  }
+  return
+}
+
 // -----
 
 func.func @omp_loop_invalid_nesting(%lb : index, %ub : index, %step : index) {

``````````

</details>


https://github.com/llvm/llvm-project/pull/190827