[GlobalISel] correctly record PHI sources coming from switch instructions

[t.p.northover]

Summary

GlobalISel lowers switches to a sequence of conditional branches spanning multiple basic blocks. This means that a single CFG edge at the IR level can change when the MIR has been generated. Successors are still well-behaved, but when querying the MachineBasicBlock predecessors (e.g. to translate a PHI), a more sophisticated analysis than getOrCreateBB is needed.

This adds a callback to IRTranslator that every translation which modifies the CFG (currently only switch translation) should be aware of. finishPendingPHIs then uses this extra information to create a correct PHI at the end of translation. If a switch instruction isn't involved in any of the PHI's predecessors then we fall back to the usually sufficient bijective map from BasicBlocks to MachineBasicBlocks.

This should fix at least one failure in the test-suite bot under GlobalISel (halide).

[rovka]

include/llvm/CodeGen/GlobalISel/IRTranslator.h

408

Typo: predecessor

[qcolombet]

Hi Tim,

Looks mostly good to me, the only thing I think we miss is clearing the map for the edges at the end of the function.
Nothing incorrect, just wasted memory :).

Cheers,
-Quentin

include/llvm/CodeGen/GlobalISel/IRTranslator.h

405

const CFGEdge &, MachineBasicBlock &

408

typo: predecessor

411

const CFGEdge &?

[t.p.northover]

include/llvm/CodeGen/GlobalISel/IRTranslator.h

405

I did think about this at the time, but decided a pair of pointers was simple enough to pass by value. I can change it, but I'm not convinced there's a strong argument either way.

[t.p.northover]

Fix typo & clear data between functions. Uploading while I've got it handy.

[qcolombet]

Hi Tim,

LGTM, one nitpick below.

Cheers,
-Quentin

include/llvm/CodeGen/GlobalISel/IRTranslator.h

405

Fine by me. I would still turn the pointer to NewPred into a reference though. That makes it obvious nullptr is not a valid argument :).

[t.p.northover]

include/llvm/CodeGen/GlobalISel/IRTranslator.h

405

I'm not so sure about that one either. The problem is that maps can't contain references, which means the querying interface essentially has to use pointers. Having those out of sync seems a bit worse than passing a raw pointer to me.

[qcolombet]

include/llvm/CodeGen/GlobalISel/IRTranslator.h

405

Fair enough.
I would add an assert that NewPred is not null in that method.

[t.p.northover]

Thanks Quentin. Committed as r291973. Hopefully the bot should be green soonish.

include/llvm/CodeGen/GlobalISel/IRTranslator.h

405

Good idea.

include/llvm/CodeGen/GlobalISel/IRTranslator.h

private:
   // do not appear in that map.
   SmallSetVector<const Constant *, 8> Constants;
 
   // N.b. it's not completely obvious that this will be sufficient for every
   // LLVM IR construct (with "invoke" being the obvious candidate to mess up our
   // lives.
   DenseMap<const BasicBlock *, MachineBasicBlock *> BBToMBB;
 
+  typedef std::pair<const BasicBlock *, const BasicBlock *> CFGEdge;
+  DenseMap<CFGEdge, SmallVector<MachineBasicBlock *, 1>> MachinePreds;
+
   // List of stubbed PHI instructions, for values and basic blocks to be filled
   // in once all MachineBasicBlocks have been created.
   SmallVector<std::pair<const PHINode *, MachineInstr *>, 4> PendingPHIs;
 
   /// Record of what frame index has been allocated to specified allocas for
   /// this function.
   DenseMap<const AllocaInst *, int> FrameIndices;
 
private:
   /// If such VReg does not exist, it is created.
   int getOrCreateFrameIndex(const AllocaInst &AI);
 
   /// Get the alignment of the given memory operation instruction. This will
   /// either be the explicitly specified value or the ABI-required alignment for
   /// the type being accessed (according to the Module's DataLayout).
   unsigned getMemOpAlignment(const Instruction &I);
 
-  /// Get the MachineBasicBlock that represents \p BB.
-  /// If such basic block does not exist, it is created.
+  /// Get the MachineBasicBlock that represents \p BB. Specifically, the block
+  /// returned will be the head of the translated block (suitable for branch
+  /// destinations). If such basic block does not exist, it is created.
   MachineBasicBlock &getOrCreateBB(const BasicBlock &BB);
 
