Skip to content

Commit c7e7901

Browse files
jamie-osecclubby789
jamie-osec
authored and
clubby789
committed
Support non-malloc functions in malloc+memset->calloc fold
1 parent 3a8b488 commit c7e7901

File tree

8 files changed

+117
-8
lines changed

8 files changed

+117
-8
lines changed

‎llvm/docs/LangRef.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1954,6 +1954,10 @@ For example:
19541954
The first three options are mutually exclusive, and the remaining options
19551955
describe more details of how the function behaves. The remaining options
19561956
are invalid for "free"-type functions.
1957+
``"alloc-variant-zeroed"="FUNCTION"``
1958+
This attribute indicates that another function is equivalent to an allocator function,
1959+
but returns zeroed memory. The function must have "zeroed" allocation behavior,
1960+
the same ``alloc-family``, and take exactly the same arguments.
19571961
``allocsize(<EltSizeParam>[, <NumEltsParam>])``
19581962
This attribute indicates that the annotated function will always return at
19591963
least a given number of bytes (or null). Its arguments are zero-indexed

‎llvm/include/llvm/IR/Attributes.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,8 @@ class Attribute {
160160
LLVM_ABI static Attribute
161161
getWithAllocSizeArgs(LLVMContext &Context, unsigned ElemSizeArg,
162162
const std::optional<unsigned> &NumElemsArg);
163+
LLVM_ABI static Attribute getWithAllocKind(LLVMContext &Context,
164+
AllocFnKind Kind);
163165
LLVM_ABI static Attribute getWithVScaleRangeArgs(LLVMContext &Context,
164166
unsigned MinValue,
165167
unsigned MaxValue);

‎llvm/lib/IR/Attributes.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,10 @@ Attribute::getWithAllocSizeArgs(LLVMContext &Context, unsigned ElemSizeArg,
300300
return get(Context, AllocSize, packAllocSizeArgs(ElemSizeArg, NumElemsArg));
301301
}
302302

303+
Attribute Attribute::getWithAllocKind(LLVMContext &Context, AllocFnKind Kind) {
304+
return get(Context, AllocKind, static_cast<uint64_t>(Kind));
305+
}
306+
303307
Attribute Attribute::getWithVScaleRangeArgs(LLVMContext &Context,
304308
unsigned MinValue,
305309
unsigned MaxValue) {

‎llvm/lib/IR/Verifier.cpp‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2382,6 +2382,31 @@ void Verifier::verifyFunctionAttrs(FunctionType *FT, AttributeList Attrs,
23822382
CheckFailed("'allockind()' can't be both zeroed and uninitialized");
23832383
}
23842384

2385+
if (Attribute A = Attrs.getFnAttr("alloc-variant-zeroed"); A.isValid()) {
2386+
StringRef S = A.getValueAsString();
2387+
Check(!S.empty(), "'alloc-variant-zeroed' must not be empty");
2388+
Function *Variant = M.getFunction(S);
2389+
if (Variant) {
2390+
Attribute Family = Attrs.getFnAttr("alloc-family");
2391+
Attribute VariantFamily = Variant->getFnAttribute("alloc-family");
2392+
if (Family.isValid())
2393+
Check(VariantFamily.isValid() &&
2394+
VariantFamily.getValueAsString() == Family.getValueAsString(),
2395+
"'alloc-variant-zeroed' must name a function belonging to the "
2396+
"same 'alloc-family'");
2397+
2398+
Check(Variant->hasFnAttribute(Attribute::AllocKind) &&
2399+
(Variant->getFnAttribute(Attribute::AllocKind).getAllocKind() &
2400+
AllocFnKind::Zeroed) != AllocFnKind::Unknown,
2401+
"'alloc-variant-zeroed' must name a function with "
2402+
"'allockind(\"zeroed\")'");
2403+
2404+
Check(FT == Variant->getFunctionType(),
2405+
"'alloc-variant-zeroed' must name a function with the same "
2406+
"signature");
2407+
}
2408+
}
2409+
23852410
if (Attrs.hasFnAttr(Attribute::VScaleRange)) {
23862411
unsigned VScaleMin = Attrs.getFnAttrs().getVScaleRangeMin();
23872412
if (VScaleMin == 0)

‎llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp‎

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2030,9 +2030,17 @@ struct DSEState {
20302030
if (!InnerCallee)
20312031
return false;
20322032
LibFunc Func;
2033+
StringRef ZeroedVariantName;
20332034
if (!TLI.getLibFunc(*InnerCallee, Func) || !TLI.has(Func) ||
2034-
Func != LibFunc_malloc)
2035-
return false;
2035+
Func != LibFunc_malloc) {
2036+
Attribute Attr = Malloc->getFnAttr("alloc-variant-zeroed");
2037+
if (!Attr.isValid())
2038+
return false;
2039+
ZeroedVariantName = Attr.getValueAsString();
2040+
if (ZeroedVariantName.empty())
2041+
return false;
2042+
}
2043+
20362044
// Gracefully handle malloc with unexpected memory attributes.
20372045
auto *MallocDef = dyn_cast_or_null<MemoryDef>(MSSA.getMemoryAccess(Malloc));
20382046
if (!MallocDef)
@@ -2059,15 +2067,32 @@ struct DSEState {
20592067

20602068
if (Malloc->getOperand(0) != MemSet->getLength())
20612069
return false;
2062-
if (!shouldCreateCalloc(Malloc, MemSet) ||
2063-
!DT.dominates(Malloc, MemSet) ||
2070+
if (!shouldCreateCalloc(Malloc, MemSet) || !DT.dominates(Malloc, MemSet) ||
20642071
!memoryIsNotModifiedBetween(Malloc, MemSet, BatchAA, DL, &DT))
20652072
return false;
20662073
IRBuilder<> IRB(Malloc);
2067-
Type *SizeTTy = Malloc->getArgOperand(0)->getType();
2068-
auto *Calloc =
2069-
emitCalloc(ConstantInt::get(SizeTTy, 1), Malloc->getArgOperand(0), IRB,
2070-
TLI, Malloc->getType()->getPointerAddressSpace());
2074+
assert(Func == LibFunc_malloc || !ZeroedVariantName.empty());
2075+
Value *Calloc = nullptr;
2076+
if (!ZeroedVariantName.empty()) {
2077+
LLVMContext &Ctx = Malloc->getContext();
2078+
AttributeList Attrs = InnerCallee->getAttributes();
2079+
AllocFnKind AllocKind =
2080+
Attrs.getFnAttr(Attribute::AllocKind).getAllocKind() |
2081+
AllocFnKind::Zeroed;
2082+
Attrs =
2083+
Attrs.addFnAttribute(Ctx, Attribute::getWithAllocKind(Ctx, AllocKind))
2084+
.removeFnAttribute(Ctx, "alloc-variant-zeroed");
2085+
FunctionCallee ZeroedVariant = Malloc->getModule()->getOrInsertFunction(
2086+
ZeroedVariantName, InnerCallee->getFunctionType(), Attrs);
2087+
SmallVector<Value *, 3> Args;
2088+
Args.append(Malloc->arg_begin(), Malloc->arg_end());
2089+
Calloc = IRB.CreateCall(ZeroedVariant, Args, ZeroedVariantName);
2090+
} else {
2091+
Type *SizeTTy = Malloc->getArgOperand(0)->getType();
2092+
Calloc =
2093+
emitCalloc(ConstantInt::get(SizeTTy, 1), Malloc->getArgOperand(0),
2094+
IRB, TLI, Malloc->getType()->getPointerAddressSpace());
2095+
}
20712096
if (!Calloc)
20722097
return false;
20732098

‎llvm/test/Transforms/DeadStoreElimination/noop-stores.ll‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -374,6 +374,19 @@ define ptr @notmalloc_memset(i64 %size, ptr %notmalloc) {
374374
ret ptr %call1
375375
}
376376

377+
; This should create a customalloc_zeroed call and eliminate the memset
378+
define ptr @customalloc_memset(i64 %size, i64 %align) {
379+
; CHECK-LABEL: @customalloc_memset
380+
; CHECK-NEXT: [[CALL:%.*]] = call ptr @customalloc_zeroed(i64 [[SIZE:%.*]], i64 [[ALIGN:%.*]])
381+
; CHECK-NEXT: ret ptr [[CALL]]
382+
%call = call ptr @customalloc(i64 %size, i64 %align)
383+
call void @llvm.memset.p0.i64(ptr %call, i8 0, i64 %size, i1 false)
384+
ret ptr %call
385+
}
386+
387+
declare ptr @customalloc(i64, i64) allockind("alloc") "alloc-family"="customalloc" "alloc-variant-zeroed"="customalloc_zeroed"
388+
declare ptr @customalloc_zeroed(i64, i64) allockind("alloc,zeroed") "alloc-family"="customalloc"
389+
377390
; This should not create recursive call to calloc.
378391
define ptr @calloc(i64 %nmemb, i64 %size) inaccessiblememonly {
379392
; CHECK-LABEL: @calloc(

‎llvm/test/Transforms/DeadStoreElimination/zeroed-missing.ll‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
; RUN: opt < %s -passes='dse,verify<memoryssa>' -S | FileCheck %s
2+
3+
declare void @llvm.memset.p0.i64(ptr nocapture, i8, i64, i1) nounwind
4+
5+
; This should create a declaration for the named variant
6+
define ptr @undeclared_customalloc(i64 %size, i64 %align) {
7+
; CHECK-LABEL: @undeclared_customalloc
8+
; CHECK-NEXT: [[CALL:%.*]] = call ptr @customalloc2_zeroed(i64 [[SIZE:%.*]], i64 [[ALIGN:%.*]])
9+
; CHECK-NEXT: ret ptr [[CALL]]
10+
%call = call ptr @customalloc2(i64 %size, i64 %align)
11+
call void @llvm.memset.p0.i64(ptr %call, i8 0, i64 %size, i1 false)
12+
ret ptr %call
13+
}
14+
15+
declare ptr @customalloc2(i64, i64) allockind("alloc") "alloc-family"="customalloc2" "alloc-variant-zeroed"="customalloc2_zeroed"
16+
; CHECK-DAG: declare ptr @customalloc2_zeroed(i64, i64) #[[CA2ATTR:[0-9]+]]
17+
; CHECK-DAG: attributes #[[CA2ATTR]] = { allockind("alloc,zeroed") "alloc-family"="customalloc2" }

‎llvm/test/Verifier/alloc-variant-zeroed.ll‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
; RUN: not llvm-as %s -o /dev/null 2>&1 | FileCheck %s
2+
3+
; CHECK: 'alloc-variant-zeroed' must not be empty
4+
declare ptr @a(i64) "alloc-variant-zeroed"=""
5+
6+
; CHECK: 'alloc-variant-zeroed' must not be empty
7+
declare ptr @b(i64) "alloc-variant-zeroed"=""
8+
9+
; CHECK: 'alloc-variant-zeroed' must name a function belonging to the same 'alloc-family'
10+
declare ptr @c(i64) "alloc-variant-zeroed"="c_zeroed" "alloc-family"="C"
11+
declare ptr @c_zeroed(i64)
12+
13+
; CHECK: 'alloc-variant-zeroed' must name a function with 'allockind("zeroed")'
14+
declare ptr @d(i64) "alloc-variant-zeroed"="d_zeroed" "alloc-family"="D"
15+
declare ptr @d_zeroed(i64) "alloc-family"="D"
16+
17+
; CHECK: 'alloc-variant-zeroed' must name a function with the same signature
18+
declare ptr @e(i64) "alloc-variant-zeroed"="e_zeroed" "alloc-family"="E"
19+
declare ptr @e_zeroed(i64, i64) "alloc-family"="E" allockind("zeroed")

0 commit comments

Comments
 (0)