[CodeGen] Postprocess PHI nodes for callbr

[void]

Summary

When processing PHI nodes after a callbr, we need to make sure that the
PHI nodes on the default branch are resolved after the callbr (inserted
after INLINEASM_BR). The PHI node values on the indirect branches are
processed before the INLINEASM_BR.

[nickdesaulniers]

test?

[void]

Add testcase.

[void]

In D88813#2312195, @nickdesaulniers wrote:

test?

Done.

[void]

Update with missing elements so that it works.

[void]

@jyknight This doesn't have the loop refactored like you asked for. I'll work on that today.

[nickdesaulniers]

test case looks reasonable (I haven't run it pre-patch to see what goes wrong yet). Just curious, is this CL related to D88823?

[void]

In D88813#2312627, @nickdesaulniers wrote:

test case looks reasonable (I haven't run it pre-patch to see what goes wrong yet). Just curious, is this CL related to D88823?

No. I thought so, but they're different issues.

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.h

private:
   }
   void visitUserOp2(const Instruction &I) {
     llvm_unreachable("UserOp2 should not exist at instruction selection time!");
   }
 
   void processIntegerCallValue(const Instruction &I,
                                SDValue Value, bool IsSigned);
 
-  void HandlePHINodesInSuccessorBlocks(const BasicBlock *LLVMBB);
+  void HandlePHINodesInSuccessorBlocks(const BasicBlock *LLVMBB,
+                                       bool FirstIter);
 
   void emitInlineAsmError(const CallBase &Call, const Twine &Message);
 
   /// If V is an function argument then create corresponding DBG_VALUE machine
   /// instruction for it now. At the end of instruction selection, they will be
   /// inserted to the entry BB.
   bool EmitFuncArgumentDbgValue(const Value *V, DILocalVariable *Variable,
                                 DIExpression *Expr, DILocation *DL,

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

SDValue SelectionDAGBuilder::getControlRoot() {
   PendingExports.append(PendingConstrainedFPStrict.begin(),
                         PendingConstrainedFPStrict.end());
   PendingConstrainedFPStrict.clear();
   return updateRoot(PendingExports);
 }
 
 void SelectionDAGBuilder::visit(const Instruction &I) {
   // Set up outgoing PHI node register values before emitting the terminator.
-  if (I.isTerminator()) {
-    HandlePHINodesInSuccessorBlocks(I.getParent());
-  }
+  if (I.isTerminator())
+    HandlePHINodesInSuccessorBlocks(I.getParent(), true);
 
   // Increase the SDNodeOrder if dealing with a non-debug instruction.
   if (!isa<DbgInfoIntrinsic>(I))
     ++SDNodeOrder;
 
   CurInst = &I;
 
   visit(I.getOpcode(), I);
 
   if (!I.isTerminator() && !HasTailCall &&
       !isa<GCStatepointInst>(I)) // statepoints handle their exports internally
     CopyToExportRegsIfNeeded(&I);
 
   CurInst = nullptr;
+  ConstantsOut.clear();
 }
 
 void SelectionDAGBuilder::visitPHI(const PHINode &) {
   llvm_unreachable("SelectionDAGBuilder shouldn't visit PHI nodes!");
 }
 
 void SelectionDAGBuilder::visit(unsigned Opcode, const User &I) {
   // Note: this doesn't use InstVisitor, because it has to work with
void SelectionDAGBuilder::visitCallBr(const CallBrInst &I) {
   assert(!I.hasOperandBundlesOtherThan(
              {LLVMContext::OB_deopt, LLVMContext::OB_funclet}) &&
          "Cannot lower callbrs with arbitrary operand bundles yet!");
 
   assert(I.isInlineAsm() && "Only know how to handle inlineasm callbr");
   visitInlineAsm(I);
   CopyToExportRegsIfNeeded(&I);
 
+  // Set up outgoing PHI node register values before emitting the branch.
+  HandlePHINodesInSuccessorBlocks(I.getParent(), false);
+
   // Retrieve successors.
   MachineBasicBlock *Return = FuncInfo.MBBMap[I.getDefaultDest()];
 
   // Update successor info.
   addSuccessorWithProb(CallBrMBB, Return, BranchProbability::getOne());
   for (unsigned i = 0, e = I.getNumIndirectDests(); i < e; ++i) {
     MachineBasicBlock *Target = FuncInfo.MBBMap[I.getIndirectDest(i)];
     addSuccessorWithProb(CallBrMBB, Target, BranchProbability::getZero());

 }
 
 /// Handle PHI nodes in successor blocks.  Emit code into the SelectionDAG to
 /// ensure constants are generated when needed.  Remember the virtual registers
 /// that need to be added to the Machine PHI nodes as input.  We cannot just
 /// directly add them, because expansion might result in multiple MBB's for one
 /// BB.  As such, the start of the BB might correspond to a different MBB than
 /// the end.
-void
-SelectionDAGBuilder::HandlePHINodesInSuccessorBlocks(const BasicBlock *LLVMBB) {
+void SelectionDAGBuilder::HandlePHINodesInSuccessorBlocks(
+    const BasicBlock *LLVMBB, bool FirstIter) {
   const Instruction *TI = LLVMBB->getTerminator();
+  const CallBrInst *CI = dyn_cast<CallBrInst>(TI);
 
   SmallPtrSet<MachineBasicBlock *, 4> SuccsHandled;
 
   // Check PHI nodes in successors that expect a value to be available from this
   // block.
   for (unsigned succ = 0, e = TI->getNumSuccessors(); succ != e; ++succ) {
     const BasicBlock *SuccBB = TI->getSuccessor(succ);
     if (!isa<PHINode>(SuccBB->begin())) continue;
+
+    if (CI) {
+      if (FirstIter) {
+        // Don't push PHI node values back before an INLINEASM_BR instruction on
+        // the default branch.
+        if (SuccBB == CI->getDefaultDest())
+          continue;
+      } else {
+        // Don't push PHI node values back after an INLINEASM_BR instruction on
+        // the indirect branch.
+        if (SuccBB != CI->getDefaultDest())
+          continue;
+      }
+    }
+
     MachineBasicBlock *SuccMBB = FuncInfo.MBBMap[SuccBB];
 
     // If this terminator has multiple identical successors (common for
     // switches), only handle each succ once.
     if (!SuccsHandled.insert(SuccMBB).second)
       continue;
 
     MachineBasicBlock::iterator MBBI = SuccMBB->begin();
for (const PHINode &PN : SuccBB->phis()) {
         unsigned NumRegisters = TLI.getNumRegisters(*DAG.getContext(), VT);
         for (unsigned i = 0, e = NumRegisters; i != e; ++i)
           FuncInfo.PHINodesToUpdate.push_back(
               std::make_pair(&*MBBI++, Reg + i));
         Reg += NumRegisters;
       }
     }
   }
-
-  ConstantsOut.clear();
 }
 
 /// Add a successor MBB to ParentMBB< creating a new MachineBB for BB if SuccMBB
 /// is 0.
 MachineBasicBlock *
 SelectionDAGBuilder::StackProtectorDescriptor::
 AddSuccessorMBB(const BasicBlock *BB,
                 MachineBasicBlock *ParentMBB,

llvm/test/CodeGen/X86/callbr-asm-phi-nodes.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=i686-- -verify-machineinstrs < %s | FileCheck %s
+
+@var = global i32 0, align 4
+
+define i32 @test1() {
+; CHECK-LABEL: test1:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    #APP
+; CHECK-NEXT:    jmp .Ltmp0
+; CHECK-NEXT:    #NO_APP
+; CHECK-NEXT:    xorl %eax, %eax
+; CHECK-NEXT:  # %bb.2:
+; CHECK-NEXT:    retl
+; CHECK-NEXT:  .Ltmp0: # Block address taken
+; CHECK-NEXT:  .LBB0_1:
+; CHECK-NEXT:    movl $1, %eax
+; CHECK-NEXT:    retl
+  callbr void asm sideeffect "jmp ${1:l}", "*m,X,~{dirflag},~{fpsr},~{flags}"(i32* nonnull @var, i8* blockaddress(@test1, %1))
+          to label %2 [label %1]
+
+1:                                                ; preds = %0
+  br label %2
+
+2:                                                ; preds = %0, %1
+  %3 = phi i32 [ 1, %1 ], [ 0, %0 ]
+  ret i32 %3
+}