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

[Kaviya Rajendiran]

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

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



>From e6da8660ae26360d0a15d115e818addbd6ee1b4d Mon Sep 17 00:00:00 2001
From: Kaviya Rajendiran <kaviyara2000 at gmail.com>
Date: Tue, 6 May 2025 11:10:27 +0530
Subject: [PATCH] [mlir][OpenMP] Allow 'cancel taskgroup' inside taskloop
 region

---
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp |  3 ++-
 mlir/test/Dialect/OpenMP/ops.mlir            | 12 ++++++++++++
 2 files changed, 14 insertions(+), 1 deletion(-)

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 {