[Mlir-commits] [mlir] [mlir][OpenMP] Allow 'cancel taskgroup' inside taskloop region (PR #138634)

[llvmlistbot at llvm.org]

llvmbot wrote:


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

@llvm/pr-subscribers-flang-openmp

Author: Kaviya Rajendiran (kaviya2510)

<details>
<summary>Changes</summary>

Previously, `cancel taskgroup` was only permitted within `omp.task` region.This change allows `cancel taskgroup` inside  `taskloop` region as well.



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


2 Files Affected:

- (modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+2-1) 
- (modified) mlir/test/Dialect/OpenMP/ops.mlir (+12) 


``````````diff
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<omp::TaskOp>(structuralParent)) {
+      (!mlir::isa<omp::TaskOp>(structuralParent) &&
+       !mlir::isa<omp::TaskloopOp>(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 {

``````````

</details>


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