[flang] Store KindMapping by value in FirOpBuilder

[tblah]

Summary

Previously only a constant reference was stored in the FirOpBuilder.
However, a lot of code was merged using

FirOpBuilder builder{rewriter, getKindMapping(mod)};

This is incorrect because the KindMapping returned will go out of scope
as soon as FirOpBuilder's constructor had run. This led to an infinite
loop running some tests using HLFIR (because the stack space containing
the kind mapping was re-used and corrupted).

One solution would have just been to fix the incorrect call sites,
however, as a large number of these had already made it past review, I
decided to instead change FirOpBuilder to store its own copy of the
KindMapping. This is not costly because nearly every time we construct a
KindMapping is exclusively to construct a FirOpBuilder. To make this
common pattern simpler, I added a new constructor to FirOpBuilder which
calls getKindMapping().

[tblah]

Changes: clang-format

[vzakhari]

flang/include/flang/Optimizer/Builder/FIRBuilder.h

59

The copy construction was cheap before this change, so it did not make sense to have a move constructor. Does it make sense to add it now?

[tblah]

Thanks for review

Changes: add a move constructor to FirOpBuilder

[vzakhari]

Thank you!

flang/include/flang/Optimizer/Builder/FIRBuilder.h

 
 #ifndef FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H
 #define FORTRAN_OPTIMIZER_BUILDER_FIRBUILDER_H
 
 #include "flang/Common/MathOptionsBase.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/FIRContext.h"
 #include "flang/Optimizer/Dialect/Support/KindMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "llvm/ADT/DenseMap.h"
 #include <optional>
