[BasicAA] Fix AA bug on dynamic allocas and stackrestore

[rnk]

Summary

BasicAA has special logic for unescaped allocas, which normally applies
equally well to dynamic and static allocas. However, llvm.stackrestore
has the power to end the lifetime of dynamic allocas, without referring
to them directly.

stackrestore is already marked with the most conservative memory
modification attributes, but because the alloca is not escaped, the
normal logic produces incorrect results. I think BasicAA needs a special
case here to teach it about the relationship between dynamic allocas and
stackrestore.

Fixes PR40118

[efriedma]

I guess stackrestore is special because it's the only way to clobber an unescaped alloca; I think this is fine, in that sense.

That said, I wish IR was defined in a more consistent way; I'm concerned that there are probably other places in LLVM that are making similar bad assumption about dynamic allocas. Since VLAs are essentially banned from most codebases, we get very little test coverage on Linux. I'm particularly concerned about places using PointerMayBeCaptured and similar logic. Is LICM safe? Is DSE safe? Is AAResults::callCapturesBefore safe? Is InstCombiner::foldAllocaCmp safe? Is TailRecursionElimination safe?

[rnk]

My experience so far with inalloca has been that, no, those passes are not safe. This is only the latest in a long line of bugs in LLVM around handling dynamic allocas. I think conflating static and dynamic allocas was a mistake, and we would've been better served with a separate construct for allocating stack memory for local variables, but at this stage, we'd be better off standardizing on the terminology of "alloca" for standard locals and pushing dynamic allocas into intrinsic instructions.

I'd like to do the same for inalloca, and use tokens to establish a region around the call site. I started a design doc for it, but it looks like I deleted it at some point. :(

[efriedma]

llvm/lib/Transforms/Scalar/MemCpyOptimizer.cpp

1064 ↗

(On Diff #179198)

(Please commit this separately.)

llvm/test/Transforms/MemCpyOpt/stackrestore.ll

31

Please add a corresponding testcase without the stacksave/stackrestore so the test doesn't bitrot.

[rnk]

• add positive test for dynamic allocas with intervening call

[efriedma]

LGTM

[hfinkel]

In D55969#1338569, @efriedma wrote:

LGTM

I agree.

llvm/lib/Analysis/BasicAliasAnalysis.cpp

ModRefInfo BasicAAResult::getModRefInfo(ImmutableCallSite CS,
   // contents of the alloca into argument registers or stack slots, so there is
   // no lifetime issue.
   if (isa<AllocaInst>(Object))
     if (const CallInst *CI = dyn_cast<CallInst>(CS.getInstruction()))
       if (CI->isTailCall() &&
           !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(CS, Intrinsic::stackrestore))
+      return ModRefInfo::Mod;
+
   // If the pointer is to a locally allocated object that does not escape,
   // then the call can not mod/ref the pointer unless the call takes the pointer
   // as an argument, and itself doesn't capture it.
   if (!isa<Constant>(Object) && CS.getInstruction() != Object &&
       isNonEscapingLocalObject(Object)) {
 
     // Optimistically assume that call doesn't touch Object and check this
     // assumption in the following loop.

llvm/test/Transforms/MemCpyOpt/stackrestore.ll

+; RUN: opt -S -memcpyopt < %s | FileCheck %s
+
+; PR40118: BasicAA didn't realize that stackrestore ends the lifetime of
+; unescaped dynamic allocas, such as those that might come from inalloca.
+
+source_filename = "t.cpp"
+target datalayout = "e-m:x-p:32:32-i64:64-f80:32-n8:16:32-a:0:32-S32"
+target triple = "i686-unknown-windows-msvc19.14.26433"
+
+@str = internal constant [9 x i8] c"abcdxxxxx"
+
+
+; Test that we can propagate memcpy through an unescaped dynamic alloca across
+; a call to @external.
+
+define i32 @test_norestore(i32 %n) {
+  %tmpmem = alloca [10 x i8], align 4
+  %tmp = getelementptr inbounds [10 x i8], [10 x i8]* %tmpmem, i32 0, i32 0
+
+  ; Make a dynamic alloca, initialize it.
+  %p = alloca i8, i32 %n, align 4
+  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %p, i8* align 1 getelementptr inbounds ([9 x i8], [9 x i8]* @str, i32 0, i32 0), i32 9, i1 false)
+
+  ; This extra byte exists to prevent memcpyopt from propagating @str.
+  %p10 = getelementptr inbounds i8, i8* %p, i32 9
+  store i8 0, i8* %p10
+
+  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %tmp, i8* %p, i32 10, i1 false)
+  call void @external()
+  %heap = call i8* @malloc(i32 9)
+  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %heap, i8* %tmp, i32 9, i1 false)
+  call void @useit(i8* %heap)
+  ret i32 0
+}
+
+; CHECK-LABEL: define i32 @test_norestore(i32 %n)
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i32(i8* %p, i8* align 1 getelementptr inbounds ([9 x i8], [9 x i8]* @str, i32 0, i32 0), i32 9, i1 false)
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i32(i8* %tmp, i8* %p, i32 10, i1 false)
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i32(i8* %heap, i8* %p, i32 9, i1 false)
+
+
+; Do not propagate memcpy from %p across the stackrestore.
+
+define i32 @test_stackrestore() {
+  %tmpmem = alloca [10 x i8], align 4
+  %tmp = getelementptr inbounds [10 x i8], [10 x i8]* %tmpmem, i32 0, i32 0
+  %inalloca.save = tail call i8* @llvm.stacksave()
+  %argmem = alloca inalloca [10 x i8], align 4
+  %p = getelementptr inbounds [10 x i8], [10 x i8]* %argmem, i32 0, i32 0
+  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %p, i8* align 1 getelementptr inbounds ([9 x i8], [9 x i8]* @str, i32 0, i32 0), i32 9, i1 false)
+
+  ; This extra byte exists to prevent memcpyopt from propagating @str.
+  %p10 = getelementptr inbounds [10 x i8], [10 x i8]* %argmem, i32 0, i32 9
+  store i8 0, i8* %p10
+
+  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %tmp, i8* %p, i32 10, i1 false)
+  call void @llvm.stackrestore(i8* %inalloca.save)
+  %heap = call i8* @malloc(i32 9)
+  call void @llvm.memcpy.p0i8.p0i8.i32(i8* %heap, i8* %tmp, i32 9, i1 false)
+  call void @useit(i8* %heap)
+  ret i32 0
+}
+
+; CHECK-LABEL: define i32 @test_stackrestore()
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i32(i8* %p, i8* align 1 getelementptr inbounds ([9 x i8], [9 x i8]* @str, i32 0, i32 0), i32 9, i1 false)
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i32(i8* %tmp, i8* %p, i32 10, i1 false)
+; CHECK: call void @llvm.memcpy.p0i8.p0i8.i32(i8* %heap, i8* %tmp, i32 9, i1 false)
+
+declare void @llvm.memcpy.p0i8.p0i8.i32(i8* nocapture writeonly, i8* nocapture readonly, i32, i1)
+declare i8* @llvm.stacksave()
+declare void @llvm.stackrestore(i8*)
+declare i8* @malloc(i32)
+declare void @useit(i8*)
+declare void @external()