+  /// Record \p NewPred as a Machine predecessor to `Edge.second`, corresponding
+  /// to `Edge.first` at the IR level. This is used when IRTranslation creates
+  /// multiple MachineBasicBlocks for a given IR block and the CFG is no longer
+  /// represented simply by the IR-level CFG.
+  void addMachineCFGPred(CFGEdge Edge, MachineBasicBlock *NewPred);
+
+  /// Returns the Machine IR predecessors for the given IR CFG edge. Usually
+  /// this is just the single MachineBasicBlock corresponding to the predecessor
+  /// in the IR. More complex lowering can result in multiple MachineBasicBlocks
+  /// preceding the original though (e.g. switch instructions).
+  ArrayRef<MachineBasicBlock *> getMachinePredBBs(CFGEdge Edge) {
+    auto RemappedEdge = MachinePreds.find(Edge);
+    if (RemappedEdge != MachinePreds.end())
+      return RemappedEdge->second;
+    return &getOrCreateBB(*Edge.first);
+  }
 
 public:
   // Ctor, nothing fancy.
   IRTranslator();
 
   StringRef getPassName() const override { return "IRTranslator"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override;

lib/CodeGen/GlobalISel/IRTranslator.cpp

 //===-- llvm/CodeGen/GlobalISel/IRTranslator.cpp - IRTranslator --*- C++ -*-==//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 /// \file
 /// This file implements the IRTranslator class.
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/GlobalISel/IRTranslator.h"
 
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/GlobalISel/CallLowering.h"
 #include "llvm/CodeGen/Analysis.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
if (!MBB) {
     MF->push_back(MBB);
 
     if (BB.hasAddressTaken())
       MBB->setHasAddressTaken();
   }
   return *MBB;
 }
 
