[mlir][openacc] Add legalize data pass for compute operation (#80351)

This patch adds a simple pass to replace the uses inside compute operation. It replaces the `varPtr` values with their corresponding `accPtr` values gathered through the dataClauseOperands. private and reductions variables are not included in this pass since they will normally be replace when they are materialized. --------- Co-authored-by: Slava Zakharin <szakharin@nvidia.com>

Author: clementval

flang/include/flang/Optimizer/Support/InitFIR.h

@@ -19,6 +19,7 @@
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
+#include "mlir/Dialect/OpenACC/Transforms/Passes.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassRegistry.h"
@@ -74,6 +75,7 @@ inline void loadDialects(mlir::MLIRContext &context) {
 /// Register the standard passes we use. This comes from registerAllPasses(),
 /// but is a smaller set since we aren't using many of the passes found there.
 inline void registerMLIRPassesForFortranTools() {
+ mlir::acc::registerOpenACCPasses();
  mlir::registerCanonicalizerPass();
  mlir::registerCSEPass();
  mlir::affine::registerAffineLoopFusionPass();

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

@@ -0,0 +1,24 @@
+// RUN: fir-opt -split-input-file --openacc-legalize-data %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>)
+ %1 = acc.copyin varPtr(%0#0 : !fir.ref<i32>) -> !fir.ref<i32> {dataClause = #acc<data_clause acc_copy>, name = "i"}
+ acc.parallel dataOperands(%1 : !fir.ref<i32>) {
+ %c0_i32 = arith.constant 0 : i32
+ hlfir.assign %c0_i32 to %0#0 : i32, !fir.ref<i32>
+ acc.yield
+ }
+ 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.parallel dataOperands(%[[COPYIN]] : !fir.ref<i32>) {
+// CHECK: %c0_i32 = arith.constant 0 : i32
+// CHECK: hlfir.assign %c0{{.*}} to %[[COPYIN]] : i32, !fir.ref<i32>
+// 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"}

mlir/include/mlir/Dialect/OpenACC/CMakeLists.txt

@@ -1,3 +1,5 @@
+add_subdirectory(Transforms)
+
 set(LLVM_TARGET_DEFINITIONS ${LLVM_MAIN_INCLUDE_DIR}/llvm/Frontend/OpenACC/ACC.td)
 mlir_tablegen(AccCommon.td --gen-directive-decl --directives-dialect=OpenACC)
 add_public_tablegen_target(acc_common_td)

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

@@ -0,0 +1,5 @@
+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name OpenACC)
+add_public_tablegen_target(MLIROpenACCPassIncGen)
+
+add_mlir_doc(Passes OpenACCPasses ./ -gen-pass-doc)

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

@@ -0,0 +1,40 @@
+//===- Passes.h - OpenACC Passes Construction and Registration ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_OPENACC_TRANSFORMS_PASSES_H
+#define MLIR_DIALECT_OPENACC_TRANSFORMS_PASSES_H
+
+#include "mlir/Dialect/LLVMIR/Transforms/AddComdats.h"
+#include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
+#include "mlir/Dialect/LLVMIR/Transforms/OptimizeForNVVM.h"
+#include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"
+#include "mlir/Dialect/LLVMIR/Transforms/TypeConsistency.h"
+#include "mlir/Pass/Pass.h"
+
+#define GEN_PASS_DECL
+#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
+
+namespace mlir {
+
+namespace func {
+class FuncOp;
+} // namespace func
+
+namespace acc {
+
+/// Create a pass to replace ssa values in region with device/host values.
+std::unique_ptr<OperationPass<func::FuncOp>> createLegalizeDataInRegion();
+
+/// Generate the code for registering conversion passes.
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/OpenACC/Transforms/Passes.h.inc"
+
+} // namespace acc
+} // namespace mlir
+
+#endif // MLIR_DIALECT_OPENACC_TRANSFORMS_PASSES_H

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

@@ -0,0 +1,28 @@
+//===-- Passes.td - OpenACC pass definition file -----------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_OPENACC_TRANSFORMS_PASSES
+#define MLIR_DIALECT_OPENACC_TRANSFORMS_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def LegalizeDataInRegion : Pass<"openacc-legalize-data", "mlir::func::FuncOp"> {
+ let summary = "Legalize the data in the compute region";
+ let description = [{
+ This pass replaces uses of varPtr in the compute region with their accPtr
+ gathered from the data clause operands.
+ }];
+ let options = [
+ Option<"hostToDevice", "host-to-device", "bool", "true",
+ "Replace varPtr uses with accPtr if true. Replace accPtr uses with "
+ "varPtr if false">
+ ];
+ let constructor = "::mlir::acc::createLegalizeDataInRegion()";
+}
+
+#endif // MLIR_DIALECT_OPENACC_TRANSFORMS_PASSES

mlir/include/mlir/InitAllPasses.h

@@ -34,6 +34,7 @@
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/Mesh/Transforms/Passes.h"
 #include "mlir/Dialect/NVGPU/Transforms/Passes.h"
+#include "mlir/Dialect/OpenACC/Transforms/Passes.h"
 #include "mlir/Dialect/SCF/Transforms/Passes.h"
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
 #include "mlir/Dialect/Shape/Transforms/Passes.h"
@@ -64,6 +65,7 @@ inline void registerAllPasses() {
  registerConversionPasses();
 
  // Dialect passes
+ acc::registerOpenACCPasses();
  affine::registerAffinePasses();
  amdgpu::registerAMDGPUPasses();
  registerAsyncPasses();

mlir/lib/Dialect/OpenACC/CMakeLists.txt

@@ -1,20 +1,2 @@
-add_mlir_dialect_library(MLIROpenACCDialect
- IR/OpenACC.cpp
-
- ADDITIONAL_HEADER_DIRS
- ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/OpenACC
-
- DEPENDS
- MLIROpenACCOpsIncGen
- MLIROpenACCEnumsIncGen
- MLIROpenACCAttributesIncGen
- MLIROpenACCOpsInterfacesIncGen
- MLIROpenACCTypeInterfacesIncGen
-
- LINK_LIBS PUBLIC
- MLIRIR
- MLIRLLVMDialect
- MLIRMemRefDialect
- MLIROpenACCMPCommon
- )
-
+add_subdirectory(IR)
+add_subdirectory(Transforms)

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

@@ -0,0 +1,20 @@
+add_mlir_dialect_library(MLIROpenACCDialect
+ OpenACC.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/OpenACC
+
+ DEPENDS
+ MLIROpenACCOpsIncGen
+ MLIROpenACCEnumsIncGen
+ MLIROpenACCAttributesIncGen
+ MLIROpenACCOpsInterfacesIncGen
+ MLIROpenACCTypeInterfacesIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRIR
+ MLIRLLVMDialect
+ MLIRMemRefDialect
+ MLIROpenACCMPCommon
+ )
+

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

@@ -0,0 +1,17 @@
+add_mlir_dialect_library(MLIROpenACCTransforms
+ LegalizeData.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/OpenACC
+
+ DEPENDS
+ MLIROpenACCPassIncGen
+ MLIROpenACCOpsIncGen
+ MLIROpenACCEnumsIncGen
+ MLIROpenACCAttributesIncGen
+ MLIROpenACCOpsInterfacesIncGen
+ MLIROpenACCTypeInterfacesIncGen
+
+ LINK_LIBS PUBLIC
+ MLIROpenACCDialect
+)