Refine the set of UniformAfterVectorization instructions

[wmi]

Summary

For %0 in the loop below, it will have both a scalar version and a vector version after vectorization, and it shouldn't be added into uniform set.

define void @test(i32* %a, i64 %n) {
entry:

br label %for.body

for.body:

%i = phi i64 [ %i.next, %for.body ], [ 0, %entry ]
%0 = trunc i64 %i to i32
%1 = getelementptr inbounds i32, i32* %a, i32 %0    ====> %0 used in uniform instruction.
store i32 %0, i32* %1, align 4                                      ====> %0 used in non-uniform instruction.
%i.next = add nuw nsw i64 %i, 1
%cond = icmp eq i64 %i.next, %n
br i1 %cond, label %for.end, label %for.body

for.end:

ret void

}

Before the patch all dependencies of conditional branch or consecutive ptrs will be added into uniform set. The patch changes that to: except seed uniform instructions (conditional branch and consecutive ptr instructions), dependencies to be added into uniform set should only be used by existing uniform instructions or intructions outside of current loop.

This patch uses parts in http://reviews.llvm.org/D21631 from mkuper.

[mkuper]

lib/Transforms/Vectorize/LoopVectorize.cpp

4947 ↗

(On Diff #61979)

Did you mean "PHIs will not be added"? Also, can you explain here why they're not added?
(I mean, this doesn't even apply to all PHIs, IIRC? Some PHIs will be added, it depends on whether the branch depends on the PHI itself, or on UpdateV, right?)

4962 ↗

(On Diff #61979)

This looks a bit weird. If the PHI has a non-uniform use, but the update doesn't, then do we expect to have a vector version of the update?
More generally, it's a bit odd that we look at exactly one instruction backwards past the Phi. But this looks like it's a net improvement over how it worked before.

test/Transforms/LoopVectorize/X86/uniform-phi.ll

31 ↗

(On Diff #61979)

You probably want a CHECK-NOT for %i being found uniform, right?

[wmi]

lib/Transforms/Vectorize/LoopVectorize.cpp

4947 ↗

(On Diff #61979)

Yes, all PHI will not be added in the above loop. That is because PHI must exist in a dep loop, when we expand Worklist in topological order, all the nodes inside dep loop are impossible to be added into Worklist.

That is why we want to handle induction PHI specially below.

4962 ↗

(On Diff #61979)

If the PHI has a non-uniform use, but the update doesn't, neither of them will be added into uniform set because we will have vector versions for both pre-increment-induction var (%vec.ind) and post-increment-induction var (%step.add).

I only look at one instruction backwords past PHI because I think it covers most of the cases.

for.body:

%indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
 ......
%indvars.iv.next = add nuw nsw i64 %indvars.iv, %step
br i1 %exitcond, label %for.end, label %for.body

Like the loop here, all the recursive dependents of induction PHI only includes %indvars.iv.next, %indvars.iv and %step. Because it is a vectorizable loop, %step is a loop invariant and should be defined out of loop. Then we only care whether %indvars.iv and %indvars.iv.next (i.e., PN and UpdateV) are uniform.

test/Transforms/LoopVectorize/X86/uniform-phi.ll

31 ↗

(On Diff #61979)

You are right. I will add it.

[wmi]

Update the testcase according to mkuper's comments.

[wmi]

Add more comments to explain why PHI should be processed separately in collectLoopUniform.

[mkuper]

The logic LGTM, now that I understand it. :-)
A few nits about the code inline.

lib/Transforms/Vectorize/LoopVectorize.cpp

4900 ↗

(On Diff #62323)

"inside current" -> "inside the current"

4901 ↗

(On Diff #62323)

"scope where" -> "the scope which"

4904 ↗

(On Diff #62323)

Why not just "return (!I || !TheLoop->contains(I))"?
But I'm ok with this if you think it's clearer.

4909 ↗

(On Diff #62323)

Can this be a SetVector<Instruction *>?
I think all of your inserts are protected by isOutOfScope(), so you can never have a non-instruction Value there now, right?

4917 ↗

(On Diff #62323)

Can you use a range for here, iterating over TheLoop->blocks()?

4920 ↗

(On Diff #62323)

Range for here too?

4971 ↗

(On Diff #62323)

Can you do something like Uniforms.insert(Worklist.begin(), Worklist.end()) instead?
We can move the debug print to where things are added to Worklist - I think it may be more useful there regardless, since it will then print the uniforms in the order that they are found.

[wmi]

lib/Transforms/Vectorize/LoopVectorize.cpp

4900 ↗

(On Diff #62323)

Fixed.

4901 ↗

(On Diff #62323)

Fixed.

4904 ↗

(On Diff #62323)

What You suggested is better. Fixed.

4909 ↗

(On Diff #62323)

It can be SetVector<Instruction *>, but then we need to use cast<instruction> before adding values into Worklist or checking whether users are already in Worklist.

4917 ↗

(On Diff #62323)

Fixed.

4920 ↗

(On Diff #62323)

Fixed.

4971 ↗

(On Diff #62323)

If I use Uniforms.insert(Worklist.begin(), Worklist.end()), The elem type of Worklist should be "Instruction *" instead of "Value *". I feel that will increase the number of cast<Instruction> needed. I am not sure which way is better. Either way is fine with me. I updated the patch according to suggestion here.

[wmi]

Address mkuper's comments.

[mkuper]

LGTM.
(If you strongly prefer the version where Worklist is a SetVector<Value *>, I'm ok with that too.)

llvm/trunk/lib/Transforms/Vectorize/LoopVectorize.cpp

if (Induction && WidestIndTy != Induction->getType())
     Induction = nullptr;
 
   return true;
 }
 
 void LoopVectorizationLegality::collectLoopUniforms() {
   // We now know that the loop is vectorizable!
   // Collect variables that will remain uniform after vectorization.
-  std::vector<Value *> Worklist;
-  BasicBlock *Latch = TheLoop->getLoopLatch();
 
-  // Start with the conditional branch and walk up the block.
-  Worklist.push_back(Latch->getTerminator()->getOperand(0));
+  // If V is not an instruction inside the current loop, it is a Value
+  // outside of the scope which we are interesting in.
+  auto isOutOfScope = [&](Value *V) -> bool {
+    Instruction *I = dyn_cast<Instruction>(V);
+    return (!I || !TheLoop->contains(I));
+  };
+
+  SetVector<Instruction *> Worklist;
+  BasicBlock *Latch = TheLoop->getLoopLatch();
+  // Start with the conditional branch.
+  if (!isOutOfScope(Latch->getTerminator()->getOperand(0))) {
+    Instruction *Cmp = cast<Instruction>(Latch->getTerminator()->getOperand(0));
+    Worklist.insert(Cmp);
+    DEBUG(dbgs() << "LV: Found uniform instruction: " << *Cmp << "\n");
+  }
 
   // Also add all consecutive pointer values; these values will be uniform
-  // after vectorization (and subsequent cleanup) and, until revectorization is
-  // supported, all dependencies must also be uniform.
-  for (Loop::block_iterator B = TheLoop->block_begin(),
-                            BE = TheLoop->block_end();
-       B != BE; ++B)
-    for (Instruction &I : **B)
-      if (I.getType()->isPointerTy() && isConsecutivePtr(&I))
-        Worklist.insert(Worklist.end(), I.op_begin(), I.op_end());
-
-  while (!Worklist.empty()) {
-    Instruction *I = dyn_cast<Instruction>(Worklist.back());
-    Worklist.pop_back();
-
-    // Look at instructions inside this loop.
-    // Stop when reaching PHI nodes.
-    // TODO: we need to follow values all over the loop, not only in this block.
-    if (!I || !TheLoop->contains(I) || isa<PHINode>(I))
-      continue;
+  // after vectorization (and subsequent cleanup).
+  for (auto *BB : TheLoop->getBlocks()) {
+    for (auto &I : *BB) {
+      if (I.getType()->isPointerTy() && isConsecutivePtr(&I)) {
+        Worklist.insert(&I);
+        DEBUG(dbgs() << "LV: Found uniform instruction: " << I << "\n");
+      }
+    }
+  }
 
-    // This is a known uniform.
-    Uniforms.insert(I);
+  // Expand Worklist in topological order: whenever a new instruction
+  // is added , its users should be either already inside Worklist, or
+  // out of scope. It ensures a uniform instruction will only be used
+  // by uniform instructions or out of scope instructions.
+  unsigned idx = 0;
+  do {
+    Instruction *I = Worklist[idx++];
 
-    // Insert all operands.
-    Worklist.insert(Worklist.end(), I->op_begin(), I->op_end());
+    for (auto OV : I->operand_values()) {
+      if (isOutOfScope(OV))
+        continue;
+      Instruction *OI = cast<Instruction>(OV);
+      if (std::all_of(OI->user_begin(), OI->user_end(), [&](User *U) -> bool {
+            return isOutOfScope(U) || Worklist.count(cast<Instruction>(U));
+          })) {
+        Worklist.insert(OI);
+        DEBUG(dbgs() << "LV: Found uniform instruction: " << *OI << "\n");
+      }
+    }
+  } while (idx != Worklist.size());
+
+  // For an instruction to be added into Worklist above, all its users inside
+  // the current loop should be already added into Worklist. This condition
+  // cannot be true for phi instructions which is always in a dependence loop.
+  // Because any instruction in the dependence cycle always depends on others
+  // in the cycle to be added into Worklist first, the result is no ones in
+  // the cycle will be added into Worklist in the end.
+  // That is why we process PHI separately.
+  for (auto &Induction : *getInductionVars()) {
+    auto *PN = Induction.first;
+    auto *UpdateV = PN->getIncomingValueForBlock(TheLoop->getLoopLatch());
+    if (std::all_of(PN->user_begin(), PN->user_end(),
+                    [&](User *U) -> bool {
+                      return U == UpdateV || isOutOfScope(U) ||
+                             Worklist.count(cast<Instruction>(U));
+                    }) &&
+        std::all_of(UpdateV->user_begin(), UpdateV->user_end(),
+                    [&](User *U) -> bool {
+                      return U == PN || isOutOfScope(U) ||
+                             Worklist.count(cast<Instruction>(U));
+                    })) {
+      Worklist.insert(cast<Instruction>(PN));
+      Worklist.insert(cast<Instruction>(UpdateV));
+      DEBUG(dbgs() << "LV: Found uniform instruction: " << *PN << "\n");
+      DEBUG(dbgs() << "LV: Found uniform instruction: " << *UpdateV << "\n");
+    }
   }
+
+  Uniforms.insert(Worklist.begin(), Worklist.end());
 }
 
 bool LoopVectorizationLegality::canVectorizeMemory() {
   LAI = &LAA->getInfo(TheLoop);
   InterleaveInfo.setLAI(LAI);
   auto &OptionalReport = LAI->getReport();
   if (OptionalReport)
     emitAnalysis(VectorizationReport(*OptionalReport));

llvm/trunk/test/Transforms/LoopVectorize/X86/uniform-phi.ll

+; RUN: opt < %s  -loop-vectorize -mtriple=x86_64-apple-macosx10.8.0 -mcpu=corei7 -debug-only=loop-vectorize -S 2>&1 | FileCheck %s
+; REQUIRES: asserts
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK-LABEL: test
+; CHECK-DAG: LV: Found uniform instruction:   %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+; CHECK-DAG: LV: Found uniform instruction:   %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+; CHECK-DAG: LV: Found uniform instruction:   %exitcond = icmp eq i64 %indvars.iv, 1599
+
+define void @test(float* noalias nocapture %a, float* noalias nocapture readonly %b) #0 {
+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 float, float* %b, i64 %indvars.iv
+  %tmp0 = load float, float* %arrayidx, align 4
+  %add = fadd float %tmp0, 1.000000e+00
+  %arrayidx5 = getelementptr inbounds float, float* %a, i64 %indvars.iv
+  store float %add, float* %arrayidx5, align 4
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp eq i64 %indvars.iv, 1599
+  br i1 %exitcond, label %for.end, label %for.body
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+
+; CHECK-LABEL: foo
+; CHECK-DAG: LV: Found uniform instruction:   %cond = icmp eq i64 %i.next, %n
+; CHECK-DAG: LV: Found uniform instruction:   %tmp1 = getelementptr inbounds i32, i32* %a, i32 %tmp0
+; CHECK-NOT: LV: Found uniform instruction:   %i = phi i64 [ %i.next, %for.body ], [ 0, %entry ]
+
+define void @foo(i32* %a, i64 %n) {
+entry:
+  br label %for.body
+
+for.body:
+  %i = phi i64 [ %i.next, %for.body ], [ 0, %entry ]
+  %tmp0 = trunc i64 %i to i32
+  %tmp1 = getelementptr inbounds i32, i32* %a, i32 %tmp0
+  store i32 %tmp0, i32* %tmp1, align 4
+  %i.next = add nuw nsw i64 %i, 1
+  %cond = icmp eq i64 %i.next, %n
+  br i1 %cond, label %for.end, label %for.body
+
+for.end:
+  ret void
+}