[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 d1f2d200b78e343af54901de5aa5506fb45572e7 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/3] 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 70597b85902c5..4cc322849591c 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 95ccb53bce3d4..2d73a073205bd 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 8ab20817c4c87cba0aa81b1cd851810e017460bb 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/3] 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 2d73a073205bd..6290aeef0bca1 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()));

>From 9e094fb66903a25d52e7258b2251a9f286256c6f Mon Sep 17 00:00:00 2001
From: Muyao Xiao <monetshaw at outlook.com>
Date: Wed, 15 Jul 2026 19:44:22 +0900
Subject: [PATCH 3/3] fix format

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

diff --git a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
index 6290aeef0bca1..b3bbf441e5e75 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -2573,8 +2573,8 @@ LogicalResult TargetDataOp::verify() {
 
 // Adapted from fir.if implementation.
 void TargetDataOp::getSuccessorRegions(
-  mlir::RegionBranchPoint point, 
-  llvm::SmallVectorImpl<::mlir::RegionSuccessor>& regions) {
+    mlir::RegionBranchPoint point,
+    llvm::SmallVectorImpl<::mlir::RegionSuccessor> &regions) {
   if (!point.isParent()) {
     regions.push_back(mlir::RegionSuccessor(getOperation()));
     return;
@@ -2866,8 +2866,8 @@ LogicalResult TargetOp::verifyRegions() {
 
 // Copy from `TargetDataOp::getSuccessorRegions`
 void TargetOp::getSuccessorRegions(
-  mlir::RegionBranchPoint point, 
-  llvm::SmallVectorImpl<::mlir::RegionSuccessor>& regions) {
+    mlir::RegionBranchPoint point,
+    llvm::SmallVectorImpl<::mlir::RegionSuccessor> &regions) {
   if (!point.isParent()) {
     regions.push_back(mlir::RegionSuccessor(getOperation()));
     return;