[Attributor] Make AANonNull perform context sensitive queries

[jdoerfert]

Summary

In order to get context sensitivity from isKnownNonZero we need to
provide a context instruction *and* a dominator tree. The latter is
passed now to which actually allows to remove some initialization code.

Tests taken from PR43833.

[uenoku]

llvm/test/Transforms/FunctionAttrs/nonnull.ll

598

Please add FIXME.

llvm/include/llvm/Transforms/IPO/Attributor.h

TargetLibraryInfo *getTargetLibraryInfoForFunction(const Function &F) {
     return AG.getAnalysis<TargetLibraryAnalysis>(F);
   }
 
   /// Return AliasAnalysis Result for function \p F.
   AAResults *getAAResultsForFunction(const Function &F) {
     return AG.getAnalysis<AAManager>(F);
   }
 
+  /// Return the analysis result from a pass \p AP for function \p F.
+  template <typename AP>
+  typename AP::Result *getAnalysisResultForFunction(const Function &F) {
+    return AG.getAnalysis<AP>(F);
+  }
+
   /// Return SCC size on call graph for function \p F.
   unsigned getSccSize(const Function &F) {
     if (!SccSizeOpt.hasValue())
       return 0;
     return (SccSizeOpt.getValue())[&F];
   }
 
   /// Return datalayout used in the module.

