[InlineFunction] add nonnull assumptions based on argument attributes

[spatel]

Summary

This was suggested in D27855: have the inliner add assumptions, so we don't lose nonnull info provided by argument attributes.

This still doesn't solve PR28430 (dyn_cast), but this gets us closer...

From the use of 'dereferenceable' in the test cases, you can guess that I was drafting an alternate patch that would provide a stronger transfer of the attribute from the callee to the caller, but that would be limited in the same way that D27855 was.

Please let me know if this needs to be disabled by default pending clang fixes.

[spatel]

Ping.

[mehdi_amini]

LGTM. Thanks.

lib/Transforms/Utils/InlineFunction.cpp

1108

could be a static variable inside the lambda, that'd reduce the scope.

[spatel]

Committed at:
rL296366

[spatel]

This one got reverted shortly after:
https://reviews.llvm.org/rL296398 / https://bugs.llvm.org/show_bug.cgi?id=32082

I was going to investigate the failure, but I got massively distracted. Reopening this review to decrease my procrastination.

[spatel]

Abandoning - this seems likely to be handled by the Attributor pass. ( cc @jdoerfert @uenoku )

lib/Transforms/Utils/InlineFunction.cpp

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)));
     }
   }
 }
 
-/// 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)
+/// Add @llvm.assume-based assumptions to preserve information supplied by
+/// argument attributes because the attributes will disappear after inlining.
+static void addAssumptions(CallSite CS, InlineFunctionInfo &IFI) {
+  if (!IFI.GetAssumptionCache)
     return;
 
   AssumptionCache *AC = &(*IFI.GetAssumptionCache)(*CS.getCaller());
   auto &DL = CS.getCaller()->getParent()->getDataLayout();
 
-  // To avoid inserting redundant assumptions, we should check for assumptions
-  // already in the caller. To do this, we might need a DT of the caller.
+  // To avoid inserting redundant assumptions, check that an assumption provides
+  // new information in the caller. This might require a dominator tree.
   DominatorTree DT;
   bool DTCalculated = false;
-
-  Function *CalledFunc = CS.getCalledFunction();
-  for (Argument &Arg : CalledFunc->args()) {
-    unsigned Align = Arg.getType()->isPointerTy() ? Arg.getParamAlignment() : 0;
-    if (Align && !Arg.hasByValOrInAllocaAttr() && !Arg.hasNUses(0)) {
-      if (!DTCalculated) {
-        DT.recalculate(*CS.getCaller());
-        DTCalculated = true;
-      }
+  auto calcDomTreeIfNeeded = [&]() {
+  };
 
+  Function *CalledFunc = CS.getCalledFunction();
+  IRBuilder<> Builder(CS.getInstruction());
+  for (Argument &Arg : CalledFunc->args()) {
+    Value *ArgVal = CS.getArgument(Arg.getArgNo());
+
+    unsigned Align = Arg.getType()->isPointerTy() ? Arg.getParamAlignment() : 0;
+    if (PreserveAlignmentAssumptions && Align &&
+        !Arg.hasByValOrInAllocaAttr() && !Arg.hasNUses(0)) {
       // If we can already prove the asserted alignment in the context of the
       // caller, then don't bother inserting the assumption.
-      Value *ArgVal = CS.getArgument(Arg.getArgNo());
-      if (getKnownAlignment(ArgVal, DL, CS.getInstruction(), AC, &DT) >= Align)
-        continue;
+      calcDomTreeIfNeeded();
+      if (getKnownAlignment(ArgVal, DL, CS.getInstruction(), AC, &DT) < Align) {
+        CallInst *Asmp = Builder.CreateAlignmentAssumption(DL, ArgVal, Align);
+        AC->registerAssumption(Asmp);
+      }
+    }
 
-      CallInst *NewAsmp = IRBuilder<>(CS.getInstruction())
-                              .CreateAlignmentAssumption(DL, ArgVal, Align);
-      AC->registerAssumption(NewAsmp);
+    if (Arg.hasNonNullAttr()) {
+      // If we can already prove nonnull in the context of the caller, then
+      // don't bother inserting the assumption.
+      calcDomTreeIfNeeded();
+      if (!isKnownNonNullAt(ArgVal, CS.getInstruction(), &DT)) {
+        Value *NotNull = Builder.CreateIsNotNull(ArgVal);
+        CallInst *Asmp = Builder.CreateAssumption(NotNull);
+        AC->registerAssumption(Asmp);
+      }
     }
   }
 }
 
 /// Once we have cloned code over from a callee into the caller,
 /// update the specified callgraph to reflect the changes we made.
 /// Note that it's possible that not all code was copied over, so only
 /// some edges of the callgraph may remain.
for (Function::const_arg_iterator I = CalledFunc->arg_begin(),
                                         CalledFunc->getParamAlignment(ArgNo+1));
         if (ActualArg != *AI)
           ByValInit.push_back(std::make_pair(ActualArg, (Value*) *AI));
       }
 
       VMap[&*I] = ActualArg;
     }
 
