MemoryBuiltins: also check function definition for allocalign

[durin42]

Summary

This got changed to use hasAttrSomewhere() during review, and I didn't
notice until today when I was writing some tests for another part of
this system that using hasAttrSomewhere only checked the callsite for
allocalign, rather than both the callsite and the definition. This fixes
that by introducing a helper method.

[nikic]

This should have a test, e.g. in InstCombine.

[durin42]

In D121641#3380875, @nikic wrote:

This should have a test, e.g. in InstCombine.

That happens as part of D121642 - is that okay?

[nikic]

In D121641#3380900, @durin42 wrote:

In D121641#3380875, @nikic wrote:

This should have a test, e.g. in InstCombine.

That happens as part of D121642 - is that okay?

I'd suggest pre-committing the tests, so that the test diff then shows up in the individual patches.

[nikic]

llvm/lib/Analysis/MemoryBuiltins.cpp

343

Maybe we should have a helper function on CallBase for this whole logic? See CallBase::getReturnedArgOperand() which implements exactly the same but with Attribute::Returned instead of Attribute::AllocAlign.

[nikic]

llvm/lib/IR/Instructions.cpp

340

That wasn't quite what I had in mind: I don't think we'll want to fetch allocalign beyond that one callsite. Rather, we could have CallBase::getArgOperandWithAttribute(AttrKind) and then use that both in getReturnedArgOperand() and in getAllocAlignment().

[nikic]

LGTM

llvm/include/llvm/IR/InstrTypes.h

public:
   bool returnDoesNotAlias() const {
     return Attrs.hasRetAttr(Attribute::NoAlias);
   }
 
   /// If one of the arguments has the 'returned' attribute, returns its
   /// operand value. Otherwise, return nullptr.
   Value *getReturnedArgOperand() const;
 
+  /// If one of the arguments has the 'allocalign' attribute, returns its
+  /// operand value. Otherwise, return nullptr.
+  Value *getAllocAlignArgOperand() const;
+ 
   /// Return true if the call should not be treated as a call to a
   /// builtin.
   bool isNoBuiltin() const {
     return hasFnAttrImpl(Attribute::NoBuiltin) &&
            !hasFnAttrImpl(Attribute::Builtin);
   }
 
   /// Determine if the call requires strict floating point semantics.

llvm/lib/Analysis/MemoryBuiltins.cpp

 }
 
 Value *llvm::getAllocAlignment(const CallBase *V,
                                const TargetLibraryInfo *TLI) {
   const Optional<AllocFnsTy> FnData = getAllocationData(V, AnyAlloc, TLI);
   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->getAllocAlignArgOperand();
 }
 
 /// When we're compiling N-bit code, and the user uses parameters that are
 /// greater than N bits (e.g. uint64_t on a 32-bit build), we can run into
 /// trouble with APInt size issues. This function handles resizing + overflow
 /// checks for us. Check and zext or trunc \p I depending on IntTyBits and
 /// I's value.
 static bool CheckedZextOrTrunc(APInt &I, unsigned IntTyBits) {
   // More bits than we can handle. Checking the bit width isn't necessary, but
   // it's faster than checking active bits, and should give `false` in the

llvm/lib/IR/Instructions.cpp

if (Attrs.hasAttrSomewhere(Attribute::Returned, &Index))
     return getArgOperand(Index - AttributeList::FirstArgIndex);
   if (const Function *F = getCalledFunction())
     if (F->getAttributes().hasAttrSomewhere(Attribute::Returned, &Index))
       return getArgOperand(Index - AttributeList::FirstArgIndex);
 
   return nullptr;
 }
 
+Value *CallBase::getAllocAlignArgOperand() const {
+  unsigned Index;
+
+  if (Attrs.hasAttrSomewhere(Attribute::AllocAlign, &Index))
+    return getArgOperand(Index - AttributeList::FirstArgIndex);
+  if (const Function *F = getCalledFunction())
+    if (F->getAttributes().hasAttrSomewhere(Attribute::AllocAlign, &Index))
+      return getArgOperand(Index - AttributeList::FirstArgIndex);
+
+  return nullptr;
+}
+
 /// Determine whether the argument or parameter has the given attribute.
 bool CallBase::paramHasAttr(unsigned ArgNo, Attribute::AttrKind Kind) const {
   assert(ArgNo < arg_size() && "Param index out of bounds!");
 
   if (Attrs.hasParamAttr(ArgNo, Kind))
     return true;
   if (const Function *F = getCalledFunction())
     return F->getAttributes().hasParamAttr(ArgNo, Kind);

llvm/test/Transforms/InstCombine/InferAlignAttribute.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -inline -instcombine %s | FileCheck %s --check-prefixes=CHECK,CHECK-INLINE
 ; RUN: opt -S -instcombine %s | FileCheck %s --check-prefixes=CHECK,CHECK-NOINLINE
 
 define i8* @widen_align_from_allocalign_callsite() {
 ; CHECK-LABEL: @widen_align_from_allocalign_callsite(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[CALL:%.*]] = tail call align 64 i8* @my_aligned_alloc_2(i32 noundef 320, i32 allocalign noundef 64)
 ; CHECK-NEXT:    ret i8* [[CALL]]
 ;
 entry:
 
   %call = tail call align 16 i8* @my_aligned_alloc_2(i32 noundef 320, i32 allocalign noundef 64)
   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:
 
   %call = tail call align 16 i8* @my_aligned_alloc(i32 noundef 320, i32 noundef 64)
   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:
   %call = tail call align 16 i8* @my_aligned_alloc(i32 noundef 320, i32 noundef 8)
   ret i8* %call
 }
 
 define i8* @my_aligned_alloc_3(i32 noundef %foo, i32 allocalign %alignment) {

 ; my_aligned_alloc_2. Note that this is a simplified version of what happens
 ; with _mm_malloc which calls posix_memalign.
 define i8* @allocalign_disappears() {
 ; CHECK-INLINE-LABEL: @allocalign_disappears(
 ; CHECK-INLINE-NEXT:    [[CALL_I:%.*]] = tail call i8* @my_aligned_alloc_2(i32 noundef 42, i32 noundef 128)
 ; 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)
   ret i8* %call
 }
 
 declare i8* @my_aligned_alloc(i32 noundef, i32 allocalign noundef)
 declare i8* @my_aligned_alloc_2(i32 noundef, i32 noundef)