[SPIR-V] Remove switch G_ICMP+G_BRCOND+G_BR before ISel

[mpaszkowski]

Summary

IRTranslator lowers switches to [G_SUB] + G_ICMP + G_BRCOND + G_BR sequences. Since values and destination MBBs are included in the spv_switch intrinsics, the sequences are not needed for ISel.

Before this patch, the information decoded by these sequences were added to spv_switch intrinsics in SPIRVPreLegalizer and the sequences were kept until SPIRVModuleAnalysis where they were marked skipped for emission.

After this patch, the [G_SUB] + G_ICMP + G_BRCOND + G_BR sequences and MBBs containing only these MIs are erased in SPIRVPreLegalizer.

[mpaszkowski]

Updated the commit with a fix to the failing OpenCL CTS test and added one new LIT test.

[mpaszkowski]

Added assert(!MBB.empty() && "MBB is empty!"). This should make it easier to spot similar issues in the future.

[andreytr]

Approve for lit tests.

[iliya-diyachkov]

The patch looks good to me, thanks Michal.

llvm/lib/Target/SPIRV/SPIRVAsmPrinter.cpp

 
 // Emit OpFunctionEnd at the end of MF and clear BBNumToRegMap.
 void SPIRVAsmPrinter::emitFunctionBodyEnd() {
   outputOpFunctionEnd();
   MAI->BBNumToRegMap.clear();
 }
 
 void SPIRVAsmPrinter::emitOpLabel(const MachineBasicBlock &MBB) {
-  if (MAI->MBBsToSkip.contains(&MBB))
-    return;
   MCInst LabelInst;
   LabelInst.setOpcode(SPIRV::OpLabel);
   LabelInst.addOperand(MCOperand::createReg(MAI->getOrCreateMBBRegister(MBB)));
   outputMCInst(LabelInst);
 }
 
 void SPIRVAsmPrinter::emitBasicBlockStart(const MachineBasicBlock &MBB) {
+  assert(!MBB.empty() && "MBB is empty!");
+
   // If it's the first MBB in MF, it has OpFunction and OpFunctionParameter, so
   // OpLabel should be output after them.
   if (MBB.getNumber() == MF->front().getNumber()) {
     for (const MachineInstr &MI : MBB)
       if (MI.getOpcode() == SPIRV::OpFunction)
         return;
     // TODO: this case should be checked by the verifier.
     report_fatal_error("OpFunction is expected in the front MBB of MF");

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.h

struct ModuleAnalysisInfo {
   DenseMap<unsigned, Register> ExtInstSetMap;
   // Contains the list of all global OpVariables in the module.
   SmallVector<MachineInstr *, 4> GlobalVarList;
   // Maps functions to corresponding function ID registers.
   DenseMap<const Function *, Register> FuncMap;
   // The set contains machine instructions which are necessary
   // for correct MIR but will not be emitted in function bodies.
   DenseSet<MachineInstr *> InstrsToDelete;
-  // The set contains machine basic blocks which are necessary
-  // for correct MIR but will not be emitted.
-  DenseSet<MachineBasicBlock *> MBBsToSkip;
   // The table contains global aliases of local registers for each machine
   // function. The aliases are used to substitute local registers during
   // code emission.
   RegisterAliasMapTy RegisterAliasTable;
   // The counter holds the maximum ID we have in the module.
   unsigned MaxID;
   // The array contains lists of MIs for each module section.
   InstrList MS[NUM_MODULE_SECTIONS];

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

for (MachineBasicBlock &MBB : *MF) {
         auto Set = MI.getOperand(2).getImm();
         if (MAI.ExtInstSetMap.find(Set) == MAI.ExtInstSetMap.end())
           MAI.ExtInstSetMap[Set] = Register::index2VirtReg(MAI.getNextID());
       }
     }
   }
 }
 
