[llvm] [InstCombine] Don't consider aligned_alloc removable if icmp uses result (PR #69474)

[Florian Hahn via llvm-commits]

https://github.com/fhahn created https://github.com/llvm/llvm-project/pull/69474

At the moment, all alloc-like functions are assumed to return non-null pointers, if their return value is only used in a compare. This is based on being allowed to substitute the allocation function with one that doesn't fail to allocate the required memory.

aligned_alloc however must also return null if the required alignment cannot be satisfied, so I don't think the same reasoning as above can be applied to it.

This patch adds a bail-out for aligned allocation functions to isAllocSiteRemovable.

>From 76741a799627e8bbc4d76e03ab51cd1ccf791da4 Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Wed, 18 Oct 2023 16:12:35 +0100
Subject: [PATCH] [InstCombine] Don't consider aligned_alloc removable if icmp
 uses result

At the moment, all alloc-like functions are assumed to return non-null
pointers, if their return value is only used in a compare. This is
based on being allowed to substitute the allocation function with one
that doesn't fail to allocate the required memory.

aligned_alloc however must also return null if the required alignment
cannot be satisfied, so I don't think the same reasoning as above can be
applied to it.

This patch adds a bail-out for aligned allocation functions to
isAllocSiteRemovable.
---
 llvm/include/llvm/Analysis/MemoryBuiltins.h              | 5 +++++
 llvm/lib/Analysis/MemoryBuiltins.cpp                     | 2 +-
 llvm/lib/Transforms/InstCombine/InstructionCombining.cpp | 5 +++++
 llvm/test/Transforms/InstCombine/malloc-free.ll          | 4 +++-
 4 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
index 827b5081b2ce755..cf18d9fcd95869d 100644
--- a/llvm/include/llvm/Analysis/MemoryBuiltins.h
+++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -131,6 +131,11 @@ Constant *getInitialValueOfAllocation(const Value *V,
 std::optional<StringRef> getAllocationFamily(const Value *I,
                                              const TargetLibraryInfo *TLI);
 
+/// If \p V is a call of a function part of an allocation family e.g.
+/// malloc/realloc/calloc/free), return the allocation function kind the called
+/// function.
+AllocFnKind getAllocFnKind(const Value *V);
+
 //===----------------------------------------------------------------------===//
 //  Utility functions to compute size of objects.
 //
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index 6ee1984f908b8a1..e52572dbbe5461a 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -272,7 +272,7 @@ getAllocationSize(const Value *V, const TargetLibraryInfo *TLI) {
   return Result;
 }
 
-static AllocFnKind getAllocFnKind(const Value *V) {
+AllocFnKind llvm::getAllocFnKind(const Value *V) {
   if (const auto *CB = dyn_cast<CallBase>(V)) {
     Attribute Attr = CB->getFnAttr(Attribute::AllocKind);
     if (Attr.isValid())
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 8a6f66e36bd80e9..fca253d67fd19b2 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -2430,6 +2430,11 @@ static bool isAllocSiteRemovable(Instruction *AI,
         unsigned OtherIndex = (ICI->getOperand(0) == PI) ? 1 : 0;
         if (!isNeverEqualToUnescapedAlloc(ICI->getOperand(OtherIndex), TLI, AI))
           return false;
+        // Do not fold compares to alligned allocation functions, as they may
+        // have to return null in case the required alignment cannot be
+        // satisified.
+        if ((getAllocFnKind(AI) & AllocFnKind::Aligned) != AllocFnKind::Unknown)
+          return false;
         Users.emplace_back(I);
         continue;
       }
diff --git a/llvm/test/Transforms/InstCombine/malloc-free.ll b/llvm/test/Transforms/InstCombine/malloc-free.ll
index 29c757f82564ad8..3522c9f3befcf4d 100644
--- a/llvm/test/Transforms/InstCombine/malloc-free.ll
+++ b/llvm/test/Transforms/InstCombine/malloc-free.ll
@@ -28,7 +28,9 @@ define i32 @dead_aligned_alloc(i32 %size, i32 %alignment, i8 %value) {
 
 define i1 @aligned_alloc_only_pointe(i32 %size, i32 %alignment, i8 %value) {
 ; CHECK-LABEL: @aligned_alloc_only_pointe(
-; CHECK-NEXT:    ret i1 true
+; CHECK-NEXT:    [[ALIGNED_ALLOCATION:%.*]] = tail call ptr @aligned_alloc(i32 [[ALIGNMENT:%.*]], i32 [[SIZE:%.*]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne ptr [[ALIGNED_ALLOCATION]], null
+; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %aligned_allocation = tail call ptr @aligned_alloc(i32 %alignment, i32 %size)
   %cmp = icmp ne ptr %aligned_allocation, null