[AMDGPU] Respect pragma unroll when loop contains convergent instructions

[yaxunl]

Summary

Currently loop unroll is conservative about loops containing convergent instructions.
It does not allow remainder for such loops, which essentially disables unroll count
requested by pragma and results in fully unrolled loop in many cases.

As such a user may specify pragma unroll 32 but instead gets the loop unrolled 512
and results in extremely long compilation time.

For some target, e.g. AMDGPU, the remainder does not cause extra divergence and
should be allowed.

This patch introduces AllowRemainderForConvergentLoop in
TargetTransformInfo::UnrollingPreferences and allows each target to specify
whether unrolling convergent loop with remainder is allowed. By default it is
false therefore no functional change for other targets.

This patch fixes shmembench-ocl compilation time issue on amdpu.

[rampitec]

LGTM

[efriedma]

include/llvm/Analysis/TargetTransformInfo.h

426

I don't like sticking this here.

From your description, it sounds like it's a *correctness* property of the target, whether or not certain transforms which duplicate convergent operations are allowed. In that case, it's not really about unrolling at all; it could apply to other transforms which clone code. So at the very least, this should be a separate hook, with a clear explanation of exactly which transforms this allows.

[yaxunl]

include/llvm/Analysis/TargetTransformInfo.h

426

For this specific transform (adding remainder for unrolling loop) we know that it will not cause extra divergence on amdgcn target. However this is something not easily applied to other cases. So far I do not see how it can be applied to other transformations.

[efriedma]

include/llvm/Analysis/TargetTransformInfo.h

426

The allowed transform can't be specifically "remainder loops produced by lib/Transforms/Utils/LoopUnrollRuntime.cpp in LLVM r324285"; there must be some set of similar transforms which are allowed (whether or not they're currently implemented in the current LLVM codebase).

test/Transforms/LoopUnroll/convergent.ll

100

I'm not sure this testcase really demonstrates what you want it to demonstrate... a trip count of 4 is divisible by an unroll count of 2, so you don't need a remainder loop anyway.

108

Are you sure this metadata is correct? It would be nice to have a separate test without the "convergent" marking to show it's making a difference.

[jlebar]

The allowed transform can't be specifically "remainder loops produced by lib/Transforms/Utils/LoopUnrollRuntime.cpp in LLVM r324285"; there must be some set of similar transforms which are allowed (whether or not they're currently implemented in the current LLVM codebase).

Agree.

If we phrase this in terms of a specific set of transformations that are/aren't allowed, we may even be able to say that a remainder loop containing convergent functions is in fact safe on all platforms. I'm not sure, need to think about it...

[nhaehnle]

test/Transforms/LoopUnroll/convergent.ll

100

Agreed. I have to say, it looks to me like the loop unroll is simply overly conservative here, and should stick with the requested 2x unroll on all targets, despite the convergent function call.

There really shouldn't be a difference between AMDGPU and NVPTX at this point.

[yaxunl]

test/Transforms/LoopUnroll/convergent.ll

100

It seems there is a separate but related bug in loop unroll: basically if AllowRemainder is disabled, pragma unroll count is not respected even though there is no remainder. I created another patch for that issue https://reviews.llvm.org/D43826

include/llvm/Analysis/TargetTransformInfo.h

struct UnrollingPreferences {
     /// (mainly to loops that fail runtime unrolling).
     bool Force;
     /// Allow using trip count upper bound to unroll loops.
     bool UpperBound;
     /// Allow peeling off loop iterations for loops with low dynamic tripcount.
     bool AllowPeeling;
     /// Allow unrolling of all the iterations of the runtime loop remainder.
     bool UnrollRemainder;
+    /// Allow unrolling convergent loop with remainder.
+    bool AllowRemainderForConvergentLoop;
   };
 
   /// \brief Get target-customized preferences for the generic loop unrolling
   /// transformation. The caller will initialize UP with the current
   /// target-independent defaults.
   void getUnrollingPreferences(Loop *L, ScalarEvolution &,
                                UnrollingPreferences &UP) const;
 

