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

[Muyao Xiao]

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

This adds RegionBranchOpInterface support to `omp.target` and `omp.target_data`, and adds RegionBranchTerminatorOpInterface to `omp.terminator`.

Both `omp.target` and `omp.target_data` contain structured regions whose control flow enters the region and eventually returns to the parent operation. Without RegionBranchOpInterface, generic region-aware data-flow analyses cannot propagate lattice state across these region boundaries.

For example, an analysis that tracks the host and device states of mapped values cannot propagate its state into an `omp.target` region and back to the parent operation after a mapped value is modified and copied out.

Changes:
- Add RegionBranchOpInterface to omp.target.
- Add RegionBranchOpInterface to omp.target_data.
- Add RegionBranchTerminatorOpInterface to omp.terminator.
- Implement the corresponding region successor methods.

This was previously raised on [LLVM Discussion](https://discourse.llvm.org/t/mlir-openmp-should-omp-target-target-data-implement-regionbranchopinterface/91227).

Similar PR: https://github.com/llvm/llvm-project/pull/202418

>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] 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
 //===----------------------------------------------------------------------===//