+#include <utility>
 
 namespace fir {
 class AbstractArrayBox;
 class ExtendedValue;
 class MutableBoxValue;
 class BoxValue;
 
 //===----------------------------------------------------------------------===//
 // FirOpBuilder
 //===----------------------------------------------------------------------===//
 
 /// Extends the MLIR OpBuilder to provide methods for building common FIR
 /// patterns.
 class FirOpBuilder : public mlir::OpBuilder, public mlir::OpBuilder::Listener {
 public:
-  explicit FirOpBuilder(mlir::Operation *op, const fir::KindMapping &kindMap)
-      : OpBuilder{op, /*listener=*/this}, kindMap{kindMap} {}
-  explicit FirOpBuilder(mlir::OpBuilder &builder,
-                        const fir::KindMapping &kindMap)
-      : OpBuilder(builder), OpBuilder::Listener(), kindMap{kindMap} {
+  explicit FirOpBuilder(mlir::Operation *op, fir::KindMapping kindMap)
+      : OpBuilder{op, /*listener=*/this}, kindMap{std::move(kindMap)} {}
+  explicit FirOpBuilder(mlir::OpBuilder &builder, fir::KindMapping kindMap)
+      : OpBuilder(builder), OpBuilder::Listener(), kindMap{std::move(kindMap)} {
+    setListener(this);
+  }
+  explicit FirOpBuilder(mlir::OpBuilder &builder, mlir::ModuleOp mod)
+      : OpBuilder(builder), OpBuilder::Listener(),
+        kindMap{getKindMapping(mod)} {
     setListener(this);
   }
 
   // The listener self-reference has to be updated in case of copy-construction.
   FirOpBuilder(const FirOpBuilder &other)
       : OpBuilder(other), OpBuilder::Listener(), kindMap{other.kindMap},
         fastMathFlags{other.fastMathFlags} {
     setListener(this);
   }
 
+  FirOpBuilder(FirOpBuilder &&other)
+      : OpBuilder(other), OpBuilder::Listener(),
+        kindMap{std::move(other.kindMap)}, fastMathFlags{other.fastMathFlags} {
+    setListener(this);
+  }
+
   /// Get the current Region of the insertion point.
   mlir::Region &getRegion() { return *getBlock()->getParent(); }
 
   /// Get the current Module
   mlir::ModuleOp getModule() {
     return getRegion().getParentOfType<mlir::ModuleOp>();
   }
 
void notifyOperationInserted(mlir::Operation *op) override {
     setCommonAttributes(op);
   }
 
 private:
   /// Set attributes (e.g. FastMathAttr) to \p op operation
   /// based on the current attributes setting.
   void setCommonAttributes(mlir::Operation *op) const;
 
-  const KindMapping &kindMap;
+  KindMapping kindMap;
 
   /// FastMathFlags that need to be set for operations that support
   /// mlir::arith::FastMathAttr.
   mlir::arith::FastMathFlags fastMathFlags{};
 };
 
 } // namespace fir
 

flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp

if (options.useThunks) {
         } else if (auto embox = mlir::dyn_cast<EmboxProcOp>(op)) {
           // Rewrite all `fir.emboxproc` ops to either `fir.convert` or a thunk
           // as required.
           mlir::Type toTy = embox.getType().cast<BoxProcType>().getEleTy();
           rewriter.setInsertionPoint(embox);
           if (embox.getHost()) {
             // Create the thunk.
             auto module = embox->getParentOfType<mlir::ModuleOp>();
-            fir::KindMapping kindMap = getKindMapping(module);
-            FirOpBuilder builder(rewriter, kindMap);
+            FirOpBuilder builder(rewriter, module);
             auto loc = embox.getLoc();
             mlir::Type i8Ty = builder.getI8Type();
             mlir::Type i8Ptr = builder.getRefType(i8Ty);
             mlir::Type buffTy = SequenceType::get({32}, i8Ty);
             auto buffer = builder.create<AllocaOp>(loc, buffTy);
             mlir::Value closure =
                 builder.createConvert(loc, i8Ptr, embox.getHost());
             mlir::Value tramp = builder.createConvert(loc, i8Ptr, buffer);

flang/lib/Optimizer/CodeGen/TargetRewrite.cpp

if (!forcedTargetTriple.empty())
       fir::setTargetTriple(mod, forcedTargetTriple);
 
     auto specifics = fir::CodeGenSpecifics::get(
         mod.getContext(), fir::getTargetTriple(mod), fir::getKindMapping(mod));
     setMembers(specifics.get(), &rewriter);
 
     // We may need to call stacksave/stackrestore later, so
     // create the FuncOps beforehand.
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(mod));
+    fir::FirOpBuilder builder(rewriter, mod);
     builder.setInsertionPointToStart(mod.getBody());
     stackSaveFn = fir::factory::getLlvmStackSave(builder);
     stackRestoreFn = fir::factory::getLlvmStackRestore(builder);
 
     // Perform type conversion on signatures and call sites.
     if (mlir::failed(convertTypes(mod))) {
       mlir::emitError(mlir::UnknownLoc::get(&context),
                       "error in converting types to target abi");
for (auto e : llvm::enumerate(
                           index, charProcAttr))
                     mlir::emitError(loc, "tuple argument will be split even "
                                          "though it does not have the `" +
                                              charProcAttr + "` attribute");
                 }
               }
               mlir::Type funcPointerType = tuple.getType(0);
               mlir::Type lenType = tuple.getType(1);
-              fir::KindMapping kindMap = fir::getKindMapping(module);
-              fir::FirOpBuilder builder(*rewriter, kindMap);
+              fir::FirOpBuilder builder(*rewriter, module);
               auto [funcPointer, len] =
                   fir::factory::extractCharacterProcedureTuple(builder, loc,
                                                                oper);
               newInTys.push_back(funcPointerType);
               newOpers.push_back(funcPointer);
               trailingInTys.push_back(lenType);
               trailingOpers.push_back(len);
             } else {
if (!func.empty()) {
           // the original argument position. The second part of the pair is
           // appended after all the original arguments.
           auto newProcPointerArg = func.front().insertArgument(
               fixup.index, newInTys[fixup.index], loc);
           auto newLenArg =
               func.front().addArgument(trailingTys[fixup.second], loc);
           auto tupleType = oldArgTys[fixup.index - offset];
           rewriter->setInsertionPointToStart(&func.front());
-          fir::KindMapping kindMap = fir::getKindMapping(getModule());
-          fir::FirOpBuilder builder(*rewriter, kindMap);
+          fir::FirOpBuilder builder(*rewriter, getModule());
           auto tuple = fir::factory::createCharacterProcedureTuple(
               builder, loc, tupleType, newProcPointerArg, newLenArg);
           func.getArgument(fixup.index + 1).replaceAllUsesWith(tuple);
           func.front().eraseArgument(fixup.index + 1);
         } break;
         }
       }
     }

flang/lib/Optimizer/HLFIR/Transforms/BufferizeHLFIR.cpp

