Support non-malloc functions in malloc+memset->calloc fold

Author: jamie-osec

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

@@ -2027,15 +2027,45 @@ struct DSEState {
     auto *InnerCallee = Malloc->getCalledFunction();
     if (!InnerCallee)
       return false;
-    LibFunc Func;
+    LibFunc Func = NotLibFunc;
     if (!TLI.getLibFunc(*InnerCallee, Func) || !TLI.has(Func) ||
-        Func != LibFunc_malloc)
-      return false;
+        Func != LibFunc_malloc) {
+      if (!isAllocLikeFn(Malloc, &TLI))
+        return false;
+    }
     // Gracefully handle malloc with unexpected memory attributes.
     auto *MallocDef = dyn_cast_or_null<MemoryDef>(MSSA.getMemoryAccess(Malloc));
     if (!MallocDef)
       return false;
 
+    // If the allocation is not a malloc, search for a matching zero-alloc
+    // function
+    Function *ZeroedAlloc = nullptr;
+
+    if (Func != LibFunc_malloc) {
+      auto OptFamily = getAllocationFamily(Malloc, &TLI);
+      if (!OptFamily)
+        return false;
+      auto &Family = *OptFamily;
+      auto *Module = InnerCallee->getParent();
+      for (auto &ModuleFunc : Module->functions()) {
+        auto Attributes = ModuleFunc.getAttributes();
+        if (!Attributes.hasFnAttr("alloc-family") ||
+            Attributes.getFnAttr("alloc-family").getValueAsString() != Family ||
+            !Attributes.hasFnAttr(llvm::Attribute::AllocKind) ||
+            (Attributes.getAllocKind() & AllocFnKind::Zeroed) ==
+                AllocFnKind::Unknown)
+          continue;
+        ZeroedAlloc = &ModuleFunc;
+        break;
+      }
+      if (!ZeroedAlloc)
+        return false;
+    }
+
+    assert((Func == LibFunc_malloc || ZeroedAlloc) &&
+           "didn't resolve allocation functions");
+
     auto shouldCreateCalloc = [](CallInst *Malloc, CallInst *Memset) {
       // Check for br(icmp ptr, null), truebb, falsebb) pattern at the end
       // of malloc block
@@ -2054,6 +2084,18 @@ struct DSEState {
         return false;
       return true;
     };
+    auto RemapArgs =
+        [&ZeroedAlloc](CallInst *Malloc) -> SmallVector<Value *, 3> {
+      // If the arg counts match, use the existing args
+      if (Malloc->arg_size() == ZeroedAlloc->arg_size()) {
+        SmallVector<Value *, 3> SV;
+        for (unsigned I = 0; I < Malloc->arg_size(); I++)
+          SV.push_back(Malloc->getArgOperand(I));
+        return SV;
+      };
+      // Unknown mapping
+      return {};
+    };
 
     if (Malloc->getOperand(0) != MemSet->getLength())
       return false;
@@ -2062,10 +2104,19 @@ struct DSEState {
         !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());
+    Value *Calloc;
+    if (Func == LibFunc_malloc) {
+      Type *SizeTTy = Malloc->getArgOperand(0)->getType();
+      Calloc =
+          emitCalloc(ConstantInt::get(SizeTTy, 1), Malloc->getArgOperand(0),
+                     IRB, TLI, Malloc->getType()->getPointerAddressSpace());
+    } else {
+      auto Args = RemapArgs(Malloc);
+      if (Args.empty())
+        return false;
+      Calloc = IRB.CreateCall(FunctionCallee(ZeroedAlloc), Args,
+                              ZeroedAlloc->getName());
+    }
     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
+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"
+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(