[InlineFunction] add nonnull assumptions based on argument attributes

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.

https://reviews.llvm.org/D29999

llvm-svn: 296366

Author: rotateright

llvm/lib/Transforms/Utils/InlineFunction.cpp

@@ -1093,38 +1093,52 @@ static void AddAliasScopeMetadata(CallSite CS, ValueToValueMapTy &VMap,
   }
 }
 
-/// 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;
+  auto calcDomTreeIfNeeded = [&]() {
+    if (!DTCalculated) {
+      DT.recalculate(*CS.getCaller());
+      DTCalculated = true;
+    }
+  };
 
   Function *CalledFunc = CS.getCalledFunction();
+  IRBuilder<> Builder(CS.getInstruction());
   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;
-      }
+    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);
+      }
     }
   }
 }
@@ -1621,10 +1635,10 @@ bool llvm::InlineFunction(CallSite CS, InlineFunctionInfo &IFI,
       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

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

@@ -12,18 +12,21 @@ define i32 @callee(i32* dereferenceable(32) %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.
 
@@ -36,6 +39,7 @@ define i32 @caller2(i32* dereferenceable(31) %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.