[LoadStoreVectorizer] Enable vectorization of stores in the presence of an aliasing load

[asbirlea]

Summary

The "getVectorizablePrefix" method would give up if it found an aliasing load for a store chain.
In practice, the aliasing load can be treated as a memory barrier and all stores that precede it
are a valid vectorizable prefix.
Issue found by volkan in D26962. Testcase is a pruned version of the one in the original patch.

[jlebar]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

486 ↗

(On Diff #78967)

Nit, pls end sentences with periods.

490 ↗

(On Diff #78967)

Can we add a comment on this loop to the effect of:

Find the first instruction (in BB order) which "conflicts with" any of ChainInstrs[0..ChainInstrIdx].

529 ↗

(On Diff #78967)

As discussed IRL, not sure why we need this early-break at all.

[asbirlea]

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

529 ↗

(On Diff #78967)

This a legality condition. It relates to the above comment:
"We can ignore the alias as long as the load comes before the store, because that means we won't be moving the load past the store [...]"

We consider loads before stores safe. This patch will somewhat relax this condition, but only for stores. First, store chains:
If a load comes before the store it's safe per above comment/condition. If a load comes after the store, we allow some leniency and don't break right away to enable vectorization of all stores that precede the "guilty" load.
Now, for load chains:
If the load comes before, again it's safe. But if we find a store that aliases, we _must_ stop processing and not take any loads that follow that store. The vectorized load will be moved at the location of the first load found, pulling up a load past an aliasing store. This produces incorrect code if the load it pulls up aliases the store(s) above it, a check that we don't make right now.
There may be room for further improvement with additional checks, but for now, don't do this.

[asbirlea]

Add more comments.

[jlebar]

As we discussed, I want to try splitting out this loop into a load and a store version to see if it clarifies things, but I'm happy to do that as a separate patch. This lgtm.

lib/Transforms/Vectorize/LoadStoreVectorizer.cpp

535 ↗

(On Diff #79003)

I'd either get rid of the second paragraph here or encode it as an assert.

[asbirlea]

Added assert and comment to justify it.

[asbirlea]

Period.

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

if (isa<LoadInst>(I) || isa<StoreInst>(I)) {
       break;
     }
   }
 
   OrderedBasicBlock OBB(Chain[0]->getParent());
 
   // Loop until we find an instruction in ChainInstrs that we can't vectorize.
   unsigned ChainInstrIdx = 0;
+  Instruction *BarrierMemoryInstr = nullptr;
+
   for (unsigned E = ChainInstrs.size(); ChainInstrIdx < E; ++ChainInstrIdx) {
     Instruction *ChainInstr = ChainInstrs[ChainInstrIdx];
-    bool AliasFound = false;
+
+    // If a barrier memory instruction was found, chain instructions that follow
+    // will not be added to the valid prefix.
+    if (BarrierMemoryInstr && OBB.dominates(BarrierMemoryInstr, ChainInstr))
+      break;
+
+    // Check (in BB order) if any instruction prevents ChainInstr from being
+    // vectorized. Find and store the first such "conflicting" instruction.
     for (Instruction *MemInstr : MemoryInstrs) {
+      // If a barrier memory instruction was found, do not check past it.
+      if (BarrierMemoryInstr && OBB.dominates(BarrierMemoryInstr, MemInstr))
+        break;
+
       if (isa<LoadInst>(MemInstr) && isa<LoadInst>(ChainInstr))
         continue;
 
       // We can ignore the alias as long as the load comes before the store,
       // because that means we won't be moving the load past the store to
       // vectorize it (the vectorized load is inserted at the location of the
       // first load in the chain).
       if (isa<StoreInst>(MemInstr) && isa<LoadInst>(ChainInstr) &&
for (Instruction *MemInstr : MemoryInstrs) {
           dbgs() << "LSV: Found alias:\n"
                     "  Aliasing instruction and pointer:\n"
                  << "  " << *MemInstr << '\n'
                  << "  " << *getPointerOperand(MemInstr) << '\n'
                  << "  Aliased instruction and pointer:\n"
                  << "  " << *ChainInstr << '\n'
                  << "  " << *getPointerOperand(ChainInstr) << '\n';
         });
-        AliasFound = true;
+        // Save this aliasing memory instruction as a barrier, but allow other
+        // instructions that precede the barrier to be vectorized with this one.
+        BarrierMemoryInstr = MemInstr;
         break;
       }
     }
-    if (AliasFound)
+    // Continue the search only for store chains, since vectorizing stores that
+    // precede an aliasing load is valid. Conversely, vectorizing loads is valid
+    // up to an aliasing store, but should not pull loads from further down in
+    // the basic block.
+    if (IsLoadChain && BarrierMemoryInstr) {
+      // The BarrierMemoryInstr is a store that precedes ChainInstr.
+      assert(OBB.dominates(BarrierMemoryInstr, ChainInstr));
       break;
     }
+  }
 
   // Find the largest prefix of Chain whose elements are all in
   // ChainInstrs[0, ChainInstrIdx).  This is the largest vectorizable prefix of
   // Chain.  (Recall that Chain is in address order, but ChainInstrs is in BB
   // order.)
   SmallPtrSet<Instruction *, 8> VectorizableChainInstrs(
       ChainInstrs.begin(), ChainInstrs.begin() + ChainInstrIdx);
   unsigned ChainIdx = 0;

llvm/trunk/test/Transforms/LoadStoreVectorizer/AMDGPU/insertion-point.ll

entry:
 
   %add = fadd float %ld.c, %ld.c.idx.1
   store float %add, float addrspace(1)* %b, align 4
   store i32 %foo, i32 addrspace(3)* null, align 4
   ret void
 }
 
 ; Here we have four stores, with an aliasing load before the last one.  We can
-; vectorize the first two stores as <2 x float>, but this vectorized store must
-; be inserted at the location of the second scalar store, not the fourth one.
+; vectorize the first three stores as <3 x float>, but this vectorized store must
+; be inserted at the location of the third scalar store, not the fourth one.
 ;
 ; CHECK-LABEL: @insert_store_point_alias
-; CHECK: store <2 x float>
-; CHECK: store float
-; CHECK-SAME: %a.idx.2
+; CHECK: store <3 x float>
 ; CHECK: load float, float addrspace(1)* %a.idx.2
 ; CHECK: store float
 ; CHECK-SAME: %a.idx.3
 define float @insert_store_point_alias(float addrspace(1)* nocapture %a, i64 %idx) {
   %a.idx = getelementptr inbounds float, float addrspace(1)* %a, i64 %idx
   %a.idx.1 = getelementptr inbounds float, float addrspace(1)* %a.idx, i64 1
   %a.idx.2 = getelementptr inbounds float, float addrspace(1)* %a.idx.1, i64 1
   %a.idx.3 = getelementptr inbounds float, float addrspace(1)* %a.idx.2, i64 1

llvm/trunk/test/Transforms/LoadStoreVectorizer/AMDGPU/store_with_aliasing_load.ll

+; RUN: opt -mtriple=amdgcn-amd-amdhsa -load-store-vectorizer -S -o - %s | FileCheck %s
+
+; Check that, in the presence of an aliasing load, the stores preceding the
+; aliasing load are safe to vectorize.
+
+; CHECK-LABEL: store_vectorize_with_alias
+; CHECK: store <4 x float>
+; CHECK: load <4 x float>
+; CHECK: store <4 x float>
+
+; Function Attrs: nounwind
+define void @store_vectorize_with_alias(i8 addrspace(1)* %a, i8 addrspace(1)* %b) #0 {
+bb:
+  %tmp = bitcast i8 addrspace(1)* %b to float addrspace(1)*
+  %tmp1 = load float, float addrspace(1)* %tmp, align 4
+
+  %tmp2 = bitcast i8 addrspace(1)* %a to float addrspace(1)*
+  store float %tmp1, float addrspace(1)* %tmp2, align 4
+  %tmp3 = getelementptr i8, i8 addrspace(1)* %a, i64 4
+  %tmp4 = bitcast i8 addrspace(1)* %tmp3 to float addrspace(1)*
+  store float %tmp1, float addrspace(1)* %tmp4, align 4
+  %tmp5 = getelementptr i8, i8 addrspace(1)* %a, i64 8
+  %tmp6 = bitcast i8 addrspace(1)* %tmp5 to float addrspace(1)*
+  store float %tmp1, float addrspace(1)* %tmp6, align 4
+  %tmp7 = getelementptr i8, i8 addrspace(1)* %a, i64 12
+  %tmp8 = bitcast i8 addrspace(1)* %tmp7 to float addrspace(1)*
+  store float %tmp1, float addrspace(1)* %tmp8, align 4
+
+  %tmp9 = getelementptr i8, i8 addrspace(1)* %b, i64 16
+  %tmp10 = bitcast i8 addrspace(1)* %tmp9 to float addrspace(1)*
+  %tmp11 = load float, float addrspace(1)* %tmp10, align 4
+  %tmp12 = getelementptr i8, i8 addrspace(1)* %b, i64 20
+  %tmp13 = bitcast i8 addrspace(1)* %tmp12 to float addrspace(1)*
+  %tmp14 = load float, float addrspace(1)* %tmp13, align 4
+  %tmp15 = getelementptr i8, i8 addrspace(1)* %b, i64 24
+  %tmp16 = bitcast i8 addrspace(1)* %tmp15 to float addrspace(1)*
+  %tmp17 = load float, float addrspace(1)* %tmp16, align 4
+  %tmp18 = getelementptr i8, i8 addrspace(1)* %b, i64 28
+  %tmp19 = bitcast i8 addrspace(1)* %tmp18 to float addrspace(1)*
+  %tmp20 = load float, float addrspace(1)* %tmp19, align 4
+
+  %tmp21 = getelementptr i8, i8 addrspace(1)* %a, i64 16
+  %tmp22 = bitcast i8 addrspace(1)* %tmp21 to float addrspace(1)*
+  store float %tmp11, float addrspace(1)* %tmp22, align 4
+  %tmp23 = getelementptr i8, i8 addrspace(1)* %a, i64 20
+  %tmp24 = bitcast i8 addrspace(1)* %tmp23 to float addrspace(1)*
+  store float %tmp14, float addrspace(1)* %tmp24, align 4
+  %tmp25 = getelementptr i8, i8 addrspace(1)* %a, i64 24
+  %tmp26 = bitcast i8 addrspace(1)* %tmp25 to float addrspace(1)*
+  store float %tmp17, float addrspace(1)* %tmp26, align 4
+  %tmp27 = getelementptr i8, i8 addrspace(1)* %a, i64 28
+  %tmp28 = bitcast i8 addrspace(1)* %tmp27 to float addrspace(1)*
+  store float %tmp20, float addrspace(1)* %tmp28, align 4
+
+  ret void
+}
+
+attributes #0 = { argmemonly nounwind }