[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 5073a5a4246d363f8444b97709ba03f68b1c0a11 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 | 73 ++++++++++++++++++++
 2 files changed, 79 insertions(+), 1 deletion(-)

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..c717d0c30a6f3d5 100644
--- a/llvm/test/Transforms/SROA/invariant-group.ll
+++ b/llvm/test/Transforms/SROA/invariant-group.ll
@@ -79,6 +79,79 @@ 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 42, ptr [[VALPTR]], align 4
+; CHECK-NEXT:    [[BARR:%.*]] = call ptr @llvm.launder.invariant.group.p0(ptr [[VALPTR]])
+; CHECK-NEXT:    [[PART_VAL:%.*]] = load i32, ptr [[VALPTR]], align 4
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[PART_VAL]], 42
+; 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 42, ptr %valptr, align 4
+  %barr = call ptr @llvm.launder.invariant.group.p0(ptr %valptr)
+  %part_val = load i32, ptr %valptr, align 4
+  %cond = icmp eq i32 %part_val, 42
+  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_use_of_struct() {
+; CHECK-LABEL: @partial_use_of_struct(
+; CHECK-NEXT:    ret void
+;
+  %struct_ptr = alloca %t, align 4
+  %field_ptr = getelementptr inbounds %t, ptr %struct_ptr, i32 0, i32 1
+  store i32 123, ptr %field_ptr, align 4
+  %barr = call ptr @llvm.launder.invariant.group.p0(ptr %struct_ptr)
+  %loaded_val = load i32, ptr %field_ptr, align 4
+  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: {{.*}}