[llvm-branch-commits] [llvm] [MemoryBuiltins] Capture more information for alloc/free from attributes (PR #198049)

[Johannes Doerfert via llvm-branch-commits]

https://github.com/jdoerfert created https://github.com/llvm/llvm-project/pull/198049

We now read the `alloc_align` attribute to provide better alignment information to users. `alloc-family` should be used as well, as described in the LangRef. Two new helpers provide argument numbers, rather than values.

>From ddc6819eeae82c457150b357c8c48df4e8e6a1c5 Mon Sep 17 00:00:00 2001
From: Johannes Doerfert <jdoerfert.llvm at gmail.com>
Date: Fri, 15 May 2026 16:50:59 -0700
Subject: [PATCH] [MemoryBuiltins] Capture more information for alloc/free from
 attributes

We now read the `alloc_align` attribute to provide better alignment
information to users. `alloc-family` should be used as well, as
described in the LangRef. Two new helpers provide argument numbers,
rather than values.
---
 llvm/include/llvm/Analysis/MemoryBuiltins.h |  4 +++
 llvm/include/llvm/IR/InstrTypes.h           |  4 +++
 llvm/lib/Analysis/MemoryBuiltins.cpp        | 33 +++++++++------------
 llvm/lib/IR/Instructions.cpp                | 13 ++++++--
 4 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/llvm/include/llvm/Analysis/MemoryBuiltins.h b/llvm/include/llvm/Analysis/MemoryBuiltins.h
index 3f34d99380512..aca1ee80d4829 100644
--- a/llvm/include/llvm/Analysis/MemoryBuiltins.h
+++ b/llvm/include/llvm/Analysis/MemoryBuiltins.h
@@ -85,6 +85,10 @@ LLVM_ABI bool isLibFreeFunction(const Function *F, const LibFunc TLIFn);
 LLVM_ABI Value *getFreedOperand(const CallBase *CB,
                                 const TargetLibraryInfo *TLI);
 
+/// If this if a call to a free function, return the freed operand number, or -1.
+LLVM_ABI int getFreedOperandNo(const CallBase *CB,
+                               const TargetLibraryInfo *TLI);
+
 //===----------------------------------------------------------------------===//
 //  Properties of allocation functions
 //
diff --git a/llvm/include/llvm/IR/InstrTypes.h b/llvm/include/llvm/IR/InstrTypes.h
index 61dc5ebef1b1d..c5821f92050dd 100644
--- a/llvm/include/llvm/IR/InstrTypes.h
+++ b/llvm/include/llvm/IR/InstrTypes.h
@@ -1902,6 +1902,10 @@ class CallBase : public Instruction {
   /// operand value. Otherwise, return nullptr.
   LLVM_ABI Value *getArgOperandWithAttribute(Attribute::AttrKind Kind) const;
 
+  /// If one of the arguments has the specified attribute, returns its
+  /// operand number. Otherwise, return -1.
+  LLVM_ABI int getArgOperandNoWithAttribute(Attribute::AttrKind Kind) const;
+
   /// Return true if the call should not be treated as a call to a
   /// builtin.
   bool isNoBuiltin() const {
diff --git a/llvm/lib/Analysis/MemoryBuiltins.cpp b/llvm/lib/Analysis/MemoryBuiltins.cpp
index e8eca14217db0..5b52211fd0097 100644
--- a/llvm/lib/Analysis/MemoryBuiltins.cpp
+++ b/llvm/lib/Analysis/MemoryBuiltins.cpp
@@ -247,12 +247,13 @@ getAllocationSize(const CallBase *CB, const TargetLibraryInfo *TLI) {
   AllocFnsTy Result;
   // Because allocsize only tells us how many bytes are allocated, we're not
   // really allowed to assume anything, so we use MallocLike.
+  // TODO: `alloc-family should be consulted here.`
   Result.AllocTy = MallocLike;
   Result.NumParams = CB->arg_size();
   Result.FstParam = Args.first;
   Result.SndParam = Args.second.value_or(-1);
-  // Allocsize has no way to specify an alignment argument
-  Result.AlignParam = -1;
+  // AllocAlign defines the alignment parameter.
+  Result.AlignParam = CB->getArgOperandNoWithAttribute(Attribute::AllocAlign);
   return Result;
 }
 
@@ -554,19 +555,26 @@ bool llvm::isLibFreeFunction(const Function *F, const LibFunc TLIFn) {
   return true;
 }
 
-Value *llvm::getFreedOperand(const CallBase *CB, const TargetLibraryInfo *TLI) {
+int llvm::getFreedOperandNo(const CallBase *CB, const TargetLibraryInfo *TLI) {
   if (const Function *Callee = getCalledFunction(CB)) {
     LibFunc TLIFn;
     if (TLI && TLI->getLibFunc(*Callee, TLIFn) && TLI->has(TLIFn) &&
         isLibFreeFunction(Callee, TLIFn)) {
       // All currently supported free functions free the first argument.
-      return CB->getArgOperand(0);
+      return 0;
     }
   }
 
   if (checkFnAllocKind(CB, AllocFnKind::Free))
-    return CB->getArgOperandWithAttribute(Attribute::AllocatedPointer);
+    return CB->getArgOperandNoWithAttribute(Attribute::AllocatedPointer);
+
+  return -1;
+}
 
+Value *llvm::getFreedOperand(const CallBase *CB, const TargetLibraryInfo *TLI) {
+  int Index = getFreedOperandNo(CB, TLI);
+  if (Index >= 0)
+    return CB->getArgOperand(Index);
   return nullptr;
 }
 
@@ -592,20 +600,7 @@ llvm::getDeallocationCallInfo(const CallBase *CB,
     }
   }
 
-  Value *FO = getFreedOperand(CB, TLI);
-  if (!FO)
-    return std::nullopt;
-
-  for (const auto &[Idx, Param] : enumerate(CB->args())) {
-    if (Param != FO)
-      continue;
-    if (DCI.FreedOperandArgNo >= 0) {
-      DCI.FreedOperandArgNo = -1;
-      break;
-    }
-    DCI.FreedOperandArgNo = Idx;
-  }
-
+  DCI.FreedOperandArgNo = getFreedOperandNo(CB, TLI);
   DCI.Family = getAllocationFamily(CB, TLI);
   return DCI;
 }
diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
index 93ec59846d360..236d828414e2d 100644
--- a/llvm/lib/IR/Instructions.cpp
+++ b/llvm/lib/IR/Instructions.cpp
@@ -397,15 +397,22 @@ bool CallBase::isReturnNonNull() const {
   return false;
 }
 
-Value *CallBase::getArgOperandWithAttribute(Attribute::AttrKind Kind) const {
+int CallBase::getArgOperandNoWithAttribute(Attribute::AttrKind Kind) const {
   unsigned Index;
 
   if (Attrs.hasAttrSomewhere(Kind, &Index))
-    return getArgOperand(Index - AttributeList::FirstArgIndex);
+    return (Index - AttributeList::FirstArgIndex);
   if (const Function *F = getCalledFunction())
     if (F->getAttributes().hasAttrSomewhere(Kind, &Index))
-      return getArgOperand(Index - AttributeList::FirstArgIndex);
+      return (Index - AttributeList::FirstArgIndex);
+
+  return -1;
+}
 
+Value *CallBase::getArgOperandWithAttribute(Attribute::AttrKind Kind) const {
+  int Index = getArgOperandNoWithAttribute(Kind);
+  if (Index >= 0)
+    return getArgOperand(Index);
   return nullptr;
 }