[llvm] [MemCpyOpt] Allow memcpy elision for non-noalias arguments (PR #107860)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/107860

We currently elide memcpys for readonly nocapture noalias arguments. noalias is checked to make sure that there are no other ways to write the memory, e.g. through a different argument or an escaped pointer.

In addition to the current noalias check, also query alias analysis, in case it can prove that modification is not possible through other means. 

This fixes the problem reported in https://discourse.llvm.org/t/problem-about-memcpy-elimination/81121.

>From 436bfbe6201c3433a171b63a2f8f1b8488c24295 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Mon, 9 Sep 2024 12:18:26 +0200
Subject: [PATCH] [MemCpyOpt] Allow memcpy elision for non-noalias arguments

We currently elide memcpys for readonly nocapture noalias
arguments. noalias is checked to make sure that there are no
other ways to write the memory, e.g. through a different argument
or an escaped pointer.

In addition to the current noalias check, also query alias
analysis, in case it can prove that modification is not possible
through other means.
---
 llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp | 17 +++++++++++++----
 llvm/test/Transforms/MemCpyOpt/memcpy.ll       |  7 ++-----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
index 3f15fa2163d270..f81ee23a0d2d9c 100644
--- a/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
+++ b/llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp
@@ -1961,12 +1961,22 @@ bool MemCpyOptPass::processByValArgument(CallBase &CB, unsigned ArgNo) {
 /// 4. The memcpy src is not modified during the call. (ModRef check shows no
 /// Mod.)
 bool MemCpyOptPass::processImmutArgument(CallBase &CB, unsigned ArgNo) {
+  BatchAAResults BAA(*AA);
+  Value *ImmutArg = CB.getArgOperand(ArgNo);
+
   // 1. Ensure passed argument is immutable during call.
-  if (!(CB.paramHasAttr(ArgNo, Attribute::NoAlias) &&
-        CB.paramHasAttr(ArgNo, Attribute::NoCapture)))
+  if (!CB.paramHasAttr(ArgNo, Attribute::NoCapture))
+    return false;
+
+  // We know that the argument is readonly at this point, but the function
+  // might still modify the same memory through a different pointer. Exclude
+  // this either via noalias, or alias analysis.
+  if (!CB.paramHasAttr(ArgNo, Attribute::NoAlias) &&
+      isModSet(
+          BAA.getModRefInfo(&CB, MemoryLocation::getBeforeOrAfter(ImmutArg))))
     return false;
+
   const DataLayout &DL = CB.getDataLayout();
-  Value *ImmutArg = CB.getArgOperand(ArgNo);
 
   // 2. Check that arg is alloca
   // TODO: Even if the arg gets back to branches, we can remove memcpy if all
@@ -1986,7 +1996,6 @@ bool MemCpyOptPass::processImmutArgument(CallBase &CB, unsigned ArgNo) {
     return false;
 
   MemCpyInst *MDep = nullptr;
-  BatchAAResults BAA(*AA);
   MemoryAccess *Clobber = MSSA->getWalker()->getClobberingMemoryAccess(
       CallAccess->getDefiningAccess(), Loc, BAA);
   if (auto *MD = dyn_cast<MemoryDef>(Clobber))
diff --git a/llvm/test/Transforms/MemCpyOpt/memcpy.ll b/llvm/test/Transforms/MemCpyOpt/memcpy.ll
index a28b0542a7c59f..ba260752ce4b56 100644
--- a/llvm/test/Transforms/MemCpyOpt/memcpy.ll
+++ b/llvm/test/Transforms/MemCpyOpt/memcpy.ll
@@ -472,9 +472,7 @@ define void @immut_param_mayalias(ptr align 4 noalias %val) {
 ; argument doesn't matter.
 define void @immut_param_unescaped_alloca(ptr align 4 noalias %val) {
 ; CHECK-LABEL: @immut_param_unescaped_alloca(
-; CHECK-NEXT:    [[VAL1:%.*]] = alloca i8, align 4
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[VAL1]], ptr align 4 [[VAL:%.*]], i64 1, i1 false)
-; CHECK-NEXT:    call void @f(ptr nocapture readonly align 4 [[VAL1]])
+; CHECK-NEXT:    call void @f(ptr nocapture readonly align 4 [[VAL:%.*]])
 ; CHECK-NEXT:    ret void
 ;
   %val1 = alloca i8, align 4
@@ -489,8 +487,7 @@ define void @immut_param_memory_argmem_read(ptr align 4 noalias %val) {
 ; CHECK-LABEL: @immut_param_memory_argmem_read(
 ; CHECK-NEXT:    [[VAL1:%.*]] = alloca i8, align 4
 ; CHECK-NEXT:    call void @f(ptr [[VAL1]])
-; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i64(ptr align 4 [[VAL1]], ptr align 4 [[VAL:%.*]], i64 1, i1 false)
-; CHECK-NEXT:    call void @f(ptr nocapture readonly align 4 [[VAL1]]) #[[ATTR6:[0-9]+]]
+; CHECK-NEXT:    call void @f(ptr nocapture readonly align 4 [[VAL:%.*]]) #[[ATTR6:[0-9]+]]
 ; CHECK-NEXT:    ret void
 ;
   %val1 = alloca i8, align 4