[llvm] [IR][DSE] Support non-malloc functions in malloc+memset->calloc fold (PR #138299)
via llvm-commits
llvm-commits at lists.llvm.org
Sat May 3 06:34:00 PDT 2025
https://github.com/clubby789 updated https://github.com/llvm/llvm-project/pull/138299
>From 910c21e93451b58268f244d36d9cdc161ec454f9 Mon Sep 17 00:00:00 2001
From: Jamie <jamie at osec.io>
Date: Fri, 2 May 2025 19:00:13 +0100
Subject: [PATCH] Support non-malloc functions in `malloc`+`memset`->`calloc`
fold
---
llvm/docs/LangRef.rst | 4 ++
llvm/include/llvm/IR/Attributes.h | 1 +
llvm/lib/IR/Attributes.cpp | 4 ++
llvm/lib/IR/Verifier.cpp | 27 ++++++++++++
.../Scalar/DeadStoreElimination.cpp | 42 +++++++++++++++----
.../DeadStoreElimination/noop-stores.ll | 27 ++++++++++++
llvm/test/Verifier/alloc-variant-zeroed.ll | 19 +++++++++
7 files changed, 116 insertions(+), 8 deletions(-)
create mode 100644 llvm/test/Verifier/alloc-variant-zeroed.ll
diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index 568843a4486e5..78718f102e19c 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -1954,6 +1954,10 @@ For example:
The first three options are mutually exclusive, and the remaining options
describe more details of how the function behaves. The remaining options
are invalid for "free"-type functions.
+``"alloc-variant-zeroed"="FUNCTION"``
+ This attribute indicates that another function is equivalent to an allocator function,
+ but returns zeroed memory. The function must have "zeroed" allocation behavior,
+ the same ``alloc-family``, and take exactly the same arguments.
``allocsize(<EltSizeParam>[, <NumEltsParam>])``
This attribute indicates that the annotated function will always return at
least a given number of bytes (or null). Its arguments are zero-indexed
diff --git a/llvm/include/llvm/IR/Attributes.h b/llvm/include/llvm/IR/Attributes.h
index 0e8e31715acd7..b7a00ad93a7e8 100644
--- a/llvm/include/llvm/IR/Attributes.h
+++ b/llvm/include/llvm/IR/Attributes.h
@@ -156,6 +156,7 @@ class Attribute {
static Attribute getWithAllocSizeArgs(
LLVMContext &Context, unsigned ElemSizeArg,
const std::optional<unsigned> &NumElemsArg);
+ static Attribute getWithAllocKind(LLVMContext &Context, AllocFnKind Kind);
static Attribute getWithVScaleRangeArgs(LLVMContext &Context,
unsigned MinValue, unsigned MaxValue);
static Attribute getWithByValType(LLVMContext &Context, Type *Ty);
diff --git a/llvm/lib/IR/Attributes.cpp b/llvm/lib/IR/Attributes.cpp
index 33ac8bfaf4e7c..d76bdcfebcb9a 100644
--- a/llvm/lib/IR/Attributes.cpp
+++ b/llvm/lib/IR/Attributes.cpp
@@ -300,6 +300,10 @@ Attribute::getWithAllocSizeArgs(LLVMContext &Context, unsigned ElemSizeArg,
return get(Context, AllocSize, packAllocSizeArgs(ElemSizeArg, NumElemsArg));
}
+Attribute Attribute::getWithAllocKind(LLVMContext &Context, AllocFnKind Kind) {
+ return get(Context, AllocKind, static_cast<uint64_t>(Kind));
+}
+
Attribute Attribute::getWithVScaleRangeArgs(LLVMContext &Context,
unsigned MinValue,
unsigned MaxValue) {
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index a798808d79656..bb0a2db67c8ac 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -2377,6 +2377,33 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
CheckFailed("'allockind()' can't be both zeroed and uninitialized");
}
+ if (Attribute A = Attrs.getFnAttr("alloc-variant-zeroed"); A.isValid()) {
+ StringRef S = A.getValueAsString();
+ Check(!S.empty(), "'alloc-variant-zeroed' must not be empty");
+ if (const Function *F = dyn_cast<Function>(V)) {
+ if (Function *Variant = F->getParent()->getFunction(S)) {
+ Attribute Family = Attrs.getFnAttr("alloc-family");
+ Attribute VariantFamily = Variant->getFnAttribute("alloc-family");
+ if (Family.isValid())
+ Check(VariantFamily.isValid() && VariantFamily.getValueAsString() ==
+ Family.getValueAsString(),
+ "'alloc-variant-zeroed' must name a function belonging to the "
+ "same 'alloc-family'");
+
+ Check(
+ Variant->hasFnAttribute(Attribute::AllocKind) &&
+ (Variant->getFnAttribute(Attribute::AllocKind).getAllocKind() &
+ AllocFnKind::Zeroed) != AllocFnKind::Unknown,
+ "'alloc-variant-zeroed' must name a function with "
+ "'allockind(\"zeroed\")'");
+
+ Check(F->getFunctionType() == Variant->getFunctionType(),
+ "'alloc-variant-zeroed' must name a function with the same "
+ "signature");
+ }
+ }
+ }
+
if (Attrs.hasFnAttr(Attribute::VScaleRange)) {
unsigned VScaleMin = Attrs.getFnAttrs().getVScaleRangeMin();
if (VScaleMin == 0)
diff --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index e318ec94db4c3..cf3069edd81e0 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -2028,9 +2028,18 @@ struct DSEState {
if (!InnerCallee)
return false;
LibFunc Func;
+ std::optional<StringRef> ZeroedVariantName = std::nullopt;
if (!TLI.getLibFunc(*InnerCallee, Func) || !TLI.has(Func) ||
- Func != LibFunc_malloc)
- return false;
+ Func != LibFunc_malloc) {
+ Attribute Attr = Malloc->getFnAttr("alloc-variant-zeroed");
+ if (!Attr.isValid()) {
+ return false;
+ }
+ ZeroedVariantName = Attr.getValueAsString();
+ if (ZeroedVariantName->empty())
+ return false;
+ }
+
// Gracefully handle malloc with unexpected memory attributes.
auto *MallocDef = dyn_cast_or_null<MemoryDef>(MSSA.getMemoryAccess(Malloc));
if (!MallocDef)
@@ -2057,15 +2066,32 @@ struct DSEState {
if (Malloc->getOperand(0) != MemSet->getLength())
return false;
- if (!shouldCreateCalloc(Malloc, MemSet) ||
- !DT.dominates(Malloc, MemSet) ||
+ if (!shouldCreateCalloc(Malloc, MemSet) || !DT.dominates(Malloc, MemSet) ||
!memoryIsNotModifiedBetween(Malloc, MemSet, BatchAA, DL, &DT))
return false;
IRBuilder<> IRB(Malloc);
- Type *SizeTTy = Malloc->getArgOperand(0)->getType();
- auto *Calloc =
- emitCalloc(ConstantInt::get(SizeTTy, 1), Malloc->getArgOperand(0), IRB,
- TLI, Malloc->getType()->getPointerAddressSpace());
+ assert(Func == LibFunc_malloc || ZeroedVariantName.has_value());
+ Value *Calloc = nullptr;
+ if (ZeroedVariantName.has_value()) {
+ LLVMContext &Ctx = Malloc->getContext();
+ AttributeList Attr = InnerCallee->getAttributes();
+ AllocFnKind AllocKind =
+ Attr.getFnAttr(Attribute::AllocKind).getAllocKind() |
+ AllocFnKind::Zeroed;
+ Attr =
+ Attr.addFnAttribute(Ctx, Attribute::getWithAllocKind(Ctx, AllocKind))
+ .removeFnAttribute(Ctx, "alloc-variant-zeroed");
+ FunctionCallee ZeroedVariant = Malloc->getModule()->getOrInsertFunction(
+ *ZeroedVariantName, InnerCallee->getFunctionType(), Attr);
+ SmallVector<Value *, 3> Args;
+ Args.append(Malloc->arg_begin(), Malloc->arg_end());
+ Calloc = IRB.CreateCall(ZeroedVariant, Args, *ZeroedVariantName);
+ } else {
+ Type *SizeTTy = Malloc->getArgOperand(0)->getType();
+ Calloc =
+ emitCalloc(ConstantInt::get(SizeTTy, 1), Malloc->getArgOperand(0),
+ IRB, TLI, Malloc->getType()->getPointerAddressSpace());
+ }
if (!Calloc)
return false;
diff --git a/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll b/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
index 9fc20d76da5eb..377eaacc14af2 100644
--- a/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
+++ b/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
@@ -374,6 +374,33 @@ define ptr @notmalloc_memset(i64 %size, ptr %notmalloc) {
ret ptr %call1
}
+; This should create a customalloc_zeroed call and eliminate the memset
+define ptr @customalloc_memset(i64 %size, i64 %align) {
+; CHECK-LABEL: @customalloc_memset
+; CHECK-NEXT: [[CALL:%.*]] = call ptr @customalloc_zeroed(i64 [[SIZE:%.*]], i64 [[ALIGN:%.*]])
+; CHECK-NEXT: ret ptr [[CALL]]
+ %call = call ptr @customalloc(i64 %size, i64 %align)
+ call void @llvm.memset.p0.i64(ptr %call, i8 0, i64 %size, i1 false)
+ ret ptr %call
+}
+
+declare ptr @customalloc(i64, i64) allockind("alloc") "alloc-family"="customalloc" "alloc-variant-zeroed"="customalloc_zeroed"
+declare ptr @customalloc_zeroed(i64, i64) allockind("alloc,zeroed") "alloc-family"="customalloc"
+
+; This should create a declaration for the named variant
+define ptr @undeclared_customalloc(i64 %size, i64 %align) {
+; CHECK-LABEL: @undeclared_customalloc
+; CHECK-NEXT: [[CALL:%.*]] = call ptr @customalloc2_zeroed(i64 [[SIZE:%.*]], i64 [[ALIGN:%.*]])
+; CHECK-NEXT: ret ptr [[CALL]]
+ %call = call ptr @customalloc2(i64 %size, i64 %align)
+ call void @llvm.memset.p0.i64(ptr %call, i8 0, i64 %size, i1 false)
+ ret ptr %call
+}
+
+declare ptr @customalloc2(i64, i64) allockind("alloc") "alloc-family"="customalloc2" "alloc-variant-zeroed"="customalloc2_zeroed"
+; CHECK-DAG: declare ptr @customalloc2_zeroed(i64, i64) #[[CA2ATTR:[0-9]+]]
+; CHECK-DAG: attributes #[[CA2ATTR]] = { allockind("alloc,zeroed") "alloc-family"="customalloc2" }
+
; This should not create recursive call to calloc.
define ptr @calloc(i64 %nmemb, i64 %size) inaccessiblememonly {
; CHECK-LABEL: @calloc(
diff --git a/llvm/test/Verifier/alloc-variant-zeroed.ll b/llvm/test/Verifier/alloc-variant-zeroed.ll
new file mode 100644
index 0000000000000..03cd7be4d9acf
--- /dev/null
+++ b/llvm/test/Verifier/alloc-variant-zeroed.ll
@@ -0,0 +1,19 @@
+; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
+
+; CHECK: 'alloc-variant-zeroed' must not be empty
+declare ptr @a(i64) "alloc-variant-zeroed"=""
+
+; CHECK: 'alloc-variant-zeroed' must not be empty
+declare ptr @b(i64) "alloc-variant-zeroed"=""
+
+; CHECK: 'alloc-variant-zeroed' must name a function belonging to the same 'alloc-family'
+declare ptr @c(i64) "alloc-variant-zeroed"="c_zeroed" "alloc-family"="C"
+declare ptr @c_zeroed(i64)
+
+; CHECK: 'alloc-variant-zeroed' must name a function with 'allockind("zeroed")'
+declare ptr @d(i64) "alloc-variant-zeroed"="d_zeroed" "alloc-family"="D"
+declare ptr @d_zeroed(i64) "alloc-family"="D"
+
+; CHECK: 'alloc-variant-zeroed' must name a function with the same signature
+declare ptr @e(i64) "alloc-variant-zeroed"="e_zeroed" "alloc-family"="E"
+declare ptr @e_zeroed(i64, i64) "alloc-family"="E" allockind("zeroed")
More information about the llvm-commits
mailing list