[InlineFunction] Update metadata on loads that are return values

[anna]

Summary

This patch builds upon D76140 by updating metadata on pointer typed loads in inlined functions, when the load is the return value, and the callsite contains return attributes which can be updated as metadata on the load.
Added test cases show this for nonnull, dereferenceable, dereferenceable_or_null

[jdoerfert]

llvm/lib/Transforms/Utils/InlineFunction.cpp

1197–1212

Isn't this just VMap.lookup(V)?

1210

Num

llvm/test/Transforms/Inline/ret_load_metadata.ll

70

Is there a test where the attribute is part of the callee and not the calls site?

[anna]

llvm/lib/Transforms/Utils/InlineFunction.cpp

1197–1212

you're right. We don't need a lambda for this.

llvm/test/Transforms/Inline/ret_load_metadata.ll

70

Will add one.

[anna]

addressed review comments

[jdoerfert]

Minor comments, otherwise LGTM.

llvm/lib/Transforms/Utils/InlineFunction.cpp

1203

Nit: Cheap condition first: UpdateXXX && isa<...>

1251–1271

if (auto *CB = dyn_cast<CallBase>(...))
should be simpler.

1261

Nit: I really hope that'll change, for non-null at least. Though, then we might also want to allow it in the metadata so all will be fine except the comment ;)

[anna]

thanks for the review @jdoerfert . Will incorporate these changes and rebase on landed D76140.

llvm/lib/Transforms/Utils/InlineFunction.cpp

1264

We'll need to avoid updating dereferenceable, dereferenceable_or_null when these metadata is already present on the load. This will avoid updating the metadata with a different value.

In effect, we will have the same behaviour as Attribute::merge, where we avoid changing the attribute value when the call already has the same attribute with a value (see test6 in D76140). Will add analogous test case here.

[nikic]

For the record, this causes a 0.6-0.8% compile-time regression on the tramp3d-v4 benchmark (there is also a 0.03% code-size improvement, so at least something is happening there codegen-wise). Can't say whether this is because adding the metadata itself has some cost, or because the presence of the metadata makes further optimization more expensive.

[hjyamauchi]

Hi, some of our internal tests seems to have started crashing due to this patch.

Here's a bug report with a reduced test: https://bugs.llvm.org/show_bug.cgi?id=45590

Please take a look.

[anna]

In D76792#1989623, @yamauchi wrote:

Hi, some of our internal tests seems to have started crashing due to this patch.

Here's a bug report with a reduced test: https://bugs.llvm.org/show_bug.cgi?id=45590

Please take a look.

thanks for the test case, I've reverted the patch for now. Will take a look.

llvm/lib/Transforms/Utils/InlineFunction.cpp

 PreserveAlignmentAssumptions("preserve-alignment-assumptions-during-inlining",
   cl::init(true), cl::Hidden,
   cl::desc("Convert align attributes to assumptions during inlining."));
 
 static cl::opt<bool> UpdateReturnAttributes(
         "update-return-attrs", cl::init(true), cl::Hidden,
             cl::desc("Update return attributes on calls within inlined body"));
 
