diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp index 0f65ace0186db..186f14de22f70 100644 --- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp +++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp @@ -3217,7 +3217,8 @@ LogicalResult CancelOp::verify() { } } if ((cct == ClauseCancellationConstructType::Taskgroup) && - !mlir::isa(structuralParent)) { + (!mlir::isa(structuralParent) && + !mlir::isa(structuralParent->getParentOp()))) { return emitOpError() << "cancel taskgroup must appear " << "inside a task region"; } diff --git a/mlir/test/Dialect/OpenMP/ops.mlir b/mlir/test/Dialect/OpenMP/ops.mlir index a3c591799b27e..b7e16b7ec35e2 100644 --- a/mlir/test/Dialect/OpenMP/ops.mlir +++ b/mlir/test/Dialect/OpenMP/ops.mlir @@ -2215,6 +2215,18 @@ func.func @omp_cancel_taskgroup() -> () { return } +func.func @omp_taskloop_cancel_taskgroup(%lb : index, %ub : index, %step : index) { + omp.taskloop { + omp.loop_nest (%iv) : index = (%lb) to (%ub) step (%step) { + // CHECK: omp.cancel cancellation_construct_type(taskgroup) + omp.cancel cancellation_construct_type(taskgroup) + // CHECK: omp.yield + omp.yield + } + } + return +} + func.func @omp_cancel_parallel_nested(%if_cond : i1) -> () { omp.parallel { scf.if %if_cond {