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

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.

Author: nikic

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;

llvm/test/Transforms/SROA/readonlynocapture.ll

@@ -406,4 +406,54 @@ 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 @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)