[llvm] 4852374 - [llvm][opt][Transforms] Replacement `calloc` should match replaced `malloc` (#110524)

[via llvm-commits]

Author: Alex Voicu
Date: 2024-10-01T02:05:28+01:00
New Revision: 4852374135773b03c14ba2003be99ed1169dedf4

URL: https://github.com/llvm/llvm-project/commit/4852374135773b03c14ba2003be99ed1169dedf4
DIFF: https://github.com/llvm/llvm-project/commit/4852374135773b03c14ba2003be99ed1169dedf4.diff

LOG: [llvm][opt][Transforms] Replacement `calloc` should match replaced `malloc` (#110524)

Currently DSE unconditionally emits `calloc` as returning a pointer to
AS0. However, this is incorrect for targets that have a non-zero default
AS, as it'd not match the `malloc` signature. This patch addresses that
by piping through the AS for the pointer returned by `malloc` into the
`calloc` insertion call.

Added: 
    llvm/test/Transforms/DeadStoreElimination/malloc-to-calloc-with-nonzero-default-as.ll

Modified: 
    llvm/include/llvm/Transforms/Utils/BuildLibCalls.h
    llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
    llvm/lib/Transforms/Utils/BuildLibCalls.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Transforms/Utils/BuildLibCalls.h b/llvm/include/llvm/Transforms/Utils/BuildLibCalls.h
index 1979c4af770b02..a8fb38e7260043 100644
--- a/llvm/include/llvm/Transforms/Utils/BuildLibCalls.h
+++ b/llvm/include/llvm/Transforms/Utils/BuildLibCalls.h
@@ -251,7 +251,7 @@ namespace llvm {
 
   /// Emit a call to the calloc function.
   Value *emitCalloc(Value *Num, Value *Size, IRBuilderBase &B,
-                    const TargetLibraryInfo &TLI);
+                    const TargetLibraryInfo &TLI, unsigned AddrSpace);
 
   /// Emit a call to the hot/cold operator new function.
   Value *emitHotColdNew(Value *Num, IRBuilderBase &B,

diff  --git a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
index a304f7b056f5f7..ce8c988ba531dd 100644
--- a/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp
@@ -1944,8 +1944,9 @@ struct DSEState {
       return false;
     IRBuilder<> IRB(Malloc);
     Type *SizeTTy = Malloc->getArgOperand(0)->getType();
-    auto *Calloc = emitCalloc(ConstantInt::get(SizeTTy, 1),
-                              Malloc->getArgOperand(0), IRB, TLI);
+    auto *Calloc =
+        emitCalloc(ConstantInt::get(SizeTTy, 1), Malloc->getArgOperand(0), IRB,
+                   TLI, Malloc->getType()->getPointerAddressSpace());
     if (!Calloc)
       return false;
 

diff  --git a/llvm/lib/Transforms/Utils/BuildLibCalls.cpp b/llvm/lib/Transforms/Utils/BuildLibCalls.cpp
index d4727dece19f62..7bb4b55fcb7cf2 100644
--- a/llvm/lib/Transforms/Utils/BuildLibCalls.cpp
+++ b/llvm/lib/Transforms/Utils/BuildLibCalls.cpp
@@ -1978,15 +1978,15 @@ Value *llvm::emitMalloc(Value *Num, IRBuilderBase &B, const DataLayout &DL,
 }
 
 Value *llvm::emitCalloc(Value *Num, Value *Size, IRBuilderBase &B,
-                        const TargetLibraryInfo &TLI) {
+                        const TargetLibraryInfo &TLI, unsigned AddrSpace) {
   Module *M = B.GetInsertBlock()->getModule();
   if (!isLibFuncEmittable(M, &TLI, LibFunc_calloc))
     return nullptr;
 
   StringRef CallocName = TLI.getName(LibFunc_calloc);
   Type *SizeTTy = getSizeTTy(B, &TLI);
-  FunctionCallee Calloc = getOrInsertLibFunc(M, TLI, LibFunc_calloc,
-                                             B.getPtrTy(), SizeTTy, SizeTTy);
+  FunctionCallee Calloc = getOrInsertLibFunc(
+      M, TLI, LibFunc_calloc, B.getPtrTy(AddrSpace), SizeTTy, SizeTTy);
   inferNonMandatoryLibFuncAttrs(M, CallocName, TLI);
   CallInst *CI = B.CreateCall(Calloc, {Num, Size}, CallocName);
 

diff  --git a/llvm/test/Transforms/DeadStoreElimination/malloc-to-calloc-with-nonzero-default-as.ll b/llvm/test/Transforms/DeadStoreElimination/malloc-to-calloc-with-nonzero-default-as.ll
new file mode 100644
index 00000000000000..977bf93fa856e1
--- /dev/null
+++ b/llvm/test/Transforms/DeadStoreElimination/malloc-to-calloc-with-nonzero-default-as.ll
@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=dse < %s | FileCheck %s
+
+define ptr addrspace(4) @malloc_to_calloc(i64 %size) {
+; CHECK-LABEL: define ptr addrspace(4) @malloc_to_calloc(
+; CHECK-SAME: i64 [[SIZE:%.*]]) {
+; CHECK-NEXT:    [[CALLOC:%.*]] = call ptr addrspace(4) @calloc(i64 1, i64 [[SIZE]])
+; CHECK-NEXT:    ret ptr addrspace(4) [[CALLOC]]
+;
+  %ret = call ptr addrspace(4) @malloc(i64 %size)
+  call void @llvm.memset.p4.i64(ptr addrspace(4) %ret, i8 0, i64 %size, i1 false)
+  ret ptr addrspace(4) %ret
+}
+
+declare void @llvm.memset.p4.i64(ptr addrspace(4) nocapture writeonly, i8, i64, i1 immarg)
+
+declare noalias ptr addrspace(4) @malloc(i64) willreturn allockind("alloc,uninitialized") "alloc-family"="malloc"