[llvm] d7cf7ec - [SROA] Handle over-large loads during presplitting

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2022-03-16T15:41:11+01:00
New Revision: d7cf7ec05d3570998d04b9c257362838026ae75b

URL: https://github.com/llvm/llvm-project/commit/d7cf7ec05d3570998d04b9c257362838026ae75b
DIFF: https://github.com/llvm/llvm-project/commit/d7cf7ec05d3570998d04b9c257362838026ae75b.diff

LOG: [SROA] Handle over-large loads during presplitting

When a load extends past the extent of the alloca, SROA will
restrict the slice size to extend to the end of the alloca only.
However, presplitting was asserting that the load size and the
slice size match exactly, which does not hold in this case.
Relax the assertion to only require that the load size is greater
or equal than the slice size.

Added: 
    

Modified: 
    llvm/lib/Transforms/Scalar/SROA.cpp
    llvm/test/Transforms/SROA/slice-width.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index eb56fea1a20bf..1f68c3734f649 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -3965,16 +3965,15 @@ bool SROAPass::presplitLoadsAndStores(AllocaInst &AI, AllocaSlices &AS) {
   for (LoadInst *LI : Loads) {
     SplitLoads.clear();
 
-    IntegerType *Ty = cast<IntegerType>(LI->getType());
-    assert(Ty->getBitWidth() % 8 == 0);
-    uint64_t LoadSize = Ty->getBitWidth() / 8;
-    assert(LoadSize > 0 && "Cannot have a zero-sized integer load!");
-
     auto &Offsets = SplitOffsetsMap[LI];
-    assert(LoadSize == Offsets.S->endOffset() - Offsets.S->beginOffset() &&
-           "Slice size should always match load size exactly!");
+    unsigned SliceSize = Offsets.S->endOffset() - Offsets.S->beginOffset();
+    assert(LI->getType()->getIntegerBitWidth() % 8 == 0 &&
+           "Load must have type size equal to store size");
+    assert(LI->getType()->getIntegerBitWidth() / 8 >= SliceSize &&
+           "Load must be >= slice size");
+
     uint64_t BaseOffset = Offsets.S->beginOffset();
-    assert(BaseOffset + LoadSize > BaseOffset &&
+    assert(BaseOffset + SliceSize > BaseOffset &&
            "Cannot represent alloca access size using 64-bit integers!");
 
     Instruction *BasePtr = cast<Instruction>(LI->getPointerOperand());
@@ -3985,7 +3984,7 @@ bool SROAPass::presplitLoadsAndStores(AllocaInst &AI, AllocaSlices &AS) {
     uint64_t PartOffset = 0, PartSize = Offsets.Splits.front();
     int Idx = 0, Size = Offsets.Splits.size();
     for (;;) {
-      auto *PartTy = Type::getIntNTy(Ty->getContext(), PartSize * 8);
+      auto *PartTy = Type::getIntNTy(LI->getContext(), PartSize * 8);
       auto AS = LI->getPointerAddressSpace();
       auto *PartPtrTy = PartTy->getPointerTo(AS);
       LoadInst *PLoad = IRB.CreateAlignedLoad(
@@ -4018,7 +4017,7 @@ bool SROAPass::presplitLoadsAndStores(AllocaInst &AI, AllocaSlices &AS) {
       // Setup the next partition.
       PartOffset = Offsets.Splits[Idx];
       ++Idx;
-      PartSize = (Idx < Size ? Offsets.Splits[Idx] : LoadSize) - PartOffset;
+      PartSize = (Idx < Size ? Offsets.Splits[Idx] : SliceSize) - PartOffset;
     }
 
     // Now that we have the split loads, do the slow walk over all uses of the

diff  --git a/llvm/test/Transforms/SROA/slice-width.ll b/llvm/test/Transforms/SROA/slice-width.ll
index beb26ea84011b..bcdac4b1b4d49 100644
--- a/llvm/test/Transforms/SROA/slice-width.ll
+++ b/llvm/test/Transforms/SROA/slice-width.ll
@@ -158,3 +158,19 @@ define void @PR50910() {
   call void @llvm.memset.p0i8.i64(i8* align 8 %t1, i8 0, i64 4294967296, i1 false)
   ret void
 }
+
+define i1 @presplit_overlarge_load() {
+; CHECK-LABEL: @presplit_overlarge_load(
+; CHECK-NEXT:    [[A_SROA_0:%.*]] = alloca i8, align 2
+; CHECK-NEXT:    [[A_SROA_0_0_A_SROA_0_0_L11:%.*]] = load i8, i8* [[A_SROA_0]], align 2
+; CHECK-NEXT:    [[A_SROA_0_0_A_1_SROA_CAST3:%.*]] = bitcast i8* [[A_SROA_0]] to i1*
+; CHECK-NEXT:    [[A_SROA_0_0_A_SROA_0_0_L2:%.*]] = load i1, i1* [[A_SROA_0_0_A_1_SROA_CAST3]], align 2
+; CHECK-NEXT:    ret i1 [[A_SROA_0_0_A_SROA_0_0_L2]]
+;
+  %A = alloca i16
+  %A.32 = bitcast i16* %A to i32*
+  %A.1 = bitcast i16* %A to i1*
+  %L1 = load i32, i32* %A.32
+  %L2 = load i1, i1* %A.1
+  ret i1 %L2
+}