[SROA] Allow load-only promotion with read-only captures

[nikic]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/130735.diff

2 Files Affected:

• (modified) llvm/lib/Transforms/Scalar/SROA.cpp (+3-3)
• (modified) llvm/test/Transforms/SROA/readonlynocapture.ll (+32)

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)

[davemgreen]

LGTM, thanks

[davemgreen]

Is it worth adding a test with captures(provenance)?

[nikic]

Done. We don't really expect to see provenance without address in practice, but it's not illegal, so it makes sense to test it...