Remove unroll pragma metadata after it is used

[meheff]

Summary

This change causes the loop unroller to remove unroll pragma metadata after it is used. This prevents a loop from being unrolled more than is directed by a pragma if the loop
unrolling pass is run more than once which it generally is.

[hfinkel]

lib/Transforms/Scalar/LoopUnrollPass.cpp

286

I suppose this 1 Value* is a placeholder for the self-reference loop-id. Should comment. Also, you expect this to be nullptr, right? You should explicit specify that so it is clear.

301

This does not seem to match the language reference, which says "Note that setting llvm.loop.vectorize.unroll to 1 disables unrolling of the loop." -- Is that out of date?

[meheff]

Thanks for the comments. I'll immediately follow this with the updated patch.

lib/Transforms/Scalar/LoopUnrollPass.cpp

286

Added a comment, and slightly changed the code so it is less surprising.

301

"llvm.loop.vectorize.unroll" is for the vectorizer unroller (interleaver), this is the concatenation loop unroller which uses "llvm.loop.unroll.*". For the concatenation unroller, setting "llvm.loop.unroll.enable" to 0 disables unrolling.

And now I see that "llvm.loop.unroll.*" metadata needs an entry in LangRef.rst and the bit about "llvm.loop.vectorize.unroll" should be changed to match the terminology in LanguageExtensions.rst (use "interleave" instead of "unroll" to avoid confusing it with the loop unroller). I'll fix this in a follow up change.

[eliben]

LGTM

[meheff]

Committed with r213412.

Mark

lib/Transforms/Scalar/LoopUnrollPass.cpp

static unsigned UnrollCountPragmaValue(const Loop *L) {
   if (CountValue) {
     unsigned Count = CountValue->getZExtValue();
     assert(Count >= 1 && "Unroll count must be positive.");
     return Count;
   }
   return 0;
 }
 