-    // Add alignment assumptions if necessary. We do this before the inlined
-    // instructions are actually cloned into the caller so that we can easily
-    // check what will be known at the start of the inlined code.
-    AddAlignmentAssumptions(CS, IFI);
+    // Add assumptions if necessary. We do this before the inlined instructions
+    // are actually cloned into the caller so that we can easily check what will
+    // be known at the start of the inlined code.
+    addAssumptions(CS, IFI);
 
     // We want the inliner to prune the code as it copies.  We would LOVE to
     // have no dead or constant instructions leftover after inlining occurs
     // (which can happen, e.g., because an argument was constant), but we'll be
     // happy with whatever the cloner can do.
     CloneAndPruneFunctionInto(Caller, CalledFunc, VMap,
                               /*ModuleLevelChanges=*/false, Returns, ".i",
                               &InlinedFunctionInfo, TheCall);

test/Transforms/Inline/arg-attr-propagation.ll

 ; RUN: opt -inline -S < %s | FileCheck %s
 
 ; The callee guarantees that the pointer argument is nonnull and dereferenceable.
 ; That information should transfer to the caller.
 
 define i32 @callee(i32* dereferenceable(32) %t1) {
 ; CHECK-LABEL: @callee(i32* dereferenceable(32) %t1)
 ; CHECK-NEXT:    [[T2:%.*]] = load i32, i32* %t1
 ; CHECK-NEXT:    ret i32 [[T2]]
 ;
   %t2 = load i32, i32* %t1
   ret i32 %t2
 }
 
-; FIXME: All dereferenceability information is lost.
+; Add a nonnull assumption. 
 ; The caller argument could be known nonnull and dereferenceable(32).
 
 define i32 @caller1(i32* %t1) {
 ; CHECK-LABEL: @caller1(i32* %t1)
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i32* %t1, null
+; CHECK-NEXT:    call void @llvm.assume(i1 [[TMP1]])
 ; CHECK-NEXT:    [[T2_I:%.*]] = load i32, i32* %t1
 ; CHECK-NEXT:    ret i32 [[T2_I]]
 ;
   %t2 = tail call i32 @callee(i32* dereferenceable(32) %t1)
   ret i32 %t2
 }
 
+; Don't add a nonnull assumption if it's redundant.
 ; The caller argument is nonnull, but that can be explicit.
 ; The dereferenceable amount could be increased.
 
 define i32 @caller2(i32* dereferenceable(31) %t1) {
 ; CHECK-LABEL: @caller2(i32* dereferenceable(31) %t1)
 ; CHECK-NEXT:    [[T2_I:%.*]] = load i32, i32* %t1
 ; CHECK-NEXT:    ret i32 [[T2_I]]
 ;
   %t2 = tail call i32 @callee(i32* dereferenceable(32) %t1)
   ret i32 %t2
 }
 
+; Don't add a nonnull assumption if it's redundant.
 ; The caller argument is nonnull, but that can be explicit.
 ; Make sure that we don't propagate a smaller dereferenceable amount.
 
 define i32 @caller3(i32* dereferenceable(33) %t1) {
 ; CHECK-LABEL: @caller3(i32* dereferenceable(33) %t1)
 ; CHECK-NEXT:    [[T2_I:%.*]] = load i32, i32* %t1
 ; CHECK-NEXT:    ret i32 [[T2_I]]
 ;
   %t2 = tail call i32 @callee(i32* dereferenceable(32) %t1)
   ret i32 %t2
 }