flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp - llvm-project - Git at Google


            
            //===-- 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/Builder/Todo.h" // remove when TODO's are done

            
            #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/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());

            
            }


            // 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 (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;

            
                if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp())) {

            
                  populateShape(shapeOpers, shapeOp);

            
                } else {

            
                  auto shiftOp =

            
                      mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp());

            
                  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())

            
                  if (auto sliceOp =

            
                          mlir::dyn_cast_or_null<fir::SliceOp>(s.getDefiningOp())) {

            
                    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 (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 {

            
                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 and CudaAttrs

            
                auto xDeclOp = fir::cg::XDeclareOp::create(

            
                    rewriter, loc, declareOp.getType(), declareOp.getMemref(), shapeOpers,

            
                    shiftOpers, declareOp.getTypeparams(), declareOp.getDummyScope(),

            
                    declareOp.getStorage(), declareOp.getStorageOffset(),

            
                    declareOp.getUniqName(), 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);

            
                  });

            
            }


            class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {

            
            public:

            
              using CodeGenRewriteBase<CodeGenRewrite>::CodeGenRewriteBase;


              void runOnOperation() override final {

            
                mlir::ModuleOp 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);

            
            }
	//===-- 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/Builder/Todo.h" // remove when TODO's are done
	#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/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());
	}

	// 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 (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;
	if (auto shapeOp = mlir::dyn_cast<fir::ShapeOp>(shapeVal.getDefiningOp())) {
	populateShape(shapeOpers, shapeOp);
	} else {
	auto shiftOp =
	mlir::dyn_cast<fir::ShapeShiftOp>(shapeVal.getDefiningOp());
	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())
	if (auto sliceOp =
	mlir::dyn_cast_or_null<fir::SliceOp>(s.getDefiningOp())) {
	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 (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 {
	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 and CudaAttrs
	auto xDeclOp = fir::cg::XDeclareOp::create(
	rewriter, loc, declareOp.getType(), declareOp.getMemref(), shapeOpers,
	shiftOpers, declareOp.getTypeparams(), declareOp.getDummyScope(),
	declareOp.getStorage(), declareOp.getStorageOffset(),
	declareOp.getUniqName(), 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);
	});
	}

	class CodeGenRewrite : public fir::impl::CodeGenRewriteBase<CodeGenRewrite> {
	public:
	using CodeGenRewriteBase<CodeGenRewrite>::CodeGenRewriteBase;

	void runOnOperation() override final {
	mlir::ModuleOp 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);
	}