[InstCombine] Enhance deduction of alignment for aligned_alloc

[xbolva00]

Summary

This patch improves https://reviews.llvm.org/D76971 (Deduce attributes for aligned_alloc in InstCombine) and implements "TODO" item mentioned in the review of that patch.

The function aligned_alloc() is the same as memalign(), except for the added restriction that size should be a multiple of alignment.

Currently, we simply bail out if we see a non-constant size - change that.

[jdoerfert]

The function aligned_alloc() is the same as memalign(), except for the added restriction that size should be a multiple of alignment. So for a non-zero alignment size need to be non-zero.

I don't think the last part is true, 0 x Alignment = 0. That said, for non-zero sizes we know alignment holds, so this is fine.

[xbolva00]

In D100785#2699460, @jdoerfert wrote:

The function aligned_alloc() is the same as memalign(), except for the added restriction that size should be a multiple of alignment. So for a non-zero alignment size need to be non-zero.

I don't think the last part is true, 0 x Alignment = 0. That said, for non-zero sizes we know alignment holds, so this is fine.

So should we just drop isKnownNonZero check? To allow annotation for case like aligned_alloc(32, dynamicsize)?

[jdoerfert]

In D100785#2699469, @xbolva00 wrote:

In D100785#2699460, @jdoerfert wrote:

The function aligned_alloc() is the same as memalign(), except for the added restriction that size should be a multiple of alignment. So for a non-zero alignment size need to be non-zero.

I don't think the last part is true, 0 x Alignment = 0. That said, for non-zero sizes we know alignment holds, so this is fine.

So should we just drop isKnownNonZero check? To allow annotation for case like aligned_alloc(32, dynamicsize)?

No, isKnownNonZero is the important part.

p = aligned_alloc(32, dynamicsize)
===>
dynamicsize == 0 || p % 32 == 0

(if p would be null for size = 0 we could avoid the check but the standard allows for a special pointer return I think)

[jdoerfert]

llvm/test/Transforms/InstCombine/deref-alloc-fns.ll

55

This is now wrong, IMHO.

[xbolva00]

Revert changes back.

[xbolva00]

In D100785#2700023, @jdoerfert wrote:

In D100785#2699469, @xbolva00 wrote:

In D100785#2699460, @jdoerfert wrote:

The function aligned_alloc() is the same as memalign(), except for the added restriction that size should be a multiple of alignment. So for a non-zero alignment size need to be non-zero.

I don't think the last part is true, 0 x Alignment = 0. That said, for non-zero sizes we know alignment holds, so this is fine.

So should we just drop isKnownNonZero check? To allow annotation for case like aligned_alloc(32, dynamicsize)?

No, isKnownNonZero is the important part.

p = aligned_alloc(32, dynamicsize)
===>
dynamicsize == 0 || p % 32 == 0

(if p would be null for size = 0 we could avoid the check but the standard allows for a special pointer return I think)

Ah, right. Reverted to previous version which was OK. Thanks.

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

static IntrinsicInst *findInitTrampoline(Value *Callee) {
 
   if (IntrinsicInst *IT = findInitTrampolineFromAlloca(TrampMem))
     return IT;
   if (IntrinsicInst *IT = findInitTrampolineFromBB(AdjustTramp, TrampMem))
     return IT;
   return nullptr;
 }
 
