[flang][hlfir] Fixed KindMapping for HLFIR intrinsics lowering.

[vzakhari]

Summary

hlfir.count lowering was using incorrect default integer kind
by ignoring the kind specified in the ModuleOp.

[tblah]

LGTM. Thanks for this!

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

explicit FirOpBuilder(mlir::OpBuilder &builder, fir::KindMapping kindMap,
       : OpBuilder(builder), OpBuilder::Listener(), kindMap{std::move(kindMap)} {
     setListener(this);
     auto fmi = mlir::dyn_cast<mlir::arith::ArithFastMathInterface>(*op);
     if (fmi) {
       // Set the builder with FastMathFlags attached to the operation.
       setFastMathFlags(fmi.getFastMathFlagsAttr().getValue());
     }
   }
+  FirOpBuilder(mlir::OpBuilder &builder, mlir::Operation *op)
+      : FirOpBuilder(builder, fir::getKindMapping(op), op) {}
 
   // 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);
   }
 

flang/include/flang/Optimizer/Dialect/Support/FIRContext.h

 #ifndef FORTRAN_OPTIMIZER_SUPPORT_FIRCONTEXT_H
 #define FORTRAN_OPTIMIZER_SUPPORT_FIRCONTEXT_H
 
 #include "llvm/ADT/StringRef.h"
 #include "llvm/TargetParser/Triple.h"
 
 namespace mlir {
 class ModuleOp;
+class Operation;
 } // namespace mlir
 
 namespace fir {
 class KindMapping;
 struct NameUniquer;
 
 /// Set the target triple for the module. `triple` must not be deallocated while
 /// module `mod` is still live.
 void setTargetTriple(mlir::ModuleOp mod, llvm::StringRef triple);
 
 /// Get the Triple instance from the Module or return the default Triple.
 llvm::Triple getTargetTriple(mlir::ModuleOp mod);
 
 /// Set the kind mapping for the module. `kindMap` must not be deallocated while
 /// module `mod` is still live.
 void setKindMapping(mlir::ModuleOp mod, KindMapping &kindMap);
 
 /// Get the KindMapping instance from the Module. If none was set, returns a
 /// default.
 KindMapping getKindMapping(mlir::ModuleOp mod);
 
+/// Get the KindMapping instance that is in effect for the specified
+/// operation. The KindMapping is taken from the operation itself,
+/// if the operation is a ModuleOp, or from its parent ModuleOp.
+/// If a ModuleOp cannot be reached, the function returns default KindMapping.
+KindMapping getKindMapping(mlir::Operation *op);
+
 /// Helper for determining the target from the host, etc. Tools may use this
 /// function to provide a consistent interpretation of the `--target=<string>`
 /// command-line option.
 /// An empty string ("") or "default" will specify that the default triple
 /// should be used. "native" will specify that the host machine be used to
 /// construct the triple.
 std::string determineTargetTriple(llvm::StringRef triple);
 
 } // namespace fir
 
 #endif // FORTRAN_OPTIMIZER_SUPPORT_FIRCONTEXT_H

flang/lib/Optimizer/Dialect/Support/FIRContext.cpp

if (auto defs = mod->getAttrOfType<mlir::StringAttr>(defKindName)) {
     auto defVals = fir::KindMapping::toDefaultKinds(defs.getValue());
     if (auto maps = mod->getAttrOfType<mlir::StringAttr>(kindMapName))
       return fir::KindMapping(ctx, maps.getValue(), defVals);
     return fir::KindMapping(ctx, defVals);
   }
   return fir::KindMapping(ctx);
 }
 
+fir::KindMapping fir::getKindMapping(mlir::Operation *op) {
+  auto moduleOp = mlir::dyn_cast<mlir::ModuleOp>(op);
+  if (moduleOp)
+    return getKindMapping(moduleOp);
+
+  moduleOp = op->getParentOfType<mlir::ModuleOp>();
+  return getKindMapping(moduleOp);
+}
+
 std::string fir::determineTargetTriple(llvm::StringRef triple) {
   // Treat "" or "default" as stand-ins for the default machine.
   if (triple.empty() || triple == "default")
     return llvm::sys::getDefaultTargetTriple();
   // Treat "native" as stand-in for the host machine.
   if (triple == "native")
     return llvm::sys::getProcessTriple();
   // TODO: normalize the triple?
   return triple.str();
 }

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

