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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-openmp

Author: Muyao Xiao (mooxiu)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/209710.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td (+6-3) 
- (modified) mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp (+39) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td b/mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td
index 70597b85902c5..1f3d2f0d9c153 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,9 @@ def TargetUpdateOp: OpenMP_Op<"target_update", traits = [
 def TargetOp : OpenMP_Op<"target", traits = [
     AttrSizedOperandSegments, BlockArgOpenMPOpInterface,
     DeclareOpInterfaceMethods<ComposableOpInterface>, IsolatedFromAbove,
-    OutlineableOpenMPOpInterface
+    OutlineableOpenMPOpInterface, 
+    DeclareOpInterfaceMethods<RegionBranchOpInterface,
+      ["getEntrySuccessorOperands", "getSuccessorInputs"]>
   ], 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..e17bb44695a99 100644
--- a/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
+++ b/mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp
@@ -15,6 +15,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/Dialect/OpenMP/OpenMPClauseOperands.h"
+#include "mlir/Dialect/OpenMP/OpenMPInterfaces.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/DialectImplementation.h"
@@ -22,6 +23,8 @@
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/SymbolTable.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/FoldInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
@@ -44,6 +47,7 @@
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc"
 #include "mlir/Dialect/OpenMP/OpenMPOpsInterfaces.cpp.inc"
 #include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.cpp.inc"
+#include "mlir/Support/LLVM.h"
 
 using namespace mlir;
 using namespace mlir::omp;
@@ -2571,6 +2575,17 @@ LogicalResult TargetDataOp::verify() {
   return verifyMapClause(*this, getMapVars(), getMapIterated());
 }
 
+// Adapted from scf.if implementation.
+void TargetDataOp::getSuccessorRegions(
+    mlir::RegionBranchPoint point,
+    llvm::SmallVectorImpl<::mlir::RegionSuccessor> &regions) {
+  if (!point.isParent()) {
+    regions.push_back(mlir::RegionSuccessor(getOperation()));
+    return;
+  }
+  regions.push_back(mlir::RegionSuccessor(&getRegion()));
+}
+
 //===----------------------------------------------------------------------===//
 // TargetEnterDataOp
 //===----------------------------------------------------------------------===//
@@ -2853,6 +2868,30 @@ LogicalResult TargetOp::verifyRegions() {
   return success();
 }
 
+void TargetOp::getSuccessorRegions(
+    mlir::RegionBranchPoint point,
+    llvm::SmallVectorImpl<::mlir::RegionSuccessor> &regions) {
+  if (!point.isParent()) {
+    regions.push_back(mlir::RegionSuccessor(getOperation()));
+    return;
+  }
+  regions.push_back(mlir::RegionSuccessor(&getRegion()));
+}
+
+OperandRange
+TargetOp::getEntrySuccessorOperands(mlir::RegionSuccessor successor) {
+  assert(successor.getSuccessor() == &getRegion());
+  return getHostEvalVars();
+}
+
+mlir::ValueRange TargetOp::getSuccessorInputs(mlir::RegionSuccessor successor) {
+  if (successor.isOperation())
+    return {};
+  assert(successor == &getRegion());
+  return mlir::cast<BlockArgOpenMPOpInterface>(getOperation())
+      .getHostEvalBlockArgs();
+}
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//

``````````

</details>


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