[ArgPromotion] Check dereferenceability on argument as well

Before walking all the callers, check whether we have a
dereferenceable attribute directly on the argument.

Also make it clearer that the code currently does not treat
alignment correctly.

Author: nikic

llvm/lib/Transforms/IPO/ArgumentPromotion.cpp

@@ -477,18 +477,20 @@ doPromotion(Function *F, SmallPtrSetImpl<Argument *> &ArgsToPromote,
 static bool allCallersPassValidPointerForArgument(Argument *Arg, Type *Ty) {
   Function *Callee = Arg->getParent();
   const DataLayout &DL = Callee->getParent()->getDataLayout();
+  Align NeededAlign(1); // TODO: This is incorrect!
+  APInt Bytes(64, DL.getTypeStoreSize(Ty));
 
-  unsigned ArgNo = Arg->getArgNo();
+  // Check if the argument itself is marked dereferenceable and aligned.
+  if (isDereferenceableAndAlignedPointer(Arg, NeededAlign, Bytes, DL))
+    return true;
 
   // Look at all call sites of the function.  At this point we know we only have
   // direct callees.
-  for (User *U : Callee->users()) {
+  return all_of(Callee->users(), [&](User *U) {
     CallBase &CB = cast<CallBase>(*U);
-
-    if (!isDereferenceablePointer(CB.getArgOperand(ArgNo), Ty, DL))
-      return false;
-  }
-  return true;
+    return isDereferenceableAndAlignedPointer(
+        CB.getArgOperand(Arg->getArgNo()), NeededAlign, Bytes, DL);
+  });
 }
 
 /// Returns true if Prefix is a prefix of longer. That means, Longer has a size

llvm/test/Transforms/ArgumentPromotion/align.ll

@@ -53,11 +53,10 @@ define void @caller_guaranteed_aligned_1(i1 %c, i32* align 16 dereferenceable(4)
 
 define internal i32 @callee_guaranteed_aligned_2(i1 %c, i32* align 16 dereferenceable(4) %p) {
 ; CHECK-LABEL: define {{[^@]+}}@callee_guaranteed_aligned_2
-; CHECK-SAME: (i1 [[C:%.*]], i32* align 16 dereferenceable(4) [[P:%.*]]) {
+; CHECK-SAME: (i1 [[C:%.*]], i32 [[P_VAL:%.*]]) {
 ; CHECK-NEXT:    br i1 [[C]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
-; CHECK-NEXT:    [[X:%.*]] = load i32, i32* [[P]], align 16
-; CHECK-NEXT:    ret i32 [[X]]
+; CHECK-NEXT:    ret i32 [[P_VAL]]
 ; CHECK:       else:
 ; CHECK-NEXT:    ret i32 -1
 ;
@@ -74,7 +73,8 @@ else:
 define void @caller_guaranteed_aligned_2(i1 %c, i32* %p) {
 ; CHECK-LABEL: define {{[^@]+}}@caller_guaranteed_aligned_2
 ; CHECK-SAME: (i1 [[C:%.*]], i32* [[P:%.*]]) {
-; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @callee_guaranteed_aligned_2(i1 [[C]], i32* [[P]])
+; CHECK-NEXT:    [[P_VAL:%.*]] = load i32, i32* [[P]], align 16
+; CHECK-NEXT:    [[TMP1:%.*]] = call i32 @callee_guaranteed_aligned_2(i1 [[C]], i32 [[P_VAL]])
 ; CHECK-NEXT:    ret void
 ;
   call i32 @callee_guaranteed_aligned_2(i1 %c, i32* %p)