[SROA] Avoid postponing rewriting load/store by ignoring lifetime intrinsics in partition's promotability checking

[slydiman]

Summary

Performance drift caused by generated ptx between slight different code calling flow.
This patch fixes a bug that generates unnecessary packing/unpacking structure code because of incorrectly handling lifetime intrinsic.

For example, a partition of an alloca may contain many slices:

Partition [0, 4):
  Slice0: [0, 4) used by: load i32 addr;
  Slice1: [0, 4) used by: store i32 v, addr;
  Slice2: [0, 16) used by lifetime.start(16, addr);

When SROA determines if the partition can be promoted, lifetime.start is currently treated as a whole alloca load/store, so Slice0 and Slice1 cannot be promoted at this attempt, but the packing/unpacking code for Slice0 and Slice1 has been generated. After rewrite lifetime.start/end intrinsic, SROA tries again with Slice0 and Slice1 and finally promotes them, but redundant packing/unpacking code remaining in the IRs.

This patch changes promotability checking to ignore lifetime intrinsic (they will be rewritten to correct sizes later), so we can promote the real users (load/store) at the first attempt with optimal code.

[tra]

llvm/test/Transforms/SROA/lifetime-intrinsic.ll

1

I'd add a comment describing what is it we're looking for here.
It's not clear from the test itself why we're checking for those phys. To make sure that we don't do excessive shifting and masking of i32->i16->i32 ?

26

These labels look overly specific. You should probably capture the actual names where the values are set.

[nikic]

llvm/test/Transforms/SROA/lifetime-intrinsic.ll

1

Generally, please reduce this test to the minimal instructions that show the problem.

[slydiman]

I have updated the test. Now it is generic, not NVPTX specific.

[nikic]

llvm/lib/Transforms/Scalar/SROA.cpp

1990

You can use isLifetimeStartOrEnd() here.

llvm/test/Transforms/SROA/lifetime-intrinsic.ll

1

Please use update_test_checks.py to generate test checks.

21

Is the loop required to demonstrate the problem?

[nikic]

If I take your example with and without lifetime intrinsics, then with current SROA, the variant without intrinsics produces a worse result: https://llvm.godbolt.org/z/oPMKWzaMq

Results of SROA should be the same with an without lifetime intrinsics, but I'm not sure your patch achieves that -- it looks like it produces a third variant that is better than both?

[slydiman]

In D124967#3496117, @nikic wrote:

If I take your example with and without lifetime intrinsics, then with current SROA, the variant without intrinsics produces a worse result: https://llvm.godbolt.org/z/oPMKWzaMq
Results of SROA should be the same with an without lifetime intrinsics, but I'm not sure your patch achieves that -- it looks like it produces a third variant that is better than both?

Here is the result after this patch. Compare the while_body section.

define i16 @foo(i32* nocapture readonly %loop) #1 {
entry:
  br label %while_cond

while_cond:                                       ; preds = %while_body, %entry
  %arr.sroa.6.0 = phi i32 [ 0, %entry ], [ %res1, %while_body ]
  %arr.sroa.0.0 = phi i32 [ 0, %entry ], [ %res0, %while_body ]
  %loopi = load i32, i32* %loop, align 4
  %loopb = icmp eq i32 %loopi, 0
  br i1 %loopb, label %while_end, label %while_body

while_body:                                       ; preds = %while_cond
  %x = call { i32, i32 } @bar(i32 %arr.sroa.0.0, i32 %arr.sroa.6.0) #0
  %res0 = extractvalue { i32, i32 } %x, 0
  %res1 = extractvalue { i32, i32 } %x, 1
  br label %while_cond

while_end:                                        ; preds = %while_cond
  %arr.sroa.0.0.extract.trunc = trunc i32 %arr.sroa.0.0 to i16
  %arr.sroa.6.4.extract.trunc = trunc i32 %arr.sroa.6.0 to i16
  %ret = add i16 %arr.sroa.0.0.extract.trunc, %arr.sroa.6.4.extract.trunc
  ret i16 %ret
}

[slydiman]

I have updated the code and removed the loop from the test.

[nikic]