lib/Target/AMDGPU/AMDGPUTargetTransformInfo.cpp

static bool dependsOnLocalPhi(const Loop *L, const Value *Cond,
   return false;
 }
 
 void AMDGPUTTIImpl::getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
                                             TTI::UnrollingPreferences &UP) {
   UP.Threshold = 300; // Twice the default.
   UP.MaxCount = std::numeric_limits<unsigned>::max();
   UP.Partial = true;
+  UP.AllowRemainderForConvergentLoop = true;
 
   // TODO: Do we want runtime unrolling?
 
   // Maximum alloca size than can fit registers. Reserve 16 registers.
   const unsigned MaxAlloca = (256 - 16) * 4;
   unsigned ThresholdPrivate = UnrollThresholdPrivate;
   unsigned ThresholdLocal = UnrollThresholdLocal;
   unsigned MaxBoost = std::max(ThresholdPrivate, ThresholdLocal);

lib/Transforms/Scalar/LoopUnrollPass.cpp

static TargetTransformInfo::UnrollingPreferences gatherUnrollingPreferences(
   UP.Partial = false;
   UP.Runtime = false;
   UP.AllowRemainder = true;
   UP.UnrollRemainder = false;
   UP.AllowExpensiveTripCount = false;
   UP.Force = false;
   UP.UpperBound = false;
   UP.AllowPeeling = true;
+  UP.AllowRemainderForConvergentLoop = false;
 
   // Override with any target specific settings
   TTI.getUnrollingPreferences(L, SE, UP);
 
   // Apply size attributes
   if (L->getHeader()->getParent()->optForSize()) {
     UP.Threshold = UP.OptSizeThreshold;
     UP.PartialThreshold = UP.PartialOptSizeThreshold;
static LoopUnrollResult tryToUnrollLoop(
   //
   // TODO: This is quite conservative.  In practice, convergent_op()
   // is likely to be called unconditionally in the loop.  In this
   // case, the program would be ill-formed (on most architectures)
   // unless n were the same on all threads in a thread group.
   // Assuming n is the same on all threads, any kind of unrolling is
   // safe.  But currently llvm's notion of convergence isn't powerful
   // enough to express this.
-  if (Convergent)
+  if (Convergent && !UP.AllowRemainderForConvergentLoop)
     UP.AllowRemainder = false;
 
   // Try to find the trip count upper bound if we cannot find the exact trip
   // count.
   bool MaxOrZero = false;
   if (!TripCount) {
     MaxTripCount = SE.getSmallConstantMaxTripCount(L);
     MaxOrZero = SE.isBackedgeTakenCountMaxOrZero(L);

test/Transforms/LoopUnroll/AMDGPU/convergent.ll

+; RUN: opt < %s -loop-unroll -S | FileCheck %s
+
+target triple = "amdgcn"
+
+declare void @f() convergent
+
+; This loop contains a convergent instruction. Since AMDGPU target transform
+; info allows unrolling loop with convergent instruction with remainders,
+; the loop is unrolled following its pragma unroll value 2, instead of
+; doing a full unroll.
+
+define void @pragma_unroll() {
+entry:
+  br label %l3, !llvm.loop !0
+
+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 !0
+
+exit:
+  ret void
+
+}
+
+!0 = !{!0, !{!"llvm.loop.unroll.count", i32 2}}

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, so allow remainder
+; is disabled. This overrides its unroll pragma -- we unroll 4 times,
+; even though 2 is requested.
+; CHECK-LABEL: @pragma_unroll2
+define void @pragma_unroll2() {
+entry:
+  br label %l3, !llvm.loop !0
+
+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()
+  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
+
+}
+
 !0 = !{!0, !{!"llvm.loop.unroll.count", i32 16}}
+!1 = !{!0, !{!"llvm.loop.unroll.count", i32 2}}