[LV] Allow interleaved accesses in loops with predicated blocks

[mssimpso]

Summary

With this patch, we allow interleaved accesses in loops containing predicated blocks. If an interleaved access is contained in a predicated block, it's only allowed to form an interleaved group with accesses in the same block.

[mssimpso]

Simplified test case

[mcrosier]

lib/Transforms/Vectorize/LoopVectorize.cpp

5091–5092

Does the Comment need updating?

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5139

With that change, analyzeInterleaving won't see all the loads/stores (as the predicated ones get skipped). Would this have any impact on the correctness of the analysis (it might be ok, but it's worth thinking about it).

[mssimpso]

lib/Transforms/Vectorize/LoopVectorize.cpp

5091–5092

I think the comment should probably stay the same. We still can't handle predicated accesses. The current patch just prevents us from giving up on the non-predicated accesses in the loop, if the loop contains a predicated block.

5139

Hey Silviu, thanks for commenting! Right, this prevents accesses from being placed into an interleaved group if they are in a predicated block, which seems like the right thing to do to me. This is what the test case demonstrates. Currently, no accesses are recognized as interleaved if the loop contains a single predicated block. If an access is non-consecutive and not in an interleaved group, it should be scalarized (unless gather/scatter is supported).

[mcrosier]

lib/Transforms/Vectorize/LoopVectorize.cpp

5091–5092

Thanks for the clarification, Matt.

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5139

The idea generally makes sense to me.

However we also have some logic here to handle the Write-After-Write dependences (in fact more than just these) which seems to rely on seeing all accesses. For example, we seen to be able to get the vectorizer confused by doing something like (i was able to modify your example to do this):

1. load a[i]
2. load a[i+1]
3. conditional: store a[i]
4. load a[i] -> it sees this as the first load, because we've ignored 3

Here is my modification of the example:

define void @load_gap_with_pred_store(%pair *%p, i64 %x, i64 %n) {

entry:
  br label %for.body

for.body:
  %i  = phi i64 [ %i.next, %if.merge ], [ 0, %entry ]
  %f1 = getelementptr inbounds %pair, %pair* %p, i64 %i, i32 1
  %f2 = getelementptr inbounds %pair, %pair* %p, i64 %i, i32 0

  %r0 = load i64, i64* %f1, align 8
  %r1 = and i64 %r0, %x
  %r2 = icmp eq i64 %r1, %x
  %r4 = add i64 %r1, 1

  br i1 %r2, label %if.then, label %if.merge

if.then:
  store i64 0, i64* %f2, align 8
  br label %if.merge

if.merge:
  %r3 = load i64, i64* %f2, align 8

  store i64 %r3, i64 *%f1, align 8 // Here we end up storing an element from the wide load, instead of having to reload

  %i.next = add nuw nsw i64 %i, 1
  %cond = icmp slt i64 %i.next, %n
  br i1 %cond, label %for.body, label %for.end

for.end:
  ret void
}

[mssimpso]

lib/Transforms/Vectorize/LoopVectorize.cpp

5139

Thanks for the counterexample, Silviu. I see what you're getting at. Let me think over this change a little more carefully. We will definitely need to do a little more work here.

[mssimpso]

lib/Transforms/Vectorize/LoopVectorize.cpp

5139

Hi Silviu,

I think there may be a bug in the interleaved access analysis, independent of this patch, that your test case led me to. It seems to be related to the same issue of reusing a loaded value instead of reloading due to an intervening store. I've created PR27626 to track the issue.

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5139

Thanks for digging further! Yes, this appears to also be a problem.

[mssimpso]

Rebased on top of D19984

With this revision, we collect all accesses in the loop, even those in the predicated blocks, so we can account for any dependences (D19984). We prevent forming interleaved groups with predicated accesses unless the accesses are in the same block. I've also added Silviu's code example as a new test case.

[mssimpso]

Rebased on top of D19984.

[mssimpso]

Rebased.

With rL273687, I think we can now resume this review.

[mssimpso]

Slightly refactored code and added a new test case

[mssimpso]

Fixed typo in test case comment.

