[llvm] [opt][HIPSTDPAR] Handle `__hipstdpar_hidden_malloc` (PR #123150)

[Alex Voicu via llvm-commits]

https://github.com/AlexVlx created https://github.com/llvm/llvm-project/pull/123150

This patch adds support for replacing calls to `__hipstdpar_hidden_malloc` with calls to `__libc_malloc`, similarly to how we handle hidden `free`. A future paired change in the forwarding header will leverage this capability in order to handle certain special cases where it is not possible / desirable to allocate via the HIP runtime.

>From b60773af8dbda229b1450726c67ade113de0ae05 Mon Sep 17 00:00:00 2001
From: Alex Voicu <alexandru.voicu at amd.com>
Date: Thu, 16 Jan 2025 02:06:11 +0000
Subject: [PATCH] Handle `__hipstdpar_hidden_malloc`

---
 llvm/lib/Transforms/HipStdPar/HipStdPar.cpp               | 8 ++++++++
 .../test/Transforms/HipStdPar/allocation-interposition.ll | 2 ++
 2 files changed, 10 insertions(+)

diff --git a/llvm/lib/Transforms/HipStdPar/HipStdPar.cpp b/llvm/lib/Transforms/HipStdPar/HipStdPar.cpp
index 92042ddab38dc7..691aaecc5e006b 100644
--- a/llvm/lib/Transforms/HipStdPar/HipStdPar.cpp
+++ b/llvm/lib/Transforms/HipStdPar/HipStdPar.cpp
@@ -298,6 +298,14 @@ HipStdParAllocationInterpositionPass::run(Module &M, ModuleAnalysisManager&) {
     }
   }
 
+  if (auto F = M.getFunction("__hipstdpar_hidden_malloc")) {
+    auto LibcMalloc =
+        M.getOrInsertFunction("__libc_malloc", F->getFunctionType(),
+                              F->getAttributes());
+    F->replaceAllUsesWith(LibcMalloc.getCallee());
+
+    eraseFromModule(*F);
+  }
   if (auto F = M.getFunction("__hipstdpar_hidden_free")) {
     auto LibcFree = M.getOrInsertFunction("__libc_free", F->getFunctionType(),
                                           F->getAttributes());
diff --git a/llvm/test/Transforms/HipStdPar/allocation-interposition.ll b/llvm/test/Transforms/HipStdPar/allocation-interposition.ll
index 291b06ed0ca9ed..9ec284b1dedb7d 100644
--- a/llvm/test/Transforms/HipStdPar/allocation-interposition.ll
+++ b/llvm/test/Transforms/HipStdPar/allocation-interposition.ll
@@ -14,6 +14,8 @@ declare i32 @__hipstdpar_posix_aligned_alloc(ptr, i64, i64)
 
 declare void @__hipstdpar_hidden_free(ptr)
 
+declare ptr @__hipstdpar_hidden_malloc(i64)
+
 declare ptr @__hipstdpar_realloc(ptr, i64)
 
 declare ptr @__hipstdpar_realloc_array(ptr, i64, i64)