185 lines · cpp
1//===- ACCRecipeBufferization.cpp -----------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// Bufferize OpenACC recipes that yield fir.box<T> to operate on10// fir.ref<fir.box<T>> and update uses accordingly.11//12//===----------------------------------------------------------------------===//13 14#include "flang/Optimizer/Dialect/FIROps.h"15#include "flang/Optimizer/OpenACC/Passes.h"16#include "mlir/Dialect/OpenACC/OpenACC.h"17#include "mlir/IR/Block.h"18#include "mlir/IR/Builders.h"19#include "mlir/IR/BuiltinOps.h"20#include "mlir/IR/SymbolTable.h"21#include "mlir/IR/Value.h"22#include "mlir/IR/Visitors.h"23#include "llvm/ADT/TypeSwitch.h"24 25namespace fir::acc {26#define GEN_PASS_DEF_ACCRECIPEBUFFERIZATION27#include "flang/Optimizer/OpenACC/Passes.h.inc"28} // namespace fir::acc29 30namespace {31 32class BufferizeInterface {33public:34 static std::optional<mlir::Type> mustBufferize(mlir::Type recipeType) {35 if (auto boxTy = llvm::dyn_cast<fir::BaseBoxType>(recipeType))36 return fir::ReferenceType::get(boxTy);37 return std::nullopt;38 }39 40 static mlir::Operation *load(mlir::OpBuilder &builder, mlir::Location loc,41 mlir::Value value) {42 return fir::LoadOp::create(builder, loc, value);43 }44 45 static mlir::Value placeInMemory(mlir::OpBuilder &builder, mlir::Location loc,46 mlir::Value value) {47 auto alloca = fir::AllocaOp::create(builder, loc, value.getType());48 fir::StoreOp::create(builder, loc, value, alloca);49 return alloca;50 }51};52 53static void bufferizeRegionArgsAndYields(mlir::Region ®ion,54 mlir::Location loc, mlir::Type oldType,55 mlir::Type newType) {56 if (region.empty())57 return;58 59 mlir::OpBuilder builder(®ion);60 for (mlir::BlockArgument arg : region.getArguments()) {61 if (arg.getType() == oldType) {62 arg.setType(newType);63 if (!arg.use_empty()) {64 mlir::Operation *loadOp = BufferizeInterface::load(builder, loc, arg);65 arg.replaceAllUsesExcept(loadOp->getResult(0), loadOp);66 }67 }68 }69 if (auto yield =70 llvm::dyn_cast<mlir::acc::YieldOp>(region.back().getTerminator())) {71 llvm::SmallVector<mlir::Value> newOperands;72 newOperands.reserve(yield.getNumOperands());73 bool changed = false;74 for (mlir::Value oldYieldArg : yield.getOperands()) {75 if (oldYieldArg.getType() == oldType) {76 builder.setInsertionPoint(yield);77 mlir::Value alloca =78 BufferizeInterface::placeInMemory(builder, loc, oldYieldArg);79 newOperands.push_back(alloca);80 changed = true;81 } else {82 newOperands.push_back(oldYieldArg);83 }84 }85 if (changed)86 yield->setOperands(newOperands);87 }88}89 90template <typename OpTy>91static void updateRecipeUse(mlir::ValueRange operands,92 llvm::StringRef recipeSymName,93 mlir::Operation *computeOp) {94 for (auto operand : operands) {95 auto op = operand.getDefiningOp<OpTy>();96 if (!op || !op.getRecipe().has_value() ||97 op.getRecipeAttr().getLeafReference() != recipeSymName)98 continue;99 100 mlir::Location loc = op->getLoc();101 102 mlir::OpBuilder builder(op);103 builder.setInsertionPointAfterValue(op.getVar());104 mlir::Value alloca =105 BufferizeInterface::placeInMemory(builder, loc, op.getVar());106 op.getVarMutable().assign(alloca);107 op.getAccVar().setType(alloca.getType());108 109 mlir::Value oldRes = op.getAccVar();110 llvm::SmallVector<mlir::Operation *> users(oldRes.getUsers().begin(),111 oldRes.getUsers().end());112 for (mlir::Operation *useOp : users) {113 if (useOp == computeOp)114 continue;115 builder.setInsertionPoint(useOp);116 mlir::Operation *load = BufferizeInterface::load(builder, loc, oldRes);117 useOp->replaceUsesOfWith(oldRes, load->getResult(0));118 }119 }120}121 122class ACCRecipeBufferization123 : public fir::acc::impl::ACCRecipeBufferizationBase<124 ACCRecipeBufferization> {125public:126 void runOnOperation() override {127 mlir::ModuleOp module = getOperation();128 129 llvm::SmallVector<llvm::StringRef> recipeNames;130 module.walk([&](mlir::Operation *recipe) {131 llvm::TypeSwitch<mlir::Operation *, void>(recipe)132 .Case<mlir::acc::PrivateRecipeOp, mlir::acc::FirstprivateRecipeOp,133 mlir::acc::ReductionRecipeOp>([&](auto recipe) {134 mlir::Type oldType = recipe.getType();135 auto bufferizedType =136 BufferizeInterface::mustBufferize(recipe.getType());137 if (!bufferizedType)138 return;139 recipe.setTypeAttr(mlir::TypeAttr::get(*bufferizedType));140 mlir::Location loc = recipe.getLoc();141 using RecipeOp = decltype(recipe);142 bufferizeRegionArgsAndYields(recipe.getInitRegion(), loc, oldType,143 *bufferizedType);144 if constexpr (std::is_same_v<RecipeOp,145 mlir::acc::FirstprivateRecipeOp>)146 bufferizeRegionArgsAndYields(recipe.getCopyRegion(), loc, oldType,147 *bufferizedType);148 if constexpr (std::is_same_v<RecipeOp,149 mlir::acc::ReductionRecipeOp>)150 bufferizeRegionArgsAndYields(recipe.getCombinerRegion(), loc,151 oldType, *bufferizedType);152 bufferizeRegionArgsAndYields(recipe.getDestroyRegion(), loc,153 oldType, *bufferizedType);154 recipeNames.push_back(recipe.getSymName());155 });156 });157 if (recipeNames.empty())158 return;159 160 module.walk([&](mlir::Operation *op) {161 llvm::TypeSwitch<mlir::Operation *, void>(op)162 .Case<mlir::acc::LoopOp, mlir::acc::ParallelOp, mlir::acc::SerialOp>(163 [&](auto computeOp) {164 for (llvm::StringRef recipeName : recipeNames) {165 if (!computeOp.getPrivateOperands().empty())166 updateRecipeUse<mlir::acc::PrivateOp>(167 computeOp.getPrivateOperands(), recipeName, op);168 if (!computeOp.getFirstprivateOperands().empty())169 updateRecipeUse<mlir::acc::FirstprivateOp>(170 computeOp.getFirstprivateOperands(), recipeName, op);171 if (!computeOp.getReductionOperands().empty())172 updateRecipeUse<mlir::acc::ReductionOp>(173 computeOp.getReductionOperands(), recipeName, op);174 }175 });176 });177 }178};179 180} // namespace181 182std::unique_ptr<mlir::Pass> fir::acc::createACCRecipeBufferizationPass() {183 return std::make_unique<ACCRecipeBufferization>();184}185