diff --git a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
--- a/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
+++ b/llvm/lib/Transforms/IPO/AttributorAttributes.cpp
@@ -2019,34 +2019,35 @@
         if (idx >= Callee->arg_size())
           break;
         Value *ArgVal = CB.getArgOperand(idx);
-        if (!ArgVal || !ArgVal->getType()->isPointerTy())
+        if (!ArgVal)
           continue;
+        // Here, we handle three cases.
+        //   (1) Not having a value means it is dead. (we can replace the value
+        //       with undef)
+        //   (2) Simplified to undef. The argument violate noundef attriubte.
+        //   (3) Simplified to null pointer where known to be nonnull.
+        //       The argument is a poison value and violate noundef attribute.
         IRPosition CalleeArgumentIRP = IRPosition::callsite_argument(CB, idx);
         if (!CalleeArgumentIRP.hasAttr({Attribute::NoUndef}))
           continue;
-        auto &NonNullAA = A.getAAFor<AANonNull>(*this, CalleeArgumentIRP,
-                                                /* TrackDependence */ false);
-        if (!NonNullAA.isKnownNonNull())
-          continue;
-        const auto &ValueSimplifyAA = A.getAAFor<AAValueSimplify>(
+        auto &ValueSimplifyAA = A.getAAFor<AAValueSimplify>(
             *this, IRPosition::value(*ArgVal), /* TrackDependence */ false);
-        Optional<Value *> SimplifiedVal =
-            ValueSimplifyAA.getAssumedSimplifiedValue(A);
-
         if (!ValueSimplifyAA.isKnown())
           continue;
-        // Here, we handle three cases.
-        //   (1) Not having a value means it is dead. (we can replace the value
-        //       with undef)
-        //   (2) Simplified to null pointer. The argument is a poison value and
-        //       violate noundef attribute.
-        //   (3) Simplified to undef. The argument violate noundef attriubte.
+        Optional<Value *> SimplifiedVal =
+            ValueSimplifyAA.getAssumedSimplifiedValue(A);
         if (!SimplifiedVal.hasValue() ||
-            isa<ConstantPointerNull>(*SimplifiedVal.getValue()) ||
             isa<UndefValue>(*SimplifiedVal.getValue())) {
           KnownUBInsts.insert(&I);
-          return true;
+          continue;
         }
+        if (!ArgVal->getType()->isPointerTy() ||
+            !isa<ConstantPointerNull>(*SimplifiedVal.getValue()))
+          continue;
+        auto &NonNullAA = A.getAAFor<AANonNull>(*this, CalleeArgumentIRP,
+                                                /* TrackDependence */ false);
+        if (NonNullAA.isKnownNonNull())
+          KnownUBInsts.insert(&I);
       }
       return true;
     };
diff --git a/llvm/test/Transforms/Attributor/undefined_behavior.ll b/llvm/test/Transforms/Attributor/undefined_behavior.ll
--- a/llvm/test/Transforms/Attributor/undefined_behavior.ll
+++ b/llvm/test/Transforms/Attributor/undefined_behavior.ll
@@ -1087,3 +1087,61 @@
   call void @callee_ptr_arg(i32* noundef undef)
   ret void
 }
+
+define i32 @argument_noundef1(i32 noundef %c) {
+; IS__TUNIT____: Function Attrs: nofree nosync nounwind readnone willreturn
+; IS__TUNIT____-LABEL: define {{[^@]+}}@argument_noundef1
+; IS__TUNIT____-SAME: (i32 noundef returned [[C:%.*]]) [[ATTR0]] {
+; IS__TUNIT____-NEXT:    ret i32 [[C]]
+;
+; IS__CGSCC____: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
+; IS__CGSCC____-LABEL: define {{[^@]+}}@argument_noundef1
+; IS__CGSCC____-SAME: (i32 noundef returned [[C:%.*]]) [[ATTR0]] {
+; IS__CGSCC____-NEXT:    ret i32 [[C]]
+;
+  ret i32 %c
+}
+
+define i32 @violate_noundef_nonpointer() {
+; IS__TUNIT____: Function Attrs: nofree nosync nounwind readnone willreturn
+; IS__TUNIT____-LABEL: define {{[^@]+}}@violate_noundef_nonpointer
+; IS__TUNIT____-SAME: () [[ATTR0]] {
+; IS__TUNIT____-NEXT:    unreachable
+;
+; IS__CGSCC____: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
+; IS__CGSCC____-LABEL: define {{[^@]+}}@violate_noundef_nonpointer
+; IS__CGSCC____-SAME: () [[ATTR0]] {
+; IS__CGSCC____-NEXT:    unreachable
+;
+  %ret = call i32 @argument_noundef1(i32 undef)
+  ret i32 %ret
+}
+
+define i32* @argument_noundef2(i32* noundef %c) {
+; IS__TUNIT____: Function Attrs: nofree nosync nounwind readnone willreturn
+; IS__TUNIT____-LABEL: define {{[^@]+}}@argument_noundef2
+; IS__TUNIT____-SAME: (i32* nofree noundef readnone returned "no-capture-maybe-returned" [[C:%.*]]) [[ATTR0]] {
+; IS__TUNIT____-NEXT:    ret i32* [[C]]
+;
+; IS__CGSCC____: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
+; IS__CGSCC____-LABEL: define {{[^@]+}}@argument_noundef2
+; IS__CGSCC____-SAME: (i32* nofree noundef readnone returned "no-capture-maybe-returned" [[C:%.*]]) [[ATTR0]] {
+; IS__CGSCC____-NEXT:    ret i32* [[C]]
+;
+  ret i32* %c
+}
+
+define i32* @violate_noundef_pointer() {
+; IS__TUNIT____: Function Attrs: nofree nosync nounwind readnone willreturn
+; IS__TUNIT____-LABEL: define {{[^@]+}}@violate_noundef_pointer
+; IS__TUNIT____-SAME: () [[ATTR0]] {
+; IS__TUNIT____-NEXT:    unreachable
+;
+; IS__CGSCC____: Function Attrs: nofree norecurse nosync nounwind readnone willreturn
+; IS__CGSCC____-LABEL: define {{[^@]+}}@violate_noundef_pointer
+; IS__CGSCC____-SAME: () [[ATTR0]] {
+; IS__CGSCC____-NEXT:    unreachable
+;
+  %ret = call i32* @argument_noundef2(i32* undef)
+  ret i32* %ret
+}