[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

4990 ↗

(On Diff #55823)

Does the Comment need updating?

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5056 ↗

(On Diff #55823)

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

4990 ↗

(On Diff #55823)

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.

5056 ↗

(On Diff #55823)

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

4990 ↗

(On Diff #55823)

Thanks for the clarification, Matt.

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5056 ↗

(On Diff #55823)

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

5056 ↗

(On Diff #55823)

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

5056 ↗

(On Diff #55823)

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

5056 ↗

(On Diff #55823)

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

36 ↗

(On Diff #63252)

Is %p.0 unused here?

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5438 ↗

(On Diff #63252)

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

5438 ↗

(On Diff #63252)

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

36 ↗

(On Diff #63252)

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

[sbaranga]

lib/Transforms/Vectorize/LoopVectorize.cpp

5438 ↗

(On Diff #63252)

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

5438 ↗

(On Diff #63252)

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

[mssimpso]

Addressed comments from Adam and Silviu. Thanks!

[sbaranga]

LGTM!

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

void collectConstStridedAccesses(
       const ValueToValueMap &Strides);
 
   /// \brief Returns true if \p Stride is allowed in an interleaved group.
   static bool isStrided(int Stride) {
     unsigned Factor = std::abs(Stride);
     return Factor >= 2 && Factor <= MaxInterleaveGroupFactor;
   }
 
+  /// \brief Returns true if \p BB is a predicated block.
+  bool isPredicated(BasicBlock *BB) const {
+    return LoopAccessInfo::blockNeedsPredication(BB, TheLoop, DT);
+  }
+
   /// \brief Returns true if LoopAccessInfo can be used for dependence queries.
   bool areDependencesValid() const {
     return LAI && LAI->getDepChecker().getDependences();
   }
 
   /// \brief Returns true if memory accesses \p B and \p A can be reordered, if
   /// necessary, when constructing interleaved groups.
   ///
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;
+
+  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 (BasicBlock *BB : make_range(DFS.beginRPO(), DFS.endRPO())) {
-    bool IsPred = LoopAccessInfo::blockNeedsPredication(BB, TheLoop, DT);
-
+  for (BasicBlock *BB : make_range(DFS.beginRPO(), DFS.endRPO()))
     for (auto &I : *BB) {
-      if (!isa<LoadInst>(&I) && !isa<StoreInst>(&I))
+      auto *LI = dyn_cast<LoadInst>(&I);
+      auto *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) {
-    auto *LI = dyn_cast<LoadInst>(I);
-    auto *SI = dyn_cast<StoreInst>(I);
 
       Value *Ptr = LI ? LI->getPointerOperand() : SI->getPointerOperand();
       int64_t Stride = getPtrStride(PSE, Ptr, TheLoop, Strides);
 
       const SCEV *Scev = replaceSymbolicStrideSCEV(PSE, Strides, Ptr);
       PointerType *PtrTy = dyn_cast<PointerType>(Ptr->getType());
       uint64_t 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
for (auto BI = std::next(AI); BI != E; ++BI) {
 
       int64_t 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<int64_t>(DesA.Size))
         continue;
 
+      // If either A or B is in a predicated block, we prevent adding them to a
+      // group. We may be able to relax this limitation in the future once we
+      // handle more complicated blocks.
+      if (isPredicated(A->getParent()) || isPredicated(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<int64_t>(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');

llvm/trunk/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 }
+
+; Ensure that we vectorize the interleaved load group even though the loop
+; contains a conditional store. The store group contains gaps and is not
+; vectorized.
+;
+; CHECK-LABEL: @interleaved_with_cond_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: pred.store.if
+; CHECK:   %[[X1:.+]] = extractelement <4 x i64> %wide.vec, i32 0
+; CHECK:   store i64 %[[X1]], {{.*}}
+;
+; CHECK: pred.store.if
+; CHECK:   %[[X2:.+]] = extractelement <4 x i64> %wide.vec, i32 2
+; CHECK:   store i64 %[[X2]], {{.*}}
+
+define void @interleaved_with_cond_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 = icmp eq i64 %0, %x
+  br i1 %1, label %if.then, label %if.merge
+
+if.then:
+  store i64 %0, 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
+}
+
+; Ensure that we don't form a single interleaved group for the two loads. The
+; conditional store prevents the second load from being hoisted. The two load
+; groups are separately vectorized. The store group contains gaps and is not
+; vectorized.
+;
+; CHECK-LABEL: @interleaved_with_cond_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: pred.store.if
+; CHECK:   %[[X1:.+]] = extractelement <4 x i64> %wide.vec, i32 0
+; CHECK:   store i64 %[[X1]], {{.*}}
+;
+; CHECK: pred.store.if
+; CHECK:   %[[X2:.+]] = extractelement <4 x i64> %wide.vec, i32 2
+; 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 @interleaved_with_cond_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 = icmp eq i64 %0, %x
+  br i1 %1, label %if.then, label %if.merge
+
+if.then:
+  store i64 %0, i64* %p.0, align 8
+  br label %if.merge
+
+if.merge:
+  %2 = load i64, i64* %p.0, align 8
+  store i64 %2, 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
+}
+
+; Ensure that we don't create a single interleaved group for the two stores.
+; The second store is conditional and we can't sink the first store inside the
+; predicated block. The load group is vectorized, and the store groups contain
+; gaps and are not vectorized.
+;
+; CHECK-LABEL: @interleaved_with_cond_store_2(
+;
+; 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:   store i64 %x, {{.*}}
+; CHECK:   store i64 %x, {{.*}}
+;
+; CHECK: pred.store.if
+; CHECK:   %[[X1:.+]] = extractelement <4 x i64> %wide.vec, i32 0
+; CHECK:   store i64 %[[X1]], {{.*}}
+;
+; CHECK: pred.store.if
+; CHECK:   %[[X2:.+]] = extractelement <4 x i64> %wide.vec, i32 2
+; CHECK:   store i64 %[[X2]], {{.*}}
+
+define void @interleaved_with_cond_store_2(%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
+  store i64 %x, i64* %p.0, align 8
+  %1 = icmp eq i64 %0, %x
+  br i1 %1, label %if.then, label %if.merge
+
+if.then:
+  store i64 %0, 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
+}