+static cl::opt<bool> UpdateLoadMetadataDuringInlining(
+        "update-load-metadata-during-inlining", cl::init(true), cl::Hidden,
+            cl::desc("Update metadata on loads within inlined body"));
+
 static cl::opt<unsigned> InlinerAttributeWindow(
     "max-inst-checked-for-throw-during-inlining", cl::Hidden,
     cl::desc("the maximum number of instructions analyzed for may throw during "
              "attribute inference in inlined body"),
     cl::init(4));
 
 llvm::InlineResult llvm::InlineFunction(CallBase *CB, InlineFunctionInfo &IFI,
                                         AAResults *CalleeAAR,
static AttrBuilder IdentifyValidAttributes(CallSite CS) {
   if (AB.contains(Attribute::NoAlias))
     Valid.addAttribute(Attribute::NoAlias);
   if (AB.contains(Attribute::NonNull))
     Valid.addAttribute(Attribute::NonNull);
   return Valid;
 }
 
 static void AddReturnAttributes(CallSite CS, ValueToValueMapTy &VMap) {
-  if (!UpdateReturnAttributes)
+  if (!UpdateReturnAttributes && !UpdateLoadMetadataDuringInlining)
     return;
 
   AttrBuilder Valid = IdentifyValidAttributes(CS);
   if (Valid.empty())
     return;
   auto *CalledFunction = CS.getCalledFunction();
   auto &Context = CalledFunction->getContext();
 
+  auto getExpectedRV = [&](Value *V) -> Instruction * {
+    if (UpdateReturnAttributes && isa<CallBase>(V))
+      return dyn_cast_or_null<CallBase>(VMap.lookup(V));
+    if (UpdateLoadMetadataDuringInlining && isa<LoadInst>(V))
+      return dyn_cast_or_null<LoadInst>(VMap.lookup(V));
+    return nullptr;
+  };
+
+ MDBuilder MDB(Context);
+  auto CreateMDNode = [&](uint64_t Num) -> MDNode * {
+    auto *Int = ConstantInt::get(Type::getInt64Ty(Context), Num);
+    return MDNode::get(Context, MDB.createConstant(Int));
+  };
+
   for (auto &BB : *CalledFunction) {
     auto *RI = dyn_cast<ReturnInst>(BB.getTerminator());
-    if (!RI || !isa<CallBase>(RI->getOperand(0)))
+    if (!RI)
       continue;
-    auto *RetVal = cast<CallBase>(RI->getOperand(0));
     // Sanity check that the cloned RetVal exists and is a call, otherwise we
     // cannot add the attributes on the cloned RetVal.
     // Simplification during inlining could have transformed the cloned
     // instruction.
-    auto *NewRetVal = dyn_cast_or_null<CallBase>(VMap.lookup(RetVal));
+    auto *NewRetVal = getExpectedRV(RI->getOperand(0));
     if (!NewRetVal)
       continue;
+    auto *RetVal = cast<Instruction>(RI->getOperand(0));
     // Backward propagation of attributes to the returned value may be incorrect
     // if it is control flow dependent.
     // Consider:
     // @callee {
     //  %rv = call @foo()
     //  %rv2 = call @bar()
     //  if (%rv2 != null)
     //    return %rv2
if (RI->getParent() != RetVal->getParent() ||
         MayContainThrowingOrExitingCall(RetVal, RI))
       continue;
     // Add to the existing attributes of NewRetVal, i.e. the cloned call
     // instruction.
     // NB! When we have the same attribute already existing on NewRetVal, but
     // with a differing value, the AttributeList's merge API honours the already
     // existing attribute value (i.e. attributes such as dereferenceable,
     // dereferenceable_or_null etc). See AttrBuilder::merge for more details.
-    AttributeList AL = NewRetVal->getAttributes();
+    if (auto *CB = dyn_cast<CallBase>(NewRetVal)) {
+      AttributeList AL = CB->getAttributes();
       AttributeList NewAL =
           AL.addAttributes(Context, AttributeList::ReturnIndex, Valid);
-    NewRetVal->setAttributes(NewAL);
+      CB->setAttributes(NewAL);
+    } else {
+      auto *NewLI = cast<LoadInst>(NewRetVal);
+      if (CS.isReturnNonNull())
+        NewLI->setMetadata(LLVMContext::MD_nonnull, CreateMDNode(1));
+      // If the load already has a dereferenceable/dereferenceable_or_null
+      // metadata, we should honour it.
+      if (uint64_t DerefBytes = Valid.getDereferenceableBytes())
+       if(!NewLI->getMetadata(LLVMContext::MD_dereferenceable))
+         NewLI->setMetadata(LLVMContext::MD_dereferenceable,
+                            CreateMDNode(DerefBytes));
+      if (uint64_t DerefOrNullBytes = Valid.getDereferenceableOrNullBytes())
+       if (!NewLI->getMetadata(LLVMContext::MD_dereferenceable_or_null))
+         NewLI->setMetadata(LLVMContext::MD_dereferenceable_or_null,
+                            CreateMDNode(DerefOrNullBytes));
+    }
+
   }
 }
 
 /// If the inlined function has non-byval align arguments, then
 /// add @llvm.assume-based alignment assumptions to preserve this information.
 static void AddAlignmentAssumptions(CallSite CS, InlineFunctionInfo &IFI) {
   if (!PreserveAlignmentAssumptions || !IFI.GetAssumptionCache)
     return;

llvm/test/Transforms/Inline/ret_load_metadata.ll

+; RUN: opt < %s -inline-threshold=0 -update-load-metadata-during-inlining=true -always-inline -S | FileCheck %s
+; RUN: opt < %s -passes=always-inline -update-load-metadata-during-inlining=true -S | FileCheck %s
+
+
+define i8* @callee(i8** %p) alwaysinline {
+; CHECK: @callee(
+; CHECK-NOT: nonnull
+  %r = load i8*, i8** %p, align 8
+  ret i8* %r
+}
+
+define i8* @test(i8** %ptr, i64 %x) {
+; CHECK-LABEL: @test
+; CHECK: load i8*, i8** %gep, align 8, !nonnull !0
+  %gep = getelementptr inbounds i8*, i8** %ptr, i64 %x
+  %p = call nonnull i8* @callee(i8** %gep)
+  ret i8* %p
+}
+
+declare void @does_not_return(i8*) nounwind
+define internal i8* @callee_negative(i8** %p) alwaysinline {
+; CHECK-NOT: @callee_negative(
+  %r = load i8*, i8** %p, align 8
+  call void @does_not_return(i8* %r)
+  ret i8* %r
+}
+
+define i8* @negative_test(i8** %ptr, i64 %x) {
+; CHECK-LABEL: @negative_test
+; CHECK: load i8*, i8** %gep, align 8
+; CHECK-NOT: nonnull
+  %gep = getelementptr inbounds i8*, i8** %ptr, i64 %x
+  %p = call nonnull i8* @callee_negative(i8** %gep)
+  ret i8* %p
+}
+
+
+define internal i8* @callee2(i8** %p) alwaysinline {
+; CHECK-NOT: @callee2(
+  %r = load i8*, i8** %p, align 8
+  ret i8* %r
+}
+
+; dereferenceable attribute in default addrspace implies nonnull
+define i8* @test2(i8** %ptr, i64 %x) {
+; CHECK-LABEL: @test2
+; CHECK: load i8*, i8** %gep, align 8, !nonnull !0, !dereferenceable !1
+  %gep = getelementptr inbounds i8*, i8** %ptr, i64 %x
+  %p = call dereferenceable(12) i8* @callee(i8** %gep)
+  ret i8* %p
+}
+
+declare void @bar(i8 addrspace(1)*) argmemonly nounwind
+
+define internal i8 addrspace(1)* @callee3(i8 addrspace(1)* addrspace(1)* %p) alwaysinline {
+  %r = load i8 addrspace(1)*, i8 addrspace(1)* addrspace(1)* %p, align 8
+  call void @bar(i8 addrspace(1)* %r)
+  ret i8 addrspace(1)* %r
+}
+
+; This load in callee already has a dereferenceable_or_null metadata. We should
+; honour it.
+define internal i8 addrspace(1)* @callee5(i8 addrspace(1)* addrspace(1)* %p) alwaysinline {
+  %r = load i8 addrspace(1)*, i8 addrspace(1)* addrspace(1)* %p, align 8, !dereferenceable_or_null !3
+  call void @bar(i8 addrspace(1)* %r)
+  ret i8 addrspace(1)* %r
+}
+
+define i8 addrspace(1)* @test3(i8 addrspace(1)* addrspace(1)* %ptr, i64 %x) {
+; CHECK-LABEL: @test3
+; CHECK: load i8 addrspace(1)*, i8 addrspace(1)* addrspace(1)* %gep, align 8, !dereferenceable_or_null !2
+; CHECK: load i8 addrspace(1)*, i8 addrspace(1)* addrspace(1)* %ptr, align 8, !dereferenceable_or_null !3
+  %gep = getelementptr inbounds i8 addrspace(1)*, i8 addrspace(1)* addrspace(1)* %ptr, i64 %x
+  %p = call dereferenceable_or_null(16) i8 addrspace(1)* @callee3(i8 addrspace(1)* addrspace(1)* %gep)
+  %q = call dereferenceable_or_null(20) i8 addrspace(1)* @callee5(i8 addrspace(1)* addrspace(1)* %ptr)
+  ret i8 addrspace(1)* %p
+}
+
+; attribute is part of the callee itself
+define nonnull i8* @callee4(i8** %p) alwaysinline {
+  %r = load i8*, i8** %p, align 8
+  ret i8* %r
+}
+
+; TODO : We should infer the attribute on the callee
+; and add the nonnull on the load
+define i8* @test4(i8** %ptr, i64 %x) {
+; CHECK-LABEL: @test4
+; CHECK: load i8*, i8** %gep, align 8
+; CHECK-NOT: nonnull
+  %gep = getelementptr inbounds i8*, i8** %ptr, i64 %x
+  %p = call i8* @callee(i8** %gep)
+  ret i8* %p
+}
+
+!0 = !{i64 1}
+!1 = !{i64 12}
+!2 = !{i64 16}
+!3 = !{i64 24}
+; CHECK: !0 = !{i64 1}
+; CHECK: !1 = !{i64 12}
+; CHECK: !2 = !{i64 16}
+; CHECK: !3 = !{i64 24}