+// Remove existing unroll metadata and add unroll disable metadata to
+// indicate the loop has already been unrolled.  This prevents a loop
+// from being unrolled more than is directed by a pragma if the loop
+// unrolling pass is run more than once (which it generally is).
+static void SetLoopAlreadyUnrolled(Loop *L) {
+  MDNode *LoopID = L->getLoopID();
+  if (!LoopID) return;
+
+  // First remove any existing loop unrolling metadata.
+  SmallVector<Value *, 4> Vals;
+  // Reserve first location for self reference to the LoopID metadata node.
+  Vals.push_back(nullptr);
+  for (unsigned i = 1, ie = LoopID->getNumOperands(); i < ie; ++i) {
+    bool IsUnrollMetadata = false;
+    MDNode *MD = dyn_cast<MDNode>(LoopID->getOperand(i));
+    if (MD) {
+      const MDString *S = dyn_cast<MDString>(MD->getOperand(0));
+      IsUnrollMetadata = S && S->getString().startswith("llvm.loop.unroll.");
+    }
+    if (!IsUnrollMetadata) Vals.push_back(LoopID->getOperand(i));
+  }
+
+  // Add unroll(disable) metadata to disable future unrolling.
+  LLVMContext &Context = L->getHeader()->getContext();
+  MDNode *DisableNode =
+      MDNode::get(Context, {MDString::get(Context, "llvm.loop.unroll.enable"),
+                            ConstantInt::get(Type::getInt1Ty(Context), 0)});
+  Vals.push_back(DisableNode);
+
+  MDNode *NewLoopID = MDNode::get(Context, Vals);
+  // Set operand 0 to refer to the loop id itself.
+  NewLoopID->replaceOperandWith(0, NewLoopID);
+  L->setLoopID(NewLoopID);
+  LoopID->replaceAllUsesWith(NewLoopID);
+}
+
 unsigned LoopUnroll::selectUnrollCount(
     const Loop *L, unsigned TripCount, bool HasEnablePragma,
     unsigned PragmaCount, const TargetTransformInfo::UnrollingPreferences &UP,
     bool &SetExplicitly) {
   SetExplicitly = true;
 
   // User-specified count (either as a command-line option or
   // constructor parameter) has highest precedence.
while (Count != 0 && UnrolledSize > PartialThreshold) {
       UnrolledSize = LoopSize * Count;
     }
     if (Count > UP.MaxCount)
       Count = UP.MaxCount;
     DEBUG(dbgs() << "  partially unrolling with count: " << Count << "\n");
   }
 
   if (HasPragma) {
+    // Mark loop as unrolled to prevent unrolling beyond that
+    // requested by the pragma.
+    SetLoopAlreadyUnrolled(L);
+
     // Emit optimization remarks if we are unable to unroll the loop
     // as directed by a pragma.
     DebugLoc LoopLoc = L->getStartLoc();
     Function *F = Header->getParent();
     LLVMContext &Ctx = F->getContext();
     if (HasEnablePragma && PragmaCount == 0) {
       if (TripCount && Count != TripCount) {
         emitOptimizationRemarkMissed(

test/Transforms/LoopUnroll/unroll-pragmas-disabled.ll

+; RUN: opt < %s -loop-unroll -S | FileCheck %s
+;
+; Verify that the unrolling pass removes existing loop unrolling metadata
+; and adds a disable unrolling node after unrolling is complete.
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; #pragma clang loop  vectorize(enable) unroll(enable) unroll_count(4) vectorize_width(8)
+;
+; Unroll metadata should be replaces with unroll(disable).  Vectorize
+; metadata should be untouched.
+;
+; CHECK-LABEL: @loop1(
+; CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[LOOP_1:.*]]
+define void @loop1(i32* nocapture %a) {
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %for.body, %entry
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32* %a, i64 %indvars.iv
+  %0 = load 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
+  %exitcond = icmp eq i64 %indvars.iv.next, 64
+  br i1 %exitcond, label %for.end, label %for.body, !llvm.loop !1
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+!1 = metadata !{metadata !1, metadata !2, metadata !3, metadata !4, metadata !5}
+!2 = metadata !{metadata !"llvm.loop.vectorize.enable", i1 true}
+!3 = metadata !{metadata !"llvm.loop.unroll.enable", i1 true}
+!4 = metadata !{metadata !"llvm.loop.unroll.count", i32 4}
+!5 = metadata !{metadata !"llvm.loop.vectorize.width", i32 8}
+
+; #pragma clang loop unroll(disable)
+;
+; Unroll metadata should not change.
+;
+; CHECK-LABEL: @loop2(
+; CHECK: br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[LOOP_2:.*]]
+define void @loop2(i32* nocapture %a) {
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %for.body, %entry
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds i32* %a, i64 %indvars.iv
+  %0 = load 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
+  %exitcond = icmp eq i64 %indvars.iv.next, 64
+  br i1 %exitcond, label %for.end, label %for.body, !llvm.loop !6
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+!6 = metadata !{metadata !6, metadata !7}
+!7 = metadata !{metadata !"llvm.loop.unroll.enable", i1 false}
+
+; CHECK: ![[LOOP_1]] = metadata !{metadata ![[LOOP_1]], metadata ![[VEC_ENABLE:.*]], metadata ![[WIDTH_8:.*]], metadata ![[UNROLL_DISABLE:.*]]}
+; CHECK: ![[VEC_ENABLE]] = metadata !{metadata !"llvm.loop.vectorize.enable", i1 true}
+; CHECK: ![[WIDTH_8]] = metadata !{metadata !"llvm.loop.vectorize.width", i32 8}
+; CHECK: ![[UNROLL_DISABLE]] = metadata !{metadata !"llvm.loop.unroll.enable", i1 false}
+; CHECK: ![[LOOP_2]] = metadata !{metadata ![[LOOP_2]], metadata ![[UNROLL_DISABLE:.*]]}

test/Transforms/LoopUnroll/unroll-pragmas.ll

 ; RUN: opt < %s -loop-unroll -S | FileCheck %s
+; RUN: opt < %s -loop-unroll -loop-unroll -S | FileCheck %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
 ; the default unrolling heuristics.  It serves as a control for the
 ; unroll(disable) pragma test loop4_with_disable.
 ;