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

[via llvm-commits]

================
@@ -374,6 +374,33 @@ 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 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: declare ptr @customalloc2_zeroed(i64, i64) #[[CA2ATTR:[0-9]+]]
+; CHECK: attributes #[[CA2ATTR]] = { allockind("alloc,zeroed") "alloc-family"="customalloc2" "alloc-variant-zeroed"="customalloc2_zeroed" }
+
----------------
clubby789 wrote:

This test works in isolation but FileCheck doesn't find the declaration when used as part of the complete file. What's the correct way of expressing this?

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