[VectorCombine] Fix alignment in single element store

[qiucf]

Summary

This concern was raised in D98240. It's a miscompile and thanks for comments from @lebedev.ri.

[tschuett]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

834

This line is too cute for me, but ...

[lebedev.ri]

Hm, this first needs more radical fixes - this is currently miscompiling vector indexes: https://alive2.llvm.org/ce/z/aWtH9w
I would suggest to first simply require the index to be constant.

[lebedev.ri]

Right, canScalarizeAccess() already does that, okay.
But then, it would be good to have a positive test for @insert_store_nonconst() :)

[lebedev.ri]

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

835

(alive *seems* to be happy with this)
This is going to be precise and optimal for constant indexes,
but i think we can get at least the lower-bound estimate for variable indexes:
the new address will be offset from the base address by DL.getTypeStoreSize(NewElement->getType()),
so i think we can do

else
  NewAlignment = commonAlignment(
          NewAlignment,
          DL.getTypeStoreSize(NewElement->getType()));

Please add positive tests:

; New alignment should be 8
define void @src(<8 x i64>* %q, i64 %s, i32 %idx) {
  %cmp = icmp ult i32 %idx, 2
  call void @llvm.assume(i1 %cmp)

  %i = load <8 x i64>, <8 x i64>* %q, align 8
  %vecins = insertelement <8 x i64> %i, i64 %s, i32 %idx
  store <8 x i64> %vecins, <8 x i64>* %q, align 8
  ret void
}

; New alignment should be 4
define void @src(<8 x i64>* %q, i64 %s, i32 %idx) {
  %cmp = icmp ult i32 %idx, 2
  call void @llvm.assume(i1 %cmp)

  %i = load <8 x i64>, <8 x i64>* %q, align 4
  %vecins = insertelement <8 x i64> %i, i64 %s, i32 %idx
  store <8 x i64> %vecins, <8 x i64>* %q, align 4
  ret void
}

[qiucf]

Add another test (first..)

[lebedev.ri]

Please feel free to just directly commit new tests.
The new tests i asked for should be positive tests - they should be getting transformed (missing @llvm.assume())

[lebedev.ri]

Thanks.
This looks fine to me now.
Can anyone spot any issues with the new alignment logic? @fhahn @spatel?

llvm/test/Transforms/VectorCombine/load-insert-store.ll

142

I think we still want those two tests i suggested,
they demonstrate that we don't increase alignment from the maximal one allowed.

Please precommit the tests.

[spatel]

llvm/test/Transforms/VectorCombine/load-insert-store.ll

23

How do we justify this increase in alignment?
The original code had minimal align 1, so it could be anything. We are creating a scalar store at an address 6 bytes over that, so it could still be anything?

[lebedev.ri]

llvm/test/Transforms/VectorCombine/load-insert-store.ll

23

This change is correct.
Before store <...>, align 1, we have already established that the %q is more aligned,
as per the load <...> with an implicit alignment, which isn't 1.
https://alive2.llvm.org/ce/z/C2qnUc

[spatel]

llvm/test/Transforms/VectorCombine/load-insert-store.ll

23

Ah, thanks for explaining. IIUC, we add explicit alignment to all load/store in IR now, so we should add the align 16 to this test to avoid confusion - and a test comment would be nice too :).

142

+1 - additional tests and pre-commit will make this easier to understand.

[qiucf]

Add explicit align to affected test.
Add comment for implicit alignment.

[spatel]

llvm/test/Transforms/VectorCombine/load-insert-store.ll

142

Let me know if I'm not seeing it, but we want 1 test with nonconst index where the original alignment is less than the presumed alignment for the new scalar store:

define void @src(<8 x i64>* %q, i64 %s, i32 %idx) {
  %cmp = icmp ult i32 %idx, 2
  call void @llvm.assume(i1 %cmp)
  %i = load <8 x i64>, <8 x i64>* %q, align 4
  %vecins = insertelement <8 x i64> %i, i64 %s, i32 %idx
  store <8 x i64> %vecins, <8 x i64>* %q, align 2 ; make this different just to exercise the logic a bit more
  ret void
}

[lebedev.ri]

(better, but still not quite there, my previous comments still stand unaddressed...)

[lebedev.ri]

LGTM unless there are other comments.
Thanks.

[spatel]

LGTM

[qiucf]

Landed. Thanks for the review!

llvm/lib/Transforms/Vectorize/VectorCombine.cpp

if (!Load->isSimple() || Load->getParent() != SI->getParent() ||
                              MemoryLocation::get(SI), AA))
       return false;
 
     Value *GEP = GetElementPtrInst::CreateInBounds(
         SI->getPointerOperand(), {ConstantInt::get(Idx->getType(), 0), Idx});
     Builder.Insert(GEP);
     StoreInst *NSI = Builder.CreateStore(NewElement, GEP);
     NSI->copyMetadata(*SI);
