[FunctionAttr] Infer nonnull attributes on returns

[reames]

Summary

Teach FunctionAttr to infer the nonnull attribute on return values of functions which never return a potentially null value. This is done both via a conservative local analysis for the function itself and a optimistic per-SCC analysis. If no function in the SCC returns anything which could be null (other than values from other functions in the SCC), we can conclude no function returned a null pointer. Even if some function within the SCC returns a null pointer, we may be able to locally conclude that some don't.

This is inspired by discussion on http://reviews.llvm.org/D9132, but is mostly orthogonal to that patch except that they both relate to nonnul attributes.

[pete]

Hi Philip

I made a few comments, but mostly just about using foreach and some comments. Nothing major.

Otherwise, LGTM.

Cheers,
Pete

lib/Transforms/IPO/FunctionAttrs.cpp

73 ↗

(On Diff #25527)

I think you need \brief here

77 ↗

(On Diff #25527)

And here, with \\\ and without listing the function name.

849 ↗

(On Diff #25527)

All function types are pointer types. I think you need F->getReturnType()->isPointerTy() as you have below.

854 ↗

(On Diff #25527)

I think this can just be a foreach loop.

858 ↗

(On Diff #25527)

I think this would be better as a worklist you pop values off until its empty.

866 ↗

(On Diff #25527)

I think this can be a dyn_cast so that you can fold it with the cast<> below.

887 ↗

(On Diff #25527)

I don't know if there's a foreach on the phi values. If there is then good to use it, otherwise I guess someone else can add it later.

[nlewycky]

lib/Transforms/IPO/FunctionAttrs.cpp

858 ↗

(On Diff #25527)

Be careful not to break the case of resolving a self recursive phi. I think the normal way to handle this is with a set of "visited" elements, plus a worklist "queue" like Pete suggests.

920 ↗

(On Diff #25527)

"Speculative assume" -> "Speculatively assume" or possibly just "Speculate"

test/Transforms/FunctionAttrs/nonnull.ll

61 ↗

(On Diff #25527)

Extra newline.

[reames]

Responding to comments. Any comment not explicitly addressed with be fixed before submission.

lib/Transforms/IPO/FunctionAttrs.cpp

849 ↗

(On Diff #25527)

Thanks. Good catch.

854 ↗

(On Diff #25527)

Yep, I copied this from code above. If you don't mind, I'll leave this as is, then fix both in one patch.

858 ↗

(On Diff #25527)

As Nick pointed out, the use of the SetVector is important here. Given this is how other algorithms in the same file are structured, I'd prefer not to change this.

887 ↗

(On Diff #25527)

I copied this from code above. If you don't mind, I'll leave this as is, then fix both in one patch.

[pete]

SGTM. I have no problems with a separate patch. Thanks!

llvm/trunk/lib/Transforms/IPO/FunctionAttrs.cpp

 #include "llvm/ADT/SCCIterator.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/CaptureTracking.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "functionattrs"
 
 STATISTIC(NumReadNone, "Number of functions marked readnone");
 STATISTIC(NumReadOnly, "Number of functions marked readonly");
 STATISTIC(NumNoCapture, "Number of arguments marked nocapture");
 STATISTIC(NumReadNoneArg, "Number of arguments marked readnone");
 STATISTIC(NumReadOnlyArg, "Number of arguments marked readonly");
 STATISTIC(NumNoAlias, "Number of function returns marked noalias");
+STATISTIC(NumNonNullReturn, "Number of function returns marked nonnull");
 STATISTIC(NumAnnotated, "Number of attributes added to library functions");
 
 namespace {
   struct FunctionAttrs : public CallGraphSCCPass {
     static char ID; // Pass identification, replacement for typeid
     FunctionAttrs() : CallGraphSCCPass(ID), AA(nullptr) {
       initializeFunctionAttrsPass(*PassRegistry::getPassRegistry());
     }
struct FunctionAttrs : public CallGraphSCCPass {
 
     // IsFunctionMallocLike - Does this function allocate new memory?
     bool IsFunctionMallocLike(Function *F,
                               SmallPtrSet<Function*, 8> &) const;
 
     // AddNoAliasAttrs - Deduce noalias attributes for the SCC.
     bool AddNoAliasAttrs(const CallGraphSCC &SCC);
 
+    /// \brief Does this function return null?  
+    bool ReturnsNonNull(Function *F, SmallPtrSet<Function*, 8> &,
+                        bool &Speculative) const;
+
+    /// \brief Deduce nonnull attributes for the SCC.
+    bool AddNonNullAttrs(const CallGraphSCC &SCC);
+
     // Utility methods used by inferPrototypeAttributes to add attributes
     // and maintain annotation statistics.
 
     void setDoesNotAccessMemory(Function &F) {
       if (!F.doesNotAccessMemory()) {
         F.setDoesNotAccessMemory();
         ++NumAnnotated;
       }
for (CallGraphSCC::iterator I = SCC.begin(), E = SCC.end(); I != E; ++I) {
     F->setDoesNotAlias(0);
     ++NumNoAlias;
     MadeChange = true;
   }
 
   return MadeChange;
 }
 
+bool FunctionAttrs::ReturnsNonNull(Function *F,
+                                   SmallPtrSet<Function*, 8> &SCCNodes,
+                                   bool &Speculative) const {
+  assert(F->getReturnType()->isPointerTy() &&
+         "nonnull only meaningful on pointer types");
+  Speculative = false;
+  
+  SmallSetVector<Value *, 8> FlowsToReturn;
+  for (BasicBlock &BB : *F)
+    if (auto *Ret = dyn_cast<ReturnInst>(BB.getTerminator()))
+      FlowsToReturn.insert(Ret->getReturnValue());
+
+  for (unsigned i = 0; i != FlowsToReturn.size(); ++i) {
+    Value *RetVal = FlowsToReturn[i];
+
+    // If this value is locally known to be non-null, we're good
+    if (isKnownNonNull(RetVal, TLI))
+      continue;
+
+    // Otherwise, we need to look upwards since we can't make any local
+    // conclusions.  
+    Instruction *RVI = dyn_cast<Instruction>(RetVal);
+    if (!RVI)
+      return false;
+    switch (RVI->getOpcode()) {
+      // Extend the analysis by looking upwards.
+    case Instruction::BitCast:
+    case Instruction::GetElementPtr:
+    case Instruction::AddrSpaceCast:
+      FlowsToReturn.insert(RVI->getOperand(0));
+      continue;
+    case Instruction::Select: {
+      SelectInst *SI = cast<SelectInst>(RVI);
+      FlowsToReturn.insert(SI->getTrueValue());
+      FlowsToReturn.insert(SI->getFalseValue());
+      continue;
+    }
+    case Instruction::PHI: {
+      PHINode *PN = cast<PHINode>(RVI);
+      for (int i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
+        FlowsToReturn.insert(PN->getIncomingValue(i));
+      continue;
+    }
+    case Instruction::Call:
+    case Instruction::Invoke: {
+      CallSite CS(RVI);
+      Function *Callee = CS.getCalledFunction();
+      // A call to a node within the SCC is assumed to return null until
+      // proven otherwise
+      if (Callee && SCCNodes.count(Callee)) {
+        Speculative = true;
+        continue;
+      }
+      return false;
+    }
+    default:
+      return false;  // Unknown source, may be null
+    };
+    llvm_unreachable("should have either continued or returned");
+  }
+
+  return true;
+}
+
+bool FunctionAttrs::AddNonNullAttrs(const CallGraphSCC &SCC) {
+  SmallPtrSet<Function*, 8> SCCNodes;
+
+  // Fill SCCNodes with the elements of the SCC.  Used for quickly
+  // looking up whether a given CallGraphNode is in this SCC.
+  for (CallGraphSCC::iterator I = SCC.begin(), E = SCC.end(); I != E; ++I)
+    SCCNodes.insert((*I)->getFunction());
+
+  // Speculative that all functions in the SCC return only nonnull
+  // pointers.  We may refute this as we analyze functions.
+  bool SCCReturnsNonNull = true;
+
+  bool MadeChange = false;
+
+  // Check each function in turn, determining which functions return nonnull
+  // pointers.
+  for (CallGraphSCC::iterator I = SCC.begin(), E = SCC.end(); I != E; ++I) {
+    Function *F = (*I)->getFunction();
+
+    if (!F || F->hasFnAttribute(Attribute::OptimizeNone))
+      // External node or node we don't want to optimize - skip it;
+      return false;
+
+    // Already nonnull.
+    if (F->getAttributes().hasAttribute(AttributeSet::ReturnIndex,
+                                        Attribute::NonNull))
+      continue;
+
+    // Definitions with weak linkage may be overridden at linktime, so
+    // treat them like declarations.
+    if (F->isDeclaration() || F->mayBeOverridden())
+      return false;
+
+    // We annotate nonnull return values, which are only applicable to 
+    // pointer types.
+    if (!F->getReturnType()->isPointerTy())
+      continue;
+
+    bool Speculative = false;
+    if (ReturnsNonNull(F, SCCNodes, Speculative)) {
+      if (!Speculative) {
+        // Mark the function eagerly since we may discover a function
+        // which prevents us from speculating about the entire SCC
+        DEBUG(dbgs() << "Eagerly marking " << F->getName() << " as nonnull\n");
+        F->addAttribute(AttributeSet::ReturnIndex, Attribute::NonNull);
+        ++NumNonNullReturn;
+        MadeChange = true;
+      }
+      continue;
+    }
+    // At least one function returns something which could be null, can't
+    // speculate any more.
+    SCCReturnsNonNull = false;
+  }
+
+  if (SCCReturnsNonNull) {
+    for (CallGraphSCC::iterator I = SCC.begin(), E = SCC.end(); I != E; ++I) {
+      Function *F = (*I)->getFunction();
+      if (F->getAttributes().hasAttribute(AttributeSet::ReturnIndex,
+                                          Attribute::NonNull) ||
+          !F->getReturnType()->isPointerTy())
+        continue;
+
+      DEBUG(dbgs() << "SCC marking " << F->getName() << " as nonnull\n");
+      F->addAttribute(AttributeSet::ReturnIndex, Attribute::NonNull);
+      ++NumNonNullReturn;
+      MadeChange = true;
+    }
+  }
+
+  return MadeChange;
+}
+
 /// inferPrototypeAttributes - Analyze the name and prototype of the
 /// given function and set any applicable attributes.  Returns true
 /// if any attributes were set and false otherwise.
 bool FunctionAttrs::inferPrototypeAttributes(Function &F) {
   if (F.hasFnAttribute(Attribute::OptimizeNone))
     return false;
 
   FunctionType *FTy = F.getFunctionType();

 bool FunctionAttrs::runOnSCC(CallGraphSCC &SCC) {
   AA = &getAnalysis<AliasAnalysis>();
   TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
 
   bool Changed = annotateLibraryCalls(SCC);
   Changed |= AddReadAttrs(SCC);
   Changed |= AddArgumentAttrs(SCC);
   Changed |= AddNoAliasAttrs(SCC);
+  Changed |= AddNonNullAttrs(SCC);
   return Changed;
 }

llvm/trunk/test/Transforms/FunctionAttrs/nonnull.ll

+; RUN: opt -S -functionattrs %s | FileCheck %s
+declare nonnull i8* @ret_nonnull()
+
+; Return a pointer trivially nonnull (call return attribute)
+define i8* @test1() {
+; CHECK: define nonnull i8* @test1
+  %ret = call i8* @ret_nonnull()
+  ret i8* %ret
+}
+
+; Return a pointer trivially nonnull (argument attribute)
+define i8* @test2(i8* nonnull %p) {
+; CHECK: define nonnull i8* @test2
+  ret i8* %p
+}
+
+; Given an SCC where one of the functions can not be marked nonnull,
+; can we still mark the other one which is trivially nonnull
+define i8* @scc_binder() {
+; CHECK: define i8* @scc_binder
+  call i8* @test3()
+  ret i8* null
+}
+
+define i8* @test3() {
+; CHECK: define nonnull i8* @test3
+  call i8* @scc_binder()
+  %ret = call i8* @ret_nonnull()
+  ret i8* %ret
+}
+
+; Given a mutual recursive set of functions, we can mark them
+; nonnull if neither can ever return null.  (In this case, they
+; just never return period.)
+define i8* @test4_helper() {
+; CHECK: define noalias nonnull i8* @test4_helper
+  %ret = call i8* @test4()
+  ret i8* %ret
+}
+
+define i8* @test4() {
+; CHECK: define noalias nonnull i8* @test4
+  %ret = call i8* @test4_helper()
+  ret i8* %ret
+}
+
+; Given a mutual recursive set of functions which *can* return null
+; make sure we haven't marked them as nonnull.
+define i8* @test5_helper() {
+; CHECK: define noalias i8* @test5_helper
+  %ret = call i8* @test5()
+  ret i8* null
+}
+
+define i8* @test5() {
+; CHECK: define noalias i8* @test5
+  %ret = call i8* @test5_helper()
+  ret i8* %ret
+}
+
+; Local analysis, but going through a self recursive phi
+define i8* @test6() {
+entry:
+; CHECK: define nonnull i8* @test6
+  %ret = call i8* @ret_nonnull()
+  br label %loop
+loop:
+  %phi = phi i8* [%ret, %entry], [%phi, %loop]
+  br i1 undef, label %loop, label %exit
+exit:
+  ret i8* %phi
+}
+
+