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(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) {