[Mlir-commits] [mlir] [OpenMP][MLIR] RegionBranchOpInterface for TargetOp, TargetDataOp (PR #209710)

[Muyao Xiao]

https://github.com/mooxiu updated https://github.com/llvm/llvm-project/pull/209710

>From b874e89f382535c46ceef67f3dc4d2231520124f Mon Sep 17 00:00:00 2001
From: Muyao Xiao <monetshaw at outlook.com>
Date: Wed, 15 Jul 2026 18:19:00 +0900
Subject: [PATCH 1/2] implement regional branch interface for omp::TargetOp,
 omp::TargetDataOp, and omp::terminatorOp

---
 mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td |  7 +++---
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp  | 22 +++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 98242a66e411f..e1840a3bee317 100644
--- a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
+++ b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
@@ -222,7 +222,7 @@ def ParallelOp : OpenMP_Op<"parallel", traits = [
   let hasRegionVerifier = 1;
 }
 
-def TerminatorOp : OpenMP_Op<"terminator", [Terminator, Pure]> {
+def TerminatorOp : OpenMP_Op<"terminator", [Terminator, Pure, RegionBranchTerminatorOpInterface]> {
   let summary = "terminator for OpenMP regions";
   let description = [{
     A terminator operation for regions that appear in the body of OpenMP
@@ -1457,7 +1457,8 @@ def MapInfoOp : OpenMP_Op<"map.info", [AttrSizedOperandSegments]> {
 //===---------------------------------------------------------------------===//
 
 def TargetDataOp: OpenMP_Op<"target_data", traits = [
-    AttrSizedOperandSegments, DeclareOpInterfaceMethods<ComposableOpInterface>
+    AttrSizedOperandSegments, DeclareOpInterfaceMethods<ComposableOpInterface>,
+    DeclareOpInterfaceMethods<RegionBranchOpInterface>
   ], clauses = [
     OpenMP_DeviceClause, OpenMP_IfClause, OpenMP_MapClause,
     OpenMP_UseDeviceAddrClause, OpenMP_UseDevicePtrClause
@@ -1621,7 +1622,7 @@ def TargetUpdateOp: OpenMP_Op<"target_update", traits = [
 def TargetOp : OpenMP_Op<"target", traits = [
     AttrSizedOperandSegments, BlockArgOpenMPOpInterface,
     DeclareOpInterfaceMethods<ComposableOpInterface>, IsolatedFromAbove,
-    OutlineableOpenMPOpInterface
+    OutlineableOpenMPOpInterface, DeclareOpInterfaceMethods<RegionBranchOpInterface>
   ], clauses = [
     // TODO: Complete clause list (defaultmap, uses_allocators).
     OpenMP_AllocateClause, OpenMP_DependClause, OpenMP_DeviceClause,
diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 3a3798c3ad00b..d50f96c8b0669 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2571,6 +2571,17 @@ LogicalResult TargetDataOp::verify() {
   return verifyMapClause(*this, getMapVars(), getMapIterated());
 }
 
+// Adapted from fir.if implementation.
+void TargetDataOp::getSuccessorRegions(
+  mlir::RegionBranchPoint point, 
+  llvm::SmallVectorImpl<::mlir::RegionSuccessor>& regions) {
+  if (!point.isParent()) {
+    regions.push_back(mlir::RegionSuccessor::parent());
+    return;
+  }
+  regions.push_back(mlir::RegionSuccessor(&getRegion()));
+}
+
 //===----------------------------------------------------------------------===//
 // TargetEnterDataOp
 //===----------------------------------------------------------------------===//
@@ -2853,6 +2864,17 @@ LogicalResult TargetOp::verifyRegions() {
   return success();
 }
 
+// Copy from `TargetDataOp::getSuccessorRegions`
+void TargetOp::getSuccessorRegions(
+  mlir::RegionBranchPoint point, 
+  llvm::SmallVectorImpl<::mlir::RegionSuccessor>& regions) {
+  if (!point.isParent()) {
+    regions.push_back(mlir::RegionSuccessor::parent());
+    return;
+  }
+  regions.push_back(mlir::RegionSuccessor(&getRegion()));
+}
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//

>From b81004d4bfbce9be9742c40c7a28e884433f6669 Mon Sep 17 00:00:00 2001
From: Muyao Xiao <monetshaw at outlook.com>
Date: Wed, 15 Jul 2026 19:22:52 +0900
Subject: [PATCH 2/2] update to latest API: no parent() anymore

---
 mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index d50f96c8b0669..ca201aba48a03 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2576,7 +2576,7 @@ void TargetDataOp::getSuccessorRegions(
   mlir::RegionBranchPoint point, 
   llvm::SmallVectorImpl<::mlir::RegionSuccessor>& regions) {
   if (!point.isParent()) {
-    regions.push_back(mlir::RegionSuccessor::parent());
+    regions.push_back(mlir::RegionSuccessor(getOperation()));
     return;
   }
   regions.push_back(mlir::RegionSuccessor(&getRegion()));
@@ -2869,7 +2869,7 @@ void TargetOp::getSuccessorRegions(
   mlir::RegionBranchPoint point, 
   llvm::SmallVectorImpl<::mlir::RegionSuccessor>& regions) {
   if (!point.isParent()) {
-    regions.push_back(mlir::RegionSuccessor::parent());
+    regions.push_back(mlir::RegionSuccessor(getOperation()));
     return;
   }
   regions.push_back(mlir::RegionSuccessor(&getRegion()));