-// Find OpIEqual and OpBranchConditional instructions originating from
-// OpSwitches, mark them skipped for emission. Also mark MBB skipped if it
-// contains only these instructions.
-static void processSwitches(const Module &M, SPIRV::ModuleAnalysisInfo &MAI,
-                            MachineModuleInfo *MMI) {
-  DenseSet<Register> SwitchRegs;
-  for (auto F = M.begin(), E = M.end(); F != E; ++F) {
-    MachineFunction *MF = MMI->getMachineFunction(*F);
-    if (!MF)
-      continue;
-    for (MachineBasicBlock &MBB : *MF)
-      for (MachineInstr &MI : MBB) {
-        if (MAI.getSkipEmission(&MI))
-          continue;
-        if (MI.getOpcode() == SPIRV::OpSwitch) {
-          assert(MI.getOperand(0).isReg());
-          SwitchRegs.insert(MI.getOperand(0).getReg());
-        }
-        if (MI.getOpcode() == SPIRV::OpISubS &&
-            SwitchRegs.contains(MI.getOperand(2).getReg())) {
-          SwitchRegs.insert(MI.getOperand(0).getReg());
-          MAI.setSkipEmission(&MI);
-        }
-        if ((MI.getOpcode() != SPIRV::OpIEqual &&
-             MI.getOpcode() != SPIRV::OpULessThanEqual) ||
-            !MI.getOperand(2).isReg() ||
-            !SwitchRegs.contains(MI.getOperand(2).getReg()))
-          continue;
-        Register CmpReg = MI.getOperand(0).getReg();
-        MachineInstr *CBr = MI.getNextNode();
-        assert(CBr && CBr->getOpcode() == SPIRV::OpBranchConditional &&
-               CBr->getOperand(0).isReg() &&
-               CBr->getOperand(0).getReg() == CmpReg);
-        MAI.setSkipEmission(&MI);
-        MAI.setSkipEmission(CBr);
-        if (&MBB.front() == &MI && &MBB.back() == CBr)
-          MAI.MBBsToSkip.insert(&MBB);
-      }
-  }
-}
-
 // RequirementHandler implementations.
 void SPIRV::RequirementHandler::getAndAddRequirements(
     SPIRV::OperandCategory::OperandCategory Category, uint32_t i,
     const SPIRVSubtarget &ST) {
   addRequirements(getSymbolicOperandRequirements(Category, i, ST, *this));
 }
 
 void SPIRV::RequirementHandler::pruneCapabilities(
bool SPIRVModuleAnalysis::runOnModule(Module &M) {
   MMI = &getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
 
   setBaseInfo(M);
 
   addDecorations(M, *TII, MMI, *ST, MAI);
 
   collectReqs(M, MAI, MMI, *ST);
 
-  processSwitches(M, MAI, MMI);
-
   // Process type/const/global var/func decl instructions, number their
   // destination registers from 0 to N, collect Extensions and Capabilities.
   processDefInstrs(M);
 
   // Number rest of registers from N+1 onwards.
   numberRegistersGlobally(M);
 
   // Collect OpName, OpEntryPoint, OpDecorate etc, process other instructions.
   processOtherInstrs(M);
 
   // If there are no entry points, we need the Linkage capability.
   if (MAI.MS[SPIRV::MB_EntryPoints].empty())
     MAI.Reqs.addCapability(SPIRV::Capability::Linkage);
 
   return false;
 }

llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp

static void processSwitches(MachineFunction &MF, SPIRVGlobalRegistry *GR,
   //   ...
   // bb.3.sw.bb1:
   //   ...
   // bb.4.sw.epilog:
   //   ...
   //
   // Sometimes (in case of range-compare switches), additional G_SUBs
   // instructions are inserted before G_ICMPs. Those need to be additionally
-  // processed and require type assignment.
+  // processed.
   //
   // This function modifies spv_switch call's operands to include destination
   // MBBs (default and for each constant value).
-  // Note that this function does not remove G_ICMP + G_BRCOND + G_BR sequences,
-  // but they are marked by ModuleAnalysis as skipped and as a result AsmPrinter
-  // does not output them.
+  //
+  // At the end, the function removes redundant [G_SUB] + G_ICMP + G_BRCOND +
+  // G_BR sequences.
 
   MachineRegisterInfo &MRI = MF.getRegInfo();
 
-  // Collect all MIs relevant to switches across all MBBs in MF.
+  // Collect spv_switches and G_ICMPs across all MBBs in MF.
   std::vector<MachineInstr *> RelevantInsts;
 
+  // Collect redundant MIs from [G_SUB] + G_ICMP + G_BRCOND + G_BR sequences.
+  // After updating spv_switches, the instructions can be removed.
+  std::vector<MachineInstr *> PostUpdateArtifacts;
+
   // Temporary set of compare registers. G_SUBs and G_ICMPs relating to
   // spv_switch use these registers.
   DenseSet<Register> CompareRegs;
   for (MachineBasicBlock &MBB : MF) {
     for (MachineInstr &MI : MBB) {
       // Calls to spv_switch intrinsics representing IR switches.
       if (isSpvIntrinsic(MI, Intrinsic::spv_switch)) {
         assert(MI.getOperand(1).isReg());
         CompareRegs.insert(MI.getOperand(1).getReg());
         RelevantInsts.push_back(&MI);
       }
 
       // G_SUBs coming from range-compare switch lowering. G_SUBs are found
       // after spv_switch but before G_ICMP.
       if (MI.getOpcode() == TargetOpcode::G_SUB && MI.getOperand(1).isReg() &&
           CompareRegs.contains(MI.getOperand(1).getReg())) {
         assert(MI.getOperand(0).isReg() && MI.getOperand(1).isReg());
         Register Dst = MI.getOperand(0).getReg();
         CompareRegs.insert(Dst);
-        SPIRVType *Ty = GR->getSPIRVTypeForVReg(MI.getOperand(1).getReg());
-        insertAssignInstr(Dst, nullptr, Ty, GR, MIB, MRI);
+        PostUpdateArtifacts.push_back(&MI);
       }
 
       // G_ICMPs relating to switches.
       if (MI.getOpcode() == TargetOpcode::G_ICMP && MI.getOperand(2).isReg() &&
           CompareRegs.contains(MI.getOperand(2).getReg())) {
         Register Dst = MI.getOperand(0).getReg();
-        // Set type info for destination register of switch's ICMP instruction.
-        if (GR->getSPIRVTypeForVReg(Dst) == nullptr) {
-          MIB.setInsertPt(*MI.getParent(), MI);
-          Type *LLVMTy = IntegerType::get(MF.getFunction().getContext(), 1);
-          SPIRVType *SpirvTy = GR->getOrCreateSPIRVType(LLVMTy, MIB);
-          MRI.setRegClass(Dst, &SPIRV::IDRegClass);
-          GR->assignSPIRVTypeToVReg(SpirvTy, Dst, MIB.getMF());
-        }
         RelevantInsts.push_back(&MI);
+        PostUpdateArtifacts.push_back(&MI);
+        MachineInstr *CBr = MRI.use_begin(Dst)->getParent();
+        assert(CBr->getOpcode() == SPIRV::G_BRCOND);
+        PostUpdateArtifacts.push_back(CBr);
+        MachineInstr *Br = CBr->getNextNode();
+        assert(Br->getOpcode() == SPIRV::G_BR);
+        PostUpdateArtifacts.push_back(Br);
       }
     }
   }
 
   // Update each spv_switch with destination MBBs.
   for (auto i = RelevantInsts.begin(); i != RelevantInsts.end(); i++) {
     if (!isSpvIntrinsic(**i, Intrinsic::spv_switch))
       continue;
for (auto j = i + 1; j != RelevantInsts.end(); j++) {
       uint64_t Value = getIConstVal(ICMP->getOperand(3).getReg(), &MRI);
       MachineInstr *CBr = MRI.use_begin(Dst)->getParent();
       assert(CBr->getOpcode() == SPIRV::G_BRCOND && CBr->getOperand(1).isMBB());
       MachineBasicBlock *MBB = CBr->getOperand(1).getMBB();
 
       // Map switch case Value to target MBB.
       ValuesToMBBs[Value] = MBB;
 
+      // Add target MBB as successor to the switch's MBB.
+      Switch->getParent()->addSuccessor(MBB);
+
       // The next MI is always G_BR to either the next case or the default.
       MachineInstr *NextMI = CBr->getNextNode();
       assert(NextMI->getOpcode() == SPIRV::G_BR &&
              NextMI->getOperand(0).isMBB());
       MachineBasicBlock *NextMBB = NextMI->getOperand(0).getMBB();
       // Default MBB does not begin with G_ICMP using spv_switch compare
       // register.
       if (NextMBB->front().getOpcode() != SPIRV::G_ICMP ||
           (NextMBB->front().getOperand(2).isReg() &&
-           NextMBB->front().getOperand(2).getReg() != CompareReg))
+           NextMBB->front().getOperand(2).getReg() != CompareReg)) {
+        // Set default MBB and add it as successor to the switch's MBB.
         DefaultMBB = NextMBB;
+        Switch->getParent()->addSuccessor(DefaultMBB);
+      }
     }
 
     // Modify considered spv_switch operands using collected Values and
     // MBBs.
     SmallVector<const ConstantInt *, 3> Values;
     SmallVector<MachineBasicBlock *, 3> MBBs;
     for (unsigned k = 2; k < Switch->getNumExplicitOperands(); k++) {
       Register CReg = Switch->getOperand(k).getReg();
for (unsigned k = Switch->getNumExplicitOperands() - 1; k > 1; k--)
       Switch->removeOperand(k);
 
     Switch->addOperand(MachineOperand::CreateMBB(DefaultMBB));
     for (unsigned k = 0; k < Values.size(); k++) {
       Switch->addOperand(MachineOperand::CreateCImm(Values[k]));
       Switch->addOperand(MachineOperand::CreateMBB(MBBs[k]));
     }
   }
+
+  for (MachineInstr *MI : PostUpdateArtifacts) {
+    MachineBasicBlock *ParentMBB = MI->getParent();
+    MI->eraseFromParent();
+    // If G_ICMP + G_BRCOND + G_BR were the only MIs in MBB, erase this MBB. It
+    // can be safely assumed, there are no breaks or phis directing into this
+    // MBB. However, we need to remove this MBB from the CFG graph. MBBs must be
+    // erased top-down.
+    if (ParentMBB->empty()) {
+      while (!ParentMBB->pred_empty())
+        (*ParentMBB->pred_begin())->removeSuccessor(ParentMBB);
+
+      while (!ParentMBB->succ_empty())
+        ParentMBB->removeSuccessor(ParentMBB->succ_begin());
+
+      ParentMBB->eraseFromParent();
+    }
+  }
 }
 
 bool SPIRVPreLegalizer::runOnMachineFunction(MachineFunction &MF) {
   // Initialize the type registry.
   const SPIRVSubtarget &ST = MF.getSubtarget<SPIRVSubtarget>();
   SPIRVGlobalRegistry *GR = ST.getSPIRVGlobalRegistry();
   GR->setCurrentFunc(MF);
   MachineIRBuilder MIB(MF);

llvm/test/CodeGen/SPIRV/branching/OpSwitchBranches.ll

+; RUN: llc -O0 -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+
+define i32 @test_switch_branches(i32 %a) {
+entry:
+  %alloc = alloca i32
+; CHECK-SPIRV:      OpSwitch %[[#]] %[[#DEFAULT:]] 1 %[[#CASE1:]] 2 %[[#CASE2:]] 3 %[[#CASE3:]]
+  switch i32 %a, label %default [
+    i32 1, label %case1
+    i32 2, label %case2
+    i32 3, label %case3
+  ]
+
+; CHECK-SPIRV:      %[[#CASE1]] = OpLabel
+case1:
+  store i32 1, ptr %alloc
+; CHECK-SPIRV:      OpBranch %[[#END:]]
+  br label %end
+
+; CHECK-SPIRV:      %[[#CASE2]] = OpLabel
+case2:
+  store i32 2, ptr %alloc
+; CHECK-SPIRV:      OpBranch %[[#END]]
+  br label %end
+
+; CHECK-SPIRV:      %[[#CASE3]] = OpLabel
+case3:
+  store i32 3, ptr %alloc
+; CHECK-SPIRV:      OpBranch %[[#END]]
+  br label %end
+
+; CHECK-SPIRV:      %[[#DEFAULT]] = OpLabel
+default:
+  store i32 0, ptr %alloc
+; CHECK-SPIRV:      OpBranch %[[#END]]
+  br label %end
+
+; CHECK-SPIRV:      %[[#END]] = OpLabel
+end:
+  %result = load i32, ptr %alloc
+  ret i32 %result
+}