[llvm] [SROA] Allow `llvm.launder.invariant.group` intrinsic to be splittable (PR #72056)

[Antonio Frighetto via llvm-commits]

https://github.com/antoniofrighetto updated https://github.com/llvm/llvm-project/pull/72056

>From 9ac0bbea770c5c0f554b49ede76136bde56bdf83 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Sun, 12 Nov 2023 17:08:20 +0100
Subject: [PATCH] [SROA] Allow `llvm.launder.invariant.group` intrinsic to be
 splittable

Fixes: https://github.com/llvm/llvm-project/issues/72035.
---
 llvm/lib/Transforms/Scalar/SROA.cpp          |  7 +-
 llvm/test/Transforms/SROA/invariant-group.ll | 82 +++++++++++++++++++-
 2 files changed, 86 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 321525d3741d0a2..095f12cd414b803 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -1144,6 +1144,7 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
     }
 
     if (II.isLaunderOrStripInvariantGroup()) {
+      insertUse(II, Offset, AllocSize, true);
       enqueueUsers(II);
       return;
     }
@@ -3328,7 +3329,8 @@ class llvm::sroa::AllocaSliceRewriter
   }
 
   bool visitIntrinsicInst(IntrinsicInst &II) {
-    assert((II.isLifetimeStartOrEnd() || II.isDroppable()) &&
+    assert((II.isLifetimeStartOrEnd() || II.isLaunderOrStripInvariantGroup() ||
+            II.isDroppable()) &&
            "Unexpected intrinsic!");
     LLVM_DEBUG(dbgs() << "    original: " << II << "\n");
 
@@ -3342,6 +3344,9 @@ class llvm::sroa::AllocaSliceRewriter
       return true;
     }
 
+    if (II.isLaunderOrStripInvariantGroup())
+      return true;
+
     assert(II.getArgOperand(1) == OldPtr);
     // Lifetime intrinsics are only promotable if they cover the whole alloca.
     // Therefore, we drop lifetime intrinsics which don't cover the whole
diff --git a/llvm/test/Transforms/SROA/invariant-group.ll b/llvm/test/Transforms/SROA/invariant-group.ll
index 083fa72026b85e6..b67dd0a1357ecda 100644
--- a/llvm/test/Transforms/SROA/invariant-group.ll
+++ b/llvm/test/Transforms/SROA/invariant-group.ll
@@ -44,9 +44,9 @@ define void @g() {
 ; CHECK-NEXT:    [[A2:%.*]] = getelementptr inbounds [[T]], ptr [[A]], i32 0, i32 1
 ; CHECK-NEXT:    [[SV1:%.*]] = call i32 @somevalue()
 ; CHECK-NEXT:    [[SV2:%.*]] = call i32 @somevalue()
-; CHECK-NEXT:    store i32 [[SV1]], ptr [[A1_I8_INV]], align 4, !invariant.group !0
+; CHECK-NEXT:    store i32 [[SV1]], ptr [[A1_I8_INV]], align 4, !invariant.group [[META0:![0-9]+]]
 ; CHECK-NEXT:    store i32 [[SV2]], ptr [[A2]], align 4
-; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[A1_I8_INV]], align 4, !invariant.group !0
+; CHECK-NEXT:    [[V1:%.*]] = load i32, ptr [[A1_I8_INV]], align 4, !invariant.group [[META0]]
 ; CHECK-NEXT:    [[V2:%.*]] = load i32, ptr [[A2]], align 4
 ; CHECK-NEXT:    call void @h(i32 [[V1]])
 ; CHECK-NEXT:    call void @h(i32 [[V2]])
@@ -79,6 +79,84 @@ define void @g() {
   ret void
 }
 
+define void @store_and_launder() {
+; CHECK-LABEL: @store_and_launder(
+; CHECK-NEXT:    ret void
+;
+  %valptr = alloca i32, align 4
+  store i32 0, ptr %valptr, align 4
+  %barr = call ptr @llvm.launder.invariant.group.p0(ptr %valptr)
+  ret void
+}
+
+define i32 @launder_and_load() {
+; CHECK-LABEL: @launder_and_load(
+; CHECK-NEXT:    ret i32 undef
+;
+  %valptr = alloca i32, align 4
+  %barr = call ptr @llvm.launder.invariant.group.p0(ptr %valptr)
+  %v2 = load i32, ptr %valptr
+  ret i32 %v2
+}
+
+define void @launder_and_ptr_arith() {
+; CHECK-LABEL: @launder_and_ptr_arith(
+; CHECK-NEXT:    ret void
+;
+  %valptr = alloca i32, align 4
+  %barr = call ptr @llvm.launder.invariant.group.p0(ptr %valptr)
+  %a2 = getelementptr inbounds i32, ptr %valptr, i32 0
+  ret void
+}
+
+define void @partial_use_of_alloca() {
+; CHECK-LABEL: @partial_use_of_alloca(
+; CHECK-NEXT:    [[VALPTR:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 0, ptr [[VALPTR]], align 4
+; CHECK-NEXT:    [[BARR:%.*]] = call ptr @llvm.launder.invariant.group.p0(ptr [[VALPTR]])
+; CHECK-NEXT:    [[LOAD_VAL:%.*]] = load i32, ptr [[VALPTR]], align 4
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[LOAD_VAL]], 0
+; CHECK-NEXT:    br i1 [[COND]], label [[USE_ALLOCA:%.*]], label [[END:%.*]]
+; CHECK:       use_alloca:
+; CHECK-NEXT:    call void @use(ptr nonnull [[VALPTR]])
+; CHECK-NEXT:    br label [[END]]
+; CHECK:       end:
+; CHECK-NEXT:    ret void
+;
+  %valptr = alloca i32, align 4
+  store i32 0, ptr %valptr, align 4
+  %barr = call ptr @llvm.launder.invariant.group.p0(ptr %valptr)
+  %load_val = load i32, ptr %valptr, align 4
+  %cond = icmp eq i32 %load_val, 0
+  br i1 %cond, label %use_alloca, label %end
+
+use_alloca:
+  call void @use(i32* nonnull %valptr)
+  br label %end
+
+end:
+  ret void
+}
+
+define void @partial_promotion_of_alloca() {
+; CHECK-LABEL: @partial_promotion_of_alloca(
+; CHECK-NEXT:    [[STRUCT_PTR_SROA_2:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store volatile i32 0, ptr [[STRUCT_PTR_SROA_2]], align 4
+; CHECK-NEXT:    [[STRUCT_PTR_SROA_2_0_STRUCT_PTR_SROA_2_4_LOAD_VAL:%.*]] = load volatile i32, ptr [[STRUCT_PTR_SROA_2]], align 4
+; CHECK-NEXT:    ret void
+;
+  %struct_ptr = alloca %t, align 4
+  %field_ptr = getelementptr inbounds %t, ptr %struct_ptr, i32 0, i32 0
+  store i32 0, ptr %field_ptr, align 4
+  %volatile_field_ptr = getelementptr inbounds %t, ptr %struct_ptr, i32 0, i32 1
+  store volatile i32 0, ptr %volatile_field_ptr, align 4, !invariant.group !0
+  %barr = call ptr @llvm.launder.invariant.group.p0(ptr %struct_ptr)
+  %load_val = load volatile i32, ptr %volatile_field_ptr, align 4, !invariant.group !0
+  ret void
+}
+
+declare void @use(i32*)
+
 !0 = !{}
 ;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
 ; CHECK-MODIFY-CFG: {{.*}}