[SPIR-V] Sort basic blocks to match the SPIR-V spec

llvm/lib/Target/SPIRV/SPIRVPostLegalizer.cpp

@@ -19,11 +19,13 @@
 #include "SPIRVUtils.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/IntrinsicsSPIRV.h"
 #include "llvm/Target/TargetIntrinsicInfo.h"
+#include <stack>
 
 #define DEBUG_TYPE "spirv-postlegalizer"
 
@@ -150,6 +152,53 @@ static void processNewInstrs(MachineFunction &MF, SPIRVGlobalRegistry *GR,
   }
 }
 
+// Do a preorder traversal of the CFG starting from the BB |Start|.
+// point. Calls |op| on each basic block encountered during the traversal.
+void visit(MachineFunction &MF, MachineBasicBlock &Start,
+           std::function<void(MachineBasicBlock *)> op) {
+  std::stack<MachineBasicBlock *> ToVisit;
+  SmallPtrSet<MachineBasicBlock *, 8> Seen;
+
+  ToVisit.push(&Start);
+  Seen.insert(ToVisit.top());
+  while (ToVisit.size() != 0) {
+    MachineBasicBlock *MBB = ToVisit.top();
+    ToVisit.pop();
+
+    op(MBB);
+
+    for (auto Succ : MBB->successors()) {
+      if (Seen.contains(Succ))
+        continue;
+      ToVisit.push(Succ);
+      Seen.insert(Succ);
+    }
+  }
+}
+
+// Do a preorder traversal of the CFG starting from the given function's entry
+// point. Calls |op| on each basic block encountered during the traversal.
+void visit(MachineFunction &MF, std::function<void(MachineBasicBlock *)> op) {
+  visit(MF, *MF.begin(), op);
+}
+
+// Sorts basic blocks by dominance to respect the SPIR-V spec.
+void sortBlocks(MachineFunction &MF) {
+  MachineDominatorTree MDT(MF);
+
+  std::unordered_map<MachineBasicBlock *, size_t> Order;
+  Order.reserve(MF.size());
+
+  size_t Index = 0;
+  visit(MF, [&Order, &Index](MachineBasicBlock *MBB) { Order[MBB] = Index++; });
+
+  auto Comparator = [&Order](MachineBasicBlock &LHS, MachineBasicBlock &RHS) {
+    return Order[&LHS] < Order[&RHS];
+  };
+
+  MF.sort(Comparator);
+}
+
 bool SPIRVPostLegalizer::runOnMachineFunction(MachineFunction &MF) {
   // Initialize the type registry.
   const SPIRVSubtarget &ST = MF.getSubtarget<SPIRVSubtarget>();
@@ -158,6 +207,7 @@ bool SPIRVPostLegalizer::runOnMachineFunction(MachineFunction &MF) {
   MachineIRBuilder MIB(MF);
 
   processNewInstrs(MF, GR, MIB);
+  sortBlocks(MF);
 
   return true;
 }

llvm/test/CodeGen/SPIRV/block-ordering.ll

@@ -0,0 +1,30 @@
+; RUN: llc -O0 -mtriple=spirv32-unknown-unknown %s -o - | FileCheck %s --check-prefix=CHECK
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv32-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; Checks SPIR-V blocks are correctly reordered so that dominators shows up
+; before others in the binary layout.
+
+define void @main() {
+; CHECK: OpLabel
+; CHECK: OpBranch %[[#l1:]]
+
+; CHECK: %[[#l1]] = OpLabel
+; CHECK:            OpBranch %[[#l2:]]
+
+; CHECK: %[[#l2]] = OpLabel
+; CHECK:            OpBranch %[[#end:]]
+
+; CHECK: %[[#end]] = OpLabel
+; CHECK:             OpReturn
+entry:
+  br label %l1
+
+l2:
+  br label %end
+
+l1:
+  br label %l2
+
+end:
+  ret void
+}

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

@@ -10,32 +10,38 @@ entry:
     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
+
+; CHECK-SPIRV:      %[[#CASE3]] = OpLabel
+; CHECK-SPIRV:      OpBranch %[[#END:]]
+
+; CHECK-SPIRV:      %[[#END]] = OpLabel
+; CHECK-SPIRV:                  OpReturnValue
+
+; CHECK-SPIRV:      %[[#CASE2]] = OpLabel
+; CHECK-SPIRV:      OpBranch %[[#END]]
+
+; CHECK-SPIRV:      %[[#CASE1]] = OpLabel
+; CHECK-SPIRV:      OpBranch %[[#END]]
+
+; CHECK-SPIRV:      %[[#DEFAULT]] = OpLabel
+; CHECK-SPIRV:      OpBranch %[[#END]]
 }

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

@@ -1,4 +1,5 @@
 ; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s --check-prefix=CHECK-SPIRV
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
 
 define spir_kernel void @test_two_switch_same_register(i32 %value) {
 ; CHECK-SPIRV:      OpSwitch %[[#REGISTER:]] %[[#DEFAULT1:]] 1 %[[#CASE1:]] 0 %[[#CASE2:]]
@@ -7,36 +8,42 @@ define spir_kernel void @test_two_switch_same_register(i32 %value) {
     i32 0, label %case2
   ]
 
-; CHECK-SPIRV:      %[[#CASE1]] = OpLabel
 case1:
-; CHECK-SPIRV-NEXT: OpBranch %[[#DEFAULT1]]
   br label %default1
 
-; CHECK-SPIRV:      %[[#CASE2]] = OpLabel
 case2:
-; CHECK-SPIRV-NEXT: OpBranch %[[#DEFAULT1]]
   br label %default1
 
-; CHECK-SPIRV:      %[[#DEFAULT1]] = OpLabel
 default1:
-; CHECK-SPIRV-NEXT:      OpSwitch %[[#REGISTER]] %[[#DEFAULT2:]] 0 %[[#CASE3:]] 1 %[[#CASE4:]]
   switch i32 %value, label %default2 [
     i32 0, label %case3
     i32 1, label %case4
   ]
 
-; CHECK-SPIRV:      %[[#CASE3]] = OpLabel
 case3:
-; CHECK-SPIRV-NEXT: OpBranch %[[#DEFAULT2]]
   br label %default2
 
-; CHECK-SPIRV:      %[[#CASE4]] = OpLabel
 case4:
-; CHECK-SPIRV-NEXT: OpBranch %[[#DEFAULT2]]
   br label %default2
 
-; CHECK-SPIRV:      %[[#DEFAULT2]] = OpLabel
 default2:
-; CHECK-SPIRV-NEXT: OpReturn
   ret void
+
+; CHECK-SPIRV:      %[[#CASE2]] = OpLabel
+; CHECK-SPIRV-NEXT: OpBranch %[[#DEFAULT1]]
+
+; CHECK-SPIRV:      %[[#CASE1]] = OpLabel
+; CHECK-SPIRV-NEXT: OpBranch %[[#DEFAULT1]]
+
+; CHECK-SPIRV:      %[[#DEFAULT1]] = OpLabel
+; CHECK-SPIRV-NEXT:      OpSwitch %[[#REGISTER]] %[[#DEFAULT2:]] 0 %[[#CASE3:]] 1 %[[#CASE4:]]
+
+; CHECK-SPIRV:      %[[#CASE4:]] = OpLabel
+; CHECK-SPIRV-NEXT: OpBranch %[[#DEFAULT2]]
+
+; CHECK-SPIRV:      %[[#CASE3]] = OpLabel
+; CHECK-SPIRV-NEXT: OpBranch %[[#DEFAULT2]]
+
+; CHECK-SPIRV:      %[[#DEFAULT2]] = OpLabel
+; CHECK-SPIRV-NEXT: OpReturn
 }

llvm/test/CodeGen/SPIRV/branching/if-merging.ll

@@ -37,15 +37,16 @@ merge_label:
 ; CHECK: [[COND:%.+]] = OpIEqual [[BOOL]] [[A]] [[B]]
 ; CHECK: OpBranchConditional [[COND]] [[TRUE_LABEL:%.+]] [[FALSE_LABEL:%.+]]
 
-; CHECK: [[TRUE_LABEL]] = OpLabel
-; CHECK: [[V1:%.+]] = OpFunctionCall [[I32]] [[FOO]]
-; CHECK: OpBranch [[MERGE_LABEL:%.+]]
-
 ; CHECK: [[FALSE_LABEL]] = OpLabel
 ; CHECK: [[V2:%.+]] = OpFunctionCall [[I32]] [[BAR]]
-; CHECK: OpBranch [[MERGE_LABEL]]
+; CHECK: OpBranch [[MERGE_LABEL:%.+]]
 
 ; CHECK: [[MERGE_LABEL]] = OpLabel
-; CHECK-NEXT: [[V:%.+]] = OpPhi [[I32]] [[V1]] [[TRUE_LABEL]] [[V2]] [[FALSE_LABEL]]
+; CHECK-NEXT: [[V:%.+]] = OpPhi [[I32]] [[V1:%.+]] [[TRUE_LABEL]] [[V2]] [[FALSE_LABEL]]
 ; CHECK: OpReturnValue [[V]]
+
+; CHECK: [[TRUE_LABEL]] = OpLabel
+; CHECK: [[V1]] = OpFunctionCall [[I32]] [[FOO]]
+; CHECK: OpBranch [[MERGE_LABEL]]
+
 ; CHECK-NEXT: OpFunctionEnd

llvm/test/CodeGen/SPIRV/branching/if-non-merging.ll

@@ -21,7 +21,7 @@ false_label:
 ; CHECK: [[ENTRY:%.+]] = OpLabel
 ; CHECK: [[COND:%.+]] = OpIEqual [[BOOL]] [[A]] [[B]]
 ; CHECK: OpBranchConditional [[COND]] [[TRUE_LABEL:%.+]] [[FALSE_LABEL:%.+]]
-; CHECK: [[TRUE_LABEL]] = OpLabel
-; CHECK: OpReturnValue [[TRUE]]
 ; CHECK: [[FALSE_LABEL]] = OpLabel
 ; CHECK: OpReturnValue [[FALSE]]
+; CHECK: [[TRUE_LABEL]] = OpLabel
+; CHECK: OpReturnValue [[TRUE]]

llvm/test/CodeGen/SPIRV/branching/switch-range-check.ll

@@ -2,11 +2,15 @@
 ; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv32-unknown-unknown %s -o - -filetype=obj | spirv-val %}
 
 ; CHECK: OpFunction
+; CHECK: OpBranchConditional %[[#]] %[[#if_then:]] %[[#if_end:]]
+; CHECK: %[[#if_end]] = OpLabel
 ; CHECK: %[[#Var:]] = OpPhi
 ; CHECK: OpSwitch %[[#Var]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]] [[#]] %[[#]]
-; CHECK-COUNT-11: OpBranch
+; CHECK-COUNT-11: OpLabel
 ; CHECK-NOT: OpBranch
 ; CHECK: OpReturn
+; CHECK: %[[#if_then]] = OpLabel
+; CHECK: OpBranch %[[#if_end]]
 ; CHECK-NEXT: OpFunctionEnd
 
 define spir_func void @foo(i64 noundef %addr, i64 noundef %as) {

llvm/test/CodeGen/SPIRV/phi-ptrcast-dominate.ll

@@ -6,89 +6,103 @@
 ; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
 ; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
 
+
 ; CHECK-DAG: OpName %[[#Case1:]] "case1"
 ; CHECK-DAG: OpName %[[#Case2:]] "case2"
 ; CHECK-DAG: OpName %[[#Case3:]] "case3"
-; CHECK: %[[#Case1]] = OpFunction
-; CHECK: OpBranchConditional
-; CHECK: OpPhi
-; CHECK: OpBranch
-; CHECK-COUNT-2: OpBranchConditional
-; CHECK: OpFunctionEnd
-; CHECK: %[[#Case2]] = OpFunction
-; CHECK: OpBranchConditional
-; CHECK: OpPhi
-; CHECK: OpBranch
-; CHECK-COUNT-2: OpBranchConditional
-; CHECK: OpFunctionEnd
-; CHECK: %[[#Case3]] = OpFunction
-; CHECK: OpBranchConditional
-; CHECK: OpPhi
-; CHECK: OpBranch
-; CHECK: OpInBoundsPtrAccessChain
-; CHECK: OpBranchConditional
-; CHECK: OpInBoundsPtrAccessChain
-; CHECK: OpBranchConditional
-; CHECK: OpFunctionEnd
 
 %struct1 = type { i64 }
 %struct2 = type { i64, i64 }
 
 @.str.1 = private unnamed_addr addrspace(1) constant [3 x i8] c"OK\00", align 1
 @.str.2 = private unnamed_addr addrspace(1) constant [6 x i8] c"WRONG\00", align 1
 
+; CHECK: %[[#Case1]] = OpFunction
 define spir_func void @case1(i1 %b1, i1 %b2, i1 %b3) {
 entry:
+; CHECK: OpBranchConditional %[[#]] %[[#l1:]] %[[#l2:]]
   br i1 %b1, label %l1, label %l2
 
 l1:
   %str = phi ptr addrspace(1) [ @.str.1, %entry ], [ @.str.2, %l2 ], [ @.str.2, %l3 ]
   br label %exit
 
+; CHECK: %[[#l2]] = OpLabel
+; CHECK:            OpBranchConditional %[[#]] %[[#l1:]] %[[#l3:]]
 l2:
   br i1 %b2, label %l1, label %l3
 
+; CHECK: %[[#l3]] = OpLabel
+; CHECK:            OpBranchConditional %[[#]] %[[#l1:]] %[[#exit:]]
 l3:
   br i1 %b3, label %l1, label %exit
 
+; CHECK: %[[#exit]] = OpLabel
+; CHECK:              OpReturn
 exit:
   ret void
+
+; CHECK:      %[[#l1]] = OpLabel
+; CHECK-NEXT:            OpPhi
+; CHECK:                 OpBranch %[[#exit:]]
 }
 
+; CHECK: %[[#Case2]] = OpFunction
 define spir_func void @case2(i1 %b1, i1 %b2, i1 %b3, ptr addrspace(1) byval(%struct1) %str1, ptr addrspace(1) byval(%struct2) %str2) {
 entry:
+; CHECK: OpBranchConditional %[[#]] %[[#l1:]] %[[#l2:]]
   br i1 %b1, label %l1, label %l2
 
 l1:
   %str = phi ptr addrspace(1) [ %str1, %entry ], [ %str2, %l2 ], [ %str2, %l3 ]
   br label %exit
 
+; CHECK: %[[#l2]] = OpLabel
+; CHECK:            OpBranchConditional %[[#]] %[[#l1:]] %[[#l3:]]
 l2:
   br i1 %b2, label %l1, label %l3
 
+; CHECK: %[[#l3]] = OpLabel
+; CHECK:            OpBranchConditional %[[#]] %[[#l1:]] %[[#exit:]]
 l3:
   br i1 %b3, label %l1, label %exit
 
+; CHECK: %[[#exit]] = OpLabel
+; CHECK:              OpReturn
 exit:
   ret void
 }
 
+; CHECK: %[[#Case3]] = OpFunction
 define spir_func void @case3(i1 %b1, i1 %b2, i1 %b3, ptr addrspace(1) byval(%struct1) %_arg_str1, ptr addrspace(1) byval(%struct2) %_arg_str2) {
 entry:
   br i1 %b1, label %l1, label %l2
 
+; CHECK: OpBranchConditional %[[#]] %[[#l1:]] %[[#l2:]]
 l1:
   %str = phi ptr addrspace(1) [ %_arg_str1, %entry ], [ %str2, %l2 ], [ %str3, %l3 ]
   br label %exit
 
+; CHECK: %[[#l2]] = OpLabel
+; CHECK:            OpInBoundsPtrAccessChain
+; CHECK:            OpBranchConditional %[[#]] %[[#l1:]] %[[#l3:]]
 l2:
   %str2 = getelementptr inbounds %struct2, ptr addrspace(1) %_arg_str2, i32 1
   br i1 %b2, label %l1, label %l3
 
+; CHECK: %[[#l3]] = OpLabel
+; CHECK:            OpInBoundsPtrAccessChain
+; CHECK:            OpBranchConditional %[[#]] %[[#l1:]] %[[#exit:]]
 l3:
   %str3 = getelementptr inbounds %struct2, ptr addrspace(1) %_arg_str2, i32 2
   br i1 %b3, label %l1, label %exit
 
+; CHECK: %[[#exit]] = OpLabel
+; CHECK:              OpReturn
 exit:
   ret void
+
+; CHECK:      %[[#l1]] = OpLabel
+; CHECK-NEXT:            OpPhi
+; CHECK:                 OpBranch %[[#exit:]]
 }