diff --git a/llvm/include/llvm/IR/InstrTypes.h b/llvm/include/llvm/IR/InstrTypes.h
--- a/llvm/include/llvm/IR/InstrTypes.h
+++ b/llvm/include/llvm/IR/InstrTypes.h
@@ -1826,7 +1826,13 @@
 
   /// If one of the arguments has the 'returned' attribute, returns its
   /// operand value. Otherwise, return nullptr.
-  Value *getReturnedArgOperand() const;
+  Value *getReturnedArgOperand() const {
+    return getArgOperandWithAttribute(Attribute::Returned);
+  }
+
+  /// If one of the arguments has the specified attribute, returns its
+  /// operand value. Otherwise, return nullptr.
+  Value *getArgOperandWithAttribute(Attribute::AttrKind Kind) const;
 
   /// Return true if the call should not be treated as a call to a
   /// builtin.
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -338,12 +338,7 @@
   if (FnData.hasValue() && FnData->AlignParam >= 0) {
     return V->getOperand(FnData->AlignParam);
   }
-  unsigned AllocAlignParam;
-  if (V->getAttributes().hasAttrSomewhere(Attribute::AllocAlign,
-                                          &AllocAlignParam)) {
-    return V->getOperand(AllocAlignParam-1);
-  }
-  return nullptr;
+  return V->getArgOperandWithAttribute(Attribute::AllocAlign);
 }
 
 /// When we're compiling N-bit code, and the user uses parameters that are
diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
--- a/llvm/lib/IR/Instructions.cpp
+++ b/llvm/lib/IR/Instructions.cpp
@@ -325,13 +325,13 @@
   return false;
 }
 
-Value *CallBase::getReturnedArgOperand() const {
+Value *CallBase::getArgOperandWithAttribute(Attribute::AttrKind Kind) const {
   unsigned Index;
 
-  if (Attrs.hasAttrSomewhere(Attribute::Returned, &Index))
+  if (Attrs.hasAttrSomewhere(Kind, &Index))
     return getArgOperand(Index - AttributeList::FirstArgIndex);
   if (const Function *F = getCalledFunction())
-    if (F->getAttributes().hasAttrSomewhere(Attribute::Returned, &Index))
+    if (F->getAttributes().hasAttrSomewhere(Kind, &Index))
       return getArgOperand(Index - AttributeList::FirstArgIndex);
 
   return nullptr;
diff --git a/llvm/test/Transforms/InstCombine/InferAlignAttribute.ll b/llvm/test/Transforms/InstCombine/InferAlignAttribute.ll
--- a/llvm/test/Transforms/InstCombine/InferAlignAttribute.ll
+++ b/llvm/test/Transforms/InstCombine/InferAlignAttribute.ll
@@ -14,11 +14,10 @@
   ret i8* %call
 }
 
-; BUG: we don't check the declaration, only the callsite. This will be fixed in the next change.
 define i8* @widen_align_from_allocalign() {
 ; CHECK-LABEL: @widen_align_from_allocalign(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CALL:%.*]] = tail call align 16 i8* @my_aligned_alloc(i32 noundef 320, i32 noundef 64)
+; CHECK-NEXT:    [[CALL:%.*]] = tail call align 64 i8* @my_aligned_alloc(i32 noundef 320, i32 noundef 64)
 ; CHECK-NEXT:    ret i8* [[CALL]]
 ;
 entry:
@@ -27,10 +26,12 @@
   ret i8* %call
 }
 
+; BUG: we shouldn't narrow this alignment since we already had a stronger
+; constraint, but we do.
 define i8* @dont_narrow_align_from_allocalign() {
 ; CHECK-LABEL: @dont_narrow_align_from_allocalign(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[CALL:%.*]] = tail call align 16 i8* @my_aligned_alloc(i32 noundef 320, i32 noundef 8)
+; CHECK-NEXT:    [[CALL:%.*]] = tail call align 8 i8* @my_aligned_alloc(i32 noundef 320, i32 noundef 8)
 ; CHECK-NEXT:    ret i8* [[CALL]]
 ;
 entry:
@@ -59,7 +60,7 @@
 ; CHECK-INLINE-NEXT:    ret i8* [[CALL_I]]
 ;
 ; CHECK-NOINLINE-LABEL: @allocalign_disappears(
-; CHECK-NOINLINE-NEXT:    [[CALL:%.*]] = tail call i8* @my_aligned_alloc_3(i32 42, i32 128)
+; CHECK-NOINLINE-NEXT:    [[CALL:%.*]] = tail call align 128 i8* @my_aligned_alloc_3(i32 42, i32 128)
 ; CHECK-NOINLINE-NEXT:    ret i8* [[CALL]]
 ;
   %call = tail call i8* @my_aligned_alloc_3(i32 42, i32 128)