[SROA] Check typeSizeEqualsStoreSize in isVectorPromotionViable

[bjope]

Summary

Commit de3445e0ef15c4209 (https://reviews.llvm.org/D132096) made
changes to isVectorPromotionViable basically doing

// Create Vector with size of V, and each element of type Ty
...
uint64_t ElementSize = DL.getTypeStoreSizeInBits(Ty).getFixedSize();
uint64_t VectorSize = DL.getTypeSizeInBits(V).getFixedSize();
...
VectorType *VTy = VectorType::get(Ty, VectorSize / ElementSize, false);

Not quite sure why it uses the TypeStoreSize for the ElementSize,
but the new vector would only match in size with the old vector in
situations when the TypeStoreSize equals the TypeSize for Ty.
Therefore this patch adds a typeSizeEqualsStoreSize check as yet
another condition for allowing the the new type as a promotion
candidate.

Without this fix the new @test15 test would fail with an assert
like this:

opt: ../lib/Transforms/Scalar/SROA.cpp:1966:

auto isVectorPromotionViable(llvm::sroa::Partition &,
                             const llvm::DataLayout &)
     ::(anonymous class)::operator()(llvm::VectorType *,
                                     llvm::VectorType *) const:
Assertion `DL.getTypeSizeInBits(RHSTy).getFixedSize() ==
           DL.getTypeSizeInBits(LHSTy).getFixedSize() &&
           "Cannot have vector types of different sizes!"' failed.

...
#8 isVectorPromotionViable(...)::$_10::operator()...
#9 llvm::SROAPass::rewritePartition(...)
#10 llvm::SROAPass::splitAlloca(...)
#11 llvm::SROAPass::runOnAlloca(...)
#12 llvm::SROAPass::runImpl(...)
#13 llvm::SROAPass::run(...)

[bjope]

Adding missing CHECKS to the new test case.

[A-Wadhwani]

Thanks for catching this. I would like someone with more experience in SROA than me to take a look, in case there's something I've missed, though.

[bjope]

llvm/lib/Transforms/Scalar/SROA.cpp

1936

I should probably add some explanatory code comment here (given that this solution is accepted).

Background to why I used the typeSizeEqualsStoreSize check:

The idea is to make a similar check to what is done at line 1859 ("While the definition of LLVM vectors is bitpacked, we don't support sizes that aren't byte sized."), although that isn't exactly what testing for typeSizeEqualsStoreSize is doing.

Not sure if it would be better to check for "byte sized". It might work as well (but it would complicate stuff for our out-of-target with 16 bit bytes if hard-coding another "8" here).

Anyhow, for the "VectorSize / ElementSize" calculation below to result in a vector with the same size as the original vector I think typeSizeEqualsStoreSize needs to be fulfilled (or we could change the ElementSize calculation to just use getTypeSizeInBits instead of getTypeStoreSizeInBits, but then I suspect that the transform might be unsafe in case typeSizeEqualsStoreSize isn't fulfilled). So checking for typeSizeEqualsStoreSize seemed like a defensive approach here to avoid the corner cases without sacrificing the normal cases that we want to handle (and that is regression tested).

[MatzeB]

Seems sensible. Thanks for catching and fixing. LGTM

llvm/lib/Transforms/Scalar/SROA.cpp

for (const Slice &S : P) {
     if (LoadInst *LI = dyn_cast<LoadInst>(S.getUse()->getUser()))
       Ty = LI->getType();
     else if (StoreInst *SI = dyn_cast<StoreInst>(S.getUse()->getUser()))
       Ty = SI->getValueOperand()->getType();
     else
       continue;
     if (isa<VectorType>(Ty))
       continue;
+    if (!DL.typeSizeEqualsStoreSize(Ty))
+      continue;
     // Create Vector with size of V, and each element of type Ty
     VectorType *V = CandidateTys[0];
     uint64_t ElementSize = DL.getTypeStoreSizeInBits(Ty).getFixedSize();
     uint64_t VectorSize = DL.getTypeSizeInBits(V).getFixedSize();
     if ((ElementSize != VectorSize) && (VectorSize % ElementSize == 0)) {
       VectorType *VTy = VectorType::get(Ty, VectorSize / ElementSize, false);
       CandidateTys.push_back(VTy);
       if (CommonEltTy != Ty)

llvm/test/Transforms/SROA/vector-promotion.ll

entry:
   %x.tmp4 = getelementptr inbounds i32, ptr %x.cast, i64 3
   %d = load i32, ptr %x.tmp4
 
   %add = add i32 %a, %b
   %add1 = add i32 %c, %d
   %add2 = add i32 %add, %add1
   ret i32 %add2
 }
+
+; This used to hit an assert after commit de3445e0ef15c4.
+; Added as regression test to verify that we handle this without crashing.
+define i1 @test15() {
+  %a = alloca <8 x i32>
+  store <2 x i64> <i64 0, i64 -1>, ptr %a
+  %l = load i1, ptr %a, align 1
+  ret i1 %l
+
+}