struct AsExprOpConversion : public mlir::OpConversionPattern<hlfir::AsExprOp> {
   using mlir::OpConversionPattern<hlfir::AsExprOp>::OpConversionPattern;
   explicit AsExprOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::AsExprOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::AsExprOp asExpr, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = asExpr->getLoc();
     auto module = asExpr->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, module);
     if (asExpr.isMove()) {
       // Move variable storage for the hlfir.expr buffer.
       mlir::Value bufferizedExpr = packageBufferizedExpr(
           loc, builder, hlfir::Entity{adaptor.getVar()}, adaptor.getMustFree());
       rewriter.replaceOp(asExpr, bufferizedExpr);
       return mlir::success();
     }
     // Otherwise, create a copy in a new buffer.
struct ShapeOfOpConversion
     : public mlir::OpConversionPattern<hlfir::ShapeOfOp> {
   using mlir::OpConversionPattern<hlfir::ShapeOfOp>::OpConversionPattern;
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::ShapeOfOp shapeOf, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = shapeOf.getLoc();
     mlir::ModuleOp mod = shapeOf->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(mod));
+    fir::FirOpBuilder builder(rewriter, mod);
 
     mlir::Value shape;
     hlfir::Entity bufferizedExpr{getBufferizedExprStorage(adaptor.getExpr())};
     if (bufferizedExpr.isVariable()) {
       shape = hlfir::genShape(loc, builder, bufferizedExpr);
     } else {
       // everything else failed so try to create a shape from static type info
       hlfir::ExprType exprTy =

flang/lib/Optimizer/HLFIR/Transforms/ConvertToFIR.cpp

public:
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::AssignOp assignOp,
                   mlir::PatternRewriter &rewriter) const override {
     mlir::Location loc = assignOp->getLoc();
     hlfir::Entity lhs(assignOp.getLhs());
     hlfir::Entity rhs(assignOp.getRhs());
     auto module = assignOp->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, module);
 
     if (rhs.getType().isa<hlfir::ExprType>()) {
       mlir::emitError(loc, "hlfir must be bufferized with --bufferize-hlfir "
                            "pass before being converted to FIR");
       return mlir::failure();
     }
     auto [rhsExv, rhsCleanUp] =
         hlfir::translateToExtendedValue(loc, builder, rhs);

flang/lib/Optimizer/HLFIR/Transforms/LowerHLFIROrderedAssignments.cpp

void OrderedAssignmentRewriter::generateSaveEntity(
   }
 }
 
 /// Lower an ordered assignment tree to fir.do_loop and hlfir.assign given
 /// a schedule.
 static void lower(hlfir::OrderedAssignmentTreeOpInterface root,
                   mlir::PatternRewriter &rewriter, hlfir::Schedule &schedule) {
   auto module = root->getParentOfType<mlir::ModuleOp>();
-  fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+  fir::FirOpBuilder builder(rewriter, module);
   OrderedAssignmentRewriter assignmentRewriter(builder, root);
   for (auto &run : schedule)
     assignmentRewriter.lowerRun(run);
   assignmentRewriter.cleanupSavedEntities();
 }
 
 /// Shared rewrite entry point for all the ordered assignment tree root
 /// operations. It calls the scheduler and then apply the schedule.

flang/lib/Optimizer/Transforms/AbstractResult.cpp

if constexpr (std::is_same_v<Op, fir::DispatchOp>) {
         newOp = rewriter.create<fir::DispatchOp>(
             loc, newResultTypes, rewriter.getStringAttr(op.getMethod()),
             op.getOperands()[0], newOperands, nullptr);
     }
 
     if (isResultBuiltinCPtr) {
       mlir::Value save = saveResult.getMemref();
       auto module = op->template getParentOfType<mlir::ModuleOp>();
-      fir::KindMapping kindMap = fir::getKindMapping(module);
-      FirOpBuilder builder(rewriter, kindMap);
+      FirOpBuilder builder(rewriter, module);
       mlir::Value saveAddr = fir::factory::genCPtrOrCFunptrAddr(
           builder, loc, save, result.getType());
       rewriter.create<fir::StoreOp>(loc, newOp->getResult(0), saveAddr);
     }
     op->dropAllReferences();
     rewriter.eraseOp(op);
     return mlir::success();
   }
