[llvm] Support non-malloc functions in malloc+memset->calloc fold (PR #138299)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: None (clubby789)

<details>
<summary>Changes</summary>

Add a `alloc-variant-zeroed` function attribute which can be used to inform folding allocation+memset. This addresses https://github.com/rust-lang/rust/issues/104847, where LLVM does not know how to perform this transformation for non-C languages.

---
Full diff: https://github.com/llvm/llvm-project/pull/138299.diff


3 Files Affected:

- (modified) llvm/docs/LangRef.rst (+4) 
- (modified) llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp (+41-8) 
- (modified) llvm/test/Transforms/DeadStoreElimination/noop-stores.ll (+13) 


``````````diff
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/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index e318ec94db4c3..9ba13450c1fdb 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -2028,9 +2028,19 @@ 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) {
+      if (!Malloc->hasFnAttr("alloc-variant-zeroed") ||
+          Malloc->getFnAttr("alloc-variant-zeroed")
+              .getValueAsString()
+              .empty()) {
+        return false;
+      }
+      ZeroedVariantName =
+          Malloc->getFnAttr("alloc-variant-zeroed").getValueAsString();
+    }
+
     // Gracefully handle malloc with unexpected memory attributes.
     auto *MallocDef = dyn_cast_or_null<MemoryDef>(MSSA.getMemoryAccess(Malloc));
     if (!MallocDef)
@@ -2057,15 +2067,38 @@ 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()) {
+      auto *ZeroedVariant =
+          Malloc->getModule()->getFunction(*ZeroedVariantName);
+      if (!ZeroedVariant)
+        return false;
+      auto Attributes = ZeroedVariant->getAttributes();
+      auto MallocFamily = getAllocationFamily(Malloc, &TLI);
+      if (MallocFamily &&
+          *MallocFamily !=
+              Attributes.getFnAttr("alloc-family").getValueAsString())
+        return false;
+      if (!Attributes.hasFnAttr(Attribute::AllocKind) ||
+          (Attributes.getAllocKind() & AllocFnKind::Zeroed) ==
+              AllocFnKind::Unknown)
+        return false;
+
+      SmallVector<Value *, 3> Args;
+      for (unsigned I = 0; I < Malloc->arg_size(); I++)
+        Args.push_back(Malloc->getArgOperand(I));
+      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..e54d93015d587 100644
--- a/llvm/test/Transforms/DeadStoreElimination/noop-stores.ll
+++ b/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
+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(

``````````

</details>


https://github.com/llvm/llvm-project/pull/138299