[llvm] [MemCpyOpt] Clear writeonly from dest argument after call slot optimization (PR #95179)

[Arthur Eubanks via llvm-commits]

https://github.com/aeubanks created https://github.com/llvm/llvm-project/pull/95179

After call slot optimization, the function may read from the parameter even if it was previously writeonly (due to the only use of it being a memcpy to it). This can cause miscompiles, as seen in #95152.

This is essentially intersecting the implicit alloca attributes with the argument attributes.

No need to worry about readnone/readonly since the dest already has a memcpy to it, so it shouldn't be marked readnone/readonly.

Fixes #95152

>From 336d7cc02abd18ee56c4c0010723a7eb06f061e3 Mon Sep 17 00:00:00 2001
From: Arthur Eubanks <aeubanks at google.com>
Date: Tue, 11 Jun 2024 22:55:01 +0000
Subject: [PATCH] [MemCpyOpt] Clear writeonly from dest argument after call
 slot optimization

After call slot optimization, the function may read from the parameter
even if it was previously writeonly (due to the only use of it being a
memcpy to it). This can cause miscompiles, as seen in #95152.

This is essentially intersecting the implicit alloca attributes with the
argument attributes.

No need to worry about readnone/readonly since the dest already has a
memcpy to it, so it shouldn't be marked readnone/readonly.

Fixes #95152
---
 .../lib/Transforms/Scalar/MemCpyOptimizer.cpp |  6 ++++++
 .../MemCpyOpt/callslot-writeonly.ll           | 19 +++++++++++++++++++
 2 files changed, 25 insertions(+)
 create mode 100644 llvm/test/Transforms/MemCpyOpt/callslot-writeonly.ll

diff --git a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
index 8fe3780bcf1b3..8f6d6d3460ee4 100644
--- a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
+++ b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
@@ -1115,6 +1115,12 @@ bool MemCpyOptPass::performCallSlotOptzn(Instruction *cpyLoad,
   if (cpyLoad != cpyStore)
     combineAAMetadata(C, cpyStore);
 
+  // Clear writeonly from the dest if it's an argument because we may now read
+  // from it in the call.
+  if (auto*A = dyn_cast<Argument>(cpyDest))
+    if (A->hasAttribute(Attribute::WriteOnly))
+      A->removeAttr(Attribute::WriteOnly);
+
   ++NumCallSlot;
   return true;
 }
diff --git a/llvm/test/Transforms/MemCpyOpt/callslot-writeonly.ll b/llvm/test/Transforms/MemCpyOpt/callslot-writeonly.ll
new file mode 100644
index 0000000000000..c336ccf8984f0
--- /dev/null
+++ b/llvm/test/Transforms/MemCpyOpt/callslot-writeonly.ll
@@ -0,0 +1,19 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=memcpyopt < %s | FileCheck %s
+
+; Check that we clear writeonly from the parameter when performing call slot optimization.
+
+define void @writeonly_param(ptr noalias writable writeonly dereferenceable(8) %dst) {
+; CHECK-LABEL: define void @writeonly_param(
+; CHECK-SAME: ptr noalias writable dereferenceable(8) [[DST:%.*]]) {
+; CHECK-NEXT:    [[SRC:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    call void @accept_ptr(ptr nocapture [[DST]]) #[[ATTR1:[0-9]+]]
+; CHECK-NEXT:    ret void
+;
+  %src = alloca i64, align 8
+  call void @accept_ptr(ptr nocapture %src) nounwind
+  call void @llvm.memcpy.p0.p0.i64(ptr align 8 %dst, ptr %src, i64 8, i1 false)
+  ret void
+}
+
+declare void @accept_ptr(ptr)