InstCombineCalls: infer return alignment from allocalign attributes

[durin42]

Summary

This exposes a couple of lingering bugs, which will be fixed in
the next two commits.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2848

We're missing change reporting here -- annotateAnyAllocSite should return a boolean and we should initialize Changed with it. But this is orthogonal to your change here.

[durin42]

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

2848

Sent a follow-up as D123232 (it's at the end of the stack now, but I'm happy to submit it if it's LG and this one is also ready)

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

 }
 
 void InstCombinerImpl::annotateAnyAllocSite(CallBase &Call, const TargetLibraryInfo *TLI) {
   // Note: We only handle cases which can't be driven from generic attributes
   // here.  So, for example, nonnull and noalias (which are common properties
   // of some allocation functions) are expected to be handled via annotation
   // of the respective allocator declaration with generic attributes.
 
+  if (isAllocationFn(&Call, TLI)) {
     uint64_t Size;
     ObjectSizeOpts Opts;
     if (getObjectSize(&Call, Size, DL, TLI, Opts) && Size > 0) {
       // TODO: We really should just emit deref_or_null here and then
       // let the generic inference code combine that with nonnull.
       if (Call.hasRetAttr(Attribute::NonNull))
-      Call.addRetAttr(Attribute::getWithDereferenceableBytes(
-          Call.getContext(), Size));
+        Call.addRetAttr(
+            Attribute::getWithDereferenceableBytes(Call.getContext(), Size));
       else
         Call.addRetAttr(Attribute::getWithDereferenceableOrNullBytes(
             Call.getContext(), Size));
     }
+  }
 
   // Add alignment attribute if alignment is a power of two constant.
   Value *Alignment = getAllocAlignment(&Call, TLI);
   if (!Alignment)
     return;
 
   ConstantInt *AlignOpC = dyn_cast<ConstantInt>(Alignment);
   if (AlignOpC && AlignOpC->getValue().ult(llvm::Value::MaximumAlignment)) {
     uint64_t AlignmentVal = AlignOpC->getZExtValue();
     if (llvm::isPowerOf2_64(AlignmentVal)) {
       Call.removeRetAttr(Attribute::Alignment);
       Call.addRetAttr(Attribute::getWithAlignment(Call.getContext(),
                                                   Align(AlignmentVal)));
     }
   }
 }
 
 /// Improvements for call, callbr and invoke instructions.
 Instruction *InstCombinerImpl::visitCallBase(CallBase &Call) {
-  if (isAllocationFn(&Call, &TLI))
-    annotateAnyAllocSite(Call, &TLI);
 
   bool Changed = false;
 
   // Mark any parameters that are known to be non-null with the nonnull
   // attribute.  This is helpful for inlining calls to functions with null
   // checks on their arguments.
   SmallVector<unsigned, 4> ArgNos;
   unsigned ArgNo = 0;

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:    ret i8* [[CALL]]
+;
+entry:
+
+  %call = tail call align 16 i8* @my_aligned_alloc(i32 noundef 320, i32 noundef 64)
+  ret i8* %call
+}
+
+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:    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) {
+; CHECK-LABEL: @my_aligned_alloc_3(
+; CHECK-NEXT:    [[CALL:%.*]] = tail call i8* @my_aligned_alloc_2(i32 noundef [[FOO:%.*]], i32 noundef [[ALIGNMENT:%.*]])
+; CHECK-NEXT:    ret i8* [[CALL]]
+;
+  %call = tail call i8* @my_aligned_alloc_2(i32 noundef %foo, i32 noundef %alignment)
+  ret i8* %call
+}
+
+; -inline is able to make my_aligned_alloc_3's arguments disappear and directly
+; call my_aligned_alloc_2, but the latter has no allocalign so the alignment just
+; disappears. This is conservatively correct but undesirable because we can't
+; figure out the `align 128` on the return value once the call is directly on
+; 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:    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)