[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

1163–1164

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

1188

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

1163–1164

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

1181

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

1226

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

1236

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

1239

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

 EnableNoAliasConversion("enable-noalias-to-md-conversion", cl::init(true),
   cl::Hidden,
   cl::desc("Convert noalias attributes to metadata 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<bool>
 PreserveAlignmentAssumptions("preserve-alignment-assumptions-during-inlining",
   cl::init(true), cl::Hidden,
   cl::desc("Convert align attributes to assumptions during inlining."));
 
 static cl::opt<unsigned> MaxInstCheckedForThrow(
     "max-inst-checked-for-throw-during-inlining", cl::Hidden,
     cl::desc("the maximum number of instructions analyzed for may throw during "
if (const Instruction *I = dyn_cast<Instruction>(VMI->first)) {
         NI->setMetadata(
             LLVMContext::MD_alias_scope,
             MDNode::concatenate(NI->getMetadata(LLVMContext::MD_alias_scope),
                                 MDNode::get(CalledFunc->getContext(), Scopes)));
     }
   }
 }
 
-static void AddReturnAttributes(CallSite CS, ValueToValueMapTy &VMap) {
-  if (!UpdateReturnAttributes)
+static void AddReturnAttributesAndUpdateMetadata(CallSite CS,
+                                                 ValueToValueMapTy &VMap) {
+  if (!UpdateReturnAttributes && !UpdateLoadMetadataDuringInlining)
     return;
   AttrBuilder AB(CS.getAttributes(), AttributeList::ReturnIndex);
   if (AB.empty())
     return;
 
   auto *CalledFunction = CS.getCalledFunction();
   auto &Context = CalledFunction->getContext();
 
   auto MayContainThrowingOrExitingCall = [&](Instruction *RVal,
                                              Instruction *RInst) {
     unsigned NumInstChecked = 0;
     // Check that all instructions in the range [RVal, Rinst) are
     // guaranteed to transfer execution to successor.
     for (auto &I : make_range(RVal->getIterator(), RInst->getIterator()))
       if (NumInstChecked++ > MaxInstCheckedForThrow ||
           !isGuaranteedToTransferExecutionToSuccessor(&I))
         return true;
     return false;
   };
 
+  auto getExpectedRV = [&](Value *V) -> Instruction * {
+    if (isa<CallBase>(V) && UpdateReturnAttributes)
+      return dyn_cast_or_null<CallBase>(VMap.lookup(V));
+    if (isa<LoadInst>(V) && UpdateLoadMetadataDuringInlining)
+      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;
     // Sanity check that the cloned return instruction exists and is a return
     // instruction itself.
     auto *NewRI = dyn_cast_or_null<ReturnInst>(VMap.lookup(RI));
     if (!NewRI)
       continue;
-    auto *RetVal = cast<CallBase>(RI->getOperand(0));
-    // Sanity check that the cloned RetVal exists and is a call.
+    // Sanity check that the cloned RetVal exists and is a call or load.
     // 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;
     // 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 (%rv == null)
     //    exit()
     //  return %rv
     // }
     // caller() {
     //   %val = call nonnull @callee()
     // }
     // Here we cannot add the nonnull attribute on either foo or bar. So, we
     // limit the check to both NewRetVal and NewRI are in the same basic block
     // and there are no throwing/exiting instructions between these
     // instructions.
     if (NewRI->getParent() != NewRetVal->getParent() ||
         MayContainThrowingOrExitingCall(NewRetVal, NewRI))
       continue;
+    if (isa<CallBase>(NewRetVal)) {
       // Add to the existing attributes.
-    AttributeList AL = NewRetVal->getAttributes();
+      AttributeList AL = cast<CallBase>(NewRetVal)->getAttributes();
       AttributeList NewAL =
           AL.addAttributes(Context, AttributeList::ReturnIndex, AB);
-    NewRetVal->setAttributes(NewAL);
+      cast<CallBase>(NewRetVal)->setAttributes(NewAL);
+    } else {
+      auto *NewLI = cast<LoadInst>(NewRetVal);
+      // We do not need to check that NewLI is of PointerType, since the
+      // following return attributes on callsite only applies to pointer typed
+      // returns.
+      if (CS.isReturnNonNull())
+        NewLI->setMetadata(LLVMContext::MD_nonnull, CreateMDNode(1));
+      if (uint64_t DerefBytes = AB.getDereferenceableBytes())
+        NewLI->setMetadata(LLVMContext::MD_dereferenceable,
+                           CreateMDNode(DerefBytes));
+      if (uint64_t DerefOrNullBytes = AB.getDereferenceableOrNullBytes())
+        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;
Function::iterator FirstNewBlock;
     // Clone existing noalias metadata if necessary.
     CloneAliasScopeMetadata(CS, VMap);
 
     // Add noalias metadata if necessary.
     AddAliasScopeMetadata(CS, VMap, DL, CalleeAAR);
 
     // Clone return attributes on the callsite into the calls within the inlined
     // function which feed into its return value.
-    AddReturnAttributes(CS, VMap);
+    // Add metadata on loads that feed into the return value in the inlined
+    // body, based on return attributes on the callsite.
+    AddReturnAttributesAndUpdateMetadata(CS, VMap);
 
     // Propagate llvm.mem.parallel_loop_access if necessary.
     PropagateParallelLoopAccessMetadata(CS, VMap);
 
     // Register any cloned assumptions.
     if (IFI.GetAssumptionCache)
       for (BasicBlock &NewBlock :
            make_range(FirstNewBlock->getIterator(), Caller->end()))

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 ![[NONNULL:[0-9]+]]
+  %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 ![[NONNULL:[0-9]+]], !dereferenceable ![[DEREF:[0-9]+]]
+  %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 {
+; CHECK-NOT: @callee3(
+  %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
+}
+
+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 ![[DEREF_OR_NULL:[0-9]+]]
+  %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)
+  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
+}
+
+
+; CHECK: ![[NONNULL]] = !{i64 1}
+; CHECK: ![[DEREF]] = !{i64 12}
+; CHECK: ![[DEREF_OR_NULL]] = !{i64 16}