llvm/lib/Transforms/IPO/Attributor.cpp

 };
 
 /// NonNull attribute for a floating value.
 struct AANonNullFloating
     : AAFromMustBeExecutedContext<AANonNull, AANonNullImpl> {
   using Base = AAFromMustBeExecutedContext<AANonNull, AANonNullImpl>;
   AANonNullFloating(const IRPosition &IRP) : Base(IRP) {}
 
-  /// See AbstractAttribute::initialize(...).
-  void initialize(Attributor &A) override {
-    Base::initialize(A);
-
-    if (isAtFixpoint())
-      return;
-
-    const IRPosition &IRP = getIRPosition();
-    const Value &V = IRP.getAssociatedValue();
-    const DataLayout &DL = A.getDataLayout();
-
-    // TODO: This context sensitive query should be removed once we can do
-    // context sensitive queries in the genericValueTraversal below.
-    if (isKnownNonZero(&V, DL, 0, /* TODO: AC */ nullptr, IRP.getCtxI(),
-                       /* TODO: DT */ nullptr))
-      indicateOptimisticFixpoint();
-  }
-
   /// See AbstractAttribute::updateImpl(...).
   ChangeStatus updateImpl(Attributor &A) override {
     ChangeStatus Change = Base::updateImpl(A);
     if (isKnownNonNull())
       return Change;
 
     if (!NullIsDefined) {
       const auto &DerefAA = A.getAAFor<AADereferenceable>(*this, getIRPosition());
       if (DerefAA.getAssumedDereferenceableBytes())
         return Change;
     }
 
     const DataLayout &DL = A.getDataLayout();
 
+    DominatorTree *DT = nullptr;
+    InformationCache &InfoCache = A.getInfoCache();
+    if (const Function *Fn = getAnchorScope())
+      DT = InfoCache.getAnalysisResultForFunction<DominatorTreeAnalysis>(*Fn);
+
     auto VisitValueCB = [&](Value &V, AANonNull::StateType &T,
                             bool Stripped) -> bool {
       const auto &AA = A.getAAFor<AANonNull>(*this, IRPosition::value(V));
       if (!Stripped && this == &AA) {
-        if (!isKnownNonZero(&V, DL, 0, /* TODO: AC */ nullptr,
-                            /* CtxI */ getCtxI(),
-                            /* TODO: DT */ nullptr))
+        if (!isKnownNonZero(&V, DL, 0, /* TODO: AC */ nullptr, getCtxI(), DT))
           T.indicatePessimisticFixpoint();
       } else {
         // Use abstract attribute information.
         const AANonNull::StateType &NS =
             static_cast<const AANonNull::StateType &>(AA.getState());
         T ^= NS;
       }
       return T.isValidState();

llvm/test/Transforms/FunctionAttrs/nonnull.ll

-; RUN: opt -S -functionattrs -enable-nonnull-arg-prop %s | FileCheck %s --check-prefixes=BOTH,FNATTR
-; RUN: opt -S -passes=function-attrs -enable-nonnull-arg-prop %s | FileCheck %s --check-prefixes=BOTH,FNATTR
-; RUN: opt -attributor --attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=8 -S < %s | FileCheck %s --check-prefixes=BOTH,ATTRIBUTOR
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -functionattrs -enable-nonnull-arg-prop %s | FileCheck %s --check-prefixes=BOTH,FNATTR,OLD
+; RUN: opt -S -passes=function-attrs -enable-nonnull-arg-prop %s | FileCheck %s --check-prefixes=BOTH,FNATTR,OLD
+; RUN: opt -attributor --attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=8 -S < %s | FileCheck %s --check-prefixes=BOTH,OLD,ATTRIBUTOR,ATTRIBUTOR_OPM
+; RUN: opt -passes=attributor --attributor-disable=false -attributor-max-iterations-verify -attributor-max-iterations=8 -S < %s | FileCheck %s --check-prefixes=BOTH,ATTRIBUTOR,ATTRIBUTOR_NPM
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 
 declare nonnull i8* @ret_nonnull()
 
 ; Return a pointer trivially nonnull (call return attribute)
 define i8* @test1() {
 ; BOTH: define nonnull i8* @test1

 define void @test13_helper() {
   %nonnullptr = tail call i8* @ret_nonnull()
   %maybenullptr = tail call i8* @unknown()
   tail call void @test13(i8* %nonnullptr, i8* %nonnullptr, i8* %maybenullptr)
   tail call void @test13(i8* %nonnullptr, i8* %maybenullptr, i8* %nonnullptr)
   ret void
 }
 define internal void @test13(i8* %a, i8* %b, i8* %c) {
-; ATTRIBUTOR: define internal void @test13(i8* nocapture nonnull readnone %a, i8* nocapture readnone %b, i8* nocapture readnone %c) 
+; ATTRIBUTOR: define internal void @test13(i8* nocapture nonnull readnone %a, i8* nocapture readnone %b, i8* nocapture readnone %c)
   ret void
 }
 
 declare nonnull i8* @nonnull()
 
 ; TEST 14
 ; Complex propagation
 ; Argument of f1, f2, f3 can be marked with nonnull.
 
 ; * Argument
 ; 1. In f1:bb6, %arg can be marked with nonnull because of the comparison in bb1
 ; 2. Because f2 is internal function, f2(i32* %arg) -> @f2(i32* nonnull %arg)
-; 3. In f1:bb4 %tmp5 is nonnull and f3 is internal function. 
+; 3. In f1:bb4 %tmp5 is nonnull and f3 is internal function.
 ;    Then, f3(i32* %arg) -> @f3(i32* nonnull %arg)
 ; 4. We get nonnull in whole f1 call sites so f1(i32* %arg) -> @f1(i32* nonnull %arg)
 
 
 define internal i32* @f1(i32* %arg) {
 ; FIXME: missing nonnull It should be nonnull @f1(i32* nonnull readonly %arg)
 ; ATTRIBUTOR: define internal nonnull i32* @f1(i32* readonly %arg)
 
; ATTRIBUTOR: %tmp7 = tail call nonnull i32* @f2(i32* readonly %arg)
   ret i32* %tmp7
 
 bb9:                                              ; preds = %bb4, %bb
   %tmp10 = phi i32* [ %tmp5, %bb4 ], [ inttoptr (i64 4 to i32*), %bb ]
   ret i32* %tmp10
 }
 
 define internal i32* @f2(i32* %arg) {
-; FIXME: missing nonnull. It should be nonnull @f2(i32* nonnull %arg) 
+; FIXME: missing nonnull. It should be nonnull @f2(i32* nonnull %arg)
 ; ATTRIBUTOR: define internal nonnull i32* @f2(i32* readonly %arg)
 bb:
 
-; FIXME: missing nonnull. It should be @f1(i32* nonnull readonly %arg) 
+; FIXME: missing nonnull. It should be @f1(i32* nonnull readonly %arg)
 ; ATTRIBUTOR:   %tmp = tail call nonnull i32* @f1(i32* readonly %arg)
   %tmp = tail call i32* @f1(i32* %arg)
   ret i32* %tmp
 }
 
 define dso_local noalias i32* @f3(i32* %arg) {
-; FIXME: missing nonnull. It should be nonnull @f3(i32* nonnull readonly %arg) 
+; FIXME: missing nonnull. It should be nonnull @f3(i32* nonnull readonly %arg)
 ; ATTRIBUTOR: define dso_local noalias i32* @f3(i32* nocapture readonly %arg)
 bb:
-; FIXME: missing nonnull. It should be @f1(i32* nonnull readonly %arg) 
+; FIXME: missing nonnull. It should be @f1(i32* nonnull readonly %arg)
 ; ATTRIBUTOR:   %tmp = call i32* @f1(i32* readonly %arg)
   %tmp = call i32* @f1(i32* %arg)
   ret i32* null
 }
 
 ; TEST 15
 define void @f15(i8* %arg) {
 ; ATTRIBUTOR:   tail call void @use1(i8* nonnull dereferenceable(4) %arg)
; BOTH-NEXT: ret i8 [[RET]]
   ret i8 %ret
 }
 
 ; Make sure that an invoke works similarly to a call.
 
 declare i32 @esfp(...)
 
 define i1 @parent8(i8* %a, i8* %bogus1, i8* %b) personality i8* bitcast (i32 (...)* @esfp to i8*){
-; BOTH-LABEL: @parent8(i8* nonnull %a, i8* nocapture readnone %bogus1, i8* nonnull %b) 
+; BOTH-LABEL: @parent8(i8* nonnull %a, i8* nocapture readnone %bogus1, i8* nonnull %b)
 ; BOTH-NEXT:  entry:
 ; FNATTR-NEXT:    invoke void @use2nonnull(i8* %a, i8* %b)
 ; ATTRIBUTOR-NEXT:    invoke void @use2nonnull(i8* nonnull %a, i8* nonnull %b)
 ; BOTH-NEXT:    to label %cont unwind label %exc
 ; BOTH:       cont:
 ; BOTH-NEXT:    [[NULL_CHECK:%.*]] = icmp eq i8* %b, null
 ; BOTH-NEXT:    ret i1 [[NULL_CHECK]]
 ; BOTH:       exc:

 }
 define void @make_live(i32* nonnull dereferenceable(8) %a) {
   call void @naked(i32* nonnull dereferenceable(8) align 16 %a)
   call void @control(i32* nonnull dereferenceable(8) align 16 %a)
   call void @optnone(i32* nonnull dereferenceable(8) align 16 %a)
   ret void
 }
 
+; Original from PR43833
+declare void @sink(i32*)
+define void @PR43833(i32* %0, i32 %1) {
+; BOTH-LABEL: @PR43833(
+; BOTH-NEXT:    [[TMP3:%.*]] = icmp sgt i32 [[TMP1:%.*]], 1
+; BOTH-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP7:%.*]]
+; BOTH:       4:
+; BOTH-NEXT:    [[TMP5:%.*]] = zext i32 [[TMP1]] to i64
+; BOTH-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* [[TMP0:%.*]], i64 [[TMP5]]
+; BOTH-NEXT:    br label [[TMP8:%.*]]
+; BOTH:       7:
+; BOTH-NEXT:    ret void
+; BOTH:       8:
+; BOTH-NEXT:    [[TMP9:%.*]] = phi i32 [ 1, [[TMP4]] ], [ [[TMP10:%.*]], [[TMP8]] ]
+; BOTH-NEXT:    tail call void @sink(i32* [[TMP6]])
+; BOTH-NEXT:    [[TMP10]] = add nuw nsw i32 [[TMP9]], 1
+; BOTH-NEXT:    [[TMP11:%.*]] = icmp eq i32 [[TMP10]], [[TMP1]]
+; BOTH-NEXT:    br i1 [[TMP11]], label [[TMP7]], label [[TMP8]]
+;
+  %3 = icmp sgt i32 %1, 1
+  br i1 %3, label %4, label %7
+
+4:                                                ; preds = %2
+  %5 = zext i32 %1 to i64
+  %6 = getelementptr inbounds i32, i32* %0, i64 %5
+  br label %8
+
+7:                                                ; preds = %8, %2
+  ret void
+
+8:                                                ; preds = %8, %4
+  %9 = phi i32 [ 1, %4 ], [ %10, %8 ]
+  tail call void @sink(i32* %6)
+  %10 = add nuw nsw i32 %9, 1
+  %11 = icmp eq i32 %10, %1
+  br i1 %11, label %7, label %8
+}
+
+; Adjusted from PR43833
+define void @PR43833_simple(i32* %0, i32 %1) {
+; OLD-LABEL: @PR43833_simple(
+; OLD-NEXT:    [[TMP3:%.*]] = icmp ne i32 [[TMP1:%.*]], 0
+; OLD-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP7:%.*]]
+; OLD:       4:
+; OLD-NEXT:    [[TMP5:%.*]] = zext i32 [[TMP1]] to i64
+; OLD-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* [[TMP0:%.*]], i64 [[TMP5]]
+; OLD-NEXT:    br label [[TMP8:%.*]]
+; OLD:       7:
+; OLD-NEXT:    ret void
+; OLD:       8:
+; OLD-NEXT:    [[TMP9:%.*]] = phi i32 [ 1, [[TMP4]] ], [ [[TMP10:%.*]], [[TMP8]] ]
+; OLD-NEXT:    tail call void @sink(i32* [[TMP6]])
+; OLD-NEXT:    [[TMP10]] = add nuw nsw i32 [[TMP9]], 1
+; OLD-NEXT:    [[TMP11:%.*]] = icmp eq i32 [[TMP10]], [[TMP1]]
+; OLD-NEXT:    br i1 [[TMP11]], label [[TMP7]], label [[TMP8]]
+;
+; ATTRIBUTOR_NPM-LABEL: @PR43833_simple(
+; ATTRIBUTOR_NPM-NEXT:    [[TMP3:%.*]] = icmp ne i32 [[TMP1:%.*]], 0
+; ATTRIBUTOR_NPM-NEXT:    br i1 [[TMP3]], label [[TMP4:%.*]], label [[TMP7:%.*]]
+; ATTRIBUTOR_NPM:       4:
+; ATTRIBUTOR_NPM-NEXT:    [[TMP5:%.*]] = zext i32 [[TMP1]] to i64
+; ATTRIBUTOR_NPM-NEXT:    [[TMP6:%.*]] = getelementptr inbounds i32, i32* [[TMP0:%.*]], i64 [[TMP5]]
+; ATTRIBUTOR_NPM-NEXT:    br label [[TMP8:%.*]]
+; ATTRIBUTOR_NPM:       7:
+; ATTRIBUTOR_NPM-NEXT:    ret void
+; ATTRIBUTOR_NPM:       8:
+; ATTRIBUTOR_NPM-NEXT:    [[TMP9:%.*]] = phi i32 [ 1, [[TMP4]] ], [ [[TMP10:%.*]], [[TMP8]] ]
+; ATTRIBUTOR_NPM-NEXT:    tail call void @sink(i32* nonnull [[TMP6]])
+; ATTRIBUTOR_NPM-NEXT:    [[TMP10]] = add nuw nsw i32 [[TMP9]], 1
+; ATTRIBUTOR_NPM-NEXT:    [[TMP11:%.*]] = icmp eq i32 [[TMP10]], [[TMP1]]
+; ATTRIBUTOR_NPM-NEXT:    br i1 [[TMP11]], label [[TMP7]], label [[TMP8]]
+;
+  %3 = icmp ne i32 %1, 0
+  br i1 %3, label %4, label %7
+
+4:                                                ; preds = %2
+  %5 = zext i32 %1 to i64
+  %6 = getelementptr inbounds i32, i32* %0, i64 %5
+  br label %8
+
+7:                                                ; preds = %8, %2
+  ret void
+
+8:                                                ; preds = %8, %4
+  %9 = phi i32 [ 1, %4 ], [ %10, %8 ]
+  tail call void @sink(i32* %6)
+  %10 = add nuw nsw i32 %9, 1
+  %11 = icmp eq i32 %10, %1
+  br i1 %11, label %7, label %8
+}
+
 attributes #0 = { "null-pointer-is-valid"="true" }
 attributes #1 = { nounwind willreturn}