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

[Antonio Frighetto via llvm-commits]

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

Fixes: https://github.com/llvm/llvm-project/issues/72035.

>From e2be56a5338a6cee905b092585200690171a0797 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 | 30 ++++++++++++++++++++
 2 files changed, 36 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..eceeca2c997bab9 100644
--- a/llvm/test/Transforms/SROA/invariant-group.ll
+++ b/llvm/test/Transforms/SROA/invariant-group.ll
@@ -79,6 +79,36 @@ 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
+}
+
 !0 = !{}
 ;; NOTE: These prefixes are unused and the list is autogenerated. Do not add tests below this line:
 ; CHECK-MODIFY-CFG: {{.*}}