[acc] Improve LegalizeDataValues pass to handle data constructs

flang/test/Fir/OpenACC/legalize-data.fir

@@ -1,4 +1,4 @@
-// RUN: fir-opt -split-input-file --openacc-legalize-data %s | FileCheck %s
+// RUN: fir-opt -split-input-file --openacc-legalize-data-values %s | FileCheck %s
 
 func.func @_QPsub1(%arg0: !fir.ref<i32> {fir.bindc_name = "i"}) {
   %0:2 = hlfir.declare %arg0 {uniq_name = "_QFsub1Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
@@ -22,3 +22,36 @@ func.func @_QPsub1(%arg0: !fir.ref<i32> {fir.bindc_name = "i"}) {
 // CHECK:   acc.yield
 // CHECK: }
 // CHECK: acc.copyout accPtr(%[[COPYIN]] : !fir.ref<i32>) to varPtr(%[[I]]#0 : !fir.ref<i32>) {dataClause = #acc<data_clause acc_copy>, name = "i"}
+
+// -----
+
+func.func @_QPsub1(%arg0: !fir.ref<i32> {fir.bindc_name = "i"}) {
+  %0:2 = hlfir.declare %arg0 {uniq_name = "_QFsub1Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+  %1 = acc.copyin varPtr(%0#0 : !fir.ref<i32>) -> !fir.ref<i32> {dataClause = #acc<data_clause acc_copy>, name = "i"}
+  acc.data dataOperands(%1 : !fir.ref<i32>) {
+    %c0_i32 = arith.constant 0 : i32
+    hlfir.assign %c0_i32 to %0#0 : i32, !fir.ref<i32>
+    acc.serial {
+      hlfir.assign %c0_i32 to %0#0 : i32, !fir.ref<i32>
+      acc.yield
+    }
+    acc.terminator
+  }
+  acc.copyout accPtr(%1 : !fir.ref<i32>) to varPtr(%0#0 : !fir.ref<i32>) {dataClause = #acc<data_clause acc_copy>, name = "i"}
+  return
+}
+
+// CHECK-LABEL: func.func @_QPsub1
+// CHECK-SAME: (%[[ARG0:.*]]: !fir.ref<i32> {fir.bindc_name = "i"})
+// CHECK: %[[I:.*]]:2 = hlfir.declare %[[ARG0]] {uniq_name = "_QFsub1Ei"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
+// CHECK: %[[COPYIN:.*]] = acc.copyin varPtr(%[[I]]#0 : !fir.ref<i32>) -> !fir.ref<i32> {dataClause = #acc<data_clause acc_copy>, name = "i"}
+// CHECK: acc.data dataOperands(%[[COPYIN]] : !fir.ref<i32>) {
+// CHECK:   %[[C0:.*]] = arith.constant 0 : i32
+// CHECK:   hlfir.assign %[[C0]] to %0#0 : i32, !fir.ref<i32>
+// CHECK:   acc.serial {
+// CHECK:     hlfir.assign %[[C0]] to %[[COPYIN]] : i32, !fir.ref<i32>
+// CHECK:     acc.yield
+// CHECK:   }
+// CHECK:   acc.terminator
+// CHECK: }
+// CHECK: acc.copyout accPtr(%[[COPYIN]] : !fir.ref<i32>) to varPtr(%[[I]]#0 : !fir.ref<i32>) {dataClause = #acc<data_clause acc_copy>, name = "i"}

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

@@ -56,14 +56,14 @@
   mlir::acc::ParallelOp, mlir::acc::KernelsOp, mlir::acc::SerialOp
 #define ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS                                     \
   ACC_COMPUTE_CONSTRUCT_OPS, mlir::acc::LoopOp
-#define OPENACC_DATA_CONSTRUCT_STRUCTURED_OPS                                  \
+#define ACC_DATA_CONSTRUCT_STRUCTURED_OPS                                      \
   mlir::acc::DataOp, mlir::acc::DeclareOp
 #define ACC_DATA_CONSTRUCT_UNSTRUCTURED_OPS                                    \
   mlir::acc::EnterDataOp, mlir::acc::ExitDataOp, mlir::acc::UpdateOp,          \
       mlir::acc::HostDataOp, mlir::acc::DeclareEnterOp,                        \
       mlir::acc::DeclareExitOp
 #define ACC_DATA_CONSTRUCT_OPS                                                 \
-  OPENACC_DATA_CONSTRUCT_STRUCTURED_OPS, ACC_DATA_CONSTRUCT_UNSTRUCTURED_OPS
+  ACC_DATA_CONSTRUCT_STRUCTURED_OPS, ACC_DATA_CONSTRUCT_UNSTRUCTURED_OPS
 #define ACC_COMPUTE_AND_DATA_CONSTRUCT_OPS                                     \
   ACC_COMPUTE_CONSTRUCT_OPS, ACC_DATA_CONSTRUCT_OPS
 #define ACC_COMPUTE_LOOP_AND_DATA_CONSTRUCT_OPS                                \

mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.h

@@ -11,9 +11,6 @@
 
 #include "mlir/Pass/Pass.h"
 
-#define GEN_PASS_DECL
-#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
-
 namespace mlir {
 
 namespace func {
@@ -22,8 +19,8 @@ class FuncOp;
 
 namespace acc {
 
-/// Create a pass to replace ssa values in region with device/host values.
-std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeDataInRegion();
+#define GEN_PASS_DECL
+#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
 
 /// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION

mlir/include/mlir/Dialect/OpenACC/Transforms/Passes.td

@@ -11,18 +11,20 @@
 
 include "mlir/Pass/PassBase.td"
 
-def LegalizeDataInRegion : Pass<"openacc-legalize-data", "mlir::func::FuncOp"> {
-  let summary = "Legalize the data in the compute region";
+def LegalizeDataValuesInRegion : Pass<"openacc-legalize-data-values", "mlir::func::FuncOp"> {
+  let summary = "Legalizes SSA values in compute regions with results from data clause operations";
   let description = [{
-    This pass replace uses of varPtr in the compute region with their accPtr
-    gathered from the data clause operands.
+    This pass replace uses of the `varPtr` in compute regions (kernels,
+    parallel, serial) with the result of data clause operations (`accPtr`).
   }];
   let options = [
     Option<"hostToDevice", "host-to-device", "bool", "true",
            "Replace varPtr uses with accPtr if true. Replace accPtr uses with "
-           "varPtr if false">
+           "varPtr if false">,
+    Option<"applyToAccDataConstruct", "apply-to-acc-data-construct", "bool", "true",
+           "Replaces varPtr uses with accPtr for acc compute regions contained "
+           "within acc.data or acc.declare region.">
   ];
-  let constructor = "::mlir::acc::createLegalizeDataInRegion()";
 }
 
 #endif // MLIR_DIALECT_OPENACC_TRANSFORMS_PASSES

mlir/lib/Dialect/OpenACC/Transforms/CMakeLists.txt

@@ -1,5 +1,5 @@
 add_mlir_dialect_library(MLIROpenACCTransforms
-  LegalizeData.cpp
+  LegalizeDataValues.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/OpenACC

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

@@ -1,4 +1,4 @@
-//===- LegalizeData.cpp - -------------------------------------------------===//
+//===- LegalizeDataValues.cpp - -------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -12,10 +12,11 @@
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/RegionUtils.h"
+#include "llvm/Support/ErrorHandling.h"
 
 namespace mlir {
 namespace acc {
-#define GEN_PASS_DEF_LEGALIZEDATAINREGION
+#define GEN_PASS_DEF_LEGALIZEDATAVALUESINREGION
 #include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
 } // namespace acc
 } // namespace mlir
@@ -24,6 +25,17 @@ using namespace mlir;
 
 namespace {
 
+static bool insideAccComputeRegion(mlir::Operation *op) {
+  mlir::Operation *parent{op->getParentOp()};
+  while (parent) {
+    if (isa<ACC_COMPUTE_CONSTRUCT_OPS>(parent)) {
+      return true;
+    }
+    parent = parent->getParentOp();
+  }
+  return false;
+}
+
 static void collectPtrs(mlir::ValueRange operands,
                         llvm::SmallVector<std::pair<Value, Value>> &values,
                         bool hostToDevice) {
@@ -39,6 +51,25 @@ static void collectPtrs(mlir::ValueRange operands,
   }
 }
 
+template <typename Op>
+static void replaceAllUsesInAccComputeRegionsWith(Value orig, Value replacement,
+                                                  Region &outerRegion) {
+  for (auto &use : llvm::make_early_inc_range(orig.getUses())) {
+    if (outerRegion.isAncestor(use.getOwner()->getParentRegion())) {
+      if constexpr (std::is_same_v<Op, acc::DataOp> ||
+                    std::is_same_v<Op, acc::DeclareOp>) {
+        // For data construct regions, only replace uses in contained compute
+        // regions.
+        if (insideAccComputeRegion(use.getOwner())) {
+          use.set(replacement);
+        }
+      } else {
+        use.set(replacement);
+      }
+    }
+  }
+}
+
 template <typename Op>
 static void collectAndReplaceInRegion(Op &op, bool hostToDevice) {
   llvm::SmallVector<std::pair<Value, Value>> values;
@@ -48,26 +79,35 @@ static void collectAndReplaceInRegion(Op &op, bool hostToDevice) {
     collectPtrs(op.getPrivateOperands(), values, hostToDevice);
   } else {
     collectPtrs(op.getDataClauseOperands(), values, hostToDevice);
-    if constexpr (!std::is_same_v<Op, acc::KernelsOp>) {
+    if constexpr (!std::is_same_v<Op, acc::KernelsOp> &&
+                  !std::is_same_v<Op, acc::DataOp> &&
+                  !std::is_same_v<Op, acc::DeclareOp>) {
       collectPtrs(op.getReductionOperands(), values, hostToDevice);
       collectPtrs(op.getGangPrivateOperands(), values, hostToDevice);
       collectPtrs(op.getGangFirstPrivateOperands(), values, hostToDevice);
     }
   }
 
   for (auto p : values)
-    replaceAllUsesInRegionWith(std::get<0>(p), std::get<1>(p), op.getRegion());
+    replaceAllUsesInAccComputeRegionsWith<Op>(std::get<0>(p), std::get<1>(p),
+                                              op.getRegion());
 }
 
-struct LegalizeDataInRegion
-    : public acc::impl::LegalizeDataInRegionBase<LegalizeDataInRegion> {
+class LegalizeDataValuesInRegion
+    : public acc::impl::LegalizeDataValuesInRegionBase<
+          LegalizeDataValuesInRegion> {
+public:
+  using LegalizeDataValuesInRegionBase<
+      LegalizeDataValuesInRegion>::LegalizeDataValuesInRegionBase;
 
   void runOnOperation() override {
     func::FuncOp funcOp = getOperation();
     bool replaceHostVsDevice = this->hostToDevice.getValue();
 
     funcOp.walk([&](Operation *op) {
-      if (!isa<ACC_COMPUTE_CONSTRUCT_OPS>(*op) && !isa<acc::LoopOp>(*op))
+      if (!isa<ACC_COMPUTE_CONSTRUCT_AND_LOOP_OPS>(*op) &&
+          !(isa<ACC_DATA_CONSTRUCT_STRUCTURED_OPS>(*op) &&
+            applyToAccDataConstruct))
         return;
 
       if (auto parallelOp = dyn_cast<acc::ParallelOp>(*op)) {
@@ -78,14 +118,15 @@ struct LegalizeDataInRegion
         collectAndReplaceInRegion(kernelsOp, replaceHostVsDevice);
       } else if (auto loopOp = dyn_cast<acc::LoopOp>(*op)) {
         collectAndReplaceInRegion(loopOp, replaceHostVsDevice);
+      } else if (auto dataOp = dyn_cast<acc::DataOp>(*op)) {
+        collectAndReplaceInRegion(dataOp, replaceHostVsDevice);
+      } else if (auto declareOp = dyn_cast<acc::DeclareOp>(*op)) {
+        collectAndReplaceInRegion(declareOp, replaceHostVsDevice);
+      } else {
+        llvm_unreachable("unsupported acc region op");
       }
     });
   }
 };
 
 } // end anonymous namespace
-
-std::unique_ptr<OperationPass<func::FuncOp>>
-mlir::acc::createLegalizeDataInRegion() {
-  return std::make_unique<LegalizeDataInRegion>();
-}

mlir/test/Dialect/OpenACC/legalize-data.mlir

@@ -1,5 +1,5 @@
-// RUN: mlir-opt -split-input-file --openacc-legalize-data %s | FileCheck %s --check-prefixes=CHECK,DEVICE
-// RUN: mlir-opt -split-input-file --openacc-legalize-data=host-to-device=false %s | FileCheck %s --check-prefixes=CHECK,HOST
+// RUN: mlir-opt -split-input-file --openacc-legalize-data-values %s | FileCheck %s --check-prefixes=CHECK,DEVICE
+// RUN: mlir-opt -split-input-file --openacc-legalize-data-values=host-to-device=false %s | FileCheck %s --check-prefixes=CHECK,HOST
 
 func.func @test(%a: memref<10xf32>, %i : index) {
   %create = acc.create varPtr(%a : memref<10xf32>) -> memref<10xf32>
@@ -61,6 +61,32 @@ func.func @test(%a: memref<10xf32>, %i : index) {
 
 // -----
 
+func.func @test(%a: memref<10xf32>, %i : index) {
+  %create = acc.create varPtr(%a : memref<10xf32>) -> memref<10xf32>
+  acc.data dataOperands(%create : memref<10xf32>) {
+    %c0 = arith.constant 0.000000e+00 : f32
+    memref.store %c0, %a[%i] : memref<10xf32>
+    acc.serial {
+      %cs = memref.load %a[%i] : memref<10xf32>
+      acc.yield
+    }
+    acc.terminator
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @test
+// CHECK-SAME: (%[[A:.*]]: memref<10xf32>, %[[I:.*]]: index)
+// CHECK: %[[CREATE:.*]] = acc.create varPtr(%[[A]] : memref<10xf32>) -> memref<10xf32>
+// CHECK: acc.data dataOperands(%[[CREATE]] : memref<10xf32>) {
+// CHECK: memref.store %{{.*}}, %[[A]][%[[I]]] : memref<10xf32>
+// DEVICE:   %{{.*}} = memref.load %[[CREATE]][%[[I]]] : memref<10xf32>
+// HOST:    %{{.*}} = memref.load %[[A]][%[[I]]] : memref<10xf32>
+// CHECK:   acc.terminator
+// CHECK: }
+
+// -----
+
 func.func @test(%a: memref<10xf32>) {
   %lb = arith.constant 0 : index
   %st = arith.constant 1 : index