[llvm] [BasicAA] MemCpyOpt fix on tail stackrestore (PR #101352)

[Shimin Cui via llvm-commits]

https://github.com/scui-ibm created https://github.com/llvm/llvm-project/pull/101352

 In BasicAAResult::getModRefInfo, the special handling of alloca/stackrestore is after checking whether the call is a tailcall or not. This is to move the specific stackrestore code before the tail call checking, as the call to stackrestore can be marked as a tail call by tailcallelim which is performed before MemCpyOpt in the transformation pipeline.
 
A new RUN is added to the existing case for the issue exposed by tail call of stackrestore.

>From 89cdac4d3acb52b3f0332e05de45093eb9fadd80 Mon Sep 17 00:00:00 2001
From: Shimin Cui <scui at ca.ibm.com>
Date: Wed, 31 Jul 2024 11:24:43 -0400
Subject: [PATCH] [BasicAA] MemCpyOpt fix on tail stackrestore

---
 llvm/lib/Analysis/BasicAliasAnalysis.cpp       | 12 ++++++------
 llvm/test/Transforms/MemCpyOpt/stackrestore.ll |  9 +++++----
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Analysis/BasicAliasAnalysis.cpp b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
index e474899fb548e..504026e44f531 100644
--- a/llvm/lib/Analysis/BasicAliasAnalysis.cpp
+++ b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
@@ -914,6 +914,12 @@ ModRefInfo BasicAAResult::getModRefInfo(const CallBase *Call,
 
   const Value *Object = getUnderlyingObject(Loc.Ptr);
 
+  // Stack restore is able to modify unescaped dynamic allocas. Assume it may
+  // modify them even though the alloca is not escaped.
+  if (auto *AI = dyn_cast<AllocaInst>(Object))
+    if (!AI->isStaticAlloca() && isIntrinsicCall(Call, Intrinsic::stackrestore))
+      return ModRefInfo::Mod;
+
   // Calls marked 'tail' cannot read or write allocas from the current frame
   // because the current frame might be destroyed by the time they run. However,
   // a tail call may use an alloca with byval. Calling with byval copies the
@@ -925,12 +931,6 @@ ModRefInfo BasicAAResult::getModRefInfo(const CallBase *Call,
           !CI->getAttributes().hasAttrSomewhere(Attribute::ByVal))
         return ModRefInfo::NoModRef;
 
-  // Stack restore is able to modify unescaped dynamic allocas. Assume it may
-  // modify them even though the alloca is not escaped.
-  if (auto *AI = dyn_cast<AllocaInst>(Object))
-    if (!AI->isStaticAlloca() && isIntrinsicCall(Call, Intrinsic::stackrestore))
-      return ModRefInfo::Mod;
-
   // A call can access a locally allocated object either because it is passed as
   // an argument to the call, or because it has escaped prior to the call.
   //
diff --git a/llvm/test/Transforms/MemCpyOpt/stackrestore.ll b/llvm/test/Transforms/MemCpyOpt/stackrestore.ll
index 0fc37c44fa9e8..f5feea98e21cd 100644
--- a/llvm/test/Transforms/MemCpyOpt/stackrestore.ll
+++ b/llvm/test/Transforms/MemCpyOpt/stackrestore.ll
@@ -1,5 +1,6 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=memcpyopt < %s -verify-memoryssa | FileCheck %s
+; RUN: opt -S -passes="tailcallelim,memcpyopt" < %s -verify-memoryssa | FileCheck %s
 
 ; PR40118: BasicAA didn't realize that stackrestore ends the lifetime of
 ; unescaped dynamic allocas, such as those that might come from inalloca.
@@ -23,7 +24,7 @@ define i32 @test_norestore(i32 %n) {
 ; CHECK-NEXT:    store i8 0, ptr [[P10]], align 1
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr [[TMPMEM]], ptr [[P]], i32 10, i1 false)
 ; CHECK-NEXT:    call void @external()
-; CHECK-NEXT:    [[HEAP:%.*]] = call ptr @malloc(i32 9)
+; CHECK-NEXT:    [[HEAP:%.*]] = tail call ptr @malloc(i32 9)
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr [[HEAP]], ptr align 1 @str, i32 9, i1 false)
 ; CHECK-NEXT:    call void @useit(ptr [[HEAP]])
 ; CHECK-NEXT:    ret i32 0
@@ -40,7 +41,7 @@ define i32 @test_norestore(i32 %n) {
 
   call void @llvm.memcpy.p0.p0.i32(ptr %tmpmem, ptr %p, i32 10, i1 false)
   call void @external()
-  %heap = call ptr @malloc(i32 9)
+  %heap = tail call ptr @malloc(i32 9)
   call void @llvm.memcpy.p0.p0.i32(ptr %heap, ptr %tmpmem, i32 9, i1 false)
   call void @useit(ptr %heap)
   ret i32 0
@@ -58,7 +59,7 @@ define i32 @test_stackrestore() {
 ; CHECK-NEXT:    store i8 0, ptr [[P10]], align 1
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr [[TMPMEM]], ptr [[ARGMEM]], i32 10, i1 false)
 ; CHECK-NEXT:    call void @llvm.stackrestore.p0(ptr [[INALLOCA_SAVE]])
-; CHECK-NEXT:    [[HEAP:%.*]] = call ptr @malloc(i32 9)
+; CHECK-NEXT:    [[HEAP:%.*]] = tail call ptr @malloc(i32 9)
 ; CHECK-NEXT:    call void @llvm.memcpy.p0.p0.i32(ptr [[HEAP]], ptr [[TMPMEM]], i32 9, i1 false)
 ; CHECK-NEXT:    call void @useit(ptr [[HEAP]])
 ; CHECK-NEXT:    ret i32 0
@@ -74,7 +75,7 @@ define i32 @test_stackrestore() {
 
   call void @llvm.memcpy.p0.p0.i32(ptr %tmpmem, ptr %argmem, i32 10, i1 false)
   call void @llvm.stackrestore(ptr %inalloca.save)
-  %heap = call ptr @malloc(i32 9)
+  %heap = tail call ptr @malloc(i32 9)
   call void @llvm.memcpy.p0.p0.i32(ptr %heap, ptr %tmpmem, i32 9, i1 false)
   call void @useit(ptr %heap)
   ret i32 0