[Mlir-commits] [mlir] Handle OpenACC host_data in LegalizeDataValues (PR #134767)

[llvmlistbot at llvm.org]

https://github.com/nvptm created https://github.com/llvm/llvm-project/pull/134767

`LegalizeDataValuesInRegion` is intended to replace the SSA values used in a region with the output of data operations, but misses the handling of the OpenACC `host_data` construct. As a result, currently

```
 !$acc host_data use_device(%var)
   ...%var...
 !$acc end host_data

```
is lowered to

```
 %dev_var = acc.use_device(%var)
 acc.host_data data_operands(%dev_var) {
   ...%var...
 }
```

This pull request updates the LegalizeDataValuesInRegion to handle HostDataOp such that lowering results in

```
 %dev_var = acc.use_device(%var)
 acc.host_data data_operands(%dev_var) {
   ...%dev_var...
 }
```






>From e95f2811278c5690595b1fc133459ea801892725 Mon Sep 17 00:00:00 2001
From: nvpm <pmathew at nvidia.com>
Date: Mon, 7 Apr 2025 13:48:08 -0700
Subject: [PATCH] Handle OpenACC host_data in LegalizeDataValues

---
 mlir/include/mlir/Dialect/OpenACC/OpenACC.h   |  5 ++---
 .../OpenACC/Transforms/LegalizeDataValues.cpp |  5 ++++-
 mlir/test/Dialect/OpenACC/legalize-data.mlir  | 20 +++++++++++++++++++
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACC.h b/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
index 748cb7f28fc8c..ff5845343313c 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
@@ -58,11 +58,10 @@
 #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::DataOp, mlir::acc::DeclareOp, mlir::acc::HostDataOp
 #define ACC_DATA_CONSTRUCT_UNSTRUCTURED_OPS                                    \
   mlir::acc::EnterDataOp, mlir::acc::ExitDataOp, mlir::acc::UpdateOp,          \
-      mlir::acc::HostDataOp, mlir::acc::DeclareEnterOp,                        \
-      mlir::acc::DeclareExitOp
+      mlir::acc::DeclareEnterOp, mlir::acc::DeclareExitOp
 #define ACC_DATA_CONSTRUCT_OPS                                                 \
   ACC_DATA_CONSTRUCT_STRUCTURED_OPS, ACC_DATA_CONSTRUCT_UNSTRUCTURED_OPS
 #define ACC_COMPUTE_AND_DATA_CONSTRUCT_OPS                                     \
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp b/mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp
index a553653c73479..f2abeab744d17 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/LegalizeDataValues.cpp
@@ -81,7 +81,8 @@ static void collectAndReplaceInRegion(Op &op, bool hostToDevice) {
     collectVars(op.getDataClauseOperands(), values, 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::DeclareOp> &&
+                  !std::is_same_v<Op, acc::HostDataOp>) {
       collectVars(op.getReductionOperands(), values, hostToDevice);
       collectVars(op.getPrivateOperands(), values, hostToDevice);
       collectVars(op.getFirstprivateOperands(), values, hostToDevice);
@@ -122,6 +123,8 @@ class LegalizeDataValuesInRegion
         collectAndReplaceInRegion(dataOp, replaceHostVsDevice);
       } else if (auto declareOp = dyn_cast<acc::DeclareOp>(*op)) {
         collectAndReplaceInRegion(declareOp, replaceHostVsDevice);
+      } else if (auto hostDataOp = dyn_cast<acc::HostDataOp>(*op)) {
+        collectAndReplaceInRegion(hostDataOp, replaceHostVsDevice);
       } else {
         llvm_unreachable("unsupported acc region op");
       }
diff --git a/mlir/test/Dialect/OpenACC/legalize-data.mlir b/mlir/test/Dialect/OpenACC/legalize-data.mlir
index baa72ae416c92..e192c0c55e599 100644
--- a/mlir/test/Dialect/OpenACC/legalize-data.mlir
+++ b/mlir/test/Dialect/OpenACC/legalize-data.mlir
@@ -226,3 +226,23 @@ func.func @test(%a: memref<10xf32>) {
 // CHECK:   }
 // CHECK:   acc.yield
 // CHECK: }
+
+// -----
+
+func.func @test(%a: memref<10xf32>) {
+  %devptr = acc.use_device varPtr(%a : memref<10xf32>) varType(tensor<10xf32>) -> memref<10xf32>
+  acc.host_data dataOperands(%devptr : memref<10xf32>) {
+    fir.call @foo(%a) fastmath<contract> : (memref<10xf32>) -> ()
+    acc.terminator
+  }
+  return
+}
+func.func private @foo(memref<10xf32>)
+
+// CHECK-LABEL: func.func @test
+// CHECK-SAME: (%[[A:.*]]: memref<10xf32>)
+// CHECK: %[[USE_DEVICE:.*]] = acc.use_device varPtr(%[[A]] : memref<10xf32>) varType(tensor<10xf32>) -> memref<10xf32>
+// CHECK: acc.host_data dataOperands(%[[USE_DEVICE]] : memref<10xf32>) {
+// DEVICE:   func.call @foo(%[[USE_DEVICE]]) : (memref<10xf32>) -> ()
+// CHECK:   acc.terminator
+// CHECK: }