matchAndRewrite(hlfir::ApplyOp apply, OpAdaptor adaptor,
     hlfir::Entity bufferizedExpr = getBufferizedExprStorage(adaptor.getExpr());
     mlir::Type resultType = hlfir::getVariableElementType(bufferizedExpr);
     mlir::Value result = rewriter.create<hlfir::DesignateOp>(
         loc, resultType, bufferizedExpr, adaptor.getIndices(),
         adaptor.getTypeparams());
     if (fir::isa_trivial(apply.getType())) {
       result = rewriter.create<fir::LoadOp>(loc, result);
     } else {
-      auto module = apply->getParentOfType<mlir::ModuleOp>();
-      fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+      fir::FirOpBuilder builder(rewriter, apply.getOperation());
       result =
           packageBufferizedExpr(loc, builder, hlfir::Entity{result}, false);
     }
     rewriter.replaceOp(apply, result);
     return mlir::success();
   }
 };
 

 struct ConcatOpConversion : public mlir::OpConversionPattern<hlfir::ConcatOp> {
   using mlir::OpConversionPattern<hlfir::ConcatOp>::OpConversionPattern;
   explicit ConcatOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::ConcatOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::ConcatOp concat, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = concat->getLoc();
-    auto module = concat->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, concat.getOperation());
     assert(adaptor.getStrings().size() >= 2 &&
            "must have at least two strings operands");
     if (adaptor.getStrings().size() > 2)
       TODO(loc, "codegen of optimized chained concatenation of more than two "
                 "strings");
     hlfir::Entity lhs = getBufferizedExprStorage(adaptor.getStrings()[0]);
     hlfir::Entity rhs = getBufferizedExprStorage(adaptor.getStrings()[1]);
     auto [lhsExv, c1] = hlfir::translateToExtendedValue(loc, builder, lhs);
struct SetLengthOpConversion
     : public mlir::OpConversionPattern<hlfir::SetLengthOp> {
   using mlir::OpConversionPattern<hlfir::SetLengthOp>::OpConversionPattern;
   explicit SetLengthOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::SetLengthOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::SetLengthOp setLength, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = setLength->getLoc();
-    auto module = setLength->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, setLength.getOperation());
     // Create a temp with the new length.
     hlfir::Entity string = getBufferizedExprStorage(adaptor.getString());
     auto charType = hlfir::getFortranElementType(setLength.getType());
     llvm::StringRef tmpName{".tmp"};
     llvm::SmallVector<mlir::Value, 1> lenParams{adaptor.getLength()};
     auto alloca = builder.createTemporary(loc, charType, tmpName,
                                           /*shape=*/std::nullopt, lenParams);
     auto declareOp = builder.create<hlfir::DeclareOp>(
struct GetLengthOpConversion
     : public mlir::OpConversionPattern<hlfir::GetLengthOp> {
   using mlir::OpConversionPattern<hlfir::GetLengthOp>::OpConversionPattern;
   explicit GetLengthOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::GetLengthOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::GetLengthOp getLength, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = getLength->getLoc();
-    auto module = getLength->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, getLength.getOperation());
     hlfir::Entity bufferizedExpr = getBufferizedExprStorage(adaptor.getExpr());
     mlir::Value length = hlfir::genCharLength(loc, builder, bufferizedExpr);
     if (!length)
       return rewriter.notifyMatchFailure(
           getLength, "could not deduce length from GetLengthOp operand");
     rewriter.replaceOp(getLength, length);
     return mlir::success();
   }