if (auto *op = returnedValue.getDefiningOp())
         if (auto declare = mlir::dyn_cast_or_null<fir::DeclareOp>(
                 resultStorage.getDefiningOp()))
           resultStorage = declare.getMemref();
         // TODO: This should be generalized for derived types, and it is
         // architecture and OS dependent.
         if (fir::isa_builtin_cptr_type(returnedValue.getType())) {
           rewriter.eraseOp(load);
           auto module = ret->getParentOfType<mlir::ModuleOp>();
-          fir::KindMapping kindMap = fir::getKindMapping(module);
-          FirOpBuilder builder(rewriter, kindMap);
+          FirOpBuilder builder(rewriter, module);
           mlir::Value retAddr = fir::factory::genCPtrOrCFunptrAddr(
               builder, loc, resultStorage, returnedValue.getType());
           mlir::Value retValue = rewriter.create<fir::LoadOp>(
               loc, fir::unwrapRefType(retAddr.getType()), retAddr);
           rewriter.replaceOpWithNewOp<mlir::func::ReturnOp>(
               ret, mlir::ValueRange{retValue});
           return mlir::success();
         }

flang/lib/Optimizer/Transforms/ArrayValueCopy.cpp

if (factory::isAssumedSize(result)) {
     auto one = rewriter.create<mlir::arith::ConstantIndexOp>(loc, 1);
     mlir::Value size = one;
     if (mlir::Value sliceArg = arrLoad.getSlice()) {
       if (auto sliceOp =
               mlir::dyn_cast_or_null<SliceOp>(sliceArg.getDefiningOp())) {
         auto triples = sliceOp.getTriples();
         const std::size_t tripleSize = triples.size();
         auto module = arrLoad->getParentOfType<mlir::ModuleOp>();
-        fir::KindMapping kindMap = getKindMapping(module);
-        FirOpBuilder builder(rewriter, kindMap);
+        FirOpBuilder builder(rewriter, module);
         size = builder.genExtentFromTriplet(loc, triples[tripleSize - 3],
                                             triples[tripleSize - 2],
                                             triples[tripleSize - 1], idxTy);
         copyUsingSlice = true;
       }
     }
     result[result.size() - 1] = size;
   }
if (skipOrig)
     originated.assign(indices.begin(), indices.end());
   else
     originated = factory::originateIndices(loc, rewriter, alloc.getType(),
                                            shape, indices);
   auto seqTy = dyn_cast_ptrOrBoxEleTy(alloc.getType());
   assert(seqTy && seqTy.isa<SequenceType>());
   const auto dimension = seqTy.cast<SequenceType>().getDimension();
   auto module = load->getParentOfType<mlir::ModuleOp>();
