398 lines · cpp
1//===-- PreCGRewrite.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// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/10//11//===----------------------------------------------------------------------===//12 13#include "flang/Optimizer/CodeGen/CodeGen.h"14 15#include "flang/Optimizer/Builder/Todo.h" // remove when TODO's are done16#include "flang/Optimizer/Dialect/FIRCG/CGOps.h"17#include "flang/Optimizer/Dialect/FIRDialect.h"18#include "flang/Optimizer/Dialect/FIROps.h"19#include "flang/Optimizer/Dialect/FIRType.h"20#include "flang/Optimizer/Dialect/Support/FIRContext.h"21#include "mlir/IR/Iterators.h"22#include "mlir/Transforms/DialectConversion.h"23#include "llvm/ADT/STLExtras.h"24#include "llvm/Support/Debug.h"25 26namespace fir {27#define GEN_PASS_DEF_CODEGENREWRITE28#include "flang/Optimizer/CodeGen/CGPasses.h.inc"29} // namespace fir30 31//===----------------------------------------------------------------------===//32// Codegen rewrite: rewriting of subgraphs of ops33//===----------------------------------------------------------------------===//34 35#define DEBUG_TYPE "flang-codegen-rewrite"36 37static void populateShape(llvm::SmallVectorImpl<mlir::Value> &vec,38 fir::ShapeOp shape) {39 vec.append(shape.getExtents().begin(), shape.getExtents().end());40}41 42// Operands of fir.shape_shift split into two vectors.43static void populateShapeAndShift(llvm::SmallVectorImpl<mlir::Value> &shapeVec,44 llvm::SmallVectorImpl<mlir::Value> &shiftVec,45 fir::ShapeShiftOp shift) {46 for (auto i = shift.getPairs().begin(), endIter = shift.getPairs().end();47 i != endIter;) {48 shiftVec.push_back(*i++);49 shapeVec.push_back(*i++);50 }51}52 53static void populateShift(llvm::SmallVectorImpl<mlir::Value> &vec,54 fir::ShiftOp shift) {55 vec.append(shift.getOrigins().begin(), shift.getOrigins().end());56}57 58namespace {59 60/// Convert fir.embox to the extended form where necessary.61///62/// The embox operation can take arguments that specify multidimensional array63/// properties at runtime. These properties may be shared between distinct64/// objects that have the same properties. Before we lower these small DAGs to65/// LLVM-IR, we gather all the information into a single extended operation. For66/// example,67/// ```68/// %1 = fir.shape_shift %4, %5 : (index, index) -> !fir.shapeshift<1>69/// %2 = fir.slice %6, %7, %8 : (index, index, index) -> !fir.slice<1>70/// %3 = fir.embox %0 (%1) [%2] : (!fir.ref<!fir.array<?xi32>>,71/// !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>72/// ```73/// can be rewritten as74/// ```75/// %1 = fircg.ext_embox %0(%5) origin %4[%6, %7, %8] :76/// (!fir.ref<!fir.array<?xi32>>, index, index, index, index, index) ->77/// !fir.box<!fir.array<?xi32>>78/// ```79class EmboxConversion : public mlir::OpRewritePattern<fir::EmboxOp> {80public:81 using OpRewritePattern::OpRewritePattern;82 83 llvm::LogicalResult84 matchAndRewrite(fir::EmboxOp embox,85 mlir::PatternRewriter &rewriter) const override {86 // If the embox does not include a shape, then do not convert it87 if (auto shapeVal = embox.getShape())88 return rewriteDynamicShape(embox, rewriter, shapeVal);89 if (mlir::isa<fir::ClassType>(embox.getType()))90 TODO(embox.getLoc(), "embox conversion for fir.class type");91 if (auto boxTy = mlir::dyn_cast<fir::BoxType>(embox.getType()))92 if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(boxTy.getEleTy()))93 if (!seqTy.hasDynamicExtents())94 return rewriteStaticShape(embox, rewriter, seqTy);95 return mlir::failure();96 }97 98 llvm::LogicalResult rewriteStaticShape(fir::EmboxOp embox,99 mlir::PatternRewriter &rewriter,100 fir::SequenceType seqTy) const {101 auto loc = embox.getLoc();102 llvm::SmallVector<mlir::Value> shapeOpers;103 auto idxTy = rewriter.getIndexType();104 for (auto ext : seqTy.getShape()) {105 auto iAttr = rewriter.getIndexAttr(ext);106 auto extVal =107 mlir::arith::ConstantOp::create(rewriter, loc, idxTy, iAttr);108 shapeOpers.push_back(extVal);109 }110 auto xbox = fir::cg::XEmboxOp::create(111 rewriter, loc, embox.getType(), embox.getMemref(), shapeOpers,112 mlir::ValueRange{}, mlir::ValueRange{}, mlir::ValueRange{},113 mlir::ValueRange{}, embox.getTypeparams(), embox.getSourceBox(),114 embox.getAllocatorIdxAttr());115 LLVM_DEBUG(llvm::dbgs() << "rewriting " << embox << " to " << xbox << '\n');116 rewriter.replaceOp(embox, xbox.getOperation()->getResults());117 return mlir::success();118 }119 120 llvm::LogicalResult rewriteDynamicShape(fir::EmboxOp embox,121 mlir::PatternRewriter &rewriter,122 mlir::Value shapeVal) const {123 auto loc = embox.getLoc();124 llvm::SmallVector<mlir::Value> shapeOpers;125 llvm::SmallVector<mlir::Value> shiftOpers;126 if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp())) {127 populateShape(shapeOpers, shapeOp);128 } else {129 auto shiftOp =130 mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp());131 assert(shiftOp && "shape is neither fir.shape nor fir.shape_shift");132 populateShapeAndShift(shapeOpers, shiftOpers, shiftOp);133 }134 llvm::SmallVector<mlir::Value> sliceOpers;135 llvm::SmallVector<mlir::Value> subcompOpers;136 llvm::SmallVector<mlir::Value> substrOpers;137 if (auto s = embox.getSlice())138 if (auto sliceOp =139 mlir::dyn_cast_or_null<fir::SliceOp>(s.getDefiningOp())) {140 sliceOpers.assign(sliceOp.getTriples().begin(),141 sliceOp.getTriples().end());142 subcompOpers.assign(sliceOp.getFields().begin(),143 sliceOp.getFields().end());144 substrOpers.assign(sliceOp.getSubstr().begin(),145 sliceOp.getSubstr().end());146 }147 auto xbox = fir::cg::XEmboxOp::create(148 rewriter, loc, embox.getType(), embox.getMemref(), shapeOpers,149 shiftOpers, sliceOpers, subcompOpers, substrOpers,150 embox.getTypeparams(), embox.getSourceBox(),151 embox.getAllocatorIdxAttr());152 LLVM_DEBUG(llvm::dbgs() << "rewriting " << embox << " to " << xbox << '\n');153 rewriter.replaceOp(embox, xbox.getOperation()->getResults());154 return mlir::success();155 }156};157 158/// Convert fir.rebox to the extended form where necessary.159///160/// For example,161/// ```162/// %5 = fir.rebox %3(%1) : (!fir.box<!fir.array<?xi32>>, !fir.shapeshift<1>) ->163/// !fir.box<!fir.array<?xi32>>164/// ```165/// converted to166/// ```167/// %5 = fircg.ext_rebox %3(%13) origin %12 : (!fir.box<!fir.array<?xi32>>,168/// index, index) -> !fir.box<!fir.array<?xi32>>169/// ```170class ReboxConversion : public mlir::OpRewritePattern<fir::ReboxOp> {171public:172 using OpRewritePattern::OpRewritePattern;173 174 llvm::LogicalResult175 matchAndRewrite(fir::ReboxOp rebox,176 mlir::PatternRewriter &rewriter) const override {177 auto loc = rebox.getLoc();178 llvm::SmallVector<mlir::Value> shapeOpers;179 llvm::SmallVector<mlir::Value> shiftOpers;180 if (auto shapeVal = rebox.getShape()) {181 if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp()))182 populateShape(shapeOpers, shapeOp);183 else if (auto shiftOp =184 mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp()))185 populateShapeAndShift(shapeOpers, shiftOpers, shiftOp);186 else if (auto shiftOp =187 mlir::dyn_cast<fir::ShiftOp>(shapeVal.getDefiningOp()))188 populateShift(shiftOpers, shiftOp);189 else190 return mlir::failure();191 }192 llvm::SmallVector<mlir::Value> sliceOpers;193 llvm::SmallVector<mlir::Value> subcompOpers;194 llvm::SmallVector<mlir::Value> substrOpers;195 if (auto s = rebox.getSlice())196 if (auto sliceOp =197 mlir::dyn_cast_or_null<fir::SliceOp>(s.getDefiningOp())) {198 sliceOpers.append(sliceOp.getTriples().begin(),199 sliceOp.getTriples().end());200 subcompOpers.append(sliceOp.getFields().begin(),201 sliceOp.getFields().end());202 substrOpers.append(sliceOp.getSubstr().begin(),203 sliceOp.getSubstr().end());204 }205 206 auto xRebox = fir::cg::XReboxOp::create(207 rewriter, loc, rebox.getType(), rebox.getBox(), shapeOpers, shiftOpers,208 sliceOpers, subcompOpers, substrOpers);209 LLVM_DEBUG(llvm::dbgs()210 << "rewriting " << rebox << " to " << xRebox << '\n');211 rewriter.replaceOp(rebox, xRebox.getOperation()->getResults());212 return mlir::success();213 }214};215 216/// Convert all fir.array_coor to the extended form.217///218/// For example,219/// ```220/// %4 = fir.array_coor %addr (%1) [%2] %0 : (!fir.ref<!fir.array<?xi32>>,221/// !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>222/// ```223/// converted to224/// ```225/// %40 = fircg.ext_array_coor %addr(%9) origin %8[%4, %5, %6<%39> :226/// (!fir.ref<!fir.array<?xi32>>, index, index, index, index, index, index) ->227/// !fir.ref<i32>228/// ```229class ArrayCoorConversion : public mlir::OpRewritePattern<fir::ArrayCoorOp> {230public:231 using OpRewritePattern::OpRewritePattern;232 233 llvm::LogicalResult234 matchAndRewrite(fir::ArrayCoorOp arrCoor,235 mlir::PatternRewriter &rewriter) const override {236 auto loc = arrCoor.getLoc();237 llvm::SmallVector<mlir::Value> shapeOpers;238 llvm::SmallVector<mlir::Value> shiftOpers;239 if (auto shapeVal = arrCoor.getShape()) {240 if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp()))241 populateShape(shapeOpers, shapeOp);242 else if (auto shiftOp =243 mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp()))244 populateShapeAndShift(shapeOpers, shiftOpers, shiftOp);245 else if (auto shiftOp =246 mlir::dyn_cast<fir::ShiftOp>(shapeVal.getDefiningOp()))247 populateShift(shiftOpers, shiftOp);248 else249 return mlir::failure();250 }251 llvm::SmallVector<mlir::Value> sliceOpers;252 llvm::SmallVector<mlir::Value> subcompOpers;253 if (auto s = arrCoor.getSlice())254 if (auto sliceOp =255 mlir::dyn_cast_or_null<fir::SliceOp>(s.getDefiningOp())) {256 sliceOpers.append(sliceOp.getTriples().begin(),257 sliceOp.getTriples().end());258 subcompOpers.append(sliceOp.getFields().begin(),259 sliceOp.getFields().end());260 assert(sliceOp.getSubstr().empty() &&261 "Don't allow substring operations on array_coor. This "262 "restriction may be lifted in the future.");263 }264 auto xArrCoor = fir::cg::XArrayCoorOp::create(265 rewriter, loc, arrCoor.getType(), arrCoor.getMemref(), shapeOpers,266 shiftOpers, sliceOpers, subcompOpers, arrCoor.getIndices(),267 arrCoor.getTypeparams());268 LLVM_DEBUG(llvm::dbgs()269 << "rewriting " << arrCoor << " to " << xArrCoor << '\n');270 rewriter.replaceOp(arrCoor, xArrCoor.getOperation()->getResults());271 return mlir::success();272 }273};274 275class DeclareOpConversion : public mlir::OpRewritePattern<fir::DeclareOp> {276 bool preserveDeclare;277 278public:279 using OpRewritePattern::OpRewritePattern;280 DeclareOpConversion(mlir::MLIRContext *ctx, bool preserveDecl)281 : OpRewritePattern(ctx), preserveDeclare(preserveDecl) {}282 283 llvm::LogicalResult284 matchAndRewrite(fir::DeclareOp declareOp,285 mlir::PatternRewriter &rewriter) const override {286 if (!preserveDeclare) {287 rewriter.replaceOp(declareOp, declareOp.getMemref());288 return mlir::success();289 }290 auto loc = declareOp.getLoc();291 llvm::SmallVector<mlir::Value> shapeOpers;292 llvm::SmallVector<mlir::Value> shiftOpers;293 if (auto shapeVal = declareOp.getShape()) {294 if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp()))295 populateShape(shapeOpers, shapeOp);296 else if (auto shiftOp =297 mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp()))298 populateShapeAndShift(shapeOpers, shiftOpers, shiftOp);299 else if (auto shiftOp =300 mlir::dyn_cast<fir::ShiftOp>(shapeVal.getDefiningOp()))301 populateShift(shiftOpers, shiftOp);302 else303 return mlir::failure();304 }305 // Extract dummy_arg_no attribute if present306 mlir::IntegerAttr dummyArgNoAttr;307 if (auto attr = declareOp->getAttrOfType<mlir::IntegerAttr>("dummy_arg_no"))308 dummyArgNoAttr = attr;309 // FIXME: Add FortranAttrs and CudaAttrs310 auto xDeclOp = fir::cg::XDeclareOp::create(311 rewriter, loc, declareOp.getType(), declareOp.getMemref(), shapeOpers,312 shiftOpers, declareOp.getTypeparams(), declareOp.getDummyScope(),313 declareOp.getStorage(), declareOp.getStorageOffset(),314 declareOp.getUniqName(), dummyArgNoAttr);315 LLVM_DEBUG(llvm::dbgs()316 << "rewriting " << declareOp << " to " << xDeclOp << '\n');317 rewriter.replaceOp(declareOp, xDeclOp.getOperation()->getResults());318 return mlir::success();319 }320};321 322class DummyScopeOpConversion323 : public mlir::OpRewritePattern<fir::DummyScopeOp> {324public:325 using OpRewritePattern::OpRewritePattern;326 327 llvm::LogicalResult328 matchAndRewrite(fir::DummyScopeOp dummyScopeOp,329 mlir::PatternRewriter &rewriter) const override {330 rewriter.replaceOpWithNewOp<fir::UndefOp>(dummyScopeOp,331 dummyScopeOp.getType());332 return mlir::success();333 }334};335 336/// Simple DCE to erase fir.shape/shift/slice/unused shape operands after this337/// pass (fir.shape and like have no codegen).338/// mlir::RegionDCE is expensive and requires running339/// mlir::eraseUnreachableBlocks. It does things that are not needed here, like340/// removing unused block arguments. fir.shape/shift/slice cannot be block341/// arguments.342/// This helper does a naive backward walk of the IR. It is not even guaranteed343/// to walk blocks according to backward dominance, but that is good enough for344/// what is done here, fir.shape/shift/slice have no usages anymore. The345/// backward walk allows getting rid of most of the unused operands, it is not a346/// problem to leave some in the weird cases.347static void simpleDCE(mlir::RewriterBase &rewriter, mlir::Operation *op) {348 op->walk<mlir::WalkOrder::PostOrder, mlir::ReverseIterator>(349 [&](mlir::Operation *subOp) {350 if (mlir::isOpTriviallyDead(subOp))351 rewriter.eraseOp(subOp);352 });353}354 355class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {356public:357 using CodeGenRewriteBase<CodeGenRewrite>::CodeGenRewriteBase;358 359 void runOnOperation() override final {360 mlir::ModuleOp mod = getOperation();361 362 auto &context = getContext();363 mlir::ConversionTarget target(context);364 target.addLegalDialect<mlir::arith::ArithDialect, fir::FIROpsDialect,365 fir::FIRCodeGenDialect, mlir::func::FuncDialect>();366 target.addIllegalOp<fir::ArrayCoorOp>();367 target.addIllegalOp<fir::ReboxOp>();368 target.addIllegalOp<fir::DeclareOp>();369 target.addIllegalOp<fir::DummyScopeOp>();370 target.addDynamicallyLegalOp<fir::EmboxOp>([](fir::EmboxOp embox) {371 return !(embox.getShape() ||372 mlir::isa<fir::SequenceType>(373 mlir::cast<fir::BaseBoxType>(embox.getType()).getEleTy()));374 });375 mlir::RewritePatternSet patterns(&context);376 fir::populatePreCGRewritePatterns(patterns, preserveDeclare);377 if (mlir::failed(378 mlir::applyPartialConversion(mod, target, std::move(patterns)))) {379 mlir::emitError(mlir::UnknownLoc::get(&context),380 "error in running the pre-codegen conversions");381 signalPassFailure();382 return;383 }384 // Erase any residual (fir.shape, fir.slice...).385 mlir::IRRewriter rewriter(&context);386 simpleDCE(rewriter, mod.getOperation());387 }388};389 390} // namespace391 392void fir::populatePreCGRewritePatterns(mlir::RewritePatternSet &patterns,393 bool preserveDeclare) {394 patterns.insert<EmboxConversion, ArrayCoorConversion, ReboxConversion,395 DummyScopeOpConversion>(patterns.getContext());396 patterns.add<DeclareOpConversion>(patterns.getContext(), preserveDeclare);397}398