Replace assumptions during llvm.loop translation with checks

There were 3 assumptions made during translation of llvm.loop metadata
into LoopMerge instruction:
1. A latch block for 'for' and 'while' loop ends with a unconditional
   branch instruction;
2. A latch block for 'do-while' loop ends with a conditional branch
   instruction;
3. For a 'conditional' latch basic block it assumes, that the first
   successor is the loop header and the second one is the exit block.

All three of them can be violated in a case of an optimized IR is passed
to the translator. In this case LoopMerge:
a. can be placed in a wrong basic block;
b. have continue target and merge block parameters sharing the same id.

This patch makes the translator to assume less and do more checks via
LLVM LoopInfo infrustructure.

Signed-off-by: Dmitry Sidorov <[email protected]>

Author: MrSidims

lib/SPIRV/SPIRVWriter.cpp

@@ -60,9 +60,11 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/ADT/Triple.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/InstrTypes.h"
@@ -1618,7 +1620,7 @@ LLVMToSPIRVBase::transValueWithoutDecoration(Value *V, SPIRVBasicBlock *BB,
     spv::LoopControlMask LoopControl = getLoopControl(Branch, Parameters);
 
     if (Branch->isUnconditional()) {
-      // For "for" and "while" loops llvm.loop metadata is attached to
+      // Usually for "for" and "while" loops llvm.loop metadata is attached to
       // an unconditional branch instruction.
       if (LoopControl != spv::LoopControlMaskNone) {
         // SuccessorTrue is the loop header BB.
@@ -1640,15 +1642,28 @@ LLVMToSPIRVBase::transValueWithoutDecoration(Value *V, SPIRVBasicBlock *BB,
       }
       return mapValue(V, BM->addBranchInst(SuccessorTrue, BB));
     }
-    // For "do-while" loops llvm.loop metadata is attached to a conditional
-    // branch instructions
+    // For "do-while" and in some cases for "for" and "while" loops llvm.loop
+    // metadata is attached to a conditional branch instructions
     SPIRVLabel *SuccessorFalse =
         static_cast<SPIRVLabel *>(transValue(Branch->getSuccessor(1), BB));
-    if (LoopControl != spv::LoopControlMaskNone)
-      // SuccessorTrue is the loop header BB.
-      BM->addLoopMergeInst(SuccessorFalse->getId(), // Merge Block
-                           BB->getId(),             // Continue Target
-                           LoopControl, Parameters, SuccessorTrue);
+    if (LoopControl != spv::LoopControlMaskNone) {
+      Function *Fun = Branch->getParent()->getParent();
+      DominatorTree DomTree(*Fun);
+      LoopInfo LI(DomTree);
+      for (const auto *LoopObj : LI.getLoopsInPreorder()) {
+        // Check whether SuccessorFalse or SuccessorTrue is the loop header BB
+        if (LoopObj->getHeader() == Branch->getSuccessor(1))
+          // SuccessorFalse is the loop header BB.
+          BM->addLoopMergeInst(SuccessorTrue->getId(), // Merge Block
+                               BB->getId(),            // Continue Target
+                               LoopControl, Parameters, SuccessorFalse);
+        else
+          // SuccessorTrue is the loop header BB.
+          BM->addLoopMergeInst(SuccessorFalse->getId(), // Merge Block
+                               BB->getId(),             // Continue Target
+                               LoopControl, Parameters, SuccessorTrue);
+      }
+    }
     return mapValue(
         V, BM->addBranchConditionalInst(transValue(Branch->getCondition(), BB),
                                         SuccessorTrue, SuccessorFalse, BB));

test/transcoding/LoopUnroll.ll

@@ -30,11 +30,18 @@
 ;       int x = i;
 ;    } while(i--> 0);
 ; }
+;
+; for_count_unusual() is a synthetically written function
+;
 ; Command:
 ; clang -cc1 -triple spir64 -O0 LoopUnroll.cl -emit-llvm -o /test/SPIRV/transcoding/LoopUnroll.ll
 
 ; RUN: llvm-as < %s > %t.bc
-; RUN: llvm-spirv %t.bc -o - -spirv-text | FileCheck %s --check-prefix=CHECK-SPIRV
+; RUN: llvm-spirv %t.bc -o %t.spv
+; RUN: llvm-spirv -to-text %t.spv -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+; RUN: llvm-spirv -r %t.spv -o %t.rev.bc
+; RUN: llvm-dis %t.rev.bc -o - | FileCheck %s --check-prefixes=CHECK-LLVM
 
 target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024"
 target triple = "spir64"
@@ -50,12 +57,12 @@ entry:
   br label %for.cond
 
 for.cond:                                         ; preds = %for.inc, %entry