llvm/test/Transforms/SROA/lifetime-intrinsic.ll

12

Can you please add a copy of this function that does not use lifetime.start/end, precommit the test, and then rebase this revision so it only shows the test diff?

[slydiman]

Ping

[nikic]

This still has the problem that the result differs depending on whether the lifetime intrinsic is present or not. It makes me suspicious that this is not fixing the right thing. I would at least like an explanation of where the difference comes from.

[slydiman]

Why do you think the results of SROA should be the same with an without lifetime intrinsics?

The precommited tests, you asked for above, contradict this statement, and demonstrate that lifetime intrinsics help with the optimization.

The problem with the current implementation is that the first element is not handled properly, fixing which is the subject of this patch.

[nikic]

I've done some due diligence, and the relevant difference starts with this code: https://github.com/llvm/llvm-project/blob/1ddc6ab1a9c321e02da4fab836c5f863a906108b/llvm/lib/Transforms/Scalar/SROA.cpp#L4412 If lifetime intrinsics are present, we'll mark both the [0,2) and [0,4) slices as unsplittable, and will then visit the [0,4) partition first, which has viable integer widening. Without lifetime intrinsics, we'll mark only [0,2) as unsplittable, and visit partition [0,2) first, where widening is not viable. I think something is very wrong here, but I'm not confident in which place it is and don't want to spend a lot of time on this.

In any case, I think your patch makes sense on a conceptual level, so let's go with it for now.

llvm/lib/Transforms/Scalar/SROA.cpp

1993

These already is a check for isLifetimeStartOrEnd() and isDroppable() below, only difference is that it's after the RelEnd > Size check. I think you'd want to add the isDroppable() check here as well, and then drop the IntrinsicInst handling below entirely.

You can also move the Use *U = S.getUse(); line before here.

[slydiman]

llvm/lib/Transforms/Scalar/SROA.cpp

1993

I have moved Use *U = S.getUse();.
Please note I cannot move isDroppable() bacause this block depends on the previous block

} else if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(U->getUser())) {

and such moving will break the current logic.

[nikic]

llvm/lib/Transforms/Scalar/SROA.cpp

1993

Do you mean in case a MemIntrinsic is droppable? This cannot happen ("droppable" here means "assume intrinsic" or similar).

[slydiman]

llvm/lib/Transforms/Scalar/SROA.cpp

1993

I can move the following code to the beginning

if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(U->getUser())) {
  if (II->isLifetimeStartOrEnd()) {
    return true;
  } else if (!II->isDroppable()) {
    return false;
  }
}

The returned value false here may be wrong because the following block may do not return anything, prevent executing the block if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(U->getUser())) below and finally the return value must be true

} else if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(U->getUser())) {
  if (MI->isVolatile() || !isa<Constant>(MI->getLength()))
    return false;
  if (!S.isSplittable())
    return false; // Skip any unsplittable intrinsics.
}

Actually such moving breakes a lot of tests.

[nikic]

llvm/lib/Transforms/Scalar/SROA.cpp

1993

To be clear, this is what I meant:

if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(U->getUser())) {
  if (II->isLifetimeStartOrEnd() || II->isDroppable())
    return true;
}

You are returning false for all non-droppable intrinsics, while we only want to return true for droppable intrinsics, and leave everything else to go through the code below (including MemIntrinsic).

[nikic]

LGTM

llvm/lib/Transforms/Scalar/SROA.cpp

static bool isIntegerWideningViableForSlice(const Slice &S,
                                             Type *AllocaTy,
                                             const DataLayout &DL,
                                             bool &WholeAllocaOp) {
   uint64_t Size = DL.getTypeStoreSize(AllocaTy).getFixedSize();
 
   uint64_t RelBegin = S.beginOffset() - AllocBeginOffset;
   uint64_t RelEnd = S.endOffset() - AllocBeginOffset;
 
+  Use *U = S.getUse();
+
+  // Lifetime intrinsics operate over the whole alloca whose sizes are usually
+  // larger than other load/store slices (RelEnd > Size). But lifetime are
+  // always promotable and should not impact other slices' promotability of the
+  // partition.
+  if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(U->getUser())) {
+    if (II->isLifetimeStartOrEnd() || II->isDroppable())
+      return true;
+  }
+
   // We can't reasonably handle cases where the load or store extends past
   // the end of the alloca's type and into its padding.
   if (RelEnd > Size)
     return false;
 
