Unroll once when -unroll-count is set

[evstupac]

Summary

Set disable unroll, when loop is unrolled once even if there were no metadata.
(basically for "-unroll-count" passed case)

[evstupac]

+test

[mzolotukhin]

Hi Evgeny,

The change looks good, but I'd like the test to be changed. We don't need to run entire O2 pipeline for it. If you're trying to catch some behavior of the second invocation of loop-unroll, please just add IR from right before that as another test.

Thanks,
Michael

test/Transforms/LoopUnroll/runtime-count.ll

1

Please run only -loop-unroll, not entire -O2.

[evstupac]

test fixed

[mzolotukhin]

LGTM with a minor remark below.

Michael

test/Transforms/LoopUnroll/unroll-count.ll

2 ↗

(On Diff #59451)

The test doesn't check that unrolling occurs only once, it only checks if the metadata is set after unrolling - this is totally fine, I just think the comment is slightly off.

[evstupac]

committed revision 272195

lib/Transforms/Scalar/LoopUnrollPass.cpp

 }
 
 // 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<Metadata *, 4> MDs;
   // Reserve first location for self reference to the LoopID metadata node.
   MDs.push_back(nullptr);
+
+  if (LoopID) {
     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)
         MDs.push_back(LoopID->getOperand(i));
     }
+  }
 
   // Add unroll(disable) metadata to disable future unrolling.
   LLVMContext &Context = L->getHeader()->getContext();
   SmallVector<Metadata *, 1> DisableOperands;
   DisableOperands.push_back(MDString::get(Context, "llvm.loop.unroll.disable"));
   MDNode *DisableNode = MDNode::get(Context, DisableOperands);
   MDs.push_back(DisableNode);
 

test/Transforms/LoopUnroll/runtime-count.ll

+; RUN: opt < %s -S -O2 -unroll-count=2 | FileCheck %s
+; Loop should be unrolled 2 times (not 4)
+;
+; CHECK-LABEL: @loop64_with_count4(
+; CHECK: store i32
+; CHECK: store i32
+; CHECK-NOT: store i32
+; CHECK: br i1
+define void @loop64_with_count4(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, 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
+  %exitcond = icmp eq i64 %indvars.iv.next, 64
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+