+void IRTranslator::addMachineCFGPred(CFGEdge Edge, MachineBasicBlock *NewPred) {
+  MachinePreds[Edge].push_back(NewPred);
+}
+
 bool IRTranslator::translateBinaryOp(unsigned Opcode, const User &U,
                                      MachineIRBuilder &MIRBuilder) {
   // FIXME: handle signed/unsigned wrapping flags.
 
   // Get or create a virtual register for each value.
   // Unless the value is a Constant => loadimm cst?
   // or inline constant each time?
   // Creation of a virtual register needs to have a size.
bool IRTranslator::translateSwitch(const User &U,
   // FIXME: could we share most of the logic/code in
   // SelectionDAGBuilder::visitSwitch between SelectionDAG and GlobalISel?
   // At first sight, it seems most of the logic in there is independent of
   // SelectionDAG-specifics and a lot of work went in to optimize switch
   // lowering in there.
 
   const SwitchInst &SwInst = cast<SwitchInst>(U);
   const unsigned SwCondValue = getOrCreateVReg(*SwInst.getCondition());
+  const BasicBlock *OrigBB = SwInst.getParent();
 
   LLT LLTi1 = LLT(*Type::getInt1Ty(U.getContext()), *DL);
   for (auto &CaseIt : SwInst.cases()) {
     const unsigned CaseValueReg = getOrCreateVReg(*CaseIt.getCaseValue());
     const unsigned Tst = MRI->createGenericVirtualRegister(LLTi1);
     MIRBuilder.buildICmp(CmpInst::ICMP_EQ, Tst, CaseValueReg, SwCondValue);
-    MachineBasicBlock &CurBB = MIRBuilder.getMBB();
-    MachineBasicBlock &TrueBB = getOrCreateBB(*CaseIt.getCaseSuccessor());
-
-    MIRBuilder.buildBrCond(Tst, TrueBB);
-    CurBB.addSuccessor(&TrueBB);
+    MachineBasicBlock &CurMBB = MIRBuilder.getMBB();
+    const BasicBlock *TrueBB = CaseIt.getCaseSuccessor();
+    MachineBasicBlock &TrueMBB = getOrCreateBB(*TrueBB);
+
+    MIRBuilder.buildBrCond(Tst, TrueMBB);
+    CurMBB.addSuccessor(&TrueMBB);
+    addMachineCFGPred({OrigBB, TrueBB}, &CurMBB);
 
-    MachineBasicBlock *FalseBB =
+    MachineBasicBlock *FalseMBB =
         MF->CreateMachineBasicBlock(SwInst.getParent());
-    MF->push_back(FalseBB);
-    MIRBuilder.buildBr(*FalseBB);
-    CurBB.addSuccessor(FalseBB);
+    MF->push_back(FalseMBB);
+    MIRBuilder.buildBr(*FalseMBB);
+    CurMBB.addSuccessor(FalseMBB);
 
-    MIRBuilder.setMBB(*FalseBB);
+    MIRBuilder.setMBB(*FalseMBB);
   }
   // handle default case
-  MachineBasicBlock &DefaultBB = getOrCreateBB(*SwInst.getDefaultDest());
-  MIRBuilder.buildBr(DefaultBB);
-  MIRBuilder.getMBB().addSuccessor(&DefaultBB);
+  const BasicBlock *DefaultBB = SwInst.getDefaultDest();
+  MachineBasicBlock &DefaultMBB = getOrCreateBB(*DefaultBB);
+  MIRBuilder.buildBr(DefaultMBB);
+  MachineBasicBlock &CurMBB = MIRBuilder.getMBB();
+  CurMBB.addSuccessor(&DefaultMBB);
+  addMachineCFGPred({OrigBB, DefaultBB}, &CurMBB);
 
   return true;
 }
 
 bool IRTranslator::translateLoad(const User &U, MachineIRBuilder &MIRBuilder) {
   const LoadInst &LI = cast<LoadInst>(U);
 
   if (!TPC->isGlobalISelAbortEnabled() && LI.isAtomic())
void IRTranslator::finishPendingPhis() {
   for (std::pair<const PHINode *, MachineInstr *> &Phi : PendingPHIs) {
     const PHINode *PI = Phi.first;
     MachineInstrBuilder MIB(*MF, Phi.second);
 
     // All MachineBasicBlocks exist, add them to the PHI. We assume IRTranslator
     // won't create extra control flow here, otherwise we need to find the
     // dominating predecessor here (or perhaps force the weirder IRTranslators
     // to provide a simple boundary).
+    SmallSet<const BasicBlock *, 4> HandledPreds;
+
     for (unsigned i = 0; i < PI->getNumIncomingValues(); ++i) {
-      assert(BBToMBB[PI->getIncomingBlock(i)]->isSuccessor(MIB->getParent()) &&
-             "I appear to have misunderstood Machine PHIs");
-      MIB.addUse(getOrCreateVReg(*PI->getIncomingValue(i)));
-      MIB.addMBB(BBToMBB[PI->getIncomingBlock(i)]);
+      auto IRPred = PI->getIncomingBlock(i);
+      if (HandledPreds.count(IRPred))
+        continue;
+
+      HandledPreds.insert(IRPred);
+      unsigned ValReg = getOrCreateVReg(*PI->getIncomingValue(i));
+      for (auto Pred : getMachinePredBBs({IRPred, PI->getParent()})) {
+        assert(Pred->isSuccessor(MIB->getParent()) &&
+               "incorrect CFG at MachineBasicBlock level");
+        MIB.addUse(ValReg);
+        MIB.addMBB(Pred);
+      }
     }
   }
 }
 
 bool IRTranslator::translate(const Instruction &Inst) {
   CurBuilder.setDebugLoc(Inst.getDebugLoc());
   switch(Inst.getOpcode()) {
 #define HANDLE_INST(NUM, OPCODE, CLASS) \

 
 void IRTranslator::finalizeFunction() {
   // Release the memory used by the different maps we
   // needed during the translation.
   PendingPHIs.clear();
   ValToVReg.clear();
   FrameIndices.clear();
   Constants.clear();
+  MachinePreds.clear();
 }
 
 bool IRTranslator::runOnMachineFunction(MachineFunction &CurMF) {
   MF = &CurMF;
   const Function &F = *MF->getFunction();
   if (F.empty())
     return false;
   CLI = MF->getSubtarget().getCallLowering();

test/CodeGen/AArch64/GlobalISel/arm64-irtranslator.ll

case200:
   %tmp2 = add i32 %argc, 2
   br label %return
 
 return:
   %res = phi i32 [ %tmp0, %default ], [ %tmp1, %case100 ], [ %tmp2, %case200 ]
   ret i32 %res
 }
 
+  ; The switch lowering code changes the CFG, which means that the original
+  ; %entry block is no longer a predecessor for the phi instruction. We need to
+  ; use the correct lowered MachineBasicBlock instead.
+; CHECK-LABEL: name: test_cfg_remap
+
+; CHECK: bb.5.entry:
+; CHECK-NEXT: successors: %[[PHI_BLOCK:bb.[0-9]+.phi.block]]
+; CHECK: G_BR %[[PHI_BLOCK]]
+
+; CHECK: [[PHI_BLOCK]]:
+; CHECK-NEXT: PHI %{{.*}}(s32), %bb.5.entry
+define i32 @test_cfg_remap(i32 %in) {
+entry:
+  switch i32 %in, label %phi.block [i32 1, label %next
+                                    i32 57, label %other]
+
+next:
+  br label %phi.block
+
+other:
+  ret i32 undef
+
+phi.block:
+  %res = phi i32 [1, %entry], [42, %next]
+  ret i32 %res
+}
+
+; CHECK-LABEL: name: test_cfg_remap_multiple_preds
+; CHECK: PHI [[ENTRY:%.*]](s32), %bb.{{[0-9]+}}.entry, [[ENTRY]](s32), %bb.{{[0-9]+}}.entry
+define i32 @test_cfg_remap_multiple_preds(i32 %in) {
+entry:
+  switch i32 %in, label %odd [i32 1, label %next
+                              i32 57, label %other
+                              i32 128, label %phi.block
+                              i32 256, label %phi.block]
+odd:
+  unreachable
+
+next:
+  br label %phi.block
+
+other:
+  ret i32 undef
+
+phi.block:
+  %res = phi i32 [1, %entry], [1, %entry], [42, %next]
+  ret i32 12
+}
+
 ; Tests for or.
 ; CHECK-LABEL: name: ori64
 ; CHECK: [[ARG1:%[0-9]+]](s64) = COPY %x0
 ; CHECK-NEXT: [[ARG2:%[0-9]+]](s64) = COPY %x1
 ; CHECK-NEXT: [[RES:%[0-9]+]](s64) = G_OR [[ARG1]], [[ARG2]]
 ; CHECK-NEXT: %x0 = COPY [[RES]]
 ; CHECK-NEXT: RET_ReallyLR implicit %x0
 define i64 @ori64(i64 %arg1, i64 %arg2) {