[libc][NFC] Move GPU allocator implementation to common header (#84690)

Summary: This is a NFC move preceding more radical functional changes to the allocator implementation. We just move it to a common utility so it will be easier to write these in tandem.
llvm · Mar 10, 2024 · ea697dc · ea697dc
1 parent a53401e
commit ea697dc
Show file tree

Hide file tree

Showing 6 changed files with 88 additions and 19 deletions.
diff --git a/libc/src/__support/GPU/CMakeLists.txt b/libc/src/__support/GPU/CMakeLists.txt
@@ -12,3 +12,15 @@ add_header_library(
   DEPENDS
     ${target_gpu_utils}
 )
+
+add_object_library(
+  allocator
+  SRCS
+    allocator.cpp
+  HDRS
+    allocator.h
+  DEPENDS
+    libc.src.__support.common
+    libc.src.__support.GPU.utils
+    libc.src.__support.RPC.rpc_client
+)
diff --git a/libc/src/__support/GPU/allocator.cpp b/libc/src/__support/GPU/allocator.cpp
@@ -0,0 +1,45 @@
+//===-- GPU memory allocator implementation ---------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "allocator.h"
+
+#include "src/__support/GPU/utils.h"
+#include "src/__support/RPC/rpc_client.h"
+
+namespace LIBC_NAMESPACE {
+namespace {
+
+void *rpc_allocate(uint64_t size) {
+  void *ptr = nullptr;
+  rpc::Client::Port port = rpc::client.open<RPC_MALLOC>();
+  port.send_and_recv([=](rpc::Buffer *buffer) { buffer->data[0] = size; },
+                     [&](rpc::Buffer *buffer) {
+                       ptr = reinterpret_cast<void *>(buffer->data[0]);
+                     });
+  port.close();
+  return ptr;
+}
+
+void rpc_free(void *ptr) {
+  rpc::Client::Port port = rpc::client.open<RPC_FREE>();
+  port.send([=](rpc::Buffer *buffer) {
+    buffer->data[0] = reinterpret_cast<uintptr_t>(ptr);
+  });
+  port.close();
+}
+
+} // namespace
+
+namespace gpu {
+
+void *allocate(uint64_t size) { return rpc_allocate(size); }
+
+void deallocate(void *ptr) { rpc_free(ptr); }
+
+} // namespace gpu
+} // namespace LIBC_NAMESPACE
diff --git a/libc/src/__support/GPU/allocator.h b/libc/src/__support/GPU/allocator.h
@@ -0,0 +1,23 @@
+//===-- GPU memory allocator implementation ---------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIBC_SRC___SUPPORT_GPU_ALLOCATOR_H
+#define LLVM_LIBC_SRC___SUPPORT_GPU_ALLOCATOR_H
+
+#include <stdint.h>
+
+namespace LIBC_NAMESPACE {
+namespace gpu {
+
+void *allocate(uint64_t size);
+void deallocate(void *ptr);
+
+} // namespace gpu
+} // namespace LIBC_NAMESPACE
+
+#endif // LLVM_LIBC_SRC___SUPPORT_GPU_ALLOCATOR_H
diff --git a/libc/src/stdlib/gpu/CMakeLists.txt b/libc/src/stdlib/gpu/CMakeLists.txt
@@ -6,7 +6,7 @@ add_entrypoint_object(
     ../malloc.h
   DEPENDS
     libc.include.stdlib
-    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.allocator
 )
 
 add_entrypoint_object(
@@ -28,5 +28,5 @@ add_entrypoint_object(
     ../abort.h
   DEPENDS
     libc.include.stdlib
-    libc.src.__support.RPC.rpc_client
+    libc.src.__support.GPU.allocator
 )
diff --git a/libc/src/stdlib/gpu/free.cpp b/libc/src/stdlib/gpu/free.cpp
@@ -7,17 +7,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/stdlib/free.h"
-#include "src/__support/RPC/rpc_client.h"
+
+#include "src/__support/GPU/allocator.h"
 #include "src/__support/common.h"
 
 namespace LIBC_NAMESPACE {
 
-LLVM_LIBC_FUNCTION(void, free, (void *ptr)) {
-  rpc::Client::Port port = rpc::client.open<RPC_FREE>();
-  port.send([=](rpc::Buffer *buffer) {
-    buffer->data[0] = reinterpret_cast<uintptr_t>(ptr);
-  });
-  port.close();
-}
+LLVM_LIBC_FUNCTION(void, free, (void *ptr)) { gpu::deallocate(ptr); }
 
 } // namespace LIBC_NAMESPACE
diff --git a/libc/src/stdlib/gpu/malloc.cpp b/libc/src/stdlib/gpu/malloc.cpp
@@ -7,20 +7,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "src/stdlib/malloc.h"
-#include "src/__support/RPC/rpc_client.h"
+
+#include "src/__support/GPU/allocator.h"
 #include "src/__support/common.h"
 
 namespace LIBC_NAMESPACE {
 
 LLVM_LIBC_FUNCTION(void *, malloc, (size_t size)) {
-  void *ptr = nullptr;
-  rpc::Client::Port port = rpc::client.open<RPC_MALLOC>();
-  port.send_and_recv([=](rpc::Buffer *buffer) { buffer->data[0] = size; },
-                     [&](rpc::Buffer *buffer) {
-                       ptr = reinterpret_cast<void *>(buffer->data[0]);
-                     });
-  port.close();
-  return ptr;
+  return gpu::allocate(size);
 }
 
 } // namespace LIBC_NAMESPACE