[llvm] [OpenMP] Implement 'omp_alloc' on the device (PR #102526)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-offload

Author: Joseph Huber (jhuber6)

<details>
<summary>Changes</summary>

Summary:
The 'omp_alloc' function should be callable from a target region. This
patch implemets it by simply calling `malloc` for every non-default
trait value allocator. All the special access modifiers are
unimplemented and return null. The null allocator returns null as the
spec states it should not be usable from the target.


---
Full diff: https://github.com/llvm/llvm-project/pull/102526.diff


4 Files Affected:

- (modified) offload/DeviceRTL/include/Allocator.h (+5) 
- (modified) offload/DeviceRTL/src/Misc.cpp (+28) 
- (modified) offload/DeviceRTL/src/State.cpp (+2-2) 
- (added) offload/test/api/omp_device_alloc.c (+25) 


``````````diff
diff --git a/offload/DeviceRTL/include/Allocator.h b/offload/DeviceRTL/include/Allocator.h
index a28eb0fb2977ea..23e0106c80a2c8 100644
--- a/offload/DeviceRTL/include/Allocator.h
+++ b/offload/DeviceRTL/include/Allocator.h
@@ -39,6 +39,11 @@ void free(void *Ptr);
 
 } // namespace ompx
 
+extern "C" {
+[[gnu::weak]] void *malloc(size_t Size);
+[[gnu::weak]] void free(void *Ptr);
+}
+
 #pragma omp end declare target
 
 #endif
diff --git a/offload/DeviceRTL/src/Misc.cpp b/offload/DeviceRTL/src/Misc.cpp
index c24af9442d16e3..fd587a7d73b72b 100644
--- a/offload/DeviceRTL/src/Misc.cpp
+++ b/offload/DeviceRTL/src/Misc.cpp
@@ -10,6 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Configuration.h"
+#include "Allocator.h"
 #include "Types.h"
 
 #include "Debug.h"
@@ -128,6 +129,33 @@ double omp_get_wtime(void) { return ompx::impl::getWTime(); }
 void *__llvm_omp_indirect_call_lookup(void *HstPtr) {
   return ompx::impl::indirectCallLookup(HstPtr);
 }
+
+void *omp_alloc(size_t size, omp_allocator_handle_t allocator) {
+    switch(allocator) {
+  case omp_default_mem_alloc:
+  case omp_large_cap_mem_alloc:
+  case omp_const_mem_alloc:
+  case omp_high_bw_mem_alloc:
+  case omp_low_lat_mem_alloc:
+    return malloc(size);
+  default:
+    return nullptr;
+    }
+}
+
+void omp_free(void *ptr, omp_allocator_handle_t allocator) {
+    switch(allocator) {
+  case omp_default_mem_alloc:
+  case omp_large_cap_mem_alloc:
+  case omp_const_mem_alloc:
+  case omp_high_bw_mem_alloc:
+  case omp_low_lat_mem_alloc:
+    free(ptr); case omp_null_allocator:
+  default:
+    return;
+    }
+}
+
 }
 
 ///}
diff --git a/offload/DeviceRTL/src/State.cpp b/offload/DeviceRTL/src/State.cpp
index a1e4fa2449d9a2..f43f2cedb431d0 100644
--- a/offload/DeviceRTL/src/State.cpp
+++ b/offload/DeviceRTL/src/State.cpp
@@ -53,12 +53,12 @@ namespace {
 extern "C" {
 #ifdef __AMDGPU__
 
-[[gnu::weak]] void *malloc(uint64_t Size) { return allocator::alloc(Size); }
+[[gnu::weak]] void *malloc(size_t Size) { return allocator::alloc(Size); }
 [[gnu::weak]] void free(void *Ptr) { allocator::free(Ptr); }
 
 #else
 
-[[gnu::weak, gnu::leaf]] void *malloc(uint64_t Size);
+[[gnu::weak, gnu::leaf]] void *malloc(size_t Size);
 [[gnu::weak, gnu::leaf]] void free(void *Ptr);
 
 #endif
diff --git a/offload/test/api/omp_device_alloc.c b/offload/test/api/omp_device_alloc.c
new file mode 100644
index 00000000000000..46153a30e2e304
--- /dev/null
+++ b/offload/test/api/omp_device_alloc.c
@@ -0,0 +1,25 @@
+// RUN: %libomptarget-compile-run-and-check-generic
+
+#include <assert.h>
+#include <omp.h>
+#include <stdio.h>
+
+int main() {
+#pragma omp target teams num_teams(4)
+#pragma omp parallel
+  {
+    int *ptr = (int *)omp_alloc(sizeof(int), omp_default_mem_alloc);
+    assert(ptr && "Ptr is (null)!");
+    *ptr = 1;
+    assert(*ptr == 1 && "Ptr is not 1");
+    omp_free(ptr, omp_default_mem_alloc);
+  }
+
+#pragma omp target
+  {
+    assert(!omp_alloc(sizeof(int), omp_null_allocator) && "Ptr is not (null)!");
+  }
+
+  // CHECK: PASS
+  printf("PASS\n");
+}

``````````

</details>


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