Silviu/Adam,

Do you have any additional comments for this change? It should be fairly straightforward now that the dependences and iteration order have been fixed.

[anemet]

Started with the tests, a question below.

test/Transforms/LoopVectorize/interleaved-accesses-pred-stores.ll

37

Is %p.0 unused here?

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5223

Could we also have a test where all the accesses in a group are predicated and in the same block? In this case would all the predicated accesses be scalarized as well? (there might be some issues with out of bounds accesses from the wide load/store and the scalar epilogue if that's not the case)

[mssimpso]

Wow, thanks for the quick response! Some comments inline.

lib/Transforms/Vectorize/LoopVectorize.cpp

5223

Yes, I wanted a test case like that, but I haven't been able to produce one that the vectorizer likes. It complains about being unable to if-convert the block I believe. I can look into this further, but currently I think we are limited to one store per predicated block. In which case, yes, we would always scalarize.

test/Transforms/LoopVectorize/interleaved-accesses-pred-stores.ll

37

Yes, thanks for noticing! I'll remove it.

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5223

Ok, makes sense.

Maybe we can remove the BlockA != BlockB condition for now and handle this case in a separate change when we can if-convert the block?

[mssimpso]

lib/Transforms/Vectorize/LoopVectorize.cpp

5223

That's a great idea. Thanks! I'll update the patch.

[mssimpso]

Addressed comments from Adam and Silviu. Thanks!

[sbaranga]

LGTM!

lib/Transforms/Vectorize/LoopVectorize.cpp

bool LoopVectorizationLegality::blockCanBePredicated(
   }
 
   return true;
 }
 
 void InterleavedAccessInfo::collectConstStridedAccesses(
     MapVector<Instruction *, StrideDescriptor> &StrideAccesses,
     const ValueToValueMap &Strides) {
-  // Holds load/store instructions in program order.
-  SmallVector<Instruction *, 16> AccessList;
+
+  // The module's data layout.
+  auto &DL = TheLoop->getHeader()->getModule()->getDataLayout();
 
   // Since it's desired that the load/store instructions be maintained in
   // "program order" for the interleaved access analysis, we have to visit the
   // blocks in the loop in reverse postorder (i.e., in a topological order).
   // Such an ordering will ensure that any load/store that may be executed
-  // before a second load/store will precede the second load/store in the
-  // AccessList.
+  // before a second load/store will precede the second load/store in
+  // StrideAccesses.
   LoopBlocksDFS DFS(TheLoop);
   DFS.perform(LI);
-  for (LoopBlocksDFS::RPOIterator I = DFS.beginRPO(), E = DFS.endRPO(); I != E;
-       ++I) {
-    BasicBlock *BB = *I;
-    bool IsPred = LoopAccessInfo::blockNeedsPredication(BB, TheLoop, DT);
-
-    for (auto &I : *BB) {
-      if (!isa<LoadInst>(&I) && !isa<StoreInst>(&I))
+  for (LoopBlocksDFS::RPOIterator BBI = DFS.beginRPO(), BBE = DFS.endRPO();
+       BBI != BBE; ++BBI)
+    for (auto &I : **BBI) {
+      LoadInst *LI = dyn_cast<LoadInst>(&I);
+      StoreInst *SI = dyn_cast<StoreInst>(&I);
+      if (!LI && !SI)
         continue;
-      // FIXME: Currently we can't handle mixed accesses and predicated accesses
-      if (IsPred)
-        return;
-
-      AccessList.push_back(&I);
-    }
-  }
-
-  if (AccessList.empty())
-    return;
-
-  auto &DL = TheLoop->getHeader()->getModule()->getDataLayout();
-  for (auto I : AccessList) {
-    LoadInst *LI = dyn_cast<LoadInst>(I);
-    StoreInst *SI = dyn_cast<StoreInst>(I);
 
       Value *Ptr = LI ? LI->getPointerOperand() : SI->getPointerOperand();
       int Stride = getPtrStride(PSE, Ptr, TheLoop, Strides);
 
       const SCEV *Scev = replaceSymbolicStrideSCEV(PSE, Strides, Ptr);
       PointerType *PtrTy = dyn_cast<PointerType>(Ptr->getType());
       unsigned Size = DL.getTypeAllocSize(PtrTy->getElementType());
 
       // An alignment of 0 means target ABI alignment.
       unsigned Align = LI ? LI->getAlignment() : SI->getAlignment();
       if (!Align)
         Align = DL.getABITypeAlignment(PtrTy->getElementType());
 
-    StrideAccesses[I] = StrideDescriptor(Stride, Scev, Size, Align);
+      StrideAccesses[&I] = StrideDescriptor(Stride, Scev, Size, Align);
     }
 }
 
 // Analyze interleaved accesses and collect them into interleaved load and
 // store groups.
 //
 // When generating code for an interleaved load group, we effectively hoist all
 // loads in the group to the location of the first load in program order. When
 // generating code for an interleaved store group, we sink all stores to the

 // reverse postorder when collecting the accesses.
 void InterleavedAccessInfo::analyzeInterleaving(
     const ValueToValueMap &Strides) {
   DEBUG(dbgs() << "LV: Analyzing interleaved accesses...\n");
 
   // Holds all the stride accesses.
   MapVector<Instruction *, StrideDescriptor> StrideAccesses;
   collectConstStridedAccesses(StrideAccesses, Strides);
 
   if (StrideAccesses.empty())
     return;
 
   // Collect the dependences in the loop.
   collectDependences();
 
   // Holds all interleaved store groups temporarily.
   SmallSetVector<InterleaveGroup *, 4> StoreGroups;
for (auto II = std::next(I); II != E; ++II) {
 
       int DistanceToA = DistToA->getAPInt().getSExtValue();
 
       // Skip if the distance is not multiple of size as they are not in the
       // same group.
       if (DistanceToA % static_cast<int>(DesA.Size))
         continue;
 
+      // If instruction B is in a predicated block, we can't form an
+      // interleaved group with instruction A unless A and B are in the same
+      // block.
+      if (LoopAccessInfo::blockNeedsPredication(B->getParent(), TheLoop, DT) &&
+          A->getParent() != B->getParent())
+        continue;
+
       // The index of B is the index of A plus the related index to A.
       int IndexB =
           Group->getIndex(A) + DistanceToA / static_cast<int>(DesA.Size);
 
       // Try to insert B into the group.
       if (Group->insertMember(B, IndexB, DesB.Align)) {
         DEBUG(dbgs() << "LV: Inserted:" << *B << '\n'
                      << "    into the interleave group with" << *A << '\n');

test/Transforms/LoopVectorize/interleaved-accesses-pred-stores.ll

+; RUN: opt -S -loop-vectorize -instcombine -force-vector-width=2 -force-vector-interleave=1 -enable-interleaved-mem-accesses -vectorize-num-stores-pred=1 -enable-cond-stores-vec < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
+
+%pair = type { i64, i64 }
+
+; CHECK-LABEL: @load_gap_with_pred_store_0
+;
+; CHECK: min.iters.checked
+; CHECK:   %n.mod.vf = and i64 %[[N:.+]], 1
+; CHECK:   %[[IsZero:[a-zA-Z0-9]+]] = icmp eq i64 %n.mod.vf, 0
+; CHECK:   %[[R:.+]] = select i1 %[[IsZero]], i64 2, i64 %n.mod.vf
+; CHECK:   %n.vec = sub i64 %[[N]], %[[R]]
+;
+; CHECK: vector.body:
+; CHECK:   %wide.vec = load <4 x i64>, <4 x i64>* %{{.*}}
+; CHECK:   %strided.vec = shufflevector <4 x i64> %wide.vec, <4 x i64> undef, <2 x i32> <i32 0, i32 2>
+; CHECK:   %[[Result:.+]] = and <2 x i64> %strided.vec, %{{.*}}
+;
+; CHECK: pred.store.if
+; CHECK:   %[[X1:.+]] = extractelement <2 x i64> %[[Result]], i32 0
+; CHECK:   store i64 %[[X1]], {{.*}}
+;
+; CHECK: pred.store.if
+; CHECK:   %[[X2:.+]] = extractelement <2 x i64> %[[Result]], i32 1
+; CHECK:   store i64 %[[X2]], {{.*}}
+
+define void @load_gap_with_pred_store_0(%pair *%p, i64 %x, i64 %n) {
+entry:
+  br label %for.body
+
+for.body:
+  %i  = phi i64 [ %i.next, %if.merge ], [ 0, %entry ]
+  %p.1 = getelementptr inbounds %pair, %pair* %p, i64 %i, i32 1
+  %0 = load i64, i64* %p.1, align 8
+  %1 = and i64 %0, %x
+  %2 = icmp eq i64 %1, %x
+  br i1 %2, label %if.then, label %if.merge
+
+if.then:
+  store i64 %1, i64* %p.1, align 8
+  br label %if.merge
+
+if.merge:
+  %i.next = add nuw nsw i64 %i, 1
+  %cond = icmp slt i64 %i.next, %n
+  br i1 %cond, label %for.body, label %for.end
+
+for.end:
+  ret void
+}
+
+; CHECK-LABEL: @load_gap_with_pred_store_1
+;
+; CHECK: min.iters.checked
+; CHECK:   %n.mod.vf = and i64 %[[N:.+]], 1
+; CHECK:   %[[IsZero:[a-zA-Z0-9]+]] = icmp eq i64 %n.mod.vf, 0
+; CHECK:   %[[R:.+]] = select i1 %[[IsZero]], i64 2, i64 %n.mod.vf
+; CHECK:   %n.vec = sub i64 %[[N]], %[[R]]
+;
+; CHECK: vector.body:
+; CHECK:   %[[L1:.+]] = load <4 x i64>, <4 x i64>* %{{.*}}
+; CHECK:   %strided.vec = shufflevector <4 x i64> %[[L1]], <4 x i64> undef, <2 x i32> <i32 0, i32 2>
+; CHECK:   %[[Result:.+]] = and <2 x i64> %strided.vec, %{{.*}}
+;
+; CHECK: pred.store.if
+; CHECK:   %[[X1:.+]] = extractelement <2 x i64> %[[Result]], i32 0
+; CHECK:   store i64 %[[X1]], {{.*}}
+;
+; CHECK: pred.store.if
+; CHECK:   %[[X2:.+]] = extractelement <2 x i64> %[[Result]], i32 1
+; CHECK:   store i64 %[[X2]], {{.*}}
+;
+; CHECK: pred.store.continue
+; CHECK:   %[[L2:.+]] = load <4 x i64>, <4 x i64>* {{.*}}
+; CHECK:   %[[X3:.+]] = extractelement <4 x i64> %[[L2]], i32 0
+; CHECK:   store i64 %[[X3]], {{.*}}
+; CHECK:   %[[X4:.+]] = extractelement <4 x i64> %[[L2]], i32 2
+; CHECK:   store i64 %[[X4]], {{.*}}
+
+define void @load_gap_with_pred_store_1(%pair *%p, i64 %x, i64 %n) {
+entry:
+  br label %for.body
+
+for.body:
+  %i  = phi i64 [ %i.next, %if.merge ], [ 0, %entry ]
+  %p.0 = getelementptr inbounds %pair, %pair* %p, i64 %i, i32 0
+  %p.1 = getelementptr inbounds %pair, %pair* %p, i64 %i, i32 1
+  %0 = load i64, i64* %p.1, align 8
+  %1 = and i64 %0, %x
+  %2 = icmp eq i64 %1, %x
+  br i1 %2, label %if.then, label %if.merge
+
+if.then:
+  store i64 %1, i64* %p.0, align 8
+  br label %if.merge
+
+if.merge:
+  %3 = load i64, i64* %p.0, align 8
+  store i64 %3, i64 *%p.1, align 8
+  %i.next = add nuw nsw i64 %i, 1
+  %cond = icmp slt i64 %i.next, %n
+  br i1 %cond, label %for.body, label %for.end
+
+for.end:
+  ret void
+}