[mlir][acc] Add LegalizeDataValues support for DeclareEnterOp #138008

SusanTan · 2025-04-30T18:13:56Z

The patch extends the existing LegalizeDataValues to support DeclareEnter and DeclareExit pair.
Since unlike other ops, DeclareEnter and DeclareExit don't have a region defined, we use dominance/post dominance information to ensure only the uses within the region dominated by DeclareEnter and post dominated by DeclareExit are updated with data on device.

llvmbot · 2025-04-30T18:14:32Z

@llvm/pr-subscribers-openacc
@llvm/pr-subscribers-mlir-openacc

@llvm/pr-subscribers-mlir

Author: Susan Tan (ス-ザン　タン) (SusanTan)

Changes

The patch extends the existing LegalizeDataValues to support DeclareEnter and DeclareExit pair.
Since unlike other ops, DeclareEnter and DeclareExit don't have a region defined, we use dominance/post dominance information to ensure only the uses within the region dominated by DeclareEnter and post dominated by DeclareExit are updated with data on device.

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

2 Files Affected:

(modified) mlir/include/mlir/Dialect/OpenACC/OpenACC.h (+2-1)
(modified) mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp (+51-5)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACC.h b/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
index ff5845343313c..9c141550b184e 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
@@ -58,7 +58,8 @@
 #define ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS                                     \
   ACC_COMPUTE_CONSTRUCT_OPS, mlir::acc::LoopOp
 #define ACC_DATA_CONSTRUCT_STRUCTURED_OPS                                      \
-  mlir::acc::DataOp, mlir::acc::DeclareOp, mlir::acc::HostDataOp
+  mlir::acc::DataOp, mlir::acc::DeclareOp, mlir::acc::HostDataOp,              \
+      mlir::acc::DeclareEnterOp
 #define ACC_DATA_CONSTRUCT_UNSTRUCTURED_OPS                                    \
   mlir::acc::EnterDataOp, mlir::acc::ExitDataOp, mlir::acc::UpdateOp,          \
       mlir::acc::DeclareEnterOp, mlir::acc::DeclareExitOp
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp b/mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp
index f2abeab744d17..24f0ed3d35a0e 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp
@@ -10,6 +10,7 @@
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/Dominance.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -70,8 +71,38 @@ static void replaceAllUsesInAccComputeRegionsWith(Value orig, Value replacement,
   }
 }
 
+// Helper function to process declare enter/exit pairs
+static void processDeclareEnterExit(
+    acc::DeclareEnterOp op, llvm::SmallVector<std::pair<Value, Value>> &values,
+    DominanceInfo &domInfo, PostDominanceInfo &postDomInfo) {
+  // For declare enter/exit pairs, verify there is exactly one exit op using the
+  // token
+  if (!op.getToken().hasOneUse())
+    op.emitError("declare enter token must have exactly one use");
+  Operation *user = *op.getToken().getUsers().begin();
+  auto declareExit = dyn_cast<acc::DeclareExitOp>(user);
+  if (!declareExit)
+    op.emitError("declare enter token must be used by declare exit op");
+
+  for (auto p : values) {
+    Value hostVal = std::get<0>(p);
+    Value deviceVal = std::get<1>(p);
+    for (auto &use : llvm::make_early_inc_range(hostVal.getUses())) {
+      Operation *owner = use.getOwner();
+      if (!domInfo.dominates(op.getOperation(), owner) ||
+          !postDomInfo.postDominates(declareExit.getOperation(), owner))
+        continue;
+      if (insideAccComputeRegion(owner))
+        use.set(deviceVal);
+    }
+  }
+}
+
 template <typename Op>
