| //===-- PreCGRewrite.cpp --------------------------------------------------===// |
| // |
| // 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 |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/ |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "flang/Optimizer/CodeGen/CodeGen.h" |
| |
| #include "flang/Optimizer/Dialect/FIRCG/CGOps.h" |
| #include "flang/Optimizer/Dialect/FIRDialect.h" |
| #include "flang/Optimizer/Dialect/FIROps.h" |
| #include "flang/Optimizer/Dialect/FIRType.h" |
| #include "flang/Optimizer/Dialect/Support/FIRContext.h" |
| #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" |
| #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" |
| #include "mlir/IR/Iterators.h" |
| #include "mlir/Interfaces/ControlFlowInterfaces.h" |
| #include "mlir/Transforms/DialectConversion.h" |
| #include "llvm/ADT/STLExtras.h" |
| #include "llvm/Support/Debug.h" |
| |
| namespace fir { |
| #define GEN_PASS_DEF_CODEGENREWRITE |
| #include "flang/Optimizer/CodeGen/CGPasses.h.inc" |
| } // namespace fir |
| |
| //===----------------------------------------------------------------------===// |
| // Codegen rewrite: rewriting of subgraphs of ops |
| //===----------------------------------------------------------------------===// |
| |
| #define DEBUG_TYPE "flang-codegen-rewrite" |
| |
| static void populateShape(llvm::SmallVectorImpl<mlir::Value> &vec, |
| fir::ShapeOp shape) { |
| vec.append(shape.getExtents().begin(), shape.getExtents().end()); |
| } |
| |
| // Operands of fir.shape_shift split into two vectors. |
| static void populateShapeAndShift(llvm::SmallVectorImpl<mlir::Value> &shapeVec, |
| llvm::SmallVectorImpl<mlir::Value> &shiftVec, |
| fir::ShapeShiftOp shift) { |
| for (auto i = shift.getPairs().begin(), endIter = shift.getPairs().end(); |
| i != endIter;) { |
| shiftVec.push_back(*i++); |
| shapeVec.push_back(*i++); |
| } |
| } |
| |
| static void populateShift(llvm::SmallVectorImpl<mlir::Value> &vec, |
| fir::ShiftOp shift) { |
| vec.append(shift.getOrigins().begin(), shift.getOrigins().end()); |
| } |
| |
| /// Report a shape or slice the rewrites cannot read. |
| /// |
| /// They are folded into the code-gen form through their defining op, so a |
| /// value that has none -- a block argument, or the result of a select -- |
| /// cannot be folded. Saying so is better than what the rewrites would |
| /// otherwise do: describe the whole array instead of the section it names. |
| static llvm::LogicalResult checkFoldableShapeAndSlice(mlir::Operation *op, |
| mlir::Value shape, |
| mlir::Value slice) { |
| if (shape && |
| !mlir::isa_and_nonnull<fir::ShapeOp, fir::ShapeShiftOp, fir::ShiftOp>( |
| shape.getDefiningOp())) |
| return op->emitOpError("shape operand is not defined by a fir.shape, " |
| "fir.shape_shift or fir.shift"); |
| if (slice && !mlir::isa_and_nonnull<fir::SliceOp>(slice.getDefiningOp())) |
| return op->emitOpError("slice operand is not defined by a fir.slice"); |
| return llvm::success(); |
| } |
| |
| // Helper to emit embox/rebox for OPTIONAL input inside a block |
| // guarded by a runtime presence check and to return an absent |
| // box when the input is not present. |
| template <typename OP> |
| static mlir::Value |
| emitOptionalBoxGuard(mlir::PatternRewriter &rewriter, OP op, |
| llvm::function_ref<mlir::Value()> buildPresent) { |
| mlir::Location loc = op.getLoc(); |
| mlir::Type boxType = op.getResult().getType(); |
| mlir::Value isPresent = fir::IsPresentOp::create( |
| rewriter, loc, rewriter.getI1Type(), op->getOperand(0)); |
| mlir::Block *condBlock = op->getBlock(); |
| mlir::Block *mergeBlock = rewriter.splitBlock(condBlock, op->getIterator()); |
| mergeBlock->addArgument(boxType, loc); |
| |
| mlir::Block *thenBlock = rewriter.createBlock(mergeBlock); |
| rewriter.setInsertionPointToStart(thenBlock); |
| mlir::Value present = buildPresent(); |
| mlir::cf::BranchOp::create(rewriter, loc, mergeBlock, |
| mlir::ValueRange{present}); |
| |
| mlir::Block *elseBlock = rewriter.createBlock(mergeBlock); |
| rewriter.setInsertionPointToStart(elseBlock); |
| mlir::Value absent = fir::AbsentOp::create(rewriter, loc, boxType); |
| mlir::cf::BranchOp::create(rewriter, loc, mergeBlock, |
| mlir::ValueRange{absent}); |
| |
| rewriter.setInsertionPointToEnd(condBlock); |
| mlir::cf::CondBranchOp::create(rewriter, loc, isPresent, thenBlock, |
| elseBlock); |
| rewriter.setInsertionPointToStart(mergeBlock); |
| return mergeBlock->getArgument(0); |
| } |
| |
| namespace { |
| |
| /// Convert fir.embox to the extended form where necessary. |
| /// |
| /// The embox operation can take arguments that specify multidimensional array |
| /// properties at runtime. These properties may be shared between distinct |
| /// objects that have the same properties. Before we lower these small DAGs to |
| /// LLVM-IR, we gather all the information into a single extended operation. For |
| /// example, |
| /// ``` |
| /// %1 = fir.shape_shift %4, %5 : (index, index) -> !fir.shapeshift<1> |
| /// %2 = fir.slice %6, %7, %8 : (index, index, index) -> !fir.slice<1> |
| /// %3 = fir.embox %0 (%1) [%2] : (!fir.ref<!fir.array<?xi32>>, |
| /// !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>> |
| /// ``` |
| /// can be rewritten as |
| /// ``` |
| /// %1 = fircg.ext_embox %0(%5) origin %4[%6, %7, %8] : |
| /// (!fir.ref<!fir.array<?xi32>>, index, index, index, index, index) -> |
| /// !fir.box<!fir.array<?xi32>> |
| /// ``` |
| class EmboxConversion : public mlir::OpRewritePattern<fir::EmboxOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| enum class RewriteKind { Dynamic, Static, DropOptional }; |
| |
| static llvm::FailureOr<RewriteKind> getRewriteKind(fir::EmboxOp embox) { |
| if (auto shapeVal = embox.getShape()) |
| return RewriteKind::Dynamic; |
| if (auto boxTy = mlir::dyn_cast<fir::BoxType>(embox.getType())) |
| if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(boxTy.getEleTy())) |
| if (!seqTy.hasDynamicExtents()) |
| return RewriteKind::Static; |
| if (embox.getOptional()) |
| return RewriteKind::DropOptional; |
| return llvm::failure(); |
| } |
| |
| llvm::LogicalResult |
| matchAndRewrite(fir::EmboxOp embox, |
| mlir::PatternRewriter &rewriter) const override { |
| llvm::FailureOr<RewriteKind> rewriteKind = getRewriteKind(embox); |
| if (llvm::failed(rewriteKind)) |
| return llvm::failure(); |
| if (*rewriteKind == RewriteKind::Dynamic && |
| llvm::failed(checkFoldableShapeAndSlice(embox, embox.getShape(), |
| embox.getSlice()))) |
| return llvm::failure(); |
| if (embox.getOptional()) { |
| mlir::Value newBox = emitOptionalBoxGuard(rewriter, embox, [&] { |
| return matchAndRewriteImpl(embox, rewriter, *rewriteKind)->getResult(0); |
| }); |
| rewriter.replaceOp(embox, newBox); |
| return mlir::success(); |
| } |
| mlir::Operation *newOp = matchAndRewriteImpl(embox, rewriter, *rewriteKind); |
| rewriter.replaceOp(embox, newOp); |
| return mlir::success(); |
| } |
| mlir::Operation *matchAndRewriteImpl(fir::EmboxOp embox, |
| mlir::PatternRewriter &rewriter, |
| RewriteKind rewriteKind) const { |
| switch (rewriteKind) { |
| case RewriteKind::Dynamic: |
| return rewriteDynamicShape(embox, rewriter, embox.getShape()); |
| case RewriteKind::Static: |
| return rewriteStaticShape(embox, rewriter, |
| fir::unwrapUntilSeqType(embox.getType())); |
| case RewriteKind::DropOptional: { |
| auto newEmbox = |
| llvm::cast<fir::EmboxOp>(rewriter.clone(*embox.getOperation())); |
| newEmbox.setOptional(false); |
| return newEmbox.getOperation(); |
| } |
| } |
| llvm_unreachable("all cases covered"); |
| return nullptr; |
| } |
| |
| mlir::Operation *rewriteStaticShape(fir::EmboxOp embox, |
| mlir::PatternRewriter &rewriter, |
| fir::SequenceType seqTy) const { |
| auto loc = embox.getLoc(); |
| llvm::SmallVector<mlir::Value> shapeOpers; |
| auto idxTy = rewriter.getIndexType(); |
| for (auto ext : seqTy.getShape()) { |
| auto iAttr = rewriter.getIndexAttr(ext); |
| auto extVal = |
| mlir::arith::ConstantOp::create(rewriter, loc, idxTy, iAttr); |
| shapeOpers.push_back(extVal); |
| } |
| auto xbox = fir::cg::XEmboxOp::create( |
| rewriter, loc, embox.getType(), embox.getMemref(), shapeOpers, |
| mlir::ValueRange{}, mlir::ValueRange{}, mlir::ValueRange{}, |
| mlir::ValueRange{}, embox.getTypeparams(), embox.getSourceBox(), |
| embox.getAllocatorIdxAttr()); |
| LLVM_DEBUG(llvm::dbgs() << "rewriting " << embox << " to " << xbox << '\n'); |
| return xbox.getOperation(); |
| } |
| |
| mlir::Operation *rewriteDynamicShape(fir::EmboxOp embox, |
| mlir::PatternRewriter &rewriter, |
| mlir::Value shapeVal) const { |
| auto loc = embox.getLoc(); |
| llvm::SmallVector<mlir::Value> shapeOpers; |
| llvm::SmallVector<mlir::Value> shiftOpers; |
| // matchAndRewrite has already reported a shape this cannot read. |
| mlir::Operation *shapeDef = shapeVal.getDefiningOp(); |
| if (auto shapeOp = mlir::dyn_cast_or_null<fir::ShapeOp>(shapeDef)) { |
| populateShape(shapeOpers, shapeOp); |
| } else { |
| auto shiftOp = mlir::dyn_cast_or_null<fir::ShapeShiftOp>(shapeDef); |
| assert(shiftOp && "shape is neither fir.shape nor fir.shape_shift"); |
| populateShapeAndShift(shapeOpers, shiftOpers, shiftOp); |
| } |
| llvm::SmallVector<mlir::Value> sliceOpers; |
| llvm::SmallVector<mlir::Value> subcompOpers; |
| llvm::SmallVector<mlir::Value> substrOpers; |
| if (auto s = embox.getSlice()) { |
| auto sliceOp = mlir::dyn_cast_or_null<fir::SliceOp>(s.getDefiningOp()); |
| // matchAndRewrite has already reported a slice this cannot read. |
| // Dropping it would describe the whole array instead of the section. |
| assert(sliceOp && "slice is not defined by a fir.slice"); |
| sliceOpers.assign(sliceOp.getTriples().begin(), |
| sliceOp.getTriples().end()); |
| subcompOpers.assign(sliceOp.getFields().begin(), |
| sliceOp.getFields().end()); |
| substrOpers.assign(sliceOp.getSubstr().begin(), |
| sliceOp.getSubstr().end()); |
| } |
| auto xbox = fir::cg::XEmboxOp::create( |
| rewriter, loc, embox.getType(), embox.getMemref(), shapeOpers, |
| shiftOpers, sliceOpers, subcompOpers, substrOpers, |
| embox.getTypeparams(), embox.getSourceBox(), |
| embox.getAllocatorIdxAttr()); |
| LLVM_DEBUG(llvm::dbgs() << "rewriting " << embox << " to " << xbox << '\n'); |
| return xbox.getOperation(); |
| } |
| }; |
| |
| /// Convert fir.rebox to the extended form where necessary. |
| /// |
| /// For example, |
| /// ``` |
| /// %5 = fir.rebox %3(%1) : (!fir.box<!fir.array<?xi32>>, !fir.shapeshift<1>) -> |
| /// !fir.box<!fir.array<?xi32>> |
| /// ``` |
| /// converted to |
| /// ``` |
| /// %5 = fircg.ext_rebox %3(%13) origin %12 : (!fir.box<!fir.array<?xi32>>, |
| /// index, index) -> !fir.box<!fir.array<?xi32>> |
| /// ``` |
| class ReboxConversion : public mlir::OpRewritePattern<fir::ReboxOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| llvm::LogicalResult |
| matchAndRewrite(fir::ReboxOp rebox, |
| mlir::PatternRewriter &rewriter) const override { |
| if (llvm::failed(checkFoldableShapeAndSlice(rebox, rebox.getShape(), |
| rebox.getSlice()))) |
| return llvm::failure(); |
| if (rebox.getOptional()) { |
| mlir::Value newBox = emitOptionalBoxGuard(rewriter, rebox, [&] { |
| return matchAndRewriteImpl(rebox, rewriter)->getResult(0); |
| }); |
| rewriter.replaceOp(rebox, newBox); |
| return mlir::success(); |
| } |
| mlir::Operation *newOp = matchAndRewriteImpl(rebox, rewriter); |
| rewriter.replaceOp(rebox, newOp); |
| return mlir::success(); |
| } |
| |
| mlir::Operation *matchAndRewriteImpl(fir::ReboxOp rebox, |
| mlir::PatternRewriter &rewriter) const { |
| auto loc = rebox.getLoc(); |
| llvm::SmallVector<mlir::Value> shapeOpers; |
| llvm::SmallVector<mlir::Value> shiftOpers; |
| if (auto shapeVal = rebox.getShape()) { |
| if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp())) |
| populateShape(shapeOpers, shapeOp); |
| else if (auto shapeShiftOp = |
| mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp())) |
| populateShapeAndShift(shapeOpers, shiftOpers, shapeShiftOp); |
| else { |
| auto shiftOp = mlir::dyn_cast<fir::ShiftOp>(shapeVal.getDefiningOp()); |
| assert(shiftOp && "unexpected shape operand type"); |
| populateShift(shiftOpers, shiftOp); |
| } |
| } |
| llvm::SmallVector<mlir::Value> sliceOpers; |
| llvm::SmallVector<mlir::Value> subcompOpers; |
| llvm::SmallVector<mlir::Value> substrOpers; |
| if (auto s = rebox.getSlice()) |
| if (auto sliceOp = |
| mlir::dyn_cast_or_null<fir::SliceOp>(s.getDefiningOp())) { |
| sliceOpers.append(sliceOp.getTriples().begin(), |
| sliceOp.getTriples().end()); |
| subcompOpers.append(sliceOp.getFields().begin(), |
| sliceOp.getFields().end()); |
| substrOpers.append(sliceOp.getSubstr().begin(), |
| sliceOp.getSubstr().end()); |
| } |
| |
| auto xRebox = fir::cg::XReboxOp::create( |
| rewriter, loc, rebox.getType(), rebox.getBox(), shapeOpers, shiftOpers, |
| sliceOpers, subcompOpers, substrOpers); |
| LLVM_DEBUG(llvm::dbgs() |
| << "rewriting " << rebox << " to " << xRebox << '\n'); |
| return xRebox.getOperation(); |
| } |
| }; |
| |
| /// Convert all fir.array_coor to the extended form. |
| /// |
| /// For example, |
| /// ``` |
| /// %4 = fir.array_coor %addr (%1) [%2] %0 : (!fir.ref<!fir.array<?xi32>>, |
| /// !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32> |
| /// ``` |
| /// converted to |
| /// ``` |
| /// %40 = fircg.ext_array_coor %addr(%9) origin %8[%4, %5, %6<%39> : |
| /// (!fir.ref<!fir.array<?xi32>>, index, index, index, index, index, index) -> |
| /// !fir.ref<i32> |
| /// ``` |
| class ArrayCoorConversion : public mlir::OpRewritePattern<fir::ArrayCoorOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| llvm::LogicalResult |
| matchAndRewrite(fir::ArrayCoorOp arrCoor, |
| mlir::PatternRewriter &rewriter) const override { |
| if (llvm::failed(checkFoldableShapeAndSlice(arrCoor, arrCoor.getShape(), |
| arrCoor.getSlice()))) |
| return llvm::failure(); |
| auto loc = arrCoor.getLoc(); |
| llvm::SmallVector<mlir::Value> shapeOpers; |
| llvm::SmallVector<mlir::Value> shiftOpers; |
| if (auto shapeVal = arrCoor.getShape()) { |
| if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp())) |
| populateShape(shapeOpers, shapeOp); |
| else if (auto shiftOp = |
| mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp())) |
| populateShapeAndShift(shapeOpers, shiftOpers, shiftOp); |
| else if (auto shiftOp = |
| mlir::dyn_cast<fir::ShiftOp>(shapeVal.getDefiningOp())) |
| populateShift(shiftOpers, shiftOp); |
| else |
| return mlir::failure(); |
| } |
| llvm::SmallVector<mlir::Value> sliceOpers; |
| llvm::SmallVector<mlir::Value> subcompOpers; |
| if (auto s = arrCoor.getSlice()) |
| if (auto sliceOp = |
| mlir::dyn_cast_or_null<fir::SliceOp>(s.getDefiningOp())) { |
| sliceOpers.append(sliceOp.getTriples().begin(), |
| sliceOp.getTriples().end()); |
| subcompOpers.append(sliceOp.getFields().begin(), |
| sliceOp.getFields().end()); |
| assert(sliceOp.getSubstr().empty() && |
| "Don't allow substring operations on array_coor. This " |
| "restriction may be lifted in the future."); |
| } |
| auto xArrCoor = fir::cg::XArrayCoorOp::create( |
| rewriter, loc, arrCoor.getType(), arrCoor.getMemref(), shapeOpers, |
| shiftOpers, sliceOpers, subcompOpers, arrCoor.getIndices(), |
| arrCoor.getTypeparams()); |
| LLVM_DEBUG(llvm::dbgs() |
| << "rewriting " << arrCoor << " to " << xArrCoor << '\n'); |
| rewriter.replaceOp(arrCoor, xArrCoor.getOperation()->getResults()); |
| return mlir::success(); |
| } |
| }; |
| |
| class DeclareOpConversion : public mlir::OpRewritePattern<fir::DeclareOp> { |
| bool preserveDeclare; |
| |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| DeclareOpConversion(mlir::MLIRContext *ctx, bool preserveDecl) |
| : OpRewritePattern(ctx), preserveDeclare(preserveDecl) {} |
| |
| llvm::LogicalResult |
| matchAndRewrite(fir::DeclareOp declareOp, |
| mlir::PatternRewriter &rewriter) const override { |
| if (!preserveDeclare) { |
| rewriter.replaceOp(declareOp, declareOp.getMemref()); |
| return mlir::success(); |
| } |
| auto loc = declareOp.getLoc(); |
| llvm::SmallVector<mlir::Value> shapeOpers; |
| llvm::SmallVector<mlir::Value> shiftOpers; |
| if (auto shapeVal = declareOp.getShape()) { |
| if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp())) |
| populateShape(shapeOpers, shapeOp); |
| else if (auto shiftOp = |
| mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp())) |
| populateShapeAndShift(shapeOpers, shiftOpers, shiftOp); |
| else if (auto shiftOp = |
| mlir::dyn_cast<fir::ShiftOp>(shapeVal.getDefiningOp())) |
| populateShift(shiftOpers, shiftOp); |
| else |
| return mlir::failure(); |
| } |
| // Extract dummy_arg_no attribute if present |
| mlir::IntegerAttr dummyArgNoAttr; |
| if (auto attr = declareOp->getAttrOfType<mlir::IntegerAttr>("dummy_arg_no")) |
| dummyArgNoAttr = attr; |
| // FIXME: Add FortranAttrs |
| auto xDeclOp = fir::cg::XDeclareOp::create( |
| rewriter, loc, declareOp.getType(), declareOp.getMemref(), shapeOpers, |
| shiftOpers, declareOp.getTypeparams(), declareOp.getDummyScope(), |
| declareOp.getStorage(), declareOp.getStorageOffset(), |
| declareOp.getUniqName(), declareOp.getDataAttrAttr(), dummyArgNoAttr); |
| LLVM_DEBUG(llvm::dbgs() |
| << "rewriting " << declareOp << " to " << xDeclOp << '\n'); |
| rewriter.replaceOp(declareOp, xDeclOp.getOperation()->getResults()); |
| return mlir::success(); |
| } |
| }; |
| |
| class DummyScopeOpConversion |
| : public mlir::OpRewritePattern<fir::DummyScopeOp> { |
| public: |
| using OpRewritePattern::OpRewritePattern; |
| |
| llvm::LogicalResult |
| matchAndRewrite(fir::DummyScopeOp dummyScopeOp, |
| mlir::PatternRewriter &rewriter) const override { |
| rewriter.replaceOpWithNewOp<fir::UndefOp>(dummyScopeOp, |
| dummyScopeOp.getType()); |
| return mlir::success(); |
| } |
| }; |
| |
| /// Simple DCE to erase fir.shape/shift/slice/unused shape operands after this |
| /// pass (fir.shape and like have no codegen). |
| /// mlir::RegionDCE is expensive and requires running |
| /// mlir::eraseUnreachableBlocks. It does things that are not needed here, like |
| /// removing unused block arguments. fir.shape/shift/slice cannot be block |
| /// arguments. |
| /// This helper does a naive backward walk of the IR. It is not even guaranteed |
| /// to walk blocks according to backward dominance, but that is good enough for |
| /// what is done here, fir.shape/shift/slice have no usages anymore. The |
| /// backward walk allows getting rid of most of the unused operands, it is not a |
| /// problem to leave some in the weird cases. |
| static void simpleDCE(mlir::RewriterBase &rewriter, mlir::Operation *op) { |
| op->walk<mlir::WalkOrder::PostOrder, mlir::ReverseIterator>( |
| [&](mlir::Operation *subOp) { |
| if (mlir::isOpTriviallyDead(subOp)) |
| rewriter.eraseOp(subOp); |
| }); |
| } |
| |
| /// A fir.shape, fir.shape_shift, fir.shift or fir.slice only describes an |
| /// array at compile time. The rewrites above read it through its defining op |
| /// and fold it into the code-gen form, so such a value never reaches codegen |
| /// and none of these types has an LLVM lowering. |
| static bool isCompileTimeOnly(mlir::Type type) { |
| return mlir::isa<fir::ShapeType, fir::ShapeShiftType, fir::ShiftType, |
| fir::SliceType>(type); |
| } |
| |
| /// Return the op \p value is built by, if this pass can rebuild it elsewhere |
| /// from its operands. A fir.slice naming components is excluded: its fields |
| /// are fir.field values, which are themselves compile-time only and so cannot |
| /// be passed along a branch. |
| static mlir::Operation *getRebuildableDefiningOp(mlir::Value value) { |
| mlir::Operation *def = value.getDefiningOp(); |
| if (!def) |
| return nullptr; |
| if (auto slice = mlir::dyn_cast<fir::SliceOp>(def)) |
| return slice.getFields().empty() ? def : nullptr; |
| if (mlir::isa<fir::ShapeOp, fir::ShapeShiftOp, fir::ShiftOp>(def)) |
| return def; |
| return nullptr; |
| } |
| |
| /// True if \p lhs and \p rhs can be rebuilt by one op taking one operand list. |
| static bool haveSameShape(mlir::Operation *lhs, mlir::Operation *rhs) { |
| if (lhs->getName() != rhs->getName() || |
| lhs->getNumOperands() != rhs->getNumOperands()) |
| return false; |
| return llvm::equal(lhs->getOperandTypes(), rhs->getOperandTypes()); |
| } |
| |
| /// Rebuild the compile-time-only block argument \p argIndex of \p block from |
| /// its operands, so the rewrites above find a defining op again. |
| /// |
| /// Every branch into the block passes a value of one of those types. The |
| /// operands behind it are integers, which can be block arguments, so they are |
| /// appended to the block and forwarded along each branch instead, and the |
| /// value is rebuilt from them at the top of the block. |
| static void rebuildBlockArgument(mlir::Block *block, unsigned argIndex) { |
| mlir::BlockArgument arg = block->getArgument(argIndex); |
| |
| // Collect one incoming edge per branch, keeping the operands each carries. |
| // A branch can name the block twice, so each edge is held separately. |
| struct Edge { |
| mlir::BranchOpInterface branch; |
| unsigned successorIndex; |
| }; |
| llvm::SmallVector<Edge> edges; |
| mlir::Operation *model = nullptr; |
| for (auto it = block->pred_begin(), e = block->pred_end(); it != e; ++it) { |
| auto branch = |
| mlir::dyn_cast<mlir::BranchOpInterface>((*it)->getTerminator()); |
| if (!branch) |
| return; |
| unsigned successorIndex = it.getSuccessorIndex(); |
| mlir::SuccessorOperands operands = |
| branch.getSuccessorOperands(successorIndex); |
| // An operand the branch itself produces has no value to read here. |
| if (operands.isOperandProduced(argIndex)) |
| return; |
| mlir::Operation *def = getRebuildableDefiningOp(operands[argIndex]); |
| if (!def) |
| return; |
| if (!model) |
| model = def; |
| else if (!haveSameShape(model, def)) |
| return; |
| edges.push_back({branch, successorIndex}); |
| } |
| if (edges.empty()) |
| return; |
| |
| // Take the operands as block arguments and forward them along each branch. |
| unsigned firstNewArg = block->getNumArguments(); |
| for (auto [type, value] : |
| llvm::zip_equal(model->getOperandTypes(), model->getOperands())) |
| block->addArgument(type, value.getLoc()); |
| for (Edge &edge : edges) { |
| mlir::SuccessorOperands operands = |
| edge.branch.getSuccessorOperands(edge.successorIndex); |
| operands.append( |
| getRebuildableDefiningOp(operands[argIndex])->getOperands()); |
| } |
| |
| // Rebuild the value at the top of the block and drop the old argument. |
| mlir::OpBuilder builder(block, block->begin()); |
| mlir::Operation *rebuilt = |
| builder.create(model->getLoc(), model->getName().getIdentifier(), |
| block->getArguments().drop_front(firstNewArg), |
| model->getResultTypes(), model->getAttrs()); |
| arg.replaceAllUsesWith(rebuilt->getResult(0)); |
| for (Edge &edge : edges) |
| edge.branch.getSuccessorOperands(edge.successorIndex).erase(argIndex); |
| block->eraseArgument(argIndex); |
| } |
| |
| /// Rebuild every compile-time-only block argument in \p op. |
| static void rebuildCompileTimeOnlyBlockArguments(mlir::Operation *op) { |
| llvm::SmallVector<mlir::Block *> blocks; |
| op->walk([&](mlir::Block *block) { |
| if (!block->isEntryBlock() && |
| llvm::any_of(block->getArgumentTypes(), isCompileTimeOnly)) |
| blocks.push_back(block); |
| }); |
| for (mlir::Block *block : blocks) |
| // Rebuilding drops an argument, so work back to front to keep the |
| // indices of the arguments still to be looked at. |
| for (unsigned i = block->getNumArguments(); i > 0; --i) |
| if (isCompileTimeOnly(block->getArgument(i - 1).getType())) |
| rebuildBlockArgument(block, i - 1); |
| } |
| |
| class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> { |
| public: |
| using CodeGenRewriteBase<CodeGenRewrite>::CodeGenRewriteBase; |
| |
| void runOnOperation() override final { |
| mlir::ModuleOp mod = getOperation(); |
| |
| rebuildCompileTimeOnlyBlockArguments(mod.getOperation()); |
| |
| auto &context = getContext(); |
| mlir::ConversionTarget target(context); |
| target.addLegalDialect<mlir::arith::ArithDialect, fir::FIROpsDialect, |
| fir::FIRCodeGenDialect, mlir::func::FuncDialect, |
| mlir::cf::ControlFlowDialect>(); |
| target.addIllegalOp<fir::ArrayCoorOp>(); |
| target.addIllegalOp<fir::ReboxOp>(); |
| target.addIllegalOp<fir::DeclareOp>(); |
| target.addIllegalOp<fir::DummyScopeOp>(); |
| target.addDynamicallyLegalOp<fir::EmboxOp>([](fir::EmboxOp embox) { |
| return llvm::failed(EmboxConversion::getRewriteKind(embox)); |
| }); |
| mlir::RewritePatternSet patterns(&context); |
| fir::populatePreCGRewritePatterns(patterns, preserveDeclare); |
| if (mlir::failed( |
| mlir::applyPartialConversion(mod, target, std::move(patterns)))) { |
| mlir::emitError(mlir::UnknownLoc::get(&context), |
| "error in running the pre-codegen conversions"); |
| signalPassFailure(); |
| return; |
| } |
| // Erase any residual (fir.shape, fir.slice...). |
| mlir::IRRewriter rewriter(&context); |
| simpleDCE(rewriter, mod.getOperation()); |
| } |
| }; |
| |
| } // namespace |
| |
| void fir::populatePreCGRewritePatterns(mlir::RewritePatternSet &patterns, |
| bool preserveDeclare) { |
| patterns.insert<EmboxConversion, ArrayCoorConversion, ReboxConversion, |
| DummyScopeOpConversion>(patterns.getContext()); |
| patterns.add<DeclareOpConversion>(patterns.getContext(), preserveDeclare); |
| } |
