[IR][DSE] Support non-malloc functions in malloc+memset->calloc fold

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

llvm/include/llvm/IR/Attributes.h

@@ -160,6 +160,8 @@ class Attribute {
   LLVM_ABI static Attribute
   getWithAllocSizeArgs(LLVMContext &Context, unsigned ElemSizeArg,
                        const std::optional<unsigned> &NumElemsArg);
+  LLVM_ABI static Attribute getWithAllocKind(LLVMContext &Context,
+                                             AllocFnKind Kind);
   LLVM_ABI static Attribute getWithVScaleRangeArgs(LLVMContext &Context,
                                                    unsigned MinValue,
                                                    unsigned MaxValue);

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) {

llvm/lib/IR/Verifier.cpp

@@ -2382,6 +2382,31 @@ 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");
+    Function *Variant = M.getFunction(S);
+    if (Variant) {
+      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(FT == 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)

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

@@ -2030,9 +2030,17 @@ struct DSEState {
     if (!InnerCallee)
       return false;
     LibFunc Func;
+    StringRef ZeroedVariantName;
     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)
@@ -2059,15 +2067,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.empty());
+    Value *Calloc = nullptr;
+    if (!ZeroedVariantName.empty()) {
+      LLVMContext &Ctx = Malloc->getContext();
+      AttributeList Attrs = InnerCallee->getAttributes();
+      AllocFnKind AllocKind =
+          Attrs.getFnAttr(Attribute::AllocKind).getAllocKind() |
+          AllocFnKind::Zeroed;
+      Attrs =
+          Attrs.addFnAttribute(Ctx, Attribute::getWithAllocKind(Ctx, AllocKind))
+              .removeFnAttribute(Ctx, "alloc-variant-zeroed");
+      FunctionCallee ZeroedVariant = Malloc->getModule()->getOrInsertFunction(
+          ZeroedVariantName, InnerCallee->getFunctionType(), Attrs);
+      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;
 

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

@@ -374,6 +374,19 @@ 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 not create recursive call to calloc.
 define ptr @calloc(i64 %nmemb, i64 %size) inaccessiblememonly {
 ; CHECK-LABEL: @calloc(

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

@@ -0,0 +1,17 @@
+; RUN: opt < %s -passes='dse,verify<memoryssa>' -S | FileCheck %s
+
+declare void @llvm.memset.p0.i64(ptr nocapture, i8, i64, i1) nounwind
+
+; 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" }

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")