[mlir] Explicitly track branch instructions in translation to LLVM IR

[ftynse]

Summary

The current implementation of the translation to LLVM IR relies on the
existence of a one-to-one mapping between MLIR blocks and LLVM IR basic blocks
in order to configure PHI nodes with appropriate source blocks. The one-to-one
mapping model is broken in presence of OpenMP operations that use LLVM's
OpenMPIRBuilder, which produces multiple blocks under the hood. This can lead
to invalid LLVM IR being emitted if OpenMPIRBuilder moved the branch operation
into a basic block different from the one it was originally created in;
specifically, a block that is not a direct predecessor could be used in the PHI
node. Instead, keep track of the mapping between MLIR LLVM dialect branch
operations and their LLVM IR counterparts and take the parent basic block of
the LLVM IR instruction at the moment of connecting the PHI nodes to
predecessors.

This behavior cannot be triggered as of now, but will be once we introduce the
conversion of OpenMP workshare loops.

[rriddle]

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

158

This will break when the terminator is not a BranchInst, e.g. SwitchInst or CallBr.

[ftynse]

Address review.

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

158

We don't have either of those in the LLVM dialect yet, but I updated to Instruction * for future-proofness.

[kiranchandramohan]

Thanks for fixing this. LGTM.

[ftynse]

Rebase

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

private:
   /// A stateful object used to translate types.
   TypeToLLVMIRTranslator typeTranslator;
 
 protected:
   /// 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 conneting 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 Value getPHISourceValue(Block *current, Block *pred,
 
   return condBranchOp.getSuccessor(0) == current
              ? condBranchOp.trueDestOperands()[index]
              : condBranchOp.falseDestOperands()[index];
 }
 
 /// 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) {
+static void connectPHINodes(
+    T &func, const DenseMap<Value, llvm::Value *> &valueMapping,
+    const DenseMap<Block *, llvm::BasicBlock *> &blockMapping,
+    const DenseMap<Operation *, llvm::Instruction *> &branchMapping) {
   // 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);
     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
+        // (in particular, OpenMP operations using OpenMPIRBuilder) may have
+        // split the blocks.
+        llvm::Instruction *terminator =
+            branchMapping.lookup(pred->getTerminator());
+        assert(terminator && "missing the mapping for a terminator");
         phiNode.addIncoming(valueMapping.lookup(getPHISourceValue(
                                 bb, pred, numArguments, index)),
-                            blockMapping.lookup(pred));
+                            terminator->getParent());
       }
     }
   }
 }
 
 /// Sort function blocks topologically.
 template <typename T>
 static llvm::SetVector<Block *> topologicalSort(T &f) {
for (auto indexedBB : llvm::enumerate(blocks)) {
 
     if (failed(convertBlock(*bb, /*ignoreArguments=*/indexedBB.index() == 0))) {
       bodyGenStatus = failure();
       return;
     }
   }
   // Finally, after all blocks have been traversed and values mapped,
   // connect the PHI nodes to the results of preceding blocks.
-  connectPHINodes(region, valueMapping, blockMapping);
+  connectPHINodes(region, valueMapping, blockMapping, branchMapping);
 }
 
 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();
if (auto lpOp = dyn_cast<LLVM::LandingpadOp>(opInst)) {
     }
     valueMapping[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);
     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));
     }
-    builder.CreateCondBr(valueMapping.lookup(condbrOp.getOperand(0)),
-                         blockMapping[condbrOp.getSuccessor(0)],
-                         blockMapping[condbrOp.getSuccessor(1)], branchWeights);
+    llvm::BranchInst *branch = builder.CreateCondBr(
+        valueMapping.lookup(condbrOp.getOperand(0)),
+    branchMapping.try_emplace(&opInst, branch);
     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();
for (Attribute attr : *attributes) {
     if (failed(checkedAddLLVMFnAttribute(loc, llvmFunc, keyAttr.getValue(),
                                          valueAttr.getValue())))
       return failure();
   }
   return success();
 }
 
 LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
-  // Clear the block and value mappings, they are only relevant within one
+  // 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());
 
   // 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;
LogicalResult ModuleTranslation::convertOneFunction(LLVMFuncOp func) {
   for (auto indexedBB : llvm::enumerate(blocks)) {
     auto *bb = indexedBB.value();
     if (failed(convertBlock(*bb, /*ignoreArguments=*/indexedBB.index() == 0)))
       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);
+  connectPHINodes(func, valueMapping, blockMapping, branchMapping);
   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();