-  fir::KindMapping kindMap = getKindMapping(module);
-  FirOpBuilder builder(rewriter, kindMap);
+  FirOpBuilder builder(rewriter, module);
   auto typeparams = getTypeParamsIfRawData(loc, builder, load, alloc.getType());
   mlir::Value result = rewriter.create<ArrayCoorOp>(
       loc, eleTy, alloc, shape, slice,
       llvm::ArrayRef<mlir::Value>{originated}.take_front(dimension),
       typeparams);
   if (dimension < originated.size())
     result = rewriter.create<fir::CoordinateOp>(
         loc, resTy, result,
for (auto sh : llvm::reverse(extents)) {
     auto ub = rewriter.create<mlir::arith::SubIOp>(loc, idxTy, ubi, one);
     auto loop = rewriter.create<DoLoopOp>(loc, zero, ub, one);
     rewriter.setInsertionPointToStart(loop.getBody());
     indices.push_back(loop.getInductionVar());
   }
   // Reverse the indices so they are in column-major order.
   std::reverse(indices.begin(), indices.end());
   auto module = arrLoad->getParentOfType<mlir::ModuleOp>();
-  fir::KindMapping kindMap = getKindMapping(module);
-  FirOpBuilder builder(rewriter, kindMap);
+  FirOpBuilder builder(rewriter, module);
   auto fromAddr = rewriter.create<ArrayCoorOp>(
       loc, getEleTy(src.getType()), src, shapeOp,
       CopyIn && copyUsingSlice ? sliceOp : mlir::Value{},
       factory::originateIndices(loc, rewriter, src.getType(), shapeOp, indices),
       getTypeParamsIfRawData(loc, builder, arrLoad, src.getType()));
   auto toAddr = rewriter.create<ArrayCoorOp>(
       loc, getEleTy(dst.getType()), dst, shapeOp,
       !CopyIn && copyUsingSlice ? sliceOp : mlir::Value{},
genArrayLoadTypeParameters(mlir::Location loc, mlir::PatternRewriter &rewriter,
                            ArrayLoadOp load) {
   if (load.getTypeparams().empty()) {
     auto eleTy =
         unwrapSequenceType(unwrapPassByRefType(load.getMemref().getType()));
     if (hasDynamicSize(eleTy)) {
       if (auto charTy = eleTy.dyn_cast<CharacterType>()) {
         assert(load.getMemref().getType().isa<BoxType>());
         auto module = load->getParentOfType<mlir::ModuleOp>();
-        fir::KindMapping kindMap = getKindMapping(module);
-        FirOpBuilder builder(rewriter, kindMap);
+        FirOpBuilder builder(rewriter, module);
         return {getCharacterLen(loc, builder, load, charTy)};
       }
       TODO(loc, "unhandled dynamic type parameters");
     }
     return {};
   }
   return load.getTypeparams();
 }
mlir::Type eleType =
       fir::unwrapSequenceType(fir::unwrapPassByRefType(baseType));
   if (fir::isRecordWithAllocatableMember(eleType)) {
     // The allocatable component descriptors need to be set to a clean
     // deallocated status before anything is done with them.
     mlir::Value box = rewriter.create<fir::EmboxOp>(
         loc, fir::BoxType::get(allocmem.getType()), allocmem, shape,
         /*slice=*/mlir::Value{}, typeParams);
     auto module = load->getParentOfType<mlir::ModuleOp>();
-    fir::KindMapping kindMap = getKindMapping(module);
-    FirOpBuilder builder(rewriter, kindMap);
+    FirOpBuilder builder(rewriter, module);
     runtime::genDerivedTypeInitialize(builder, loc, box);
     // Any allocatable component that may have been allocated must be
     // deallocated during the clean-up.
     auto cleanup = [=](mlir::PatternRewriter &r) {
-      fir::KindMapping kindMap = getKindMapping(module);
-      FirOpBuilder builder(r, kindMap);
+      FirOpBuilder builder(r, module);
       runtime::genDerivedTypeDestroy(builder, loc, box);
       r.create<FreeMemOp>(loc, allocmem);
     };
     return {allocmem, cleanup};
   }
   auto cleanup = [=](mlir::PatternRewriter &r) {
     r.create<FreeMemOp>(loc, allocmem);
   };

flang/lib/Optimizer/Transforms/LoopVersioning.cpp

if (!argsInLoop.empty()) {
       OpsWithArgs ops = {loop, argsInLoop};
       loopsOfInterest.push_back(ops);
     }
   });
   if (loopsOfInterest.empty())
     return;
 
   // If we get here, there are loops to process.
-  fir::FirOpBuilder builder{module, kindMap};
+  fir::FirOpBuilder builder{module, std::move(kindMap)};
   mlir::Location loc = builder.getUnknownLoc();
   mlir::IndexType idxTy = builder.getIndexType();
 
   LLVM_DEBUG(llvm::dbgs() << "Module Before transformation:");
   LLVM_DEBUG(module->dump());
 
   LLVM_DEBUG(llvm::dbgs() << "loopsOfInterest: " << loopsOfInterest.size()
                           << "\n");

flang/lib/Optimizer/Transforms/StackArrays.cpp

return rewriter.create<fir::AllocaOp>(loc, varTy, uniqName, bindcName,
                                         oldAlloc.getTypeparams(),
                                         oldAlloc.getShape());
 }
 
 void AllocMemConversion::insertStackSaveRestore(
     fir::AllocMemOp &oldAlloc, mlir::PatternRewriter &rewriter) const {
   auto oldPoint = rewriter.saveInsertionPoint();
   auto mod = oldAlloc->getParentOfType<mlir::ModuleOp>();
-  fir::KindMapping kindMap = fir::getKindMapping(mod);
-  fir::FirOpBuilder builder{rewriter, kindMap};
+  fir::FirOpBuilder builder{rewriter, mod};
 
   mlir::func::FuncOp stackSaveFn = fir::factory::getLlvmStackSave(builder);
   mlir::SymbolRefAttr stackSaveSym =
       builder.getSymbolRefAttr(stackSaveFn.getName());
 
   builder.setInsertionPoint(oldAlloc);
   mlir::Value sp =
       builder