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

[llvmlistbot at llvm.org]

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

>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 1/2] 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: }

>From 378ebb0d4f207254fdd3bff8e740e874eaa5d06c Mon Sep 17 00:00:00 2001
From: nvpm <pmathew at nvidia.com>
Date: Thu, 10 Apr 2025 13:26:37 -0700
Subject: [PATCH 2/2] Fix test to use FUNC dialect instead of FIR which is
 outside of the MLIR tree

---
 mlir/test/Dialect/OpenACC/legalize-data.mlir | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/OpenACC/legalize-data.mlir b/mlir/test/Dialect/OpenACC/legalize-data.mlir
index e192c0c55e599..ae854a42adb3f 100644
--- a/mlir/test/Dialect/OpenACC/legalize-data.mlir
+++ b/mlir/test/Dialect/OpenACC/legalize-data.mlir
@@ -232,7 +232,7 @@ func.func @test(%a: memref<10xf32>) {
 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>) -> ()
+    func.call @foo(%a) fastmath<contract> : (memref<10xf32>) -> ()
     acc.terminator
   }
   return