-; CHECK-SPIRV: Label [[Header:[0-9]+]]
+; CHECK-SPIRV: Label [[#HEADER:]]
   %0 = load i32, i32* %i, align 4
   %cmp = icmp slt i32 %0, 1024
 ; Per SPIRV spec p3.23 "DontUnroll" loop control = 0x2
-; CHECK-SPIRV: 4 LoopMerge [[MergeBlock:[0-9]+]] [[ContinueTarget:[0-9]+]] 2
-; CHECK-SPIRV: BranchConditional {{[0-9]+}} {{[0-9]+}} [[MergeBlock]]
+; CHECK-SPIRV: LoopMerge [[#MERGEBLOCK:]] [[#CONTINUE:]] 2
+; CHECK-SPIRV: BranchConditional [[#]] [[#]] [[#MERGEBLOCK]]
   br i1 %cmp, label %for.body, label %for.end
 
 for.body:                                         ; preds = %for.cond
@@ -76,15 +83,16 @@ if.end:                                           ; preds = %for.body
   br label %for.inc
 
 for.inc:                                          ; preds = %if.end, %if.then
-; CHECK-SPIRV: Label [[ContinueTarget]]
+; CHECK-SPIRV: Label [[#CONTINUE]]
   %3 = load i32, i32* %i, align 4
   %inc = add nsw i32 %3, 1
   store i32 %inc, i32* %i, align 4
   br label %for.cond, !llvm.loop !5
-; CHECK-SPIRV: Branch [[Header]]
+; CHECK-LLVM: br label %for.cond, !llvm.loop ![[#UNROLLDISABLE:]]
+; CHECK-SPIRV: Branch [[#HEADER]]
 
 for.end:                                          ; preds = %for.cond
-; CHECK-SPIRV: Label [[MergeBlock]]
+; CHECK-SPIRV: Label [[#MERGEBLOCK]]
   ret void
 }
 
@@ -99,14 +107,14 @@ entry:
   br label %while.cond
 
 while.cond:                                       ; preds = %if.end, %if.then, %entry
-; CHECK-SPIRV: Label [[Header:[0-9]+]]
+; CHECK-SPIRV: Label [[#HEADER:]]
   %0 = load i32, i32* %i, align 4
   %dec = add nsw i32 %0, -1
   store i32 %dec, i32* %i, align 4
   %cmp = icmp sgt i32 %0, 0
 ; Per SPIRV spec p3.23 "Unroll" loop control = 0x1
-; CHECK-SPIRV: 5 LoopMerge [[MergeBlock:[0-9]+]] [[ContinueTarget:[0-9]+]] 256 8
-; CHECK-SPIRV: BranchConditional {{[0-9]+}} {{[0-9]+}} [[MergeBlock]]
+; CHECK-SPIRV: LoopMerge [[#MERGEBLOCK:]] [[#CONTINUE:]] 256 8
+; CHECK-SPIRV: BranchConditional [[#]] [[#]] [[#MERGEBLOCK]]
   br i1 %cmp, label %while.body, label %while.end
 
 while.body:                                       ; preds = %while.cond
@@ -118,12 +126,13 @@ while.body:                                       ; preds = %while.cond
 
 if.then:                                          ; preds = %while.body
 ; CHECK-SPIRV: Label
+; CHECK-LLVM: br label %while.cond, !llvm.loop ![[#UNROLLCOUNT:]]
   br label %while.cond, !llvm.loop !7
 
 ; loop-simplify pass will create extra basic block which is the only one in
 ; loop having a back-edge to the header
-; CHECK-SPIRV: [[ContinueTarget]]
-; CHECK-SPIRV: Branch [[Header]]
+; CHECK-SPIRV: [[#CONTINUE]]
+; CHECK-SPIRV: Branch [[#HEADER]]
 
 if.end:                                           ; preds = %while.body
 ; CHECK-SPIRV: Label
@@ -132,7 +141,7 @@ if.end:                                           ; preds = %while.body
   br label %while.cond, !llvm.loop !7
 
 while.end:                                        ; preds = %while.cond
-; CHECK-SPIRV: [[MergeBlock]]
+; CHECK-SPIRV: [[#MERGEBLOCK]]
   ret void
 }
 
@@ -144,15 +153,15 @@ entry:
   %i = alloca i32, align 4
   %x = alloca i32, align 4
   store i32 1024, i32* %i, align 4
-  br label %do.body, !llvm.loop !9
+  br label %do.body
 
 do.body:                                          ; preds = %do.cond, %entry
-; CHECK-SPIRV: Label [[Header:[0-9]+]]
+; CHECK-SPIRV: Label [[#HEADER:]]
   %0 = load i32, i32* %i, align 4
   %rem = srem i32 %0, 2
   %tobool = icmp ne i32 %rem, 0
 ; Per SPIRV spec p3.23 "Unroll" loop control = 0x1
-; CHECK-SPIRV: 4 LoopMerge [[MergeBlock:[0-9]+]] [[ContinueTarget:[0-9]+]] 1
+; CHECK-SPIRV: LoopMerge [[#MERGEBLOCK:]] [[#CONTINUE:]] 1
 ; CHECK-SPIRV: BranchConditional
   br i1 %tobool, label %if.then, label %if.end
 
@@ -167,16 +176,62 @@ if.end:                                           ; preds = %do.body
   br label %do.cond
 
 do.cond:                                          ; preds = %if.end, %if.then
-; CHECK-SPIRV: Label [[ContinueTarget]]
+; CHECK-SPIRV: Label [[#CONTINUE]]
   %2 = load i32, i32* %i, align 4
   %dec = add nsw i32 %2, -1
   store i32 %dec, i32* %i, align 4
   %cmp = icmp sgt i32 %2, 0
-; CHECK-SPIRV: BranchConditional {{[0-9]+}} [[Header]] [[MergeBlock]]
+; CHECK-SPIRV: BranchConditional [[#]] [[#HEADER]] [[#MERGEBLOCK]]
+; CHECK-LLVM: br i1 %cmp, label %do.body, label %do.end, !llvm.loop ![[#UNROLLENABLE1:]]
   br i1 %cmp, label %do.body, label %do.end, !llvm.loop !9
 
 do.end:                                           ; preds = %do.cond
-; CHECK-SPIRV: Label [[MergeBlock]]
+; CHECK-SPIRV: Label [[#MERGEBLOCK]]
+  ret void
+}
+
+; CHECK-SPIRV: Function
+; Function Attrs: noinline nounwind optnone
+define spir_func void @for_count_unusual() #0 {
+entry:
+; CHECK-SPIRV: Label
+  %i = alloca i32, align 4
+  %x = alloca i32, align 4
+  store i32 1024, i32* %i, align 4
+  br label %for.body
+
+for.body:                                          ; preds = %for.cond, %entry
+; CHECK-SPIRV: Label [[#HEADER:]]
+  %0 = load i32, i32* %i, align 4
+  %rem = srem i32 %0, 2
+  %tobool = icmp ne i32 %rem, 0
+; Per SPIRV spec p3.23 "Unroll" loop control = 0x1
+; CHECK-SPIRV: LoopMerge [[#MERGEBLOCK:]] [[#CONTINUE:]] 1
+; CHECK-SPIRV: BranchConditional
+  br i1 %tobool, label %if.then, label %if.end
+
+if.then:                                          ; preds = %for.body
+; CHECK-SPIRV: Label
+  br label %for.cond
+
+if.end:                                           ; preds = %for.body
+; CHECK-SPIRV: Label
+  %1 = load i32, i32* %i, align 4
+  store i32 %1, i32* %x, align 4
+  br label %for.cond
+
+for.cond:                                          ; preds = %if.end, %if.then
+; CHECK-SPIRV: Label [[#CONTINUE]]
+  %2 = load i32, i32* %i, align 4
+  %dec = add nsw i32 %2, -1
+  store i32 %dec, i32* %i, align 4
+  %cmp = icmp sgt i32 %2, 0
+; CHECK-SPIRV: BranchConditional [[#]] [[#MERGEBLOCK]] [[#HEADER]]
+; CHECK-LLVM: br i1 %cmp, label %for.end, label %for.body, !llvm.loop ![[#UNROLLENABLE2:]]
+  br i1 %cmp, label %for.end, label %for.body, !llvm.loop !9
+
+for.end:                                           ; preds = %for.cond
+; CHECK-SPIRV: Label [[#MERGEBLOCK]]
   ret void
 }
 
@@ -202,3 +257,11 @@ attributes #0 = { noinline nounwind optnone "correctly-rounded-divide-sqrt-fp-ma
 !8 = !{!"llvm.loop.unroll.count", i32 8}
 !9 = distinct !{!9, !10}
 !10 = !{!"llvm.loop.unroll.enable"}
+
+; CHECK-LLVM: ![[#UNROLLDISABLE]] = distinct !{![[#UNROLLDISABLE]], ![[#DISABLE:]]}
+; CHECK-LLVM: ![[#DISABLE]] = !{!"llvm.loop.unroll.disable"}
+; CHECK-LLVM: ![[#UNROLLCOUNT]] = distinct !{![[#UNROLLCOUNT]], ![[#COUNT:]]}
+; CHECK-LLVM: ![[#COUNT]] = !{!"llvm.loop.unroll.count", i32 8}
+; CHECK-LLVM: ![[#UNROLLENABLE1]] = distinct !{![[#UNROLLENABLE1]], ![[#ENABLE:]]}
+; CHECK-LLVM: ![[#ENABLE]] = !{!"llvm.loop.unroll.enable"}
+; CHECK-LLVM: ![[#UNROLLENABLE2]] = distinct !{![[#UNROLLENABLE2]], ![[#ENABLE]]}