[llvm] [SROA] Allow load-only promotion with read-only captures (PR #130735)

[Nikita Popov via llvm-commits]

https://github.com/nikic updated https://github.com/llvm/llvm-project/pull/130735

>From 495c6d1fa9319187ea747db981b2a18bd55dbe0f Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Tue, 11 Mar 2025 09:35:13 +0100
Subject: [PATCH 1/2] [SROA] Allow load-only promotion with read-only captures

It's okay if the address or read-provenance of the pointer is
captured. We only have to make sure that there are no unanalyzable
writes to the pointer.
---
 llvm/lib/Transforms/Scalar/SROA.cpp           |  6 ++--
 .../test/Transforms/SROA/readonlynocapture.ll | 32 +++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/SROA.cpp b/llvm/lib/Transforms/Scalar/SROA.cpp
index 69e7ce83f82e4..1b82e5dc58707 100644
--- a/llvm/lib/Transforms/Scalar/SROA.cpp
+++ b/llvm/lib/Transforms/Scalar/SROA.cpp
@@ -1397,10 +1397,10 @@ class AllocaSlices::SliceBuilder : public PtrUseVisitor<SliceBuilder> {
   void visitInstruction(Instruction &I) { PI.setAborted(&I); }
 
   void visitCallBase(CallBase &CB) {
-    // If the call operand is NoCapture ReadOnly, then we mark it as
-    // EscapedReadOnly.
+    // If the call operand is read-only and only does a read-only or address
+    // capture, then we mark it as EscapedReadOnly.
     if (CB.isDataOperand(U) &&
-        CB.doesNotCapture(U->getOperandNo()) &&
+        !capturesFullProvenance(CB.getCaptureInfo(U->getOperandNo())) &&
         CB.onlyReadsMemory(U->getOperandNo())) {
       PI.setEscapedReadOnly(&CB);
       return;
diff --git a/llvm/test/Transforms/SROA/readonlynocapture.ll b/llvm/test/Transforms/SROA/readonlynocapture.ll
index 611c90ac32b5a..5ae436e52258b 100644
--- a/llvm/test/Transforms/SROA/readonlynocapture.ll
+++ b/llvm/test/Transforms/SROA/readonlynocapture.ll
@@ -406,4 +406,36 @@ define i32 @simple_byval() {
   ret i32 %l1
 }
 
+declare void @callee_address_only_capture(ptr readonly captures(address) %p)
+
+define i32 @address_only_capture() {
+; CHECK-LABEL: @address_only_capture(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 0, ptr [[A]], align 4
+; CHECK-NEXT:    call void @callee_address_only_capture(ptr [[A]])
+; CHECK-NEXT:    ret i32 0
+;
+  %a = alloca i32
+  store i32 0, ptr %a
+  call void @callee_address_only_capture(ptr %a)
+  %l1 = load i32, ptr %a
+  ret i32 %l1
+}
+
+declare void @callee_read_only_capture(ptr readonly captures(address, read_provenance) %p)
+
+define i32 @read_only_capture() {
+; CHECK-LABEL: @read_only_capture(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 0, ptr [[A]], align 4
+; CHECK-NEXT:    call void @callee_read_only_capture(ptr [[A]])
+; CHECK-NEXT:    ret i32 0
+;
+  %a = alloca i32
+  store i32 0, ptr %a
+  call void @callee_read_only_capture(ptr %a)
+  %l1 = load i32, ptr %a
+  ret i32 %l1
+}
+
 declare void @llvm.memcpy.p0.p0.i64(ptr, ptr, i64, i1)

>From 29b94f9b1494e14e6fd6f61f0818634653f25667 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Thu, 13 Mar 2025 09:51:48 +0100
Subject: [PATCH 2/2] Add extra test using captures(provenance)

---
 llvm/test/Transforms/SROA/readonlynocapture.ll | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/llvm/test/Transforms/SROA/readonlynocapture.ll b/llvm/test/Transforms/SROA/readonlynocapture.ll
index 5ae436e52258b..2ad20fcc51dc5 100644
--- a/llvm/test/Transforms/SROA/readonlynocapture.ll
+++ b/llvm/test/Transforms/SROA/readonlynocapture.ll
@@ -438,4 +438,22 @@ define i32 @read_only_capture() {
   ret i32 %l1
 }
 
+declare void @callee_provenance_only_capture(ptr readonly captures(provenance) %p)
+
+; Should not be transformed, as write-provenance is captured.
+define i32 @provenance_only_capture() {
+; CHECK-LABEL: @provenance_only_capture(
+; CHECK-NEXT:    [[A:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    store i32 0, ptr [[A]], align 4
+; CHECK-NEXT:    call void @callee_provenance_only_capture(ptr [[A]])
+; CHECK-NEXT:    [[L1:%.*]] = load i32, ptr [[A]], align 4
+; CHECK-NEXT:    ret i32 [[L1]]
+;
+  %a = alloca i32
+  store i32 0, ptr %a
+  call void @callee_provenance_only_capture(ptr %a)
+  %l1 = load i32, ptr %a
+  ret i32 %l1
+}
+
 declare void @llvm.memcpy.p0.p0.i64(ptr, ptr, i64, i1)