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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Florian Hahn (fhahn)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/69474.diff


4 Files Affected:

- (modified) llvm/include/llvm/Analysis/MemoryBuiltins.h (+5) 
- (modified) llvm/lib/Analysis/MemoryBuiltins.cpp (+1-1) 
- (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+5) 
- (modified) llvm/test/Transforms/InstCombine/malloc-free.ll (+3-1) 


``````````diff
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

``````````

</details>


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