[FunctionAttrs] Account for memory effects of inalloca/preallocated

[nikic]

Summary

The code for inferring memory attributes on arguments claims that inalloca/preallocated arguments are always clobbered: https://github.com/llvm/llvm-project/blob/d71ad4108056d685f48407447095d8d92fd7685d/llvm/lib/Transforms/IPO/FunctionAttrs.cpp#L640-L642

However, we would still infer memory attributes for the whole function without taking this into account, so we could still end up inferring readnone for the function. This adds an argument clobber if there are any inalloca/preallocated arguments.

Disclaimer: I'm completely trusting the existing comment here, I'm not familiar with inalloc/preallocated arguments and why they're special in this regard.

[rnk]

Inalloca (& preallcoated) effectively allows the caller to take the address of the argument stack memory used for the call. The call effectively deallocates the argument memory, since it can immediately be reused by the next call, so it should be modeled like free.

+@akhuang because we were just discussing documentation of this feature

[nikic]

In D135783#3853152, @rnk wrote:

Inalloca (& preallcoated) effectively allows the caller to take the address of the argument stack memory used for the call. The call effectively deallocates the argument memory, since it can immediately be reused by the next call, so it should be modeled like free.

I see, thanks! It looks like this is also mentioned in https://llvm.org/docs/InAlloca.html#callee-cleanup-calling-conventions:

Another wrinkle is the existence of callee-cleanup conventions. On Windows, all methods and many other functions adjust the stack to clear the memory used to pass their arguments. In some sense, this means that the allocas are automatically cleared by the call. However, LLVM instead models this as a write of undef to all of the inalloca values passed to the call instead of a stack adjustment. Frontends should still restore the stack pointer to avoid a stack leak.

[rnk]

In D135783#3853445, @nikic wrote:

Another wrinkle is the existence of callee-cleanup conventions. On Windows, all methods and many other functions adjust the stack to clear the memory used to pass their arguments. In some sense, this means that the allocas are automatically cleared by the call. However, LLVM instead models this as a write of undef to all of the inalloca values passed to the call instead of a stack adjustment. Frontends should still restore the stack pointer to avoid a stack leak.

I don't recall why I thought modelling it as a "write of undef" was better than saying it deallocated the memory. It's effectively the same result as a stackrestore, which can be used to deallocate dynamic allocas.

It also kind of doesn't matter whether the call is callee cleanup, because it is valid for a pass to place a side-effect free call immediately after the call that uses inalloca.

llvm/lib/Transforms/IPO/FunctionAttrs.cpp