-  Use *U = S.getUse();
-
   if (LoadInst *LI = dyn_cast<LoadInst>(U->getUser())) {
     if (LI->isVolatile())
       return false;
     // We can't handle loads that extend past the allocated memory.
     if (DL.getTypeStoreSize(LI->getType()).getFixedSize() > Size)
       return false;
     // So far, AllocaSliceRewriter does not support widening split slice tails
     // in rewriteIntegerLoad.
if (IntegerType *ITy = dyn_cast<IntegerType>(ValueTy)) {
       // they are promotable.
       return false;
     }
   } else if (MemIntrinsic *MI = dyn_cast<MemIntrinsic>(U->getUser())) {
     if (MI->isVolatile() || !isa<Constant>(MI->getLength()))
       return false;
     if (!S.isSplittable())
       return false; // Skip any unsplittable intrinsics.
-  } else if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(U->getUser())) {
-    if (!II->isLifetimeStartOrEnd() && !II->isDroppable())
-      return false;
   } else {
     return false;
   }
 
   return true;
 }
 
 /// Test whether the given alloca partition's integer operations can be

llvm/test/Transforms/SROA/lifetime-intrinsic.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=sroa -S | FileCheck %s
 
 %i32x2 = type { [2 x i32] }
 
 ; Note %arr is the union
 ; union {
 ;   int i[2];
 ;   short s[4];
 ; };
 
 define i16 @with_lifetime(i32 %a, i32 %b) #0 {
 ; CHECK-LABEL: @with_lifetime(
-; CHECK-NEXT:    [[ARR_SROA_0_SROA_0_0_EXTRACT_TRUNC:%.*]] = trunc i32 [[A:%.*]] to i16
-; CHECK-NEXT:    [[ARR_SROA_0_SROA_4_0_EXTRACT_SHIFT:%.*]] = lshr i32 [[A]], 16
-; CHECK-NEXT:    [[ARR_SROA_0_SROA_4_0_EXTRACT_TRUNC:%.*]] = trunc i32 [[ARR_SROA_0_SROA_4_0_EXTRACT_SHIFT]] to i16
+; CHECK-NEXT:    [[ARR_SROA_0_0_EXTRACT_TRUNC:%.*]] = trunc i32 [[A:%.*]] to i16
 ; CHECK-NEXT:    [[ARR_SROA_4_4_EXTRACT_TRUNC:%.*]] = trunc i32 [[B:%.*]] to i16
-; CHECK-NEXT:    [[RET:%.*]] = add i16 [[ARR_SROA_0_SROA_0_0_EXTRACT_TRUNC]], [[ARR_SROA_4_4_EXTRACT_TRUNC]]
+; CHECK-NEXT:    [[RET:%.*]] = add i16 [[ARR_SROA_0_0_EXTRACT_TRUNC]], [[ARR_SROA_4_4_EXTRACT_TRUNC]]
 ; CHECK-NEXT:    ret i16 [[RET]]
 ;
   %arr = alloca %i32x2, align 4
   %p = bitcast %i32x2* %arr to i8*
   call void @llvm.lifetime.start.p0i8(i64 8, i8* %p)
   %p0 = getelementptr inbounds %i32x2, %i32x2* %arr, i64 0, i32 0, i32 0
   %p1 = getelementptr inbounds %i32x2, %i32x2* %arr, i64 0, i32 0, i32 1
   store i32 %a, i32* %p0, align 4
   store i32 %b, i32* %p1, align 4
   %p0_i16 = bitcast %i32x2* %arr to i16*
   %s0 = load i16, i16* %p0_i16, align 4
   %p1_i16 = bitcast i32* %p1 to i16*
   %s2 = load i16, i16* %p1_i16, align 4
   %ret = add i16 %s0, %s2
   call void @llvm.lifetime.end.p0i8(i64 8, i8* %p)
   ret i16 %ret
 }