[InstCombine] Enhance deduction of alignment for aligned_alloc

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.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D100785

3 files changed, 36 insertions, 9 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 615f10e..3932446 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -2049,7 +2049,7 @@ static IntrinsicInst *findInitTrampoline(Value *Callee) {
   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 =
@@ -2068,12 +2068,14 @@ static void annotateAnyAllocSite(CallBase &Call, const TargetLibraryInfo *TLI) {
       Call.addAttribute(AttributeList::ReturnIndex,
                         Attribute::getWithDereferenceableOrNullBytes(
                             Call.getContext(), Op0C->getZExtValue()));
-  } else if (isAlignedAllocLikeFn(&Call, TLI) && Op1C) {
-    Call.addAttribute(AttributeList::ReturnIndex,
-                      Attribute::getWithDereferenceableOrNullBytes(
-                          Call.getContext(), Op1C->getZExtValue()));
+  } 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,
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 8907d4b..edf8f0f 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -187,6 +187,7 @@ public:
                                  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;
diff --git a/llvm/test/Transforms/InstCombine/deref-alloc-fns.ll b/llvm/test/Transforms/InstCombine/deref-alloc-fns.ll
index 2726a5f..b660e6b 100644
--- a/llvm/test/Transforms/InstCombine/deref-alloc-fns.ll
+++ b/llvm/test/Transforms/InstCombine/deref-alloc-fns.ll
@@ -38,13 +38,37 @@ define noalias i8* @aligned_alloc_constant_size() {
   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)