struct AssociateOpConversion
     : public mlir::OpConversionPattern<hlfir::AssociateOp> {
   using mlir::OpConversionPattern<hlfir::AssociateOp>::OpConversionPattern;
   explicit AssociateOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::AssociateOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::AssociateOp associate, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = associate->getLoc();
-    auto module = associate->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, associate.getOperation());
     mlir::Value bufferizedExpr = getBufferizedExprStorage(adaptor.getSource());
     const bool isTrivialValue = fir::isa_trivial(bufferizedExpr.getType());
 
     auto getEndAssociate =
         [](hlfir::AssociateOp associate) -> mlir::Operation * {
       for (mlir::Operation *useOp : associate->getUsers())
         if (mlir::isa<hlfir::EndAssociateOp>(useOp))
           return useOp;
struct EndAssociateOpConversion
     : public mlir::OpConversionPattern<hlfir::EndAssociateOp> {
   using mlir::OpConversionPattern<hlfir::EndAssociateOp>::OpConversionPattern;
   explicit EndAssociateOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::EndAssociateOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::EndAssociateOp endAssociate, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = endAssociate->getLoc();
-    auto module = endAssociate->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, endAssociate.getOperation());
     genFreeIfMustFree(loc, builder, adaptor.getVar(), adaptor.getMustFree());
     rewriter.eraseOp(endAssociate);
     return mlir::success();
   }
 };
 
 struct DestroyOpConversion
     : public mlir::OpConversionPattern<hlfir::DestroyOp> {
   using mlir::OpConversionPattern<hlfir::DestroyOp>::OpConversionPattern;
   explicit DestroyOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::DestroyOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::DestroyOp destroy, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     // If expr was bufferized on the heap, now is time to deallocate the buffer.
     mlir::Location loc = destroy->getLoc();
     hlfir::Entity bufferizedExpr = getBufferizedExprStorage(adaptor.getExpr());
     if (!fir::isa_trivial(bufferizedExpr.getType())) {
-      auto module = destroy->getParentOfType<mlir::ModuleOp>();
-      fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+      fir::FirOpBuilder builder(rewriter, destroy.getOperation());
       mlir::Value mustFree = getBufferizedExprMustFreeFlag(adaptor.getExpr());
       mlir::Value firBase = bufferizedExpr.getFirBase();
       genFreeIfMustFree(loc, builder, firBase, mustFree);
     }
     rewriter.eraseOp(destroy);
     return mlir::success();
   }
 };
 
 struct NoReassocOpConversion
     : public mlir::OpConversionPattern<hlfir::NoReassocOp> {
   using mlir::OpConversionPattern<hlfir::NoReassocOp>::OpConversionPattern;
   explicit NoReassocOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::NoReassocOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::NoReassocOp noreassoc, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = noreassoc->getLoc();
-    auto module = noreassoc->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, noreassoc.getOperation());
     mlir::Value bufferizedExpr = getBufferizedExprStorage(adaptor.getVal());
     mlir::Value result =
         builder.create<hlfir::NoReassocOp>(loc, bufferizedExpr);
 
     if (!fir::isa_trivial(bufferizedExpr.getType())) {
       // NoReassocOp should not be needed on the mustFree path.
       mlir::Value mustFree = getBufferizedExprMustFreeFlag(adaptor.getVal());
       result =
explicit ElementalOpConversion(mlir::MLIRContext *ctx)
     // for recursive pattern application. The recursion is bounded
     // by the nesting level of ElementalOp's.
     setHasBoundedRewriteRecursion();
   }
   mlir::LogicalResult
   matchAndRewrite(hlfir::ElementalOp elemental, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = elemental->getLoc();
-    auto module = elemental->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, elemental.getOperation());
     // The body of the elemental op may contain operation that will require
     // to be translated. Notify the rewriter about the cloned operations.
     HLFIRListener listener{builder, rewriter};
     builder.setListener(&listener);
 
     mlir::Value shape = adaptor.getShape();
     auto extents = hlfir::getIndexExtents(loc, builder, shape);
     auto [temp, cleanup] =
struct CharExtremumOpConversion
     : public mlir::OpConversionPattern<hlfir::CharExtremumOp> {
   using mlir::OpConversionPattern<hlfir::CharExtremumOp>::OpConversionPattern;
   explicit CharExtremumOpConversion(mlir::MLIRContext *ctx)
       : mlir::OpConversionPattern<hlfir::CharExtremumOp>{ctx} {}
   mlir::LogicalResult
   matchAndRewrite(hlfir::CharExtremumOp char_extremum, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     mlir::Location loc = char_extremum->getLoc();
-    auto module = char_extremum->getParentOfType<mlir::ModuleOp>();
     auto predicate = char_extremum.getPredicate();
     bool predIsMin =
         predicate == hlfir::CharExtremumPredicate::min ? true : false;
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, char_extremum.getOperation());
     assert(adaptor.getStrings().size() >= 2 &&
            "must have at least two strings operands");
     auto numOperands = adaptor.getStrings().size();
 
     std::vector<hlfir::Entity> chars;
     std::vector<
         std::pair<fir::ExtendedValue, std::optional<hlfir::CleanupFunction>>>
         pairs;

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

auto res =
             .getResults();
     return {res[0], res[1]};
   }
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::CopyInOp copyInOp,
                   mlir::PatternRewriter &rewriter) const override {
     mlir::Location loc = copyInOp.getLoc();
-    auto module = copyInOp->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, copyInOp.getOperation());
     CopyInResult result = copyInOp.getVarIsPresent()
                               ? genOptionalCopyIn(loc, builder, copyInOp)
                               : genNonOptionalCopyIn(loc, builder, copyInOp);
     rewriter.replaceOp(copyInOp, {result.addr, result.wasCopied});
     return mlir::success();
   }
 };
 
 class CopyOutOpConversion : public mlir::OpRewritePattern<hlfir::CopyOutOp> {
 public:
   explicit CopyOutOpConversion(mlir::MLIRContext *ctx)
       : OpRewritePattern{ctx} {}
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::CopyOutOp copyOutOp,
                   mlir::PatternRewriter &rewriter) const override {
     mlir::Location loc = copyOutOp.getLoc();
-    auto module = copyOutOp->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, copyOutOp.getOperation());
 
     builder.genIfThen(loc, copyOutOp.getWasCopied())
         .genThen([&]() {
           mlir::Value temp = copyOutOp.getTemp();
           if (mlir::Value var = copyOutOp.getVar()) {
             auto mutableBoxTo = builder.createTemporary(loc, var.getType());
             builder.create<fir::StoreOp>(loc, var, mutableBoxTo);
             // Generate CopyOutAssign() call to copy data from the temporary
auto firBase = rewriter
                        .create<fir::DeclareOp>(
                            loc, memref.getType(), memref, declareOp.getShape(),
                            declareOp.getTypeparams(), declareOp.getUniqName(),
                            fortranAttrs)
                        .getResult();
     mlir::Value hlfirBase;
     mlir::Type hlfirBaseType = declareOp.getBase().getType();
     if (hlfirBaseType.isa<fir::BaseBoxType>()) {
-      auto module = declareOp->getParentOfType<mlir::ModuleOp>();
-      fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+      fir::FirOpBuilder builder(rewriter, declareOp.getOperation());
       // Helper to generate the hlfir fir.box with the local lower bounds and
       // type parameters.
       auto genHlfirBox = [&]() -> mlir::Value {
         if (!firBase.getType().isa<fir::BaseBoxType>()) {
           llvm::SmallVector<mlir::Value> typeParams;
           auto maybeCharType =
               fir::unwrapSequenceType(fir::unwrapPassByRefType(hlfirBaseType))
                   .dyn_cast<fir::CharacterType>();

 public:
   explicit DesignateOpConversion(mlir::MLIRContext *ctx)
       : OpRewritePattern{ctx} {}
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::DesignateOp designate,
                   mlir::PatternRewriter &rewriter) const override {
     mlir::Location loc = designate.getLoc();
-    auto module = designate->getParentOfType<mlir::ModuleOp>();
-    fir::FirOpBuilder builder(rewriter, fir::getKindMapping(module));
+    fir::FirOpBuilder builder(rewriter, designate.getOperation());
 
     hlfir::Entity baseEntity(designate.getMemref());
 
     if (baseEntity.isMutableBox())
       TODO(loc, "hlfir::designate load of pointer or allocatable");
 
     mlir::Type designateResultType = designate.getResult().getType();
     llvm::SmallVector<mlir::Value> firBaseTypeParameters;

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

 //===- InlineElementals.cpp - Inline chained hlfir.elemental ops ----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 // Chained elemental operations like a + b + c can inline the first elemental
 // at the hlfir.apply in the body of the second one (as described in
 // docs/HighLevelFIR.md). This has to be done in a pass rather than in lowering
 // so that it happens after the HLFIR intrinsic simplification pass.
 //===----------------------------------------------------------------------===//
 
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/HLFIRTools.h"
 #include "flang/Optimizer/Dialect/Support/FIRContext.h"
-#include "flang/Optimizer/Dialect/Support/KindMapping.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/HLFIR/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/DialectConversion.h"
if (elemental.isOrdered()) {
       return rewriter.notifyMatchFailure(elemental,
                                          "hlfir.elemental is ordered");
     }
     auto [apply, destroy] = *maybeTuple;
 
     assert(elemental.getRegion().hasOneBlock() &&
            "expect elemental region to have one block");
 
-    fir::FirOpBuilder builder{rewriter,
-                              fir::KindMapping{rewriter.getContext()}};
+    fir::FirOpBuilder builder{rewriter, elemental.getOperation()};
     builder.setInsertionPointAfter(apply);
     hlfir::YieldElementOp yield = hlfir::inlineElementalOp(
         elemental.getLoc(), builder, elemental, apply.getIndices());
 
     // remove the old elemental and all of the bookkeeping
     rewriter.replaceAllUsesWith(apply.getResult(), yield.getElementValue());
     rewriter.eraseOp(yield);
     rewriter.eraseOp(apply);

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

protected:
   /// Lower the arguments to the intrinsic: adding necessary boxing and
   /// conversion to match the signature of the intrinsic in the runtime library.
   llvm::SmallVector<fir::ExtendedValue, 3>
   lowerArguments(mlir::Operation *op,
                  const llvm::ArrayRef<IntrinsicArgument> &args,
                  mlir::PatternRewriter &rewriter,
                  const fir::IntrinsicArgumentLoweringRules *argLowering) const {
     mlir::Location loc = op->getLoc();
-    fir::KindMapping kindMapping{rewriter.getContext()};
-    fir::FirOpBuilder builder{rewriter, kindMapping, op};
+    fir::FirOpBuilder builder{rewriter, op};
 
     llvm::SmallVector<fir::ExtendedValue, 3> ret;
     llvm::SmallVector<std::function<void()>, 2> cleanupFns;
 
     for (size_t i = 0; i < args.size(); ++i) {
       mlir::Value arg = args[i].val;
       mlir::Type desiredType = args[i].desiredType;
       if (!arg) {
matchAndRewrite(OP operation,
     } else if constexpr (std::is_same_v<OP, hlfir::AnyOp>) {
       opName = "any";
     } else if constexpr (std::is_same_v<OP, hlfir::AllOp>) {
       opName = "all";
     } else {
       return mlir::failure();
     }
 
-    fir::KindMapping kindMapping{rewriter.getContext()};
-    fir::FirOpBuilder builder{rewriter, kindMapping, operation};
+    fir::FirOpBuilder builder{rewriter, operation.getOperation()};
     const mlir::Location &loc = operation->getLoc();
 
     mlir::Type i32 = builder.getI32Type();
     mlir::Type logicalType = fir::LogicalType::get(
         builder.getContext(), builder.getKindMap().defaultLogicalKind());
 
     llvm::SmallVector<fir::ExtendedValue, 0> args;
 

 using AllOpConversion = HlfirReductionIntrinsicConversion<hlfir::AllOp>;
 
 struct CountOpConversion : public HlfirIntrinsicConversion<hlfir::CountOp> {
   using HlfirIntrinsicConversion<hlfir::CountOp>::HlfirIntrinsicConversion;
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::CountOp count,
                   mlir::PatternRewriter &rewriter) const override {
-    fir::KindMapping kindMapping{rewriter.getContext()};
-    fir::FirOpBuilder builder{rewriter, kindMapping, count};
+    fir::FirOpBuilder builder{rewriter, count.getOperation()};
     const mlir::Location &loc = count->getLoc();
 
     mlir::Type i32 = builder.getI32Type();
     mlir::Type logicalType = fir::LogicalType::get(
         builder.getContext(), builder.getKindMap().defaultLogicalKind());
 
     llvm::SmallVector<IntrinsicArgument, 3> inArgs;
     inArgs.push_back({count.getMask(), logicalType});

 };
 
 struct MatmulOpConversion : public HlfirIntrinsicConversion<hlfir::MatmulOp> {
   using HlfirIntrinsicConversion<hlfir::MatmulOp>::HlfirIntrinsicConversion;
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::MatmulOp matmul,
                   mlir::PatternRewriter &rewriter) const override {
-    fir::KindMapping kindMapping{rewriter.getContext()};
-    fir::FirOpBuilder builder{rewriter, kindMapping, matmul};
+    fir::FirOpBuilder builder{rewriter, matmul.getOperation()};
     const mlir::Location &loc = matmul->getLoc();
 
     mlir::Value lhs = matmul.getLhs();
     mlir::Value rhs = matmul.getRhs();
     llvm::SmallVector<IntrinsicArgument, 2> inArgs;
     inArgs.push_back({lhs, lhs.getType()});
     inArgs.push_back({rhs, rhs.getType()});
 

 
 struct DotProductOpConversion
     : public HlfirIntrinsicConversion<hlfir::DotProductOp> {
   using HlfirIntrinsicConversion<hlfir::DotProductOp>::HlfirIntrinsicConversion;
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::DotProductOp dotProduct,
                   mlir::PatternRewriter &rewriter) const override {
-    fir::KindMapping kindMapping{rewriter.getContext()};
-    fir::FirOpBuilder builder{rewriter, kindMapping, dotProduct};
+    fir::FirOpBuilder builder{rewriter, dotProduct.getOperation()};
     const mlir::Location &loc = dotProduct->getLoc();
 
     mlir::Value lhs = dotProduct.getLhs();
     mlir::Value rhs = dotProduct.getRhs();
     llvm::SmallVector<IntrinsicArgument, 2> inArgs;
     inArgs.push_back({lhs, lhs.getType()});
     inArgs.push_back({rhs, rhs.getType()});
 

 
 class TransposeOpConversion
     : public HlfirIntrinsicConversion<hlfir::TransposeOp> {
   using HlfirIntrinsicConversion<hlfir::TransposeOp>::HlfirIntrinsicConversion;
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::TransposeOp transpose,
                   mlir::PatternRewriter &rewriter) const override {
-    fir::KindMapping kindMapping{rewriter.getContext()};
-    fir::FirOpBuilder builder{rewriter, kindMapping, transpose};
+    fir::FirOpBuilder builder{rewriter, transpose.getOperation()};
     const mlir::Location &loc = transpose->getLoc();
 
     mlir::Value arg = transpose.getArray();
     llvm::SmallVector<IntrinsicArgument, 1> inArgs;
     inArgs.push_back({arg, arg.getType()});
 
     auto *argLowering = fir::getIntrinsicArgumentLowering("transpose");
     llvm::SmallVector<fir::ExtendedValue, 1> args =

 struct MatmulTransposeOpConversion
     : public HlfirIntrinsicConversion<hlfir::MatmulTransposeOp> {
   using HlfirIntrinsicConversion<
       hlfir::MatmulTransposeOp>::HlfirIntrinsicConversion;
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::MatmulTransposeOp multranspose,
                   mlir::PatternRewriter &rewriter) const override {
-    fir::KindMapping kindMapping{rewriter.getContext()};
-    fir::FirOpBuilder builder{rewriter, kindMapping, multranspose};
+    fir::FirOpBuilder builder{rewriter, multranspose.getOperation()};
     const mlir::Location &loc = multranspose->getLoc();
 
     mlir::Value lhs = multranspose.getLhs();
     mlir::Value rhs = multranspose.getRhs();
     llvm::SmallVector<IntrinsicArgument, 2> inArgs;
     inArgs.push_back({lhs, lhs.getType()});
     inArgs.push_back({rhs, rhs.getType()});
 

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

 //===- SimplifyHLFIRIntrinsics.cpp - Simplify HLFIR Intrinsics ------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 // Normally transformational intrinsics are lowered to calls to runtime
 // functions. However, some cases of the intrinsics are faster when inlined
 // into the calling function.
 //===----------------------------------------------------------------------===//
 
 #include "flang/Optimizer/Builder/FIRBuilder.h"
 #include "flang/Optimizer/Builder/HLFIRTools.h"
 #include "flang/Optimizer/Dialect/FIRDialect.h"
-#include "flang/Optimizer/Dialect/Support/KindMapping.h"
 #include "flang/Optimizer/HLFIR/HLFIRDialect.h"
 #include "flang/Optimizer/HLFIR/HLFIROps.h"
 #include "flang/Optimizer/HLFIR/Passes.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/Pass/Pass.h"
class TransposeAsElementalConversion
     : public mlir::OpRewritePattern<hlfir::TransposeOp> {
 public:
   using mlir::OpRewritePattern<hlfir::TransposeOp>::OpRewritePattern;
 
   mlir::LogicalResult
   matchAndRewrite(hlfir::TransposeOp transpose,
                   mlir::PatternRewriter &rewriter) const override {
     mlir::Location loc = transpose.getLoc();
-    fir::KindMapping kindMapping{rewriter.getContext()};
-    fir::FirOpBuilder builder{rewriter, kindMapping};
+    fir::FirOpBuilder builder{rewriter, transpose.getOperation()};
     hlfir::ExprType expr = transpose.getType();
     mlir::Type elementType = expr.getElementType();
     hlfir::Entity array = hlfir::Entity{transpose.getArray()};
     mlir::Value resultShape = genResultShape(loc, builder, array);
     llvm::SmallVector<mlir::Value, 1> typeParams;
     hlfir::genLengthParameters(loc, builder, array, typeParams);
 
     auto genKernel = [&array](mlir::Location loc, fir::FirOpBuilder &builder,

flang/test/HLFIR/count-lowering-default-int-kinds.fir

+// Test hlfir.count operation lowering with different default integer kinds.
+// RUN: fir-opt %s -lower-hlfir-intrinsics | FileCheck %s
+
+module attributes {fir.defaultkind = "a1c4d8i8l4r4", fir.kindmap = ""} {
+  func.func @test_i8(%arg0: !fir.box<!fir.array<?x?x!fir.logical<4>>> {fir.bindc_name = "x"}, %arg1: i64) {
+    %4 = hlfir.count %arg0 dim %arg1 : (!fir.box<!fir.array<?x?x!fir.logical<4>>>, i64) -> !hlfir.expr<?xi64>
+    return
+  }
+}
+// CHECK-LABEL: func.func @test_i8
+// CHECK: %[[KIND:.*]] = arith.constant 8 : index
+// CHECK: %[[KIND_ARG:.*]] = fir.convert %[[KIND]] : (index) -> i32
+// CHECK: fir.call @_FortranACountDim(%{{.*}}, %{{.*}}, %{{.*}}, %[[KIND_ARG]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i32, !fir.ref<i8>, i32) -> none
+
+module attributes {fir.defaultkind = "a1c4d8i4l4r4", fir.kindmap = ""} {
+  func.func @test_i4(%arg0: !fir.box<!fir.array<?x?x!fir.logical<4>>> {fir.bindc_name = "x"}, %arg1: i64) {
+    %4 = hlfir.count %arg0 dim %arg1 : (!fir.box<!fir.array<?x?x!fir.logical<4>>>, i64) -> !hlfir.expr<?xi32>
+    return
+  }
+}
+// CHECK-LABEL: func.func @test_i4
+// CHECK: %[[KIND:.*]] = arith.constant 4 : index
+// CHECK: %[[KIND_ARG:.*]] = fir.convert %[[KIND]] : (index) -> i32
+// CHECK: fir.call @_FortranACountDim(%{{.*}}, %{{.*}}, %{{.*}}, %[[KIND_ARG]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i32, !fir.ref<i8>, i32) -> none
+
+module attributes {fir.defaultkind = "a1c4d8i2l4r4", fir.kindmap = ""} {
+  func.func @test_i2(%arg0: !fir.box<!fir.array<?x?x!fir.logical<4>>> {fir.bindc_name = "x"}, %arg1: i64) {
+    %4 = hlfir.count %arg0 dim %arg1 : (!fir.box<!fir.array<?x?x!fir.logical<4>>>, i64) -> !hlfir.expr<?xi16>
+    return
+  }
+}
+// CHECK-LABEL: func.func @test_i2
+// CHECK: %[[KIND:.*]] = arith.constant 2 : index
+// CHECK: %[[KIND_ARG:.*]] = fir.convert %[[KIND]] : (index) -> i32
+// CHECK: fir.call @_FortranACountDim(%{{.*}}, %{{.*}}, %{{.*}}, %[[KIND_ARG]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i32, !fir.ref<i8>, i32) -> none
+
+module attributes {fir.defaultkind = "a1c4d8i1l4r4", fir.kindmap = ""} {
+  func.func @test_i1(%arg0: !fir.box<!fir.array<?x?x!fir.logical<4>>> {fir.bindc_name = "x"}, %arg1: i64) {
+    %4 = hlfir.count %arg0 dim %arg1 : (!fir.box<!fir.array<?x?x!fir.logical<4>>>, i64) -> !hlfir.expr<?xi8>
+    return
+  }
+}
+// CHECK-LABEL: func.func @test_i1
+// CHECK: arith.constant 1 : index
+// CHECK: %[[KIND:.*]] = arith.constant 1 : index
+// CHECK: %[[KIND_ARG:.*]] = fir.convert %[[KIND]] : (index) -> i32
+// CHECK: fir.call @_FortranACountDim(%{{.*}}, %{{.*}}, %{{.*}}, %[[KIND_ARG]], %{{.*}}, %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i32, !fir.ref<i8>, i32) -> none

flang/test/HLFIR/count-lowering.fir

 // CHECK:           %[[ARG0:.*]]: !fir.box<!fir.array<?x?x!fir.logical<4>>>
 // CHECK:           %[[ARG1:.*]]: !fir.box<!fir.array<?xi32>
 // CHECK:           %[[ARG2:.*]]: !fir.ref<i32>
 // CHECK-DAG:     %[[MASK:.*]]:2 = hlfir.declare %[[ARG0]]
 // CHECK-DAG:     %[[DIM_VAR:.*]]:2 = hlfir.declare %[[ARG2]]
 // CHECK-DAG:     %[[RES:.*]]:2 = hlfir.declare %[[ARG1]]
 
 // CHECK-DAG:     %[[RET_BOX:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xi32>>>
+// CHECK-DAG:     %[[KIND:.*]] = arith.constant 4 : index
 // CHECK-DAG:     %[[RET_ADDR:.*]] = fir.zero_bits !fir.heap<!fir.array<?xi32>>
 // CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[RET_SHAPE:.*]] = fir.shape %[[C0]] : (index) -> !fir.shape<1>
 // CHECK-DAG:     %[[RET_EMBOX:.*]] = fir.embox %[[RET_ADDR]](%[[RET_SHAPE]])
 // CHECK-DAG:     fir.store %[[RET_EMBOX]] to %[[RET_BOX]]
 
 // CHECK-DAG:     %[[DIM:.*]] = fir.load %[[DIM_VAR]]#0 : !fir.ref<i32>
 // CHECK-DAG:     %[[RET_ARG:.*]] = fir.convert %[[RET_BOX]]
 // CHECK-DAG:     %[[MASK_ARG:.*]] = fir.convert %[[MASK]]#1
+// CHECK-DAG:     %[[KIND_ARG:.*]] = fir.convert %[[KIND]] : (index) -> i32
 
-// CHECK:         %[[NONE:.*]] = fir.call @_FortranACountDim(%[[RET_ARG]], %[[MASK_ARG]], %[[DIM]], %[[LOC_STR:.*]], %[[LOC_N:.*]]) : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i32, !fir.ref<i8>, i32) -> none
+// CHECK:         %[[NONE:.*]] = fir.call @_FortranACountDim(%[[RET_ARG]], %[[MASK_ARG]], %[[DIM]], %[[KIND_ARG]], %[[LOC_STR:.*]], %[[LOC_N:.*]]) : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32, i32, !fir.ref<i8>, i32) -> none
 // CHECK:         %[[RET:.*]] = fir.load %[[RET_BOX]]
 // CHECK:         %[[BOX_DIMS:.*]]:3 = fir.box_dims %[[RET]]
 // CHECK-NEXT:    %[[ADDR:.*]] = fir.box_addr %[[RET]]
 // CHECK-NEXT:    %[[SHIFT:.*]] = fir.shape_shift %[[BOX_DIMS]]#0, %[[BOX_DIMS]]#1
 // CHECK-NEXT:    %[[TMP:.*]]:2 = hlfir.declare %[[ADDR]](%[[SHIFT]]) {uniq_name = ".tmp.intrinsic_result"}
 // CHECK:         %[[TRUE:.*]] = arith.constant true
 // CHECK:         %[[EXPR:.*]] = hlfir.as_expr %[[TMP]]#0 move %[[TRUE]] : (!fir.box<!fir.array<?xi32>>, i1) -> !hlfir.expr<?xi32>
 // CHECK:         hlfir.assign %[[EXPR]] to %[[RES]]#0
func.func @_QPcount3(%arg0: !fir.ref<!fir.array<2xi32>> {fir.bindc_name = "s"}) {
   %5 = hlfir.count %2#0 dim %c1_i32 {fastmath = #arith.fastmath<contract>} : (!fir.ref<!fir.array<2x2x!fir.logical<4>>>, i32) -> !hlfir.expr<2xi32>
   hlfir.assign %5 to %4#0 : !hlfir.expr<2xi32>, !fir.ref<!fir.array<2xi32>>
   hlfir.destroy %5 : !hlfir.expr<2xi32>
   return
 }
 // CHECK-LABEL:  func.func @_QPcount3(
 // CHECK:           %[[ARG0:.*]]: !fir.ref<!fir.array<2xi32>>
 // CHECK-DAG:     %[[RET_BOX:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xi32>>>
+// CHECK-DAG:     %[[KIND:.*]] = arith.constant 4 : index
 // CHECK-DAG:     %[[RET_ADDR:.*]] = fir.zero_bits !fir.heap<!fir.array<?xi32>>
 // CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[RET_SHAPE:.*]] = fir.shape %[[C0]] : (index) -> !fir.shape<1>
 // CHECK-DAG:     %[[RET_EMBOX:.*]] = fir.embox %[[RET_ADDR]](%[[RET_SHAPE]])
 // CHECK-DAG:     fir.store %[[RET_EMBOX]] to %[[RET_BOX]]
 // CHECK-DAG:     %[[RES:.*]]:2 = hlfir.declare %[[ARG0]](%[[RES_SHAPE:.*]])
 
 // CHECK-DAG:     %[[MASK_ADDR:.*]] = fir.address_of
 // CHECK-DAG:     %[[MASK_VAR:.*]]:2 = hlfir.declare %[[MASK_ADDR]](%[[MASK_SHAPE:.*]])
 // CHECK-DAG:     %[[MASK_BOX:.*]] = fir.embox %[[MASK_VAR]]#1(%[[MASK_SHAPE:.*]])
 
 // CHECK-DAG:     %[[DIM:.*]] = arith.constant 1 : i32
 
 // CHECK-DAG:     %[[RET_ARG:.*]] = fir.convert %[[RET_BOX]]
 // CHECK-DAG:     %[[MASK_ARG:.*]] = fir.convert %[[MASK_BOX]] : (!fir.box<!fir.array<2x2x!fir.logical<4>>>) -> !fir.box<none>
-// CHECK:         %[[NONE:.*]] = fir.call @_FortranACountDim(%[[RET_ARG]], %[[MASK_ARG]], %[[DIM]], %[[LOC_STR:.*]], %[[LOC_N:.*]])
+// CHECK-DAG:     %[[KIND_ARG:.*]] = fir.convert %[[KIND]] : (index) -> i32
+
+// CHECK:         %[[NONE:.*]] = fir.call @_FortranACountDim(%[[RET_ARG]], %[[MASK_ARG]], %[[DIM]], %[[KIND_ARG]], %[[LOC_STR:.*]], %[[LOC_N:.*]])
 // CHECK:         %[[RET:.*]] = fir.load %[[RET_BOX]]
 // CHECK:         %[[BOX_DIMS:.*]]:3 = fir.box_dims %[[RET]]
 // CHECK-NEXT:    %[[ADDR:.*]] = fir.box_addr %[[RET]]
 // CHECK-NEXT:    %[[SHIFT:.*]] = fir.shape_shift %[[BOX_DIMS]]#0, %[[BOX_DIMS]]#1
 // CHECK-NEXT:    %[[TMP:.*]]:2 = hlfir.declare %[[ADDR]](%[[SHIFT]]) {uniq_name = ".tmp.intrinsic_result"}
 // CHECK:         %[[TRUE:.*]] = arith.constant true
 // CHECK:         %[[EXPR:.*]] = hlfir.as_expr %[[TMP]]#0 move %[[TRUE]] : (!fir.box<!fir.array<?xi32>>, i1) -> !hlfir.expr<?xi32>
 // CHECK:         hlfir.assign %[[EXPR]] to %[[RES]]