[mlir] make ModuleTranslation mapping fields private

[ftynse]

Summary

ModuleTranslation contains multiple fields that keep track of the mappings
between various MLIR and LLVM IR components. The original ModuleTranslation
extension model was based on inheritance, with these fields being protected and
thus accessible in the ModuleTranslation and derived classes. The
inheritance-based model doesn't scale to translation of more than one derived
dialect and will be progressively replaced with a more flexible one based on
dialect interfaces and a translation state that is separate from
ModuleTranslation. This change prepares the replacement by making the mappings
private and providing public methods to access them.

Depends On D96436

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

translateModule(Operation *m, llvm::LLVMContext &llvmContext,
 
     return std::move(translator.llvmModule);
   }
 
   /// A helper method to get the single Block in an operation honoring LLVM's
   /// module requirements.
   static Block &getModuleBody(Operation *m) { return m->getRegion(0).front(); }
 
+  /// Stores the mapping between a function name and its LLVM IR representation.
+  void mapFunction(StringRef name, llvm::Function *func) {
+    auto result = functionMapping.try_emplace(name, func);
+    (void)result;
+    assert(result.second &&
+           "attempting to map a function that is already mapped");
+  }
+
+  /// Finds an LLVM IR function by its name.
+  llvm::Function *lookupFunction(StringRef name) const {
+    return functionMapping.lookup(name);
+  }
+
+  /// Stores the mapping between an MLIR value and its LLVM IR counterpart.
+  void mapValue(Value mlir, llvm::Value *llvm) { mapValue(mlir) = llvm; }
+
+  /// Provides write-once access to store the LLVM IR value corresponding to the
+  /// given MLIR value.
+  llvm::Value *&mapValue(Value value) {
+    llvm::Value *&llvm = valueMapping[value];
+    assert(llvm == nullptr &&
+           "attempting to map a value that is already mapped");
+    return llvm;
+  }
+
+  /// Finds an LLVM IR value corresponding to the given MLIR value.
+  llvm::Value *lookupValue(Value value) const {
+    return valueMapping.lookup(value);
+  }
+
+  /// Stores the mapping between an MLIR block and LLVM IR basic block.
+  void mapBlock(Block *mlir, llvm::BasicBlock *llvm) {
+    auto result = blockMapping.try_emplace(mlir, llvm);
+    (void)result;
+    assert(result.second && "attempting to map a block that is already mapped");
+  }
+
+  /// Finds an LLVM IR basic block that corresponds to the given MLIR block.
+  llvm::BasicBlock *lookupBlock(Block *block) const {
+    return blockMapping.lookup(block);
+  }
+
+  /// Stores the mapping between an MLIR operation with successors and a
+  /// corresponding LLVM IR instruction.
+  void mapBranch(Operation *mlir, llvm::Instruction *llvm) {
+    auto result = branchMapping.try_emplace(mlir, llvm);
+    (void)result;
+    assert(result.second &&
+           "attempting to map a branch that is already mapped");
+  }
+
+  /// Finds an LLVM IR instruction that corresponds to the given MLIR operation
+  /// with successors.
+  llvm::Instruction *lookupBranch(Operation *op) const {
+    return branchMapping.lookup(op);
+  }
+
 protected:
   /// Translate the given MLIR module expressed in MLIR LLVM IR dialect into an
   /// LLVM IR module. The MLIR LLVM IR dialect holds a pointer to an
   /// LLVMContext, the LLVM IR module will be created in that context.
   ModuleTranslation(Operation *module,
                     std::unique_ptr<llvm::Module> llvmModule);
   virtual ~ModuleTranslation();
 
   virtual LogicalResult convertOperation(Operation &op,
                                          llvm::IRBuilder<> &builder);
   virtual LogicalResult convertOmpOperation(Operation &op,
                                             llvm::IRBuilder<> &builder);
   virtual LogicalResult convertOmpParallel(Operation &op,
                                            llvm::IRBuilder<> &builder);
   virtual LogicalResult convertOmpMaster(Operation &op,
                                          llvm::IRBuilder<> &builder);
   void convertOmpOpRegions(Region &region, StringRef blockName,
-                           DenseMap<Value, llvm::Value *> &valueMapping,
-                           DenseMap<Block *, llvm::BasicBlock *> &blockMapping,
                            llvm::BasicBlock &sourceBlock,
                            llvm::BasicBlock &continuationBlock,
                            llvm::IRBuilder<> &builder,
                            LogicalResult &bodyGenStatus);
   virtual LogicalResult convertOmpWsLoop(Operation &opInst,
                                          llvm::IRBuilder<> &builder);
 
   /// Converts the type from MLIR LLVM dialect to LLVM.
private:
   const Dialect *ompDialect;
 
   /// Mappings between llvm.mlir.global definitions and corresponding globals.
   DenseMap<Operation *, llvm::GlobalValue *> globalsMapping;
 
   /// A stateful object used to translate types.
   TypeToLLVMIRTranslator typeTranslator;
 
