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 ↗

(On Diff #136090)

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

llvm/trunk/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 || (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;

llvm/trunk/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_unroll_divisible_trip_count
+define void @pragma_unroll_divisible_trip_count() {
+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_unroll_divisible_trip_multiple
+define i32 @pragma_unroll_divisible_trip_multiple(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
+}
+
+; This loop contains a convergent instruction. Since the pragma loop unroll
+; count 2 is unknown to divide runtime trip count, the loop is not unrolled
+; since remainder is forbidden for unrolling convergent loop.
+; ToDo: Forbidding remainder for unrolling convergent loop may be relaxed
+; in the future.
+; CHECK-LABEL: @pragma_unroll_indivisible_runtime_trip_count
+define i32 @pragma_unroll_indivisible_runtime_trip_count(i32 %n) {
+entry:
+  br label %l3, !llvm.loop !1
+
+l3:
+  %x.0 = phi i32 [ 0, %entry ], [ %inc, %l3 ]
+; 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, %n
+  br i1 %exitcond, label %exit, label %l3, !llvm.loop !1
+
+exit:
+  ret i32 0
+}
+
+; This loop contains a convergent instruction. Since the pragma loop unroll
+; count 2 does not divide trip count 5, the loop is not unrolled by 2
+; since remainder is forbidden for unrolling convergent loop. Instead, the
+; loop gets fully unrolled.
+; ToDo: Forbidding remainder for unrolling convergent loop may be relaxed
+; in the future.
+; CHECK-LABEL: @pragma_unroll_indivisible_trip_count
+define i32 @pragma_unroll_indivisible_trip_count() {
+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: call void @f()
+; 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, 5
+  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}}
+

llvm/trunk/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

 
 for.end:                                          ; preds = %for.body, %entry
   ret void
 }
 !8 = !{!8, !4}
 
 ; #pragma clang loop unroll_count(4)
 ; Loop has a runtime trip count.  Runtime unrolling should occur and loop
-; should be duplicated (original and 4x unrolled).
+; should be duplicated (original and 4x unrolled) if remainder is allowed,
+; otherwise loop should not be unrolled.
 ;
 ; CHECK-LABEL: @runtime_loop_with_count4(
 ; CHECK: for.body
 ; CHECK: store
-; CHECK: store
-; CHECK: store
-; CHECK: store
+; REM: store
+; REM: store
+; REM: store
 ; CHECK-NOT: store
 ; CHECK: br i1
-; CHECK: for.body.epil:
-; CHECK: store
+; REM: for.body.epil:
+; REM: store
+; NOREM-NOT: for.body.epil:
+; NOREM-NOT: store
 ; CHECK-NOT: store
-; CHECK: br i1
+; REM: br i1
+; NOREM-NOT: br i1
 define void @runtime_loop_with_count4(i32* nocapture %a, i32 %b) {
 entry:
   %cmp3 = icmp sgt i32 %b, 0
   br i1 %cmp3, label %for.body, label %for.end, !llvm.loop !9
 
 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

 for.end:                                          ; preds = %for.body
   ret void
 }
 !13 = !{!13, !14}
 !14 = !{!"llvm.loop.unroll.enable"}
 
 ; #pragma clang loop unroll(enable)
 ; Loop has a runtime trip count and should be runtime unrolled and duplicated
-; (original and 8x).
+; (original and 8x) if remainder is allowed, otherwise it should not be
+; unrolled.
 ;
 ; 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
   %0 = load i32, i32* %arrayidx, align 4
   %inc = add nsw i32 %0, 1
   store i32 %inc, i32* %arrayidx, align 4
   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
   %lftr.wideiv = trunc i64 %indvars.iv.next to i32
   %exitcond = icmp eq i32 %lftr.wideiv, %b
   br i1 %exitcond, label %for.end, label %for.body, !llvm.loop !15
 
 for.end:                                          ; preds = %for.body, %entry
   ret void
 }
 !15 = !{!15, !14}
 
 ; #pragma clang loop unroll_count(3)
 ; Loop has a runtime trip count.  Runtime unrolling should occur and loop
-; should be duplicated (original and 3x unrolled).
+; should be duplicated (original and 3x unrolled) if remainder is allowed,
+; otherwise it should not be unrolled.
 ;
 ; CHECK-LABEL: @runtime_loop_with_count3(
 ; CHECK: for.body
 ; CHECK: store
-; CHECK: store
-; CHECK: store
+; REM: store
+; REM: store
 ; CHECK-NOT: store
 ; CHECK: br i1
-; CHECK: for.body.epil:
-; CHECK: store
+; REM: for.body.epil:
+; REM: store
+; NOREM-NOT: for.body.epil:
+; NOREM-NOT: store
 ; CHECK-NOT: store
-; CHECK: br i1
+; REM: br i1
 define void @runtime_loop_with_count3(i32* nocapture %a, i32 %b) {
 entry:
   %cmp3 = icmp sgt i32 %b, 0
   br i1 %cmp3, label %for.body, label %for.end, !llvm.loop !16
 
 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