LoopUnroll: respect pragma unroll when AllowRemainder is disabled

[yaxunl]

Summary

Currently when AllowRemainder is disabled, pragma unroll count is not
respected even though there is no remainder. This bug causes a loop
fully unrolled in many cases even though the user specifies a unroll
count. Especially it affects OpenCL/CUDA since in many cases a loop
contains convergent instructions and currently AllowRemainder is
disabled for such loops.

This patch is related to https://reviews.llvm.org/D43594 but that one
deals with situations where remainder exists.

[nhaehnle]

Thanks, this makes more sense than the other change. LGTM, although you may want to wait a bit for other people to chime in.

[efriedma]

lib/Transforms/Scalar/LoopUnrollPass.cpp

733

I don't think the TripCount check here is right; if the trip count isn't constant, TripCount will be zero.

The check is unnecessary, anyway: if the trip count is constant, TripCount is the same as TripMultiple.

[yaxunl]

Revised by Eli's comments and added more tests for runtime trip count and indivisible trip count.

[efriedma]

LGTM

lib/Transforms/Scalar/LoopUnrollPass.cpp

static bool computeUnrollCount(
 
   // 2nd priority is unroll count set by pragma.
   unsigned PragmaCount = UnrollCountPragmaValue(L);
   if (PragmaCount > 0) {
     UP.Count = PragmaCount;
     UP.Runtime = true;
     UP.AllowExpensiveTripCount = true;
     UP.Force = true;
-    if (UP.AllowRemainder &&
+    if ((UP.AllowRemainder || (TripCount % PragmaCount == 0) ||
+         (TripMultiple % PragmaCount == 0)) &&
         getUnrolledLoopSize(LoopSize, UP) < PragmaUnrollThreshold)
       return true;
   }
   bool PragmaFullUnroll = HasUnrollFullPragma(L);
   if (PragmaFullUnroll && TripCount != 0) {
     UP.Count = TripCount;
     if (getUnrolledLoopSize(LoopSize, UP) < PragmaUnrollThreshold)
       return false;

test/Transforms/LoopUnroll/convergent.ll

; CHECK-NOT: call void @f()
   %inc = add nsw i32 %x.0, 1
   %exitcond = icmp eq i32 %inc, %loop_ctl
   br i1 %exitcond, label %exit, label %l3, !llvm.loop !0
 
 exit:
   ret i32 0
 }
 
+; This loop contains a convergent instruction. Since the pragma loop unroll
+; count 2 divides trip count 4. The loop unroll should respect the pragma.
+; CHECK-LABEL: @pragma_unroll2
+define void @pragma_unroll2() {
+entry:
+  br label %l3, !llvm.loop !1
+
+l3:
+  %x.0 = phi i32 [ 0, %entry ], [ %inc, %l3 ]
+; CHECK: call void @f()
+; CHECK: call void @f()
+; CHECK-NOT: call void @f()
+  call void @f() convergent
+  %inc = add nsw i32 %x.0, 1
+  %exitcond = icmp eq i32 %inc, 4
+  br i1 %exitcond, label %exit, label %l3, !llvm.loop !1
+
+exit:
+  ret void
+}
+
+; This loop contains a convergent instruction. Since the pragma loop unroll
+; count 2 divides trip multiple 2. The loop unroll should respect the pragma.
+; CHECK-LABEL: @pragma_unroll3
+define i32 @pragma_unroll3(i32 %n) {
+entry:
+  %loop_ctl = mul nsw i32 %n, 2
+  br label %l3, !llvm.loop !1
+
+l3:
+  %x.0 = phi i32 [ 0, %entry ], [ %inc, %l3 ]
+; CHECK: call void @f()
+; CHECK: call void @f()
+; CHECK-NOT: call void @f()
+  call void @f() convergent
+  %inc = add nsw i32 %x.0, 1
+  %exitcond = icmp eq i32 %inc, %loop_ctl
+  br i1 %exitcond, label %exit, label %l3, !llvm.loop !1
+
+exit:
+  ret i32 0
+}
+
 !0 = !{!0, !{!"llvm.loop.unroll.count", i32 16}}
+!1 = !{!1, !{!"llvm.loop.unroll.count", i32 2}}
+

test/Transforms/LoopUnroll/unroll-pragmas.ll

-; RUN: opt < %s -loop-unroll -pragma-unroll-threshold=1024 -S | FileCheck %s
-; RUN: opt < %s -loop-unroll -loop-unroll -pragma-unroll-threshold=1024 -S | FileCheck %s
+; RUN: opt < %s -loop-unroll -pragma-unroll-threshold=1024 -S | FileCheck -check-prefixes=CHECK,REM %s
+; RUN: opt < %s -loop-unroll -loop-unroll -pragma-unroll-threshold=1024 -S | FileCheck -check-prefixes=CHECK,REM %s
+; RUN: opt < %s -loop-unroll -unroll-allow-remainder=0 -pragma-unroll-threshold=1024 -S | FileCheck -check-prefixes=CHECK,NOREM %s
 ;
 ; Run loop unrolling twice to verify that loop unrolling metadata is properly
 ; removed and further unrolling is disabled after the pass is run once.
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; loop4 contains a small loop which should be completely unrolled by

 
 ; #pragma clang loop unroll(enable)
 ; Loop has a runtime trip count and should be runtime unrolled and duplicated
 ; (original and 8x).
 ;
 ; CHECK-LABEL: @runtime_loop_with_enable(
 ; CHECK: for.body:
 ; CHECK: store i32
-; CHECK: store i32
-; CHECK: store i32
-; CHECK: store i32
-; CHECK: store i32
-; CHECK: store i32
-; CHECK: store i32
-; CHECK: store i32
+; REM: store i32
+; REM: store i32
+; REM: store i32
+; REM: store i32
+; REM: store i32
+; REM: store i32
+; REM: store i32
 ; CHECK-NOT: store i32
 ; CHECK: br i1
-; CHECK: for.body.epil:
-; CHECK: store
+; REM: for.body.epil:
+; NOREM-NOT: for.body.epil:
+; REM: store
 ; CHECK-NOT: store
-; CHECK: br i1
+; REM: br i1
+; NOREM-NOT: br i1
 define void @runtime_loop_with_enable(i32* nocapture %a, i32 %b) {
 entry:
   %cmp3 = icmp sgt i32 %b, 0
   br i1 %cmp3, label %for.body, label %for.end, !llvm.loop !8
 
 for.body:                                         ; preds = %entry, %for.body
   %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
   %arrayidx = getelementptr inbounds i32, i32* %a, i64 %indvars.iv