3 files changed, 55 insertions, 12 deletions
diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.cpp b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
index e577458f8fc..c71cdbb6d6f 100644
--- a/llvm/lib/CodeGen/MIRParser/MIParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
@@ -174,9 +174,8 @@ private:
   /// Return 0 if the name isn't a subregister index class.
   unsigned getSubRegIndex(StringRef Name);
 
-  void initSlots2BasicBlocks();
-
   const BasicBlock *getIRBlock(unsigned Slot);
+  const BasicBlock *getIRBlock(unsigned Slot, const Function &F);
 
   void initNames2TargetIndices();
 
@@ -858,7 +857,7 @@ bool MIParser::parseIRBlock(BasicBlock *&BB, const Function &F) {
     unsigned SlotNumber = 0;
     if (getUnsigned(SlotNumber))
       return true;
-    BB = const_cast<BasicBlock *>(getIRBlock(SlotNumber));
+    BB = const_cast<BasicBlock *>(getIRBlock(SlotNumber, F));
     if (!BB)
       return error(Twine("use of undefined IR block '%ir-block.") +
                    Twine(SlotNumber) + "'");
@@ -1208,11 +1207,10 @@ unsigned MIParser::getSubRegIndex(StringRef Name) {
   return SubRegInfo->getValue();
 }
 
-void MIParser::initSlots2BasicBlocks() {
-  if (!Slots2BasicBlocks.empty())
-    return;
-  const auto &F = *MF.getFunction();
-  ModuleSlotTracker MST(F.getParent());
+static void initSlots2BasicBlocks(
+    const Function &F,
+    DenseMap<unsigned, const BasicBlock *> &Slots2BasicBlocks) {
+  ModuleSlotTracker MST(F.getParent(), /*ShouldInitializeAllMetadata=*/false);
   MST.incorporateFunction(F);
   for (auto &BB : F) {
     if (BB.hasName())
@@ -1224,14 +1222,29 @@ void MIParser::initSlots2BasicBlocks() {
   }
 }
 
-const BasicBlock *MIParser::getIRBlock(unsigned Slot) {
-  initSlots2BasicBlocks();
+static const BasicBlock *getIRBlockFromSlot(
+    unsigned Slot,
+    const DenseMap<unsigned, const BasicBlock *> &Slots2BasicBlocks) {
   auto BlockInfo = Slots2BasicBlocks.find(Slot);
   if (BlockInfo == Slots2BasicBlocks.end())
     return nullptr;
   return BlockInfo->second;
 }
 
+const BasicBlock *MIParser::getIRBlock(unsigned Slot) {
+  if (Slots2BasicBlocks.empty())
+    initSlots2BasicBlocks(*MF.getFunction(), Slots2BasicBlocks);
+  return getIRBlockFromSlot(Slot, Slots2BasicBlocks);
+}
+
+const BasicBlock *MIParser::getIRBlock(unsigned Slot, const Function &F) {
+  if (&F == MF.getFunction())
+    return getIRBlock(Slot);
+  DenseMap<unsigned, const BasicBlock *> CustomSlots2BasicBlocks;
+  initSlots2BasicBlocks(F, CustomSlots2BasicBlocks);
+  return getIRBlockFromSlot(Slot, CustomSlots2BasicBlocks);
+}
+
 void MIParser::initNames2TargetIndices() {
   if (!Names2TargetIndices.empty())
     return;
diff --git a/llvm/lib/CodeGen/MIRPrinter.cpp b/llvm/lib/CodeGen/MIRPrinter.cpp
index cc131bc4641..e162a8f73e2 100644
--- a/llvm/lib/CodeGen/MIRPrinter.cpp
+++ b/llvm/lib/CodeGen/MIRPrinter.cpp
@@ -476,7 +476,16 @@ void MIPrinter::printIRBlockReference(const BasicBlock &BB) {
     printLLVMNameWithoutPrefix(OS, BB.getName());
     return;
   }
-  int Slot = MST.getLocalSlot(&BB);
+  const Function *F = BB.getParent();
+  int Slot;
+  if (F == MST.getCurrentFunction()) {
+    Slot = MST.getLocalSlot(&BB);
+  } else {
+    ModuleSlotTracker CustomMST(F->getParent(),
+                                /*ShouldInitializeAllMetadata=*/false);
+    CustomMST.incorporateFunction(*F);
+    Slot = CustomMST.getLocalSlot(&BB);
+  }
   if (Slot == -1)
     OS << "<badref>";
   else
diff --git a/llvm/test/CodeGen/MIR/X86/block-address-operands.mir b/llvm/test/CodeGen/MIR/X86/block-address-operands.mir
index c96720d3559..67e298e4932 100644
--- a/llvm/test/CodeGen/MIR/X86/block-address-operands.mir
+++ b/llvm/test/CodeGen/MIR/X86/block-address-operands.mir
@@ -26,6 +26,12 @@
     ret void
   }
 
+  define void @slot_in_other_function(i8** %addr) {
+  entry:
+    store volatile i8* blockaddress(@test3, %0), i8** %addr
+    ret void
+  }
+
   define void @test3() {
   entry:
     store volatile i8* blockaddress(@test3, %0), i8** @addr
@@ -81,6 +87,20 @@ body:
       - RETQ
 ...
 ---
+name:            slot_in_other_function
+tracksRegLiveness: true
+body:
+  - id:              0
+    name:            entry
+    liveins:         [ '%rdi' ]
+    instructions:
+# CHECK: name: slot_in_other_function
+# CHECK: %rax = LEA64r %rip, 1, _, blockaddress(@test3, %ir-block.0), _
+      - '%rax = LEA64r %rip, 1, _, blockaddress(@test3, %ir-block.0), _'
+      - 'MOV64mr killed %rdi, 1, _, 0, _, killed %rax'
+      - RETQ
+...
+---
 name:            test3
 tracksRegLiveness: true
 body:
@@ -88,7 +108,8 @@ body:
     name:            entry
     successors:      [ '%bb.1' ]
     instructions:
-    # CHECK: %rax = LEA64r %rip, 1, _, blockaddress(@test3, %ir-block.0), _
+# CHECK: name: test3
+# CHECK: %rax = LEA64r %rip, 1, _, blockaddress(@test3, %ir-block.0), _
       - '%rax = LEA64r %rip, 1, _, blockaddress(@test3, %ir-block.0), _'
       - 'MOV64mr %rip, 1, _, @addr, _, killed %rax'
       - 'JMP64m %rip, 1, _, @addr, _'