[SPIR-V] Add pre-headers to loops. (#75844)

This is the first of the 7 steps outlined in #75801. This PR explicitely
calls the SimplifyLoops pass. Directly following this pass should follow
the 6 others required to structurize the IR.

Running this pass could generate empty basic-blocks, which are implicit
fallthrough to the successor BB.
There was a specific condition in the SPIR-V ISel which handled implicit
fallthrough, but it couldn't work on empty basic-blocks. This commits
removes the old logic, and adds this new logic, which checks all
basic-blocks for implicit fallthroughs, including empty ones.

---------

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

llvm/lib/Target/SPIRV/SPIRVPreLegalizer.cpp

@@ -607,6 +607,40 @@ static void processSwitches(MachineFunction &MF, SPIRVGlobalRegistry *GR,
   }
 }
 
+static bool isImplicitFallthrough(MachineBasicBlock &MBB) {
+  if (MBB.empty())
+    return true;
+
+  // Branching SPIR-V intrinsics are not detected by this generic method.
+  // Thus, we can only trust negative result.
+  if (!MBB.canFallThrough())
+    return false;
+
+  // Otherwise, we must manually check if we have a SPIR-V intrinsic which
+  // prevent an implicit fallthrough.
+  for (MachineBasicBlock::reverse_iterator It = MBB.rbegin(), E = MBB.rend();
+       It != E; ++It) {
+    if (isSpvIntrinsic(*It, Intrinsic::spv_switch))
+      return false;
+  }
+  return true;
+}
+
+static void removeImplicitFallthroughs(MachineFunction &MF,
+                                       MachineIRBuilder MIB) {
+  // It is valid for MachineBasicBlocks to not finish with a branch instruction.
+  // In such cases, they will simply fallthrough their immediate successor.
+  for (MachineBasicBlock &MBB : MF) {
+    if (!isImplicitFallthrough(MBB))
+      continue;
+
+    assert(std::distance(MBB.successors().begin(), MBB.successors().end()) ==
+           1);
+    MIB.setInsertPt(MBB, MBB.end());
+    MIB.buildBr(**MBB.successors().begin());
+  }
+}
+
 bool SPIRVPreLegalizer::runOnMachineFunction(MachineFunction &MF) {
   // Initialize the type registry.
   const SPIRVSubtarget &ST = MF.getSubtarget<SPIRVSubtarget>();
@@ -619,6 +653,7 @@ bool SPIRVPreLegalizer::runOnMachineFunction(MachineFunction &MF) {
   generateAssignInstrs(MF, GR, MIB);
   processSwitches(MF, GR, MIB);
   processInstrsWithTypeFolding(MF, GR, MIB);
+  removeImplicitFallthroughs(MF, MIB);
 
   return true;
 }

llvm/lib/Target/SPIRV/SPIRVTargetMachine.cpp

@@ -29,6 +29,7 @@
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Pass.h"
 #include "llvm/Target/TargetOptions.h"
+#include "llvm/Transforms/Utils.h"
 #include <optional>
 
 using namespace llvm;
@@ -151,6 +152,19 @@ TargetPassConfig *SPIRVTargetMachine::createPassConfig(PassManagerBase &PM) {
 }
 
 void SPIRVPassConfig::addIRPasses() {
+  if (TM.getSubtargetImpl()->isVulkanEnv()) {
+    // Once legalized, we need to structurize the CFG to follow the spec.
+    // This is done through the following 8 steps.
+    // TODO(#75801): add the remaining steps.
+
+    // 1.  Simplify loop for subsequent transformations. After this steps, loops
+    // have the following properties:
+    //  - loops have a single entry edge (pre-header to loop header).
+    //  - all loop exits are dominated by the loop pre-header.
+    //  - loops have a single back-edge.
+    addPass(createLoopSimplifyPass());
+  }
+
   TargetPassConfig::addIRPasses();
   addPass(createSPIRVRegularizerPass());
   addPass(createSPIRVPrepareFunctionsPass(TM));

llvm/lib/Target/SPIRV/SPIRVUtils.cpp

@@ -228,8 +228,8 @@ uint64_t getIConstVal(Register ConstReg, const MachineRegisterInfo *MRI) {
   return MI->getOperand(1).getCImm()->getValue().getZExtValue();
 }
 
-bool isSpvIntrinsic(MachineInstr &MI, Intrinsic::ID IntrinsicID) {
-  if (auto *GI = dyn_cast<GIntrinsic>(&MI))
+bool isSpvIntrinsic(const MachineInstr &MI, Intrinsic::ID IntrinsicID) {
+  if (const auto *GI = dyn_cast<GIntrinsic>(&MI))
     return GI->is(IntrinsicID);
   return false;
 }

llvm/lib/Target/SPIRV/SPIRVUtils.h

@@ -79,7 +79,7 @@ MachineInstr *getDefInstrMaybeConstant(Register &ConstReg,
 uint64_t getIConstVal(Register ConstReg, const MachineRegisterInfo *MRI);
 
 // Check if MI is a SPIR-V specific intrinsic call.
-bool isSpvIntrinsic(MachineInstr &MI, Intrinsic::ID IntrinsicID);
+bool isSpvIntrinsic(const MachineInstr &MI, Intrinsic::ID IntrinsicID);
 
 // Get type of i-th operand of the metadata node.
 Type *getMDOperandAsType(const MDNode *N, unsigned I);

llvm/test/CodeGen/SPIRV/scfg-add-pre-headers.ll

@@ -0,0 +1,66 @@
+; RUN: llc -mtriple=spirv-unknown-unknown -O0 %s -o - | FileCheck %s
+
+; CHECK-DAG:    %[[#bool:]] = OpTypeBool
+; CHECK-DAG:    %[[#uint:]] = OpTypeInt 32 0
+; CHECK-DAG:  %[[#uint_0:]] = OpConstant %[[#uint]] 0
+
+define void @main() #1 {
+  %1 = icmp ne i32 0, 0
+  br i1 %1, label %l1, label %l2
+
+; CHECK:        %[[#cond:]] = OpINotEqual %[[#bool]] %[[#uint_0]] %[[#uint_0]]
+; CHECK:                      OpBranchConditional %[[#cond]] %[[#l1_pre:]] %[[#l2_pre:]]
+
+; CHECK-DAG:   %[[#l2_pre]] = OpLabel
+; CHECK-NEXT:                 OpBranch %[[#l2_header:]]
+
+; CHECK-DAG:   %[[#l1_pre]] = OpLabel
+; CHECK-NEXT:                 OpBranch %[[#l1_header:]]
+
+l1:
+  br i1 %1, label %l1_body, label %l1_end
+; CHECK-DAG:    %[[#l1_header]] = OpLabel
+; CHECK-NEXT:                     OpBranchConditional %[[#cond]] %[[#l1_body:]] %[[#l1_end:]]
+
+l1_body:
+  br label %l1_continue
+; CHECK-DAG:   %[[#l1_body]] = OpLabel
+; CHECK-NEXT:                  OpBranch %[[#l1_continue:]]
+
+l1_continue:
+  br label %l1
+; CHECK-DAG:   %[[#l1_continue]] = OpLabel
+; CHECK-NEXT:                      OpBranch %[[#l1_header]]
+
+l1_end:
+  br label %end
+; CHECK-DAG:   %[[#l1_end]] = OpLabel
+; CHECK-NEXT:                 OpBranch %[[#end:]]
+
+l2:
+  br i1 %1, label %l2_body, label %l2_end
+; CHECK-DAG:    %[[#l2_header]] = OpLabel
+; CHECK-NEXT:                     OpBranchConditional %[[#cond]] %[[#l2_body:]] %[[#l2_end:]]
+
+l2_body:
+  br label %l2_continue
+; CHECK-DAG:   %[[#l2_body]] = OpLabel
+; CHECK-NEXT:                  OpBranch %[[#l2_continue:]]
+
+l2_continue:
+  br label %l2
+; CHECK-DAG:   %[[#l2_continue]] = OpLabel
+; CHECK-NEXT:                      OpBranch %[[#l2_header]]
+
+l2_end:
+  br label %end
+; CHECK-DAG:   %[[#l2_end]] = OpLabel
+; CHECK-NEXT:                 OpBranch %[[#end:]]
+
+end:
+  ret void
+; CHECK-DAG:       %[[#end]] = OpLabel
+; CHECK-NEXT:                  OpReturn
+}
+
+attributes #1 = { "hlsl.numthreads"="4,8,16" "hlsl.shader"="compute" convergent }