[BranchFolding] do not iterate the aliases of virtual registers

[jingyue]

Summary

MCRegAliasIterator only works for physical registers. So, do not run it
on virtual registers.

With this issue fixed, we can resurrect the BranchFolding pass in NVPTX
backend.

[jingyue]

Ping?

[bkramer]

This looks good with one comment below. Also the top comment in this file could use an update so it's clear that the pass has to handle vregs too.

lib/CodeGen/BranchFolding.cpp

1813–1819

The alias iterator has the IncludeSelf flag set, so shouldn't Reg be removed from the set even if it's a vreg?

[jingyue]

fixed the bug Ben pointed out

[jingyue]

Nice catch. Thank you!

[jingyue]

update header comments

lib/CodeGen/BranchFolding.cpp

 //===-- BranchFolding.cpp - Fold machine code branch instructions ---------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass forwards branches to unconditional branches to make them branch
 // directly to the target block.  This pass often results in dead MBB's, which
 // it then removes.
 //
 // Note that this pass must be run after register allocation, it cannot handle
-// SSA form.
+// SSA form. It also must handle virtual registers for targets that emit virtual
+// ISA (e.g. NVPTX).
 //
 //===----------------------------------------------------------------------===//
 
 #include "BranchFolding.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
for (MachineBasicBlock::succ_iterator SI = BB->succ_begin(),
          E = BB->succ_end(); SI != E; ++SI) {
     MachineBasicBlock *SuccBB = *SI;
     if (SuccBB != TrueBB)
       return SuccBB;
   }
   return nullptr;
 }
 
+template <class Container>
+static void addRegAndItsAliases(unsigned Reg, const TargetRegisterInfo *TRI,
+                                Container &Set) {
+  if (TargetRegisterInfo::isPhysicalRegister(Reg)) {
+    for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
+      Set.insert(*AI);
+  } else {
+    Set.insert(Reg);
+  }
+}
+
 /// findHoistingInsertPosAndDeps - Find the location to move common instructions
 /// in successors to. The location is usually just before the terminator,
 /// however if the terminator is a conditional branch and its previous
 /// instruction is the flag setting instruction, the previous instruction is
 /// the preferred location. This function also gathers uses and defs of the
 /// instructions from the insertion point to the end of the block. The data is
 /// used by HoistCommonCodeInSuccs to ensure safety.
 static
MachineBasicBlock::iterator findHoistingInsertPosAndDeps(MachineBasicBlock *MBB,
   for (unsigned i = 0, e = Loc->getNumOperands(); i != e; ++i) {
     const MachineOperand &MO = Loc->getOperand(i);
     if (!MO.isReg())
       continue;
     unsigned Reg = MO.getReg();
     if (!Reg)
       continue;
     if (MO.isUse()) {
-      for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-        Uses.insert(*AI);
+      addRegAndItsAliases(Reg, TRI, Uses);
     } else {
       if (!MO.isDead())
         // Don't try to hoist code in the rare case the terminator defines a
         // register that is later used.
         return MBB->end();
 
       // If the terminator defines a register, make sure we don't hoist
       // the instruction whose def might be clobbered by the terminator.
-      for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-        Defs.insert(*AI);
+      addRegAndItsAliases(Reg, TRI, Defs);
     }
   }
 
   if (Uses.empty())
     return Loc;
   if (Loc == MBB->begin())
     return MBB->end();
 
MachineBasicBlock::iterator findHoistingInsertPosAndDeps(MachineBasicBlock *MBB,
   for (unsigned i = 0, e = PI->getNumOperands(); i != e; ++i) {
     const MachineOperand &MO = PI->getOperand(i);
     if (!MO.isReg())
       continue;
     unsigned Reg = MO.getReg();
     if (!Reg)
       continue;
     if (MO.isUse()) {
-      for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-        Uses.insert(*AI);
+      addRegAndItsAliases(Reg, TRI, Uses);
     } else {
       if (Uses.erase(Reg)) {
+        if (TargetRegisterInfo::isPhysicalRegister(Reg)) {
           for (MCSubRegIterator SubRegs(Reg, TRI); SubRegs.isValid(); ++SubRegs)
             Uses.erase(*SubRegs); // Use sub-registers to be conservative
         }
-      for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-        Defs.insert(*AI);
+      }
+      addRegAndItsAliases(Reg, TRI, Defs);
     }
   }
 
   return PI;
 }
 
 /// HoistCommonCodeInSuccs - If the successors of MBB has common instruction
 /// sequence at the start of the function, move the instructions before MBB
while (TIB != TIE && FIB != FIE) {
     // Remove kills from LocalDefsSet, these registers had short live ranges.
     for (unsigned i = 0, e = TIB->getNumOperands(); i != e; ++i) {
       MachineOperand &MO = TIB->getOperand(i);
       if (!MO.isReg() || !MO.isUse() || !MO.isKill())
         continue;
       unsigned Reg = MO.getReg();
       if (!Reg || !LocalDefsSet.count(Reg))
         continue;
+      if (TargetRegisterInfo::isPhysicalRegister(Reg)) {
         for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
           LocalDefsSet.erase(*AI);
+      } else {
+        LocalDefsSet.erase(Reg);
+      }
     }
 
     // Track local defs so we can update liveins.
     for (unsigned i = 0, e = TIB->getNumOperands(); i != e; ++i) {
       MachineOperand &MO = TIB->getOperand(i);
       if (!MO.isReg() || !MO.isDef() || MO.isDead())
         continue;
       unsigned Reg = MO.getReg();
       if (!Reg)
         continue;
       LocalDefs.push_back(Reg);
-      for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI)
-        LocalDefsSet.insert(*AI);
+      addRegAndItsAliases(Reg, TRI, LocalDefsSet);
     }
 
     HasDups = true;
     ++TIB;
     ++FIB;
   }
 
   if (!HasDups)

lib/Target/NVPTX/NVPTXTargetMachine.cpp

 void NVPTXPassConfig::addIRPasses() {
   // The following passes are known to not play well with virtual regs hanging
   // around after register allocation (which in our case, is *all* registers).
   // We explicitly disable them here.  We do, however, need some functionality
   // of the PrologEpilogCodeInserter pass, so we emulate that behavior in the
   // NVPTXPrologEpilog pass (see NVPTXPrologEpilogPass.cpp).
   disablePass(&PrologEpilogCodeInserterID);
   disablePass(&MachineCopyPropagationID);
-  disablePass(&BranchFolderPassID);
   disablePass(&TailDuplicateID);
 
   addPass(createNVPTXImageOptimizerPass());
   TargetPassConfig::addIRPasses();
   addPass(createNVPTXAssignValidGlobalNamesPass());
   addPass(createGenericToNVVMPass());
   addPass(createNVPTXLowerKernelArgsPass(&getNVPTXTargetMachine()));
   // NVPTXLowerKernelArgs emits alloca for byval parameters which can often

test/CodeGen/NVPTX/branch-fold.ll

+; RUN: llc < %s -march=nvptx64 -mcpu=sm_35 -disable-cgp | FileCheck %s
+; Disable CGP which also folds branches, so that only BranchFolding is under
+; the spotlight.
+
+target datalayout = "e-i64:64-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-nvidia-cuda"
+
+define void @foo(i32 %x, float* %output) {
+; CHECK-LABEL: .visible .func foo(
+; CHECK-NOT: bra.uni
+; CHECK-NOT: LBB0_
+  %1 = icmp eq i32 %x, 1
+  br i1 %1, label %then, label %else
+
+then:
+  br label %merge
+
+else:
+  br label %merge
+
+merge:
+  store float 2.0, float* %output
+  ret void
+}