-    if (SI->getAlign() < NSI->getAlign())
-      NSI->setAlignment(SI->getAlign());
+    Align NewAlignment = std::max(SI->getAlign(), Load->getAlign());
+    if (auto *C = dyn_cast<ConstantInt>(Idx))
+      NewAlignment = commonAlignment(
+          NewAlignment,
+          C->getZExtValue() * DL.getTypeStoreSize(NewElement->getType()));
+    else
+      NewAlignment = commonAlignment(
+          NewAlignment, DL.getTypeStoreSize(NewElement->getType()));
+    NSI->setAlignment(NewAlignment);
     replaceValue(I, *NSI);
     // Need erasing the store manually.
     I.eraseFromParent();
     return true;
   }
 
   return false;
 }

llvm/test/Transforms/VectorCombine/load-insert-store.ll

entry:
   store <16 x i8> %vecins, <16 x i8>* %q, align 16
   ret void
 }
 
 define void @insert_store_i16_align1(<8 x i16>* %q, i16 zeroext %s) {
 ; CHECK-LABEL: @insert_store_i16_align1(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <8 x i16>, <8 x i16>* [[Q:%.*]], i32 0, i32 3
-; CHECK-NEXT:    store i16 [[S:%.*]], i16* [[TMP0]], align 1
+; CHECK-NEXT:    store i16 [[S:%.*]], i16* [[TMP0]], align 2
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %0 = load <8 x i16>, <8 x i16>* %q
   %vecins = insertelement <8 x i16> %0, i16 %s, i32 3
   store <8 x i16> %vecins, <8 x i16>* %q, align 1
   ret void
 }
 
+define void @insert_store_zero_offset(<8 x i16>* %q, i16 zeroext %s) {
+; CHECK-LABEL: @insert_store_zero_offset(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <8 x i16>, <8 x i16>* [[Q:%.*]], i32 0, i32 0
+; CHECK-NEXT:    store i16 [[S:%.*]], i16* [[TMP0]], align 128
+; CHECK-NEXT:    ret void
+;
+entry:
+  %0 = load <8 x i16>, <8 x i16>* %q, align 128
+  %vecins = insertelement <8 x i16> %0, i16 %s, i32 0
+  store <8 x i16> %vecins, <8 x i16>* %q, align 128
+  ret void
+}
+
 ; To verify case when index is out of bounds
 define void @insert_store_outofbounds(<8 x i16>* %q, i16 zeroext %s) {
 ; CHECK-LABEL: @insert_store_outofbounds(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load <8 x i16>, <8 x i16>* [[Q:%.*]], align 16
 ; CHECK-NEXT:    [[VECINS:%.*]] = insertelement <8 x i16> [[TMP0]], i16 [[S:%.*]], i32 9
 ; CHECK-NEXT:    store <8 x i16> [[VECINS]], <8 x i16>* [[Q]], align 16
 ; CHECK-NEXT:    ret void

 ;
 entry:
   %0 = load <16 x i8>, <16 x i8>* %q
   %vecins = insertelement <16 x i8> %0, i8 %s, i32 %idx
   store <16 x i8> %vecins, <16 x i8>* %q
   ret void
 }
 
+define void @insert_store_nonconst_large_alignment(<4 x i32>* %q, i32 zeroext %s, i32 %idx) {
+; CHECK-LABEL: @insert_store_nonconst_large_alignment(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[IDX:%.*]], 4
+; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <4 x i32>, <4 x i32>* [[Q:%.*]], i32 0, i32 [[IDX]]
+; CHECK-NEXT:    store i32 [[S:%.*]], i32* [[TMP0]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp = icmp ult i32 %idx, 4
+  call void @llvm.assume(i1 %cmp)
+  %i = load <4 x i32>, <4 x i32>* %q, align 128
+  %vecins = insertelement <4 x i32> %i, i32 %s, i32 %idx
+  store <4 x i32> %vecins, <4 x i32>* %q, align 128
+  ret void
+}
+
 define void @insert_store_nonconst_index_known_valid_by_assume(<16 x i8>* %q, i8 zeroext %s, i32 %idx) {
 ; CHECK-LABEL: @insert_store_nonconst_index_known_valid_by_assume(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[IDX:%.*]], 4
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[CMP]])
 ; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds <16 x i8>, <16 x i8>* [[Q:%.*]], i32 0, i32 [[IDX]]
 ; CHECK-NEXT:    store i8 [[S:%.*]], i8* [[TMP0]], align 1
 ; CHECK-NEXT:    ret void