-static void collectAndReplaceInRegion(Op &op, bool hostToDevice) {
+static void
+collectAndReplaceInRegion(Op &op, bool hostToDevice,
+                          DominanceInfo *domInfo = nullptr,
+                          PostDominanceInfo *postDomInfo = nullptr) {
   llvm::SmallVector<std::pair<Value, Value>> values;
 
   if constexpr (std::is_same_v<Op, acc::LoopOp>) {
@@ -82,16 +113,24 @@ static void collectAndReplaceInRegion(Op &op, bool hostToDevice) {
     if constexpr (!std::is_same_v<Op, acc::KernelsOp> &&
                   !std::is_same_v<Op, acc::DataOp> &&
                   !std::is_same_v<Op, acc::DeclareOp> &&
-                  !std::is_same_v<Op, acc::HostDataOp>) {
+                  !std::is_same_v<Op, acc::HostDataOp> &&
+                  !std::is_same_v<Op, acc::DeclareEnterOp>) {
       collectVars(op.getReductionOperands(), values, hostToDevice);
       collectVars(op.getPrivateOperands(), values, hostToDevice);
       collectVars(op.getFirstprivateOperands(), values, hostToDevice);
     }
   }
 
-  for (auto p : values)
-    replaceAllUsesInAccComputeRegionsWith<Op>(std::get<0>(p), std::get<1>(p),
-                                              op.getRegion());
+  if constexpr (std::is_same_v<Op, acc::DeclareEnterOp>) {
+    assert(domInfo && postDomInfo &&
+           "Dominance info required for DeclareEnterOp");
+    processDeclareEnterExit(op, values, *domInfo, *postDomInfo);
+  } else {
+    for (auto p : values) {
+      replaceAllUsesInAccComputeRegionsWith<Op>(std::get<0>(p), std::get<1>(p),
+                                                op.getRegion());
+    }
+  }
 }
 
 class LegalizeDataValuesInRegion
@@ -105,6 +144,10 @@ class LegalizeDataValuesInRegion
     func::FuncOp funcOp = getOperation();
     bool replaceHostVsDevice = this->hostToDevice.getValue();
 
+    // Get dominance info for the function
+    DominanceInfo domInfo(funcOp);
+    PostDominanceInfo postDomInfo(funcOp);
+
     funcOp.walk([&](Operation *op) {
       if (!isa<ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS>(*op) &&
           !(isa<ACC_DATA_CONSTRUCT_STRUCTURED_OPS>(*op) &&
@@ -125,6 +168,9 @@ class LegalizeDataValuesInRegion
         collectAndReplaceInRegion(declareOp, replaceHostVsDevice);
       } else if (auto hostDataOp = dyn_cast<acc::HostDataOp>(*op)) {
         collectAndReplaceInRegion(hostDataOp, replaceHostVsDevice);
+      } else if (auto declareEnterOp = dyn_cast<acc::DeclareEnterOp>(*op)) {
+        collectAndReplaceInRegion(declareEnterOp, replaceHostVsDevice, &domInfo,
+                                  &postDomInfo);
       } else {
         llvm_unreachable("unsupported acc region op");
       }

razvanlupusoru · 2025-04-30T18:48:54Z

mlir/include/mlir/Dialect/OpenACC/OpenACC.h

@@ -58,7 +58,8 @@
 #define ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS                                     \
  ACC_COMPUTE_CONSTRUCT_OPS, mlir::acc::LoopOp
 #define ACC_DATA_CONSTRUCT_STRUCTURED_OPS                                      \
-  mlir::acc::DataOp, mlir::acc::DeclareOp, mlir::acc::HostDataOp
+  mlir::acc::DataOp, mlir::acc::DeclareOp, mlir::acc::HostDataOp,              \
+      mlir::acc::DeclareEnterOp


This should not be added to this list because although it can result from a structured declare construct, the operation itself allows unstructured flow.

yep, i realized that was a mistake, changed it in the next commit

razvanlupusoru · 2025-04-30T18:50:45Z

mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp

@@ -105,6 +144,10 @@ class LegalizeDataValuesInRegion
    func::FuncOp funcOp = getOperation();
    bool replaceHostVsDevice = this->hostToDevice.getValue();

+    // Get dominance info for the function
+    DominanceInfo domInfo(funcOp);


Let's compute this lazily - only compute it first time when you encounter acc.declare_enter.

razvanlupusoru · 2025-04-30T19:38:00Z

mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp

+  if constexpr (std::is_same_v<Op, acc::DeclareEnterOp>) {
+    // For declare enter/exit pairs, verify there is exactly one exit op using
+    // the token
+    if (!op.getToken().hasOneUse())


This is too restrictive since from a language perspective multiple exits are allowed - and a dialect may choose to permit that. I think the algorithm here should collect all of the declare_exit's and ensure that the replaced use post-dominates one of them.

razvanlupusoru · 2025-04-30T19:42:01Z

mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp

+    if (!declareExit)
+      op.emitError("declare enter token must be used by declare exit op");
+    exitOp = declareExit;
+  } else if constexpr (std::is_same_v<Op, acc::EnterDataOp>) {


I am excited about the potential to support acc.enter_data. As I looked at your algorithm here though I realized that maybe it should be left for another change. This is because enter_data and exit_data don't come in pairs - so a bit more analysis is needed to get this right. More specifically, finding an exit_data does not constitute the end of the lifetime for a particular variable.

razvanlupusoru · 2025-04-30T20:51:19Z

mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp

+        exitOps.push_back(declareExit);
+    }
+    if (exitOps.empty())
+      op.emitError(


This is not an error. For example in the acc.global_ctor, an acc.declare_enter is used without an exit operation. That said, I think it seems OK to avoid doing any replacement of values in such a case - just don't make it an error.

razvanlupusoru

Nice job!

…38008) The patch extends the existing LegalizeDataValues to support DeclareEnter and DeclareExit pair. Since unlike other ops, DeclareEnter and DeclareExit don't have a region defined, we use dominance/post dominance information to ensure only the uses within the region dominated by DeclareEnter and post dominated by DeclareExit are updated with data on device.

add support for acc.declare_enter

4c458ea

llvmbot added mlir mlir:openacc openacc labels Apr 30, 2025

SusanTan marked this pull request as draft April 30, 2025 18:15

razvanlupusoru reviewed Apr 30, 2025

View reviewed changes

add support for enter data

03b89b1

SusanTan marked this pull request as ready for review April 30, 2025 19:15

SusanTan added 2 commits April 30, 2025 12:30

change how dominfo is computed

cd70987

leave out initialization

014a8b8

razvanlupusoru reviewed Apr 30, 2025

View reviewed changes

SusanTan added 3 commits April 30, 2025 12:55

add lit test

823d31d

change to support multiple exits

4814f8f

remove remaining enter_data

cb0a1b0

razvanlupusoru reviewed Apr 30, 2025

View reviewed changes

change the logic for when exitops is empty

a132a37

razvanlupusoru approved these changes May 1, 2025

View reviewed changes

razvanlupusoru merged commit a073bb5 into llvm:main May 1, 2025
11 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][acc] Add LegalizeDataValues support for DeclareEnterOp #138008

[mlir][acc] Add LegalizeDataValues support for DeclareEnterOp #138008

Uh oh!

SusanTan commented Apr 30, 2025

Uh oh!

llvmbot commented Apr 30, 2025 •

edited

Loading

Uh oh!

razvanlupusoru Apr 30, 2025

Uh oh!

SusanTan Apr 30, 2025

Uh oh!

razvanlupusoru Apr 30, 2025

Uh oh!

razvanlupusoru Apr 30, 2025

Uh oh!

razvanlupusoru Apr 30, 2025

Uh oh!

razvanlupusoru Apr 30, 2025

Uh oh!

razvanlupusoru left a comment

Uh oh!

Uh oh!

Uh oh!

[mlir][acc] Add LegalizeDataValues support for DeclareEnterOp #138008

[mlir][acc] Add LegalizeDataValues support for DeclareEnterOp #138008

Uh oh!

Conversation

SusanTan commented Apr 30, 2025

Uh oh!

llvmbot commented Apr 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

razvanlupusoru Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

SusanTan Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

razvanlupusoru Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

razvanlupusoru Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

razvanlupusoru Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

razvanlupusoru Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

razvanlupusoru left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

llvmbot commented Apr 30, 2025 •

edited

Loading