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

[Nikita Popov via llvm-commits]

================
@@ -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());
----------------
nikic wrote:

I was going to say that we should handle malloc with the generic mechanism by inferring the attribute in https://github.com/llvm/llvm-project/blob/55e39100dae2a131c7cbcdbbe628c0b82f8c12ca/llvm/lib/Transforms/Utils/BuildLibCalls.cpp#L515 But calloc has a different signature, so the generic mechanism doesn't work for it. And it's probably not worthwhile to try to extend it in a way to make it work.

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