[mlir][OpenMP] Allow 'cancel taskgroup' inside taskloop region

[kaviya2510]

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

[llvmbot]

@llvm/pr-subscribers-mlir-openmp

@llvm/pr-subscribers-flang-openmp

Author: Kaviya Rajendiran (kaviya2510)

Changes

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

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Kaviya Rajendiran (kaviya2510)

Changes

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

[kaviya2510]

This changes is required to fix the testcase taskloop-cancel.f90 , which is part of the PR #138646

[tblah]

Thanks for the fix!