From 8d58c8e57b2769e640a33d7ad4b57014c3d2df99 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Thu, 21 Jul 2022 10:05:48 +0200
Subject: [PATCH] Reapply [InstCombine] Don't check for alloc fn before
 fetching alloc size

Reapply the patch with getObjectSize() replaced by getAllocSize().
The former will also look through calls that return their argument,
and we'll end up placing dereferenceable attributes on intrinsics
like llvm.launder.invariant.group. While this isn't wrong, it also
doesn't seem to be particularly useful. For now, use getAllocSize()
instead, which sticks closer to the original behavior of this code.

-----

This code is just interested in the allocsize, not any other
allocator properties.
---
 .../Transforms/InstCombine/InstCombineCalls.cpp    | 30 +++++++++++-----------
 .../test/Transforms/InstCombine/deref-alloc-fns.ll | 19 ++++++++++++++
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 6a2cfd6..8ac9be5 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -2885,21 +2885,21 @@ bool InstCombinerImpl::annotateAnyAllocSite(CallBase &Call,
   // of the respective allocator declaration with generic attributes.
   bool Changed = false;
 
-  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)) {
-        Changed = !Call.hasRetAttr(Attribute::Dereferenceable);
-        Call.addRetAttr(
-            Attribute::getWithDereferenceableBytes(Call.getContext(), Size));
-      } else {
-        Changed = !Call.hasRetAttr(Attribute::DereferenceableOrNull);
-        Call.addRetAttr(Attribute::getWithDereferenceableOrNullBytes(
-            Call.getContext(), Size));
-      }
+  if (!Call.getType()->isPointerTy())
+    return Changed;
+
+  Optional<APInt> Size = getAllocSize(&Call, TLI);
+  if (Size && *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)) {
+      Changed = !Call.hasRetAttr(Attribute::Dereferenceable);
+      Call.addRetAttr(Attribute::getWithDereferenceableBytes(
+          Call.getContext(), Size->getLimitedValue()));
+    } else {
+      Changed = !Call.hasRetAttr(Attribute::DereferenceableOrNull);
+      Call.addRetAttr(Attribute::getWithDereferenceableOrNullBytes(
+          Call.getContext(), Size->getLimitedValue()));
     }
   }
 
diff --git a/llvm/test/Transforms/InstCombine/deref-alloc-fns.ll b/llvm/test/Transforms/InstCombine/deref-alloc-fns.ll
index 0b52935..41a2427 100644
--- a/llvm/test/Transforms/InstCombine/deref-alloc-fns.ll
+++ b/llvm/test/Transforms/InstCombine/deref-alloc-fns.ll
@@ -13,6 +13,8 @@ declare noalias i8* @aligned_alloc(i64, i64)
 declare noalias align 16 i8* @memalign(i64, i64)
 ; new[](unsigned int, align_val_t)
 declare noalias i8* @_ZnajSt11align_val_t(i64 %size, i64 %align)
+declare i8* @my_malloc(i64) allocsize(0)
+declare i8* @my_calloc(i64, i64) allocsize(0, 1)
 
 @.str = private unnamed_addr constant [6 x i8] c"hello\00", align 1
 
@@ -354,3 +356,20 @@ define noalias i8* @op_new_align() {
   ret i8* %call
 }
 
+define i8* @my_malloc_constant_size() {
+; CHECK-LABEL: @my_malloc_constant_size(
+; CHECK-NEXT:    [[CALL:%.*]] = call dereferenceable_or_null(32) i8* @my_malloc(i64 32)
+; CHECK-NEXT:    ret i8* [[CALL]]
+;
+  %call = call i8* @my_malloc(i64 32)
+  ret i8* %call
+}
+
+define i8* @my_calloc_constant_size() {
+; CHECK-LABEL: @my_calloc_constant_size(
+; CHECK-NEXT:    [[CALL:%.*]] = call dereferenceable_or_null(128) i8* @my_calloc(i64 32, i64 4)
+; CHECK-NEXT:    ret i8* [[CALL]]
+;
+  %call = call i8* @my_calloc(i64 32, i64 4)
+  ret i8* %call
+}
-- 
2.7.4