[llvm] BasicAA: handle base ptrs from alloc fn (PR #111181)

[Ramkumar Ramachandra via llvm-commits]

https://github.com/artagnon created https://github.com/llvm/llvm-project/pull/111181

Extend isBaseOfObject to also query llvm::isAllocationFn using TLI.

>From e7bd084be7bbe806b638d7c09e2b908e4c71cf69 Mon Sep 17 00:00:00 2001
From: Ramkumar Ramachandra <ramkumar.ramachandra at codasip.com>
Date: Fri, 4 Oct 2024 16:43:18 +0100
Subject: [PATCH] BasicAA: handle base ptrs from alloc fn

Extend isBaseOfObject to also query llvm::isAllocationFn using TLI.
---
 llvm/lib/Analysis/BasicAliasAnalysis.cpp | 18 ++++++++++--------
 llvm/test/Analysis/BasicAA/negoffset.ll  | 15 +++++++++++++++
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Analysis/BasicAliasAnalysis.cpp b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
index f471c0db11d3ef..d1757a63bc36b5 100644
--- a/llvm/lib/Analysis/BasicAliasAnalysis.cpp
+++ b/llvm/lib/Analysis/BasicAliasAnalysis.cpp
@@ -1073,12 +1073,14 @@ ModRefInfo BasicAAResult::getModRefInfo(const CallBase *Call1,
 /// for the underlying object.  Note that just being isIdentifiedObject() is
 /// not enough - For example, a negative offset from a noalias argument or call
 /// can be inbounds w.r.t the actual underlying object.
-static bool isBaseOfObject(const Value *V) {
-  // TODO: We can handle other cases here
-  // 1) For GC languages, arguments to functions are often required to be
-  //    base pointers.
-  // 2) Result of allocation routines are often base pointers.  Leverage TLI.
-  return (isa<AllocaInst>(V) || isa<GlobalVariable>(V));
+static bool isBaseOfObject(const Value *V, const TargetLibraryInfo *TLI) {
+  // TODO: for GC languages, arguments to functions are often required to be
+  // base pointers.
+  if (isa<AllocaInst>(V) || isa<GlobalVariable>(V))
+    return true;
+  if (auto *I = dyn_cast<Instruction>(V))
+    return isAllocationFn(I, TLI);
+  return false;
 }
 
 /// Provides a bunch of ad-hoc rules to disambiguate a GEP instruction against
@@ -1133,14 +1135,14 @@ AliasResult BasicAAResult::aliasGEP(
   if (DecompGEP1.NWFlags.isInBounds() && DecompGEP1.VarIndices.empty() &&
       V2Size.hasValue() && !V2Size.isScalable() &&
       DecompGEP1.Offset.sge(V2Size.getValue()) &&
-      isBaseOfObject(DecompGEP2.Base))
+      isBaseOfObject(DecompGEP2.Base, &TLI))
     return AliasResult::NoAlias;
 
   // Symmetric case to above.
   if (DecompGEP2.NWFlags.isInBounds() && DecompGEP1.VarIndices.empty() &&
       V1Size.hasValue() && !V1Size.isScalable() &&
       DecompGEP1.Offset.sle(-V1Size.getValue()) &&
-      isBaseOfObject(DecompGEP1.Base))
+      isBaseOfObject(DecompGEP1.Base, &TLI))
     return AliasResult::NoAlias;
 
   // For GEPs with identical offsets, we can preserve the size and AAInfo
diff --git a/llvm/test/Analysis/BasicAA/negoffset.ll b/llvm/test/Analysis/BasicAA/negoffset.ll
index addeabf17b0eb7..b4dd27056bff9f 100644
--- a/llvm/test/Analysis/BasicAA/negoffset.ll
+++ b/llvm/test/Analysis/BasicAA/negoffset.ll
@@ -5,6 +5,7 @@ target triple = "i386-unknown-linux-gnu"
 
 declare ptr @random.i32(ptr %ptr)
 declare ptr @random.i8(ptr %ptr)
+declare ptr @malloc(i64) allockind("alloc")
 
 ; CHECK-LABEL: Function: arr:
 ; CHECK-DAG: MayAlias: i32* %alloca, i32* %p0
@@ -20,6 +21,20 @@ define void @arr() {
   ret void
 }
 
+; CHECK-LABEL: Function: allocfn:
+; CHECK-DAG: MayAlias: i8* %malloc, i8* %p0
+; CHECK-DAG: NoAlias:  i8* %malloc, i8* %p1
+define void @allocfn() {
+  %malloc = call ptr @malloc(i64 4)
+  %random = call ptr @random.i8(ptr %malloc)
+  %p0 = getelementptr inbounds i32, ptr %random, i32 0
+  %p1 = getelementptr inbounds i32, ptr %random, i32 1
+  load i8, ptr %malloc
+  load i8, ptr %p0
+  load i8, ptr %p1
+  ret void
+}
+
 ; CHECK-LABEL: Function: arg:
 ; CHECK-DAG: MayAlias: i32* %arg, i32* %p0
 ; CHECK-DAG: MayAlias: i32* %arg, i32* %p1