188 lines · cpp
1//===- InlineHLFIRCopyIn.cpp - Inline hlfir.copy_in ops -------------------===//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// Transform hlfir.copy_in array operations into loop nests performing element9// per element assignments. For simplicity, the inlining is done for trivial10// data types when the copy_in does not require a corresponding copy_out and11// when the input array is not behind a pointer. This may change in the future.12//===----------------------------------------------------------------------===//13 14#include "flang/Optimizer/Builder/FIRBuilder.h"15#include "flang/Optimizer/Builder/HLFIRTools.h"16#include "flang/Optimizer/Dialect/FIRType.h"17#include "flang/Optimizer/HLFIR/HLFIROps.h"18#include "flang/Optimizer/OpenMP/Passes.h"19#include "mlir/IR/PatternMatch.h"20#include "mlir/Support/LLVM.h"21#include "mlir/Transforms/GreedyPatternRewriteDriver.h"22 23namespace hlfir {24#define GEN_PASS_DEF_INLINEHLFIRCOPYIN25#include "flang/Optimizer/HLFIR/Passes.h.inc"26} // namespace hlfir27 28#define DEBUG_TYPE "inline-hlfir-copy-in"29 30static llvm::cl::opt<bool> noInlineHLFIRCopyIn(31 "no-inline-hlfir-copy-in",32 llvm::cl::desc("Do not inline hlfir.copy_in operations"),33 llvm::cl::init(false));34 35namespace {36class InlineCopyInConversion : public mlir::OpRewritePattern<hlfir::CopyInOp> {37public:38 using mlir::OpRewritePattern<hlfir::CopyInOp>::OpRewritePattern;39 40 llvm::LogicalResult41 matchAndRewrite(hlfir::CopyInOp copyIn,42 mlir::PatternRewriter &rewriter) const override;43};44 45llvm::LogicalResult46InlineCopyInConversion::matchAndRewrite(hlfir::CopyInOp copyIn,47 mlir::PatternRewriter &rewriter) const {48 fir::FirOpBuilder builder(rewriter, copyIn.getOperation());49 mlir::Location loc = copyIn.getLoc();50 hlfir::Entity inputVariable{copyIn.getVar()};51 mlir::Type resultAddrType = copyIn.getCopiedIn().getType();52 if (!fir::isa_trivial(inputVariable.getFortranElementType()))53 return rewriter.notifyMatchFailure(copyIn,54 "CopyInOp's data type is not trivial");55 56 // There should be exactly one user of WasCopied - the corresponding57 // CopyOutOp.58 if (!copyIn.getWasCopied().hasOneUse())59 return rewriter.notifyMatchFailure(60 copyIn, "CopyInOp's WasCopied has no single user");61 // The copy out should always be present, either to actually copy or just62 // deallocate memory.63 auto copyOut = mlir::dyn_cast<hlfir::CopyOutOp>(64 copyIn.getWasCopied().user_begin().getCurrent().getUser());65 66 if (!copyOut)67 return rewriter.notifyMatchFailure(copyIn,68 "CopyInOp has no direct CopyOut");69 70 if (mlir::cast<fir::BaseBoxType>(resultAddrType).isAssumedRank())71 return rewriter.notifyMatchFailure(copyIn,72 "The result array is assumed-rank");73 74 // Only inline the copy_in when copy_out does not need to be done, i.e. in75 // case of intent(in).76 if (copyOut.getVar())77 return rewriter.notifyMatchFailure(copyIn, "CopyIn needs a copy-out");78 79 inputVariable =80 hlfir::derefPointersAndAllocatables(loc, builder, inputVariable);81 mlir::Type sequenceType =82 hlfir::getFortranElementOrSequenceType(inputVariable.getType());83 fir::BoxType resultBoxType = fir::BoxType::get(sequenceType);84 mlir::Value isContiguous =85 fir::IsContiguousBoxOp::create(builder, loc, inputVariable);86 mlir::Operation::result_range results =87 builder88 .genIfOp(loc, {resultBoxType, builder.getI1Type()}, isContiguous,89 /*withElseRegion=*/true)90 .genThen([&]() {91 mlir::Value result = inputVariable;92 if (fir::isPointerType(inputVariable.getType())) {93 result = fir::ReboxOp::create(builder, loc, resultBoxType,94 inputVariable, mlir::Value{},95 mlir::Value{});96 }97 fir::ResultOp::create(98 builder, loc,99 mlir::ValueRange{result, builder.createBool(loc, false)});100 })101 .genElse([&] {102 mlir::Value shape = hlfir::genShape(loc, builder, inputVariable);103 llvm::SmallVector<mlir::Value> extents =104 hlfir::getIndexExtents(loc, builder, shape);105 llvm::StringRef tmpName{".tmp.copy_in"};106 llvm::SmallVector<mlir::Value> lenParams;107 mlir::Value alloc = builder.createHeapTemporary(108 loc, sequenceType, tmpName, extents, lenParams);109 110 auto declareOp = hlfir::DeclareOp::create(builder, loc, alloc,111 tmpName, shape, lenParams,112 /*dummy_scope=*/nullptr,113 /*storage=*/nullptr,114 /*storage_offset=*/0);115 hlfir::Entity temp{declareOp.getBase()};116 hlfir::LoopNest loopNest =117 hlfir::genLoopNest(loc, builder, extents, /*isUnordered=*/true,118 flangomp::shouldUseWorkshareLowering(copyIn),119 /*couldVectorize=*/false);120 builder.setInsertionPointToStart(loopNest.body);121 hlfir::Entity elem = hlfir::getElementAt(122 loc, builder, inputVariable, loopNest.oneBasedIndices);123 elem = hlfir::loadTrivialScalar(loc, builder, elem);124 hlfir::Entity tempElem = hlfir::getElementAt(125 loc, builder, temp, loopNest.oneBasedIndices);126 hlfir::AssignOp::create(builder, loc, elem, tempElem);127 builder.setInsertionPointAfter(loopNest.outerOp);128 129 mlir::Value result;130 // Make sure the result is always a boxed array by boxing it131 // ourselves if need be.132 if (mlir::isa<fir::BaseBoxType>(temp.getType())) {133 result = temp;134 } else {135 fir::ReferenceType refTy =136 fir::ReferenceType::get(temp.getElementOrSequenceType());137 mlir::Value refVal = builder.createConvert(loc, refTy, temp);138 result = fir::EmboxOp::create(builder, loc, resultBoxType, refVal,139 shape);140 }141 142 fir::ResultOp::create(143 builder, loc,144 mlir::ValueRange{result, builder.createBool(loc, true)});145 })146 .getResults();147 148 mlir::OpResult resultBox = results[0];149 mlir::OpResult needsCleanup = results[1];150 151 // Prepare the corresponding copyOut to free the temporary if it is required152 auto alloca = fir::AllocaOp::create(builder, loc, resultBox.getType());153 auto store = fir::StoreOp::create(builder, loc, resultBox, alloca);154 rewriter.startOpModification(copyOut);155 copyOut->setOperand(0, store.getMemref());156 copyOut->setOperand(1, needsCleanup);157 rewriter.finalizeOpModification(copyOut);158 159 rewriter.replaceOp(copyIn, {resultBox, builder.genNot(loc, isContiguous)});160 return mlir::success();161}162 163class InlineHLFIRCopyInPass164 : public hlfir::impl::InlineHLFIRCopyInBase<InlineHLFIRCopyInPass> {165public:166 void runOnOperation() override {167 mlir::MLIRContext *context = &getContext();168 169 mlir::GreedyRewriteConfig config;170 // Prevent the pattern driver from merging blocks.171 config.setRegionSimplificationLevel(172 mlir::GreedySimplifyRegionLevel::Disabled);173 174 mlir::RewritePatternSet patterns(context);175 if (!noInlineHLFIRCopyIn) {176 patterns.insert<InlineCopyInConversion>(context);177 }178 179 if (mlir::failed(mlir::applyPatternsGreedily(180 getOperation(), std::move(patterns), config))) {181 mlir::emitError(getOperation()->getLoc(),182 "failure in hlfir.copy_in inlining");183 signalPassFailure();184 }185 }186};187} // namespace188