checkFunctionMemoryAccess(Function &F, bool ThisBody, AAResults &AAR,
   FunctionModRefBehavior OrigMRB = AAR.getModRefBehavior(&F);
   if (OrigMRB.doesNotAccessMemory())
     // Already perfect!
     return OrigMRB;
 
   if (!ThisBody)
     return OrigMRB;
 
-  // Scan the function body for instructions that may read or write memory.
   FunctionModRefBehavior MRB = FunctionModRefBehavior::none();
+  // Inalloca and preallocated arguments are always clobbered by the call.
+  if (F.getAttributes().hasAttrSomewhere(Attribute::InAlloca) ||
+      F.getAttributes().hasAttrSomewhere(Attribute::Preallocated))
+    MRB |= FunctionModRefBehavior::argMemOnly(ModRefInfo::ModRef);
+
   // Returns true if Ptr is not based on a function argument.
   auto IsArgumentOrAlloca = [](const Value *Ptr) {
     const Value *UO = getUnderlyingObject(Ptr);
     return isa<Argument>(UO) || isa<AllocaInst>(UO);
   };
+  // Scan the function body for instructions that may read or write memory.
   for (Instruction &I : instructions(F)) {
     // Some instructions can be ignored even if they read or write memory.
     // Detect these now, skipping to the next instruction if one is found.
     if (auto *Call = dyn_cast<CallBase>(&I)) {
       // Ignore calls to functions in the same SCC, as long as the call sites
       // don't have operand bundles.  Calls with operand bundles are allowed to
       // have memory effects not described by the memory effects of the call
       // target.

llvm/test/Transforms/FunctionAttrs/readattrs.ll

 ;
   store i8* %q, i8** %p
   call void @test6_1()
   ret void
 }
 
 ; inalloca parameters are always considered written
 define void @test7_1(i32* inalloca(i32) %a) {
-; CHECK: Function Attrs: mustprogress nofree norecurse nosync nounwind readnone willreturn
+; CHECK: Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind willreturn
 ; CHECK-LABEL: define {{[^@]+}}@test7_1
-; CHECK-SAME: (i32* nocapture inalloca(i32) [[A:%.*]]) #[[ATTR1]] {
+; CHECK-SAME: (i32* nocapture inalloca(i32) [[A:%.*]]) #[[ATTR5:[0-9]+]] {
 ; CHECK-NEXT:    ret void
 ;
   ret void
 }
 
 ; preallocated parameters are always considered written
 define void @test7_2(i32* preallocated(i32) %a) {
-; CHECK: Function Attrs: mustprogress nofree norecurse nosync nounwind readnone willreturn
+; CHECK: Function Attrs: argmemonly mustprogress nofree norecurse nosync nounwind willreturn
 ; CHECK-LABEL: define {{[^@]+}}@test7_2
-; CHECK-SAME: (i32* nocapture preallocated(i32) [[A:%.*]]) #[[ATTR1]] {
+; CHECK-SAME: (i32* nocapture preallocated(i32) [[A:%.*]]) #[[ATTR5]] {
 ; CHECK-NEXT:    ret void
 ;
   ret void
 }
 
 define i32* @test8_1(i32* %p) {
 ; CHECK: Function Attrs: mustprogress nofree norecurse nosync nounwind readnone willreturn
 ; CHECK-LABEL: define {{[^@]+}}@test8_1
entry:
   ret void
 }
 
 declare void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32>%val, <4 x i32*>, i32, <4 x i1>)
 
 define void @test9(<4 x i32*> %ptrs, <4 x i32>%val) {
 ; CHECK: Function Attrs: mustprogress nofree nosync nounwind willreturn writeonly
 ; CHECK-LABEL: define {{[^@]+}}@test9
-; CHECK-SAME: (<4 x i32*> [[PTRS:%.*]], <4 x i32> [[VAL:%.*]]) #[[ATTR6:[0-9]+]] {
+; CHECK-SAME: (<4 x i32*> [[PTRS:%.*]], <4 x i32> [[VAL:%.*]]) #[[ATTR7:[0-9]+]] {
 ; CHECK-NEXT:    call void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32> [[VAL]], <4 x i32*> [[PTRS]], i32 4, <4 x i1> <i1 true, i1 false, i1 true, i1 false>)
 ; CHECK-NEXT:    ret void
 ;
   call void @llvm.masked.scatter.v4i32.v4p0i32(<4 x i32>%val, <4 x i32*> %ptrs, i32 4, <4 x i1><i1 true, i1 false, i1 true, i1 false>)
   ret void
 }
 
 declare <4 x i32> @llvm.masked.gather.v4i32.v4p0i32(<4 x i32*>, i32, <4 x i1>, <4 x i32>)
 define <4 x i32> @test10(<4 x i32*> %ptrs) {
 ; CHECK: Function Attrs: mustprogress nofree nosync nounwind readonly willreturn
 ; CHECK-LABEL: define {{[^@]+}}@test10
-; CHECK-SAME: (<4 x i32*> [[PTRS:%.*]]) #[[ATTR8:[0-9]+]] {
+; CHECK-SAME: (<4 x i32*> [[PTRS:%.*]]) #[[ATTR9:[0-9]+]] {
 ; CHECK-NEXT:    [[RES:%.*]] = call <4 x i32> @llvm.masked.gather.v4i32.v4p0i32(<4 x i32*> [[PTRS]], i32 4, <4 x i1> <i1 true, i1 false, i1 true, i1 false>, <4 x i32> undef)
 ; CHECK-NEXT:    ret <4 x i32> [[RES]]
 ;
   %res = call <4 x i32> @llvm.masked.gather.v4i32.v4p0i32(<4 x i32*> %ptrs, i32 4, <4 x i1><i1 true, i1 false, i1 true, i1 false>, <4 x i32>undef)
   ret <4 x i32> %res
 }
 
 declare <4 x i32> @test11_1(<4 x i32*>) argmemonly nounwind readonly
 define <4 x i32> @test11_2(<4 x i32*> %ptrs) {
 ; CHECK: Function Attrs: argmemonly nofree nounwind readonly
 ; CHECK-LABEL: define {{[^@]+}}@test11_2
-; CHECK-SAME: (<4 x i32*> [[PTRS:%.*]]) #[[ATTR10:[0-9]+]] {
+; CHECK-SAME: (<4 x i32*> [[PTRS:%.*]]) #[[ATTR11:[0-9]+]] {
 ; CHECK-NEXT:    [[RES:%.*]] = call <4 x i32> @test11_1(<4 x i32*> [[PTRS]])
 ; CHECK-NEXT:    ret <4 x i32> [[RES]]
 ;
   %res = call <4 x i32> @test11_1(<4 x i32*> %ptrs)
   ret <4 x i32> %res
 }
 
 declare <4 x i32> @test12_1(<4 x i32*>) argmemonly nounwind
 define <4 x i32> @test12_2(<4 x i32*> %ptrs) {
 ; CHECK: Function Attrs: argmemonly nounwind
 ; CHECK-LABEL: define {{[^@]+}}@test12_2
-; CHECK-SAME: (<4 x i32*> [[PTRS:%.*]]) #[[ATTR11:[0-9]+]] {
+; CHECK-SAME: (<4 x i32*> [[PTRS:%.*]]) #[[ATTR12:[0-9]+]] {
 ; CHECK-NEXT:    [[RES:%.*]] = call <4 x i32> @test12_1(<4 x i32*> [[PTRS]])
 ; CHECK-NEXT:    ret <4 x i32> [[RES]]
 ;
   %res = call <4 x i32> @test12_1(<4 x i32*> %ptrs)
   ret <4 x i32> %res
 }
 
 define i32 @volatile_load(i32* %p) {
 ; CHECK: Function Attrs: argmemonly mustprogress nofree norecurse nounwind willreturn
 ; CHECK-LABEL: define {{[^@]+}}@volatile_load
-; CHECK-SAME: (i32* [[P:%.*]]) #[[ATTR12:[0-9]+]] {
+; CHECK-SAME: (i32* [[P:%.*]]) #[[ATTR13:[0-9]+]] {
 ; CHECK-NEXT:    [[LOAD:%.*]] = load volatile i32, i32* [[P]], align 4
 ; CHECK-NEXT:    ret i32 [[LOAD]]
 ;
   %load = load volatile i32, i32* %p
   ret i32 %load
 }
 
 declare void @escape_readnone_ptr(i8** %addr, i8* readnone %ptr)