-static void annotateAnyAllocSite(CallBase &Call, const TargetLibraryInfo *TLI) {
+void InstCombinerImpl::annotateAnyAllocSite(CallBase &Call, const TargetLibraryInfo *TLI) {
   unsigned NumArgs = Call.getNumArgOperands();
   ConstantInt *Op0C = dyn_cast<ConstantInt>(Call.getOperand(0));
   ConstantInt *Op1C =
       (NumArgs == 1) ? nullptr : dyn_cast<ConstantInt>(Call.getOperand(1));
   // Bail out if the allocation size is zero (or an invalid alignment of zero
   // with aligned_alloc).
   if ((Op0C && Op0C->isNullValue()) || (Op1C && Op1C->isNullValue()))
     return;
 
   if (isMallocLikeFn(&Call, TLI) && Op0C) {
     if (isOpNewLikeFn(&Call, TLI))
       Call.addAttribute(AttributeList::ReturnIndex,
                         Attribute::getWithDereferenceableBytes(
                             Call.getContext(), Op0C->getZExtValue()));
     else
       Call.addAttribute(AttributeList::ReturnIndex,
                         Attribute::getWithDereferenceableOrNullBytes(
                             Call.getContext(), Op0C->getZExtValue()));
-  } else if (isAlignedAllocLikeFn(&Call, TLI) && Op1C) {
+  } else if (isAlignedAllocLikeFn(&Call, TLI)) {
+    if (Op1C)
       Call.addAttribute(AttributeList::ReturnIndex,
                         Attribute::getWithDereferenceableOrNullBytes(
                             Call.getContext(), Op1C->getZExtValue()));
     // Add alignment attribute if alignment is a power of two constant.
-    if (Op0C && Op0C->getValue().ult(llvm::Value::MaximumAlignment)) {
+    if (Op0C && Op0C->getValue().ult(llvm::Value::MaximumAlignment) &&
+        isKnownNonZero(Call.getOperand(1), DL, 0, &AC, &Call, &DT)) {
       uint64_t AlignmentVal = Op0C->getZExtValue();
       if (llvm::isPowerOf2_64(AlignmentVal))
         Call.addAttribute(AttributeList::ReturnIndex,
                           Attribute::getWithAlignment(Call.getContext(),
                                                       Align(AlignmentVal)));
     }
   } else if (isReallocLikeFn(&Call, TLI) && Op1C) {
     Call.addAttribute(AttributeList::ReturnIndex,

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   /// Try to replace select with select operand SIOpd in SI-ICmp sequence.
   bool replacedSelectWithOperand(SelectInst *SI, const ICmpInst *Icmp,
                                  const unsigned SIOpd);
 
   LoadInst *combineLoadToNewType(LoadInst &LI, Type *NewTy,
                                  const Twine &Suffix = "");
 
 private:
+  void annotateAnyAllocSite(CallBase &Call, const TargetLibraryInfo *TLI);
   bool shouldChangeType(unsigned FromBitWidth, unsigned ToBitWidth) const;
   bool shouldChangeType(Type *From, Type *To) const;
   Value *dyn_castNegVal(Value *V) const;
   Type *FindElementAtOffset(PointerType *PtrTy, int64_t Offset,
                             SmallVectorImpl<Value *> &NewIndices);
 
   /// Classify whether a cast is worth optimizing.
   ///

llvm/test/Transforms/InstCombine/deref-alloc-fns.ll

 ; CHECK-LABEL: @aligned_alloc_constant_size(
 ; CHECK-NEXT:    [[CALL:%.*]] = tail call noalias align 32 dereferenceable_or_null(512) i8* @aligned_alloc(i64 32, i64 512)
 ; CHECK-NEXT:    ret i8* [[CALL]]
 ;
   %call = tail call noalias i8* @aligned_alloc(i64 32, i64 512)
   ret i8* %call
 }
 
+define noalias i8* @aligned_alloc_unknown_size_nonzero(i1 %c) {
+; CHECK-LABEL: @aligned_alloc_unknown_size_nonzero(
+; CHECK-NEXT:    [[SIZE:%.*]] = select i1 [[C:%.*]], i64 64, i64 128
+; CHECK-NEXT:    [[CALL:%.*]] = tail call noalias align 32 i8* @aligned_alloc(i64 32, i64 [[SIZE]])
+; CHECK-NEXT:    ret i8* [[CALL]]
+;
+  %size = select i1 %c, i64 64, i64 128
+  %call = tail call noalias i8* @aligned_alloc(i64 32, i64 %size)
+  ret i8* %call
+}
+
+define noalias i8* @aligned_alloc_unknown_size_possibly_zero(i1 %c) {
+; CHECK-LABEL: @aligned_alloc_unknown_size_possibly_zero(
+; CHECK-NEXT:    [[SIZE:%.*]] = select i1 [[C:%.*]], i64 64, i64 0
+; CHECK-NEXT:    [[CALL:%.*]] = tail call noalias i8* @aligned_alloc(i64 32, i64 [[SIZE]])
+; CHECK-NEXT:    ret i8* [[CALL]]
+;
+  %size = select i1 %c, i64 64, i64 0
+  %call = tail call noalias i8* @aligned_alloc(i64 32, i64 %size)
+  ret i8* %call
+}
+
 declare noalias i8* @foo(i8*, i8*, i8*)
 
 define noalias i8* @aligned_alloc_dynamic_args(i64 %align, i64 %size) {
 ; CHECK-LABEL: @aligned_alloc_dynamic_args(
-; CHECK-NEXT:    tail call noalias dereferenceable_or_null(1024) i8* @aligned_alloc(i64 %{{.*}}, i64 1024)
-; CHECK-NEXT:    tail call noalias i8* @aligned_alloc(i64 0, i64 1024)
-; CHECK-NEXT:    tail call noalias i8* @aligned_alloc(i64 32, i64 %{{.*}})
+; CHECK-NEXT:    [[CALL:%.*]] = tail call noalias dereferenceable_or_null(1024) i8* @aligned_alloc(i64 [[ALIGN:%.*]], i64 1024)
+; CHECK-NEXT:    [[CALL_1:%.*]] = tail call noalias i8* @aligned_alloc(i64 0, i64 1024)
+; CHECK-NEXT:    [[CALL_2:%.*]] = tail call noalias i8* @aligned_alloc(i64 32, i64 [[SIZE:%.*]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call i8* @foo(i8* [[CALL]], i8* [[CALL_1]], i8* [[CALL_2]])
+; CHECK-NEXT:    ret i8* [[CALL]]
 ;
   %call = tail call noalias i8* @aligned_alloc(i64 %align, i64 1024)
   %call_1 = tail call noalias i8* @aligned_alloc(i64 0, i64 1024)
   %call_2 = tail call noalias i8* @aligned_alloc(i64 32, i64 %size)
 
   call i8* @foo(i8* %call, i8* %call_1, i8* %call_2)
   ret i8* %call
 }