-protected:
+private:
   /// Mappings between original and translated values, used for lookups.
   llvm::StringMap<llvm::Function *> functionMapping;
   DenseMap<Value, llvm::Value *> valueMapping;
   DenseMap<Block *, llvm::BasicBlock *> blockMapping;
 
   /// A mapping between MLIR LLVM dialect terminators and LLVM IR terminators
   /// they are converted to. This allows for connecting PHI nodes to the source
   /// values after all operations are converted.
   DenseMap<Operation *, llvm::Instruction *> branchMapping;
 };
 
 } // namespace LLVM
 } // namespace mlir
 
 #endif // MLIR_TARGET_LLVMIR_MODULETRANSLATION_H

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

static llvm::Type *getInnermostElementType(llvm::Type *type) {
   do {
     if (auto *arrayTy = dyn_cast<llvm::ArrayType>(type)) {
       type = arrayTy->getElementType();
     } else if (auto *vectorTy = dyn_cast<llvm::VectorType>(type)) {
       type = vectorTy->getElementType();
     } else {
       return type;
     }
-  } while (1);
+  } while (true);
 }
 
 /// Create an LLVM IR constant of `llvmType` from the MLIR attribute `attr`.
 /// This currently supports integer, floating point, splat and dense element
 /// attributes and combinations thereof.  In case of error, report it to `loc`
 /// and return nullptr.
 llvm::Constant *ModuleTranslation::getLLVMConstant(llvm::Type *llvmType,
                                                    Attribute attr,
                                                    Location loc) {
   if (!attr)
     return llvm::UndefValue::get(llvmType);
   if (llvmType->isStructTy()) {
     emitError(loc, "struct types are not supported in constants");
     return nullptr;
   }
   // For integer types, we allow a mismatch in sizes as the index type in
   // MLIR might have a different size than the index type in the LLVM module.
   if (auto intAttr = attr.dyn_cast<IntegerAttr>())
     return llvm::ConstantInt::get(
         llvmType,
         intAttr.getValue().sextOrTrunc(llvmType->getIntegerBitWidth()));
   if (auto floatAttr = attr.dyn_cast<FloatAttr>())
     return llvm::ConstantFP::get(llvmType, floatAttr.getValue());
   if (auto funcAttr = attr.dyn_cast<FlatSymbolRefAttr>())
-    return llvm::ConstantExpr::getBitCast(
-        functionMapping.lookup(funcAttr.getValue()), llvmType);
+    return llvm::ConstantExpr::getBitCast(lookupFunction(funcAttr.getValue()),
+                                          llvmType);
   if (auto splatAttr = attr.dyn_cast<SplatElementsAttr>()) {
     llvm::Type *elementType;
     uint64_t numElements;
     if (auto *arrayTy = dyn_cast<llvm::ArrayType>(llvmType)) {
       elementType = arrayTy->getElementType();
       numElements = arrayTy->getNumElements();
     } else {
       auto *vectorTy = cast<llvm::FixedVectorType>(llvmType);
static Value getPHISourceValue(Block *current, Block *pred,
   // For instructions that branch based on a condition value, we need to take
   // the operands for the branch that was taken.
   if (auto condBranchOp = dyn_cast<LLVM::CondBrOp>(terminator)) {
     // For conditional branches, we take the operands from either the "true" or
     // the "false" branch.
     return condBranchOp.getSuccessor(0) == current
                ? condBranchOp.trueDestOperands()[index]
                : condBranchOp.falseDestOperands()[index];
-  } else if (auto switchOp = dyn_cast<LLVM::SwitchOp>(terminator)) {
+  }
+
+  if (auto switchOp = dyn_cast<LLVM::SwitchOp>(terminator)) {
     // For switches, we take the operands from either the default case, or from
     // the case branch that was taken.
     if (switchOp.defaultDestination() == current)
       return switchOp.defaultOperands()[index];
     for (auto i : llvm::enumerate(switchOp.caseDestinations()))
       if (i.value() == current)
         return switchOp.getCaseOperands(i.index())[index];
   }
 
   llvm_unreachable("only branch or switch operations can be terminators of a "
                    "block that has successors");
 }
 
 /// Connect the PHI nodes to the results of preceding blocks.
 template <typename T>
-static void connectPHINodes(
-    T &func, const DenseMap<Value, llvm::Value *> &valueMapping,
-    const DenseMap<Block *, llvm::BasicBlock *> &blockMapping,
-    const DenseMap<Operation *, llvm::Instruction *> &branchMapping) {
+static void connectPHINodes(T &func, const ModuleTranslation &state) {
   // Skip the first block, it cannot be branched to and its arguments correspond
   // to the arguments of the LLVM function.
   for (auto it = std::next(func.begin()), eit = func.end(); it != eit; ++it) {
     Block *bb = &*it;
-    llvm::BasicBlock *llvmBB = blockMapping.lookup(bb);
+    llvm::BasicBlock *llvmBB = state.lookupBlock(bb);
     auto phis = llvmBB->phis();
     auto numArguments = bb->getNumArguments();
     assert(numArguments == std::distance(phis.begin(), phis.end()));
     for (auto &numberedPhiNode : llvm::enumerate(phis)) {
       auto &phiNode = numberedPhiNode.value();
       unsigned index = numberedPhiNode.index();
       for (auto *pred : bb->getPredecessors()) {
         // Find the LLVM IR block that contains the converted terminator
         // instruction and use it in the PHI node. Note that this block is not
-        // necessarily the same as blockMapping.lookup(pred), some operations
+        // necessarily the same as state.lookupBlock(pred), some operations
         // (in particular, OpenMP operations using OpenMPIRBuilder) may have
         // split the blocks.
         llvm::Instruction *terminator =
-            branchMapping.lookup(pred->getTerminator());
+            state.lookupBranch(pred->getTerminator());
         assert(terminator && "missing the mapping for a terminator");
-        phiNode.addIncoming(valueMapping.lookup(getPHISourceValue(
-                                bb, pred, numArguments, index)),
-                            terminator->getParent());
+        phiNode.addIncoming(
+            state.lookupValue(getPHISourceValue(bb, pred, numArguments, index)),
       }
     }
   }
 }
 
 /// Sort function blocks topologically.
 template <typename T>
 static llvm::SetVector<Block *> topologicalSort(T &f) {
ModuleTranslation::convertOmpParallel(Operation &opInst,
   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
   // relying on captured variables.
   LogicalResult bodyGenStatus = success();
 
   auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
                        llvm::BasicBlock &continuationBlock) {
     // ParallelOp has only one region associated with it.
     auto &region = cast<omp::ParallelOp>(opInst).getRegion();
-    convertOmpOpRegions(region, "omp.par.region", valueMapping, blockMapping,
-                        *codeGenIP.getBlock(), continuationBlock, builder,
-                        bodyGenStatus);
+    convertOmpOpRegions(region, "omp.par.region", *codeGenIP.getBlock(),
+                        continuationBlock, builder, bodyGenStatus);
   };
 
   // TODO: Perform appropriate actions according to the data-sharing
   // attribute (shared, private, firstprivate, ...) of variables.
   // Currently defaults to shared.
   auto privCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
                     llvm::Value &, llvm::Value &vPtr,
                     llvm::Value *&replacementValue) -> InsertPointTy {
     replacementValue = &vPtr;
 
     return codeGenIP;
   };
 
   // TODO: Perform finalization actions for variables. This has to be
   // called for variables which have destructors/finalizers.
   auto finiCB = [&](InsertPointTy codeGenIP) {};
 
   llvm::Value *ifCond = nullptr;
   if (auto ifExprVar = cast<omp::ParallelOp>(opInst).if_expr_var())
-    ifCond = valueMapping.lookup(ifExprVar);
+    ifCond = lookupValue(ifExprVar);
   llvm::Value *numThreads = nullptr;
   if (auto numThreadsVar = cast<omp::ParallelOp>(opInst).num_threads_var())
-    numThreads = valueMapping.lookup(numThreadsVar);
+    numThreads = lookupValue(numThreadsVar);
   llvm::omp::ProcBindKind pbKind = llvm::omp::OMP_PROC_BIND_default;
   if (auto bind = cast<omp::ParallelOp>(opInst).proc_bind_val())
     pbKind = llvm::omp::getProcBindKind(bind.getValue());
   // TODO: Is the Parallel construct cancellable?
   bool isCancellable = false;
   // TODO: Determine the actual alloca insertion point, e.g., the function
   // entry or the alloca insertion point as provided by the body callback
   // above.
   llvm::OpenMPIRBuilder::InsertPointTy allocaIP(builder.saveIP());
   if (failed(bodyGenStatus))
     return failure();
   builder.restoreIP(
       ompBuilder->createParallel(builder, allocaIP, bodyGenCB, privCB, finiCB,
                                  ifCond, numThreads, pbKind, isCancellable));
   return success();
 }
 
 void ModuleTranslation::convertOmpOpRegions(
     Region &region, StringRef blockName,
-    DenseMap<Value, llvm::Value *> &valueMapping,
-    DenseMap<Block *, llvm::BasicBlock *> &blockMapping,
     llvm::BasicBlock &sourceBlock, llvm::BasicBlock &continuationBlock,
     llvm::IRBuilder<> &builder, LogicalResult &bodyGenStatus) {
   llvm::LLVMContext &llvmContext = builder.getContext();
   for (Block &bb : region) {
     llvm::BasicBlock *llvmBB = llvm::BasicBlock::Create(
         llvmContext, blockName, builder.GetInsertBlock()->getParent());
-    blockMapping[&bb] = llvmBB;
+    mapBlock(&bb, llvmBB);
   }
 
   llvm::Instruction *sourceTerminator = sourceBlock.getTerminator();
 
   // Convert blocks one by one in topological order to ensure
   // defs are converted before uses.
   llvm::SetVector<Block *> blocks = topologicalSort(region);
   for (Block *bb : blocks) {
-    llvm::BasicBlock *llvmBB = blockMapping[bb];
+    llvm::BasicBlock *llvmBB = lookupBlock(bb);
     // Retarget the branch of the entry block to the entry block of the
     // converted region (regions are single-entry).
     if (bb->isEntryBlock()) {
       assert(sourceTerminator->getNumSuccessors() == 1 &&
              "provided entry block has multiple successors");
       assert(sourceTerminator->getSuccessor(0) == &continuationBlock &&
              "ContinuationBlock is not the successor of the entry block");
       sourceTerminator->setSuccessor(0, llvmBB);
for (Block *bb : blocks) {
     // ModuleTranslation class to set up the correct insertion point. This is
     // also consistent with MLIR's idiom of handling special region terminators
     // in the same code that handles the region-owning operation.
     if (isa<omp::TerminatorOp, omp::YieldOp>(bb->getTerminator()))
       builder.CreateBr(&continuationBlock);
   }
   // Finally, after all blocks have been traversed and values mapped,
   // connect the PHI nodes to the results of preceding blocks.
-  connectPHINodes(region, valueMapping, blockMapping, branchMapping);
+  connectPHINodes(region, *this);
 }
 
 LogicalResult ModuleTranslation::convertOmpMaster(Operation &opInst,
                                                   llvm::IRBuilder<> &builder) {
   using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
   // relying on captured variables.
   LogicalResult bodyGenStatus = success();
 
   auto bodyGenCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
                        llvm::BasicBlock &continuationBlock) {
     // MasterOp has only one region associated with it.
     auto &region = cast<omp::MasterOp>(opInst).getRegion();
-    convertOmpOpRegions(region, "omp.master.region", valueMapping, blockMapping,
-                        *codeGenIP.getBlock(), continuationBlock, builder,
-                        bodyGenStatus);
+    convertOmpOpRegions(region, "omp.master.region", *codeGenIP.getBlock(),
+                        continuationBlock, builder, bodyGenStatus);
   };
 
   // TODO: Perform finalization actions for variables. This has to be
   // called for variables which have destructors/finalizers.
   auto finiCB = [&](InsertPointTy codeGenIP) {};
 
   builder.restoreIP(ompBuilder->createMaster(builder, bodyGenCB, finiCB));
   return success();
LogicalResult ModuleTranslation::convertOmpWsLoop(Operation &opInst,
 
   if (loop.schedule_val().hasValue() &&
       omp::symbolizeClauseScheduleKind(loop.schedule_val().getValue()) !=
           omp::ClauseScheduleKind::Static)
     return opInst.emitOpError(
         "only static (default) loop schedule is currently supported");
 
   // Find the loop configuration.
-  llvm::Value *lowerBound = valueMapping.lookup(loop.lowerBound()[0]);
-  llvm::Value *upperBound = valueMapping.lookup(loop.upperBound()[0]);
-  llvm::Value *step = valueMapping.lookup(loop.step()[0]);
+  llvm::Value *lowerBound = lookupValue(loop.lowerBound()[0]);
+  llvm::Value *upperBound = lookupValue(loop.upperBound()[0]);
+  llvm::Value *step = lookupValue(loop.step()[0]);
   llvm::Type *ivType = step->getType();
   llvm::Value *chunk = loop.schedule_chunk_var()
-                           ? valueMapping[loop.schedule_chunk_var()]
+                           ? lookupValue(loop.schedule_chunk_var())
                            : llvm::ConstantInt::get(ivType, 1);
 
   // Set up the source location value for OpenMP runtime.
   llvm::DISubprogram *subprogram =
       builder.GetInsertBlock()->getParent()->getSubprogram();
   const llvm::DILocation *diLoc =
       debugTranslation->translateLoc(opInst.getLoc(), subprogram);
   llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder.saveIP(),
                                                     llvm::DebugLoc(diLoc));
 
   // Generator of the canonical loop body. Produces an SESE region of basic
   // blocks.
   // TODO: support error propagation in OpenMPIRBuilder and use it instead of
   // relying on captured variables.
   LogicalResult bodyGenStatus = success();
   auto bodyGen = [&](llvm::OpenMPIRBuilder::InsertPointTy ip, llvm::Value *iv) {
     llvm::IRBuilder<>::InsertPointGuard guard(builder);
 
     // Make sure further conversions know about the induction variable.
-    valueMapping[loop.getRegion().front().getArgument(0)] = iv;
+    mapValue(loop.getRegion().front().getArgument(0), iv);
 
     llvm::BasicBlock *entryBlock = ip.getBlock();
     llvm::BasicBlock *exitBlock =
         entryBlock->splitBasicBlock(ip.getPoint(), "omp.wsloop.exit");
 
     // Convert the body of the loop.
-    convertOmpOpRegions(loop.region(), "omp.wsloop.region", valueMapping,
-                        blockMapping, *entryBlock, *exitBlock, builder,
-                        bodyGenStatus);
+    convertOmpOpRegions(loop.region(), "omp.wsloop.region", *entryBlock,
+                        *exitBlock, builder, bodyGenStatus);
   };
 
   // Delegate actual loop construction to the OpenMP IRBuilder.
   // TODO: this currently assumes WsLoop is semantically similar to SCF loop,
   // i.e. it has a positive step, uses signed integer semantics. Reconsider
   // this code when WsLoop clearly supports more cases.
   llvm::BasicBlock *insertBlock = builder.GetInsertBlock();
   llvm::CanonicalLoopInfo *loopInfo = ompBuilder->createCanonicalLoop(
#include "mlir/Dialect/LLVMIR/LLVMConversions.inc"
   // Emit function calls.  If the "callee" attribute is present, this is a
   // direct function call and we also need to look up the remapped function
   // itself.  Otherwise, this is an indirect call and the callee is the first
   // operand, look it up as a normal value.  Return the llvm::Value representing
   // the function result, which may be of llvm::VoidTy type.
   auto convertCall = [this, &builder](Operation &op) -> llvm::Value * {
     auto operands = lookupValues(op.getOperands());
     ArrayRef<llvm::Value *> operandsRef(operands);
-    if (auto attr = op.getAttrOfType<FlatSymbolRefAttr>("callee")) {
-      return builder.CreateCall(functionMapping.lookup(attr.getValue()),
-                                operandsRef);
-    } else {
-      auto *calleePtrType =
-          cast<llvm::PointerType>(operandsRef.front()->getType());
-      auto *calleeType =
-          cast<llvm::FunctionType>(calleePtrType->getElementType());
-      return builder.CreateCall(calleeType, operandsRef.front(),
-                                operandsRef.drop_front());
-    }
+    if (auto attr = op.getAttrOfType<FlatSymbolRefAttr>("callee"))
+      return builder.CreateCall(lookupFunction(attr.getValue()), operandsRef);
   };
 
   // Emit calls.  If the called function has a result, remap the corresponding
   // value.  Note that LLVM IR dialect CallOp has either 0 or 1 result.
   if (isa<LLVM::CallOp>(opInst)) {
     llvm::Value *result = convertCall(opInst);
     if (opInst.getNumResults() != 0) {
-      valueMapping[opInst.getResult(0)] = result;
+      mapValue(opInst.getResult(0), result);
       return success();
     }
     // Check that LLVM call returns void for 0-result functions.
     return success(result->getType()->isVoidTy());
   }
 
   if (auto inlineAsmOp = dyn_cast<LLVM::InlineAsmOp>(opInst)) {
     // TODO: refactor function type creation which usually occurs in std-LLVM
llvm::InlineAsm *inlineAsmInst =
                   convertAsmDialectToLLVM(*inlineAsmOp.asm_dialect()))
             : llvm::InlineAsm::get(
                   static_cast<llvm::FunctionType *>(convertType(ft)),
                   inlineAsmOp.asm_string(), inlineAsmOp.constraints(),
                   inlineAsmOp.has_side_effects(), inlineAsmOp.is_align_stack());
     llvm::Value *result =
         builder.CreateCall(inlineAsmInst, lookupValues(inlineAsmOp.operands()));
     if (opInst.getNumResults() != 0)
-      valueMapping[opInst.getResult(0)] = result;
+      mapValue(opInst.getResult(0), result);
     return success();
   }
 
   if (auto invOp = dyn_cast<LLVM::InvokeOp>(opInst)) {
     auto operands = lookupValues(opInst.getOperands());
     ArrayRef<llvm::Value *> operandsRef(operands);
     if (auto attr = opInst.getAttrOfType<FlatSymbolRefAttr>("callee")) {
-      builder.CreateInvoke(functionMapping.lookup(attr.getValue()),
-                           blockMapping[invOp.getSuccessor(0)],
-                           blockMapping[invOp.getSuccessor(1)], operandsRef);
+      builder.CreateInvoke(lookupFunction(attr.getValue()),
+                           lookupBlock(invOp.getSuccessor(0)),
+                           lookupBlock(invOp.getSuccessor(1)), operandsRef);
     } else {
       auto *calleePtrType =
           cast<llvm::PointerType>(operandsRef.front()->getType());
       auto *calleeType =
           cast<llvm::FunctionType>(calleePtrType->getElementType());
       builder.CreateInvoke(
-          calleeType, operandsRef.front(), blockMapping[invOp.getSuccessor(0)],
-          blockMapping[invOp.getSuccessor(1)], operandsRef.drop_front());
+          calleeType, operandsRef.front(), lookupBlock(invOp.getSuccessor(0)),
+          lookupBlock(invOp.getSuccessor(1)), operandsRef.drop_front());
     }
     return success();
   }
 
   if (auto lpOp = dyn_cast<LLVM::LandingpadOp>(opInst)) {
     llvm::Type *ty = convertType(lpOp.getType());
     llvm::LandingPadInst *lpi =
         builder.CreateLandingPad(ty, lpOp.getNumOperands());
 
     // Add clauses
-    for (auto operand : lookupValues(lpOp.getOperands())) {
+    for (llvm::Value *operand : lookupValues(lpOp.getOperands())) {
       // All operands should be constant - checked by verifier
-      if (auto constOperand = dyn_cast<llvm::Constant>(operand))
+      if (auto *constOperand = dyn_cast<llvm::Constant>(operand))
         lpi->addClause(constOperand);
     }
-    valueMapping[lpOp.getResult()] = lpi;
+    mapValue(lpOp.getResult(), lpi);
     return success();
   }
 
   // Emit branches.  We need to look up the remapped blocks and ignore the block
   // arguments that were transformed into PHI nodes.
   if (auto brOp = dyn_cast<LLVM::BrOp>(opInst)) {
     llvm::BranchInst *branch =
-        builder.CreateBr(blockMapping[brOp.getSuccessor()]);
-    branchMapping.try_emplace(&opInst, branch);
+        builder.CreateBr(lookupBlock(brOp.getSuccessor()));
+    mapBranch(&opInst, branch);
     return success();
   }
   if (auto condbrOp = dyn_cast<LLVM::CondBrOp>(opInst)) {
     auto weights = condbrOp.branch_weights();
     llvm::MDNode *branchWeights = nullptr;
     if (weights) {
       // Map weight attributes to LLVM metadata.
       auto trueWeight =
           weights.getValue().getValue(0).cast<IntegerAttr>().getInt();
       auto falseWeight =
           weights.getValue().getValue(1).cast<IntegerAttr>().getInt();
       branchWeights =
           llvm::MDBuilder(llvmModule->getContext())
               .createBranchWeights(static_cast<uint32_t>(trueWeight),
                                    static_cast<uint32_t>(falseWeight));
     }
     llvm::BranchInst *branch = builder.CreateCondBr(
-        valueMapping.lookup(condbrOp.getOperand(0)),
-        blockMapping[condbrOp.getSuccessor(0)],
-        blockMapping[condbrOp.getSuccessor(1)], branchWeights);
-    branchMapping.try_emplace(&opInst, branch);
+        lookupValue(condbrOp.getOperand(0)),
+        lookupBlock(condbrOp.getSuccessor(0)),
+        lookupBlock(condbrOp.getSuccessor(1)), branchWeights);
+    mapBranch(&opInst, branch);
     return success();
   }
   if (auto switchOp = dyn_cast<LLVM::SwitchOp>(opInst)) {
     llvm::MDNode *branchWeights = nullptr;
     if (auto weights = switchOp.branch_weights()) {
       llvm::SmallVector<uint32_t> weightValues;
       weightValues.reserve(weights->size());
       for (llvm::APInt weight : weights->cast<DenseIntElementsAttr>())
         weightValues.push_back(weight.getLimitedValue());
       branchWeights = llvm::MDBuilder(llvmModule->getContext())
                           .createBranchWeights(weightValues);
     }
 
     llvm::SwitchInst *switchInst =
-        builder.CreateSwitch(valueMapping[switchOp.value()],
-                             blockMapping[switchOp.defaultDestination()],
+        builder.CreateSwitch(lookupValue(switchOp.value()),
+                             lookupBlock(switchOp.defaultDestination()),
                              switchOp.caseDestinations().size(), branchWeights);
 
     auto *ty =
         llvm::cast<llvm::IntegerType>(convertType(switchOp.value().getType()));
     for (auto i :
          llvm::zip(switchOp.case_values()->cast<DenseIntElementsAttr>(),
                    switchOp.caseDestinations()))
       switchInst->addCase(
           llvm::ConstantInt::get(ty, std::get<0>(i).getLimitedValue()),
-          blockMapping[std::get<1>(i)]);
+          lookupBlock(std::get<1>(i)));
 
-    branchMapping.try_emplace(&opInst, switchInst);
+    mapBranch(&opInst, switchInst);
     return success();
   }
 
   // Emit addressof.  We need to look up the global value referenced by the
   // operation and store it in the MLIR-to-LLVM value mapping.  This does not
   // emit any LLVM instruction.
   if (auto addressOfOp = dyn_cast<LLVM::AddressOfOp>(opInst)) {
     LLVM::GlobalOp global = addressOfOp.getGlobal();
     LLVM::LLVMFuncOp function = addressOfOp.getFunction();
 
     // The verifier should not have allowed this.
     assert((global || function) &&
            "referencing an undefined global or function");
 
-    valueMapping[addressOfOp.getResult()] =
-        global ? globalsMapping.lookup(global)
-               : functionMapping.lookup(function.getName());
+    mapValue(addressOfOp.getResult(), global
+                                          ? globalsMapping.lookup(global)
+                                          : lookupFunction(function.getName()));
     return success();
   }
 
   if (ompDialect && opInst.getDialect() == ompDialect)
     return convertOmpOperation(opInst, builder);
 
   return opInst.emitError("unsupported or non-LLVM operation: ")
          << opInst.getName();
 }
 
 /// Convert block to LLVM IR.  Unless `ignoreArguments` is set, emit PHI nodes
 /// to define values corresponding to the MLIR block arguments.  These nodes
 /// are not connected to the source basic blocks, which may not exist yet.  Uses
 /// `builder` to construct the LLVM IR. Expects the LLVM IR basic block to have
 /// been created for `bb` and included in the block mapping.  Inserts new
 /// instructions at the end of the block and leaves `builder` in a state
 /// suitable for further insertion into the end of the block.
 LogicalResult ModuleTranslation::convertBlock(Block &bb, bool ignoreArguments,
                                               llvm::IRBuilder<> &builder) {
-  builder.SetInsertPoint(blockMapping[&bb]);
+  builder.SetInsertPoint(lookupBlock(&bb));
   auto *subprogram = builder.GetInsertBlock()->getParent()->getSubprogram();
 
   // Before traversing operations, make block arguments available through
   // value remapping and PHI nodes, but do not add incoming edges for the PHI
   // nodes just yet: those values may be defined by this or following blocks.
   // This step is omitted if "ignoreArguments" is set.  The arguments of the
   // first block have been already made available through the remapping of
   // LLVM function arguments.
   if (!ignoreArguments) {
     auto predecessors = bb.getPredecessors();
     unsigned numPredecessors =
         std::distance(predecessors.begin(), predecessors.end());
     for (auto arg : bb.getArguments()) {
       auto wrappedType = arg.getType();
       if (!isCompatibleType(wrappedType))
         return emitError(bb.front().getLoc(),
                          "block argument does not have an LLVM type");
       llvm::Type *type = convertType(wrappedType);
       llvm::PHINode *phi = builder.CreatePHI(type, numPredecessors);
-      valueMapping[arg] = phi;
+      mapValue(arg, phi);
     }
   }
 
   // Traverse operations.
   for (auto &op : bb) {
     // Set the current debug location within the builder.
     builder.SetCurrentDebugLocation(
         debugTranslation->translateLoc(op.getLoc(), subprogram));
if (op.getValueOrNull()) {
       } else if (!(cst = getLLVMConstant(type, op.getValueOrNull(),
                                          op.getLoc()))) {
         return failure();
       }
     } else if (Block *initializer = op.getInitializerBlock()) {
       llvm::IRBuilder<> builder(llvmModule->getContext());
       for (auto &op : initializer->without_terminator()) {
         if (failed(convertOperation(op, builder)) ||
-            !isa<llvm::Constant>(valueMapping.lookup(op.getResult(0))))
+            !isa<llvm::Constant>(lookupValue(op.getResult(0))))
           return emitError(op.getLoc(), "unemittable constant value");
       }
       ReturnOp ret = cast<ReturnOp>(initializer->getTerminator());
-      cst = cast<llvm::Constant>(valueMapping.lookup(ret.getOperand(0)));
+      cst = cast<llvm::Constant>(lookupValue(ret.getOperand(0)));
     }
 
     auto linkage = convertLinkageToLLVM(op.linkage());
     bool anyExternalLinkage =
         ((linkage == llvm::GlobalVariable::ExternalLinkage &&
           isa<llvm::UndefValue>(cst)) ||
          linkage == llvm::GlobalVariable::ExternalWeakLinkage);
     auto addrSpace = op.addr_space();

 }
 
 LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   // Clear the block, branch value mappings, they are only relevant within one
   // function.
   blockMapping.clear();
   valueMapping.clear();
   branchMapping.clear();
-  llvm::Function *llvmFunc = functionMapping.lookup(func.getName());
+  llvm::Function *llvmFunc = lookupFunction(func.getName());
 
   // Translate the debug information for this function.
   debugTranslation->translate(func, *llvmFunc);
 
   // Add function arguments to the value remapping table.
   // If there was noalias info then we decorate each argument accordingly.
   unsigned int argIdx = 0;
   for (auto kvp : llvm::zip(func.getArguments(), llvmFunc->args())) {
if (auto attr = func.getArgAttrOfType<UnitAttr>(argIdx, "llvm.byval")) {
       auto argTy = mlirArg.getType();
       if (!argTy.isa<LLVM::LLVMPointerType>())
         return func.emitError(
             "llvm.byval attribute attached to LLVM non-pointer argument");
       llvmArg.addAttrs(llvm::AttrBuilder().addByValAttr(
           llvmArg.getType()->getPointerElementType()));
     }
 
-    valueMapping[mlirArg] = &llvmArg;
+    mapValue(mlirArg, &llvmArg);
     argIdx++;
   }
 
   // Check the personality and set it.
   if (func.personality().hasValue()) {
     llvm::Type *ty = llvm::Type::getInt8PtrTy(llvmFunc->getContext());
     if (llvm::Constant *pfunc =
             getLLVMConstant(ty, func.personalityAttr(), func.getLoc()))
       llvmFunc->setPersonalityFn(pfunc);
   }
 
   // First, create all blocks so we can jump to them.
   llvm::LLVMContext &llvmContext = llvmFunc->getContext();
   for (auto &bb : func) {
     auto *llvmBB = llvm::BasicBlock::Create(llvmContext);
     llvmBB->insertInto(llvmFunc);
-    blockMapping[&bb] = llvmBB;
+    mapBlock(&bb, llvmBB);
   }
 
   // Then, convert blocks one by one in topological order to ensure defs are
   // converted before uses.
   auto blocks = topologicalSort(func);
   for (Block *bb : blocks) {
     llvm::IRBuilder<> builder(llvmContext);
     if (failed(convertBlock(*bb, bb->isEntryBlock(), builder)))
       return failure();
   }
 
   // Finally, after all blocks have been traversed and values mapped, connect
   // the PHI nodes to the results of preceding blocks.
-  connectPHINodes(func, valueMapping, blockMapping, branchMapping);
+  connectPHINodes(func, *this);
   return success();
 }
 
 LogicalResult ModuleTranslation::checkSupportedModuleOps(Operation *m) {
   for (Operation &o : getModuleBody(m).getOperations())
     if (!isa<LLVM::LLVMFuncOp, LLVM::GlobalOp>(&o) && !o.isKnownTerminator())
       return o.emitOpError("unsupported module-level operation");
   return success();
 }
 
 LogicalResult ModuleTranslation::convertFunctionSignatures() {
   // Declare all functions first because there may be function calls that form a
   // call graph with cycles, or global initializers that reference functions.
   for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     llvm::FunctionCallee llvmFuncCst = llvmModule->getOrInsertFunction(
         function.getName(),
         cast<llvm::FunctionType>(convertType(function.getType())));
     llvm::Function *llvmFunc = cast<llvm::Function>(llvmFuncCst.getCallee());
     llvmFunc->setLinkage(convertLinkageToLLVM(function.linkage()));
-    functionMapping[function.getName()] = llvmFunc;
+    mapFunction(function.getName(), llvmFunc);
 
     // Forward the pass-through attributes to LLVM.
     if (failed(forwardPassthroughAttributes(function.getLoc(),
                                             function.passthrough(), llvmFunc)))
       return failure();
   }
 
   return success();
llvm::Type *ModuleTranslation::convertType(Type type) {
   return typeTranslator.translateType(type);
 }
 
 /// A helper to look up remapped operands in the value remapping table.`
 SmallVector<llvm::Value *, 8>
 ModuleTranslation::lookupValues(ValueRange values) {
   SmallVector<llvm::Value *, 8> remapped;
   remapped.reserve(values.size());
-  for (Value v : values) {
-    assert(valueMapping.count(v) && "referencing undefined value");
-    remapped.push_back(valueMapping.lookup(v));
-  }
+  for (Value v : values)
+    remapped.push_back(lookupValue(v));
   return remapped;
 }
 
 std::unique_ptr<llvm::Module> ModuleTranslation::prepareLLVMModule(
     Operation *m, llvm::LLVMContext &llvmContext, StringRef name) {
   m->getContext()->getOrLoadDialect<LLVM::LLVMDialect>();
   auto llvmModule = std::make_unique<llvm::Module>(name, llvmContext);
   if (auto dataLayoutAttr =

mlir/tools/mlir-tblgen/LLVMIRConversionGen.cpp

static bool emitOneBuilder(const Record &record, raw_ostream &os) {
   llvm::raw_string_ostream bs(builder);
   while (auto loc = findNextVariable(builderStrRef)) {
     auto name = loc.in(builderStrRef).drop_front();
     // First, insert the non-matched part as is.
     bs << builderStrRef.substr(0, loc.pos);
     // Then, rewrite the name based on its kind.
     bool isVariadicOperand = isVariadicOperandName(op, name);
     if (isOperandName(op, name)) {
-      auto result = isVariadicOperand
-                        ? formatv("lookupValues(op.{0}())", name)
-                        : formatv("valueMapping.lookup(op.{0}())", name);
+      auto result = isVariadicOperand ? formatv("lookupValues(op.{0}())", name)
+                                      : formatv("lookupValue(op.{0}())", name);
       bs << result;
     } else if (isAttributeName(op, name)) {
       bs << formatv("op.{0}()", name);
     } else if (isResultName(op, name)) {
-      bs << formatv("valueMapping[op.{0}()]", name);
+      bs << formatv("mapValue(op.{0}())", name);
     } else if (name == "_resultType") {
       bs << "convertType(op.getResult().getType())";
     } else if (name == "_hasResult") {
       bs << "opInst.getNumResults() == 1";
     } else if (name == "_location") {
       bs << "opInst.getLoc()";
     } else if (name == "_numOperands") {
       bs << "opInst.getNumOperands()";