[OpenMP][AMDGPU] Add the support for `malloc` and `free` for AMDGPU

[tianshilei1992]

Summary

This patch adds the basic support for malloc and free for AMDGPU.
It borrows the idea from clang/lib/Headers/__clang_hip_runtime_wrapper.h but
ignores the two checks: HIP version, and whether enabling device malloc. Arguably,
this doesn't make things worse. Currently no matter which HIP version users are
using, or whether the device malloc is enabled, they will all get link error if
they use malloc and free in target region. Then for users using pre ROCm 4.5,
they still get link error. That also helps to get rid of the check whether device
malloc is enabled because it is only needed for before ROCm 4.5. The good part is,
users using ROCm 4.5+ would get the feature.

[tianshilei1992]

remove unrelated changes

[tianshilei1992]

fix rong header macro

[jdoerfert]

openmp/libomptarget/DeviceRTL/include/Memory.h

19

I don't think we can/should define it like this. Maybe using size_t = decltype(sizeof(char)); instead? Also, this should arguably be in Types.h.

openmp/libomptarget/DeviceRTL/src/Memory.cpp

32

I think we want (now or later) a second level of indirection here.
malloc -> _OMP::impl::malloc -> (__ockl_dm_alloc, ...)

Also, we need to annotate all external functions properly, e.g., see malloc and free in State.cpp.
Finally, combine this with existing malloc/free declarations and uses in the runtime.

[tianshilei1992]

fix comments

[tianshilei1992]

openmp/libomptarget/DeviceRTL/src/Memory.cpp

32

That would be our next step after we have a proper allocator.

[jdoerfert]

LG, minor nits

openmp/libomptarget/DeviceRTL/include/Types.h

45

We don't require 8 byte size_t, do we? size_t is by definition supposed to be the sizeof result type, not 8 byte.

openmp/libomptarget/DeviceRTL/src/Memory.cpp

25

Rename the arguments.

Unrelated: I really have no words for dealloc taking a unsigned long long.

[tianshilei1992]

rebase and fix comments

[JonChesterfield]

This is missing changes to the host plugin to set up the state malloc uses, and missing a host call implementation to do the allocation. So as written this isn't going to work. What's the benefit?

[tianshilei1992]

It looks like there are a couple of things still missing:

• malloc and free defined for AMDGPU using declare variant cannot be linked into the device module correctly. Their names still contain declare variant mangling such that they are dropped directly, leaving malloc and free still undefined.
• Even manually link the functions into device module, because some host setup are missing, as Jon pointed out, there is still some memory error:

[GPU Memory Error] Addr: 0x0 Reason: Page not present or supervisor privilege.
Memory access fault by GPU node-8 (Agent handle: 0x1428d10) on address (nil). Reason: Page not present or supervisor privilege.

[JonChesterfield]

Yep, the host plugin will be passing 0 at an offset where the ockl machinery expects some non-null value. I think that allocator uses some initial block allocated per kernel launch (or possibly per device image), but it should also fall back on a host call to get more memory when that is exhausted which will not work in trunk. It might be documented what setup is expected for the pointer (in particular whether it's shared across kernels and how much memory should be allocated there), otherwise it could be reverse engineered from hip.

Alternatively we could use a different implementation for malloc/free.

[tianshilei1992]

Move to a more general way to implement it w/o depending on vendor interfaces.

openmp/libomptarget/DeviceRTL/CMakeLists.txt

 set(LIBOMPTARGET_DEVICE_DEBUG FALSE CACHE BOOL
   "Activate DeviceRTL debug messages.")
 
 set(include_files
   ${include_directory}/Configuration.h
   ${include_directory}/Debug.h
   ${include_directory}/Interface.h
   ${include_directory}/Mapping.h
+  ${include_directory}/Memory.h
   ${include_directory}/State.h
   ${include_directory}/Synchronization.h
   ${include_directory}/Types.h
   ${include_directory}/Utils.h
 )
 
 set(src_files
   ${source_directory}/Configuration.cpp
   ${source_directory}/Debug.cpp
   ${source_directory}/Kernel.cpp
   ${source_directory}/Mapping.cpp
+  ${source_directory}/Memory.cpp
   ${source_directory}/Misc.cpp
   ${source_directory}/Parallelism.cpp
   ${source_directory}/Reduction.cpp
   ${source_directory}/State.cpp
   ${source_directory}/Synchronization.cpp
   ${source_directory}/Tasking.cpp
   ${source_directory}/Utils.cpp
   ${source_directory}/Workshare.cpp

openmp/libomptarget/DeviceRTL/include/Memory.h

+//===--- Memory.h - OpenMP device runtime memory allocator -------- 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 OMPTARGET_MEMORY_H
+#define OMPTARGET_MEMORY_H
+
+#include "Types.h"
+
+extern "C" {
+__attribute__((leaf)) void *malloc(uint64_t Size);
+__attribute__((leaf)) void free(void *Ptr);
+}
+
+#endif

openmp/libomptarget/DeviceRTL/include/Types.h

 using int8_t = char;
 using uint8_t = unsigned char;
 using int16_t = short;
 using uint16_t = unsigned short;
 using int32_t = int;
 using uint32_t = unsigned int;
 using int64_t = long;
 using uint64_t = unsigned long;
+using size_t = decltype(sizeof(char));
 
 static_assert(sizeof(int8_t) == 1, "type size mismatch");
 static_assert(sizeof(uint8_t) == 1, "type size mismatch");
 static_assert(sizeof(int16_t) == 2, "type size mismatch");
 static_assert(sizeof(uint16_t) == 2, "type size mismatch");
 static_assert(sizeof(int32_t) == 4, "type size mismatch");
 static_assert(sizeof(uint32_t) == 4, "type size mismatch");
 static_assert(sizeof(int64_t) == 8, "type size mismatch");
 static_assert(sizeof(uint64_t) == 8, "type size mismatch");
 ///}
 
 enum omp_proc_bind_t {
   omp_proc_bind_false = 0,
   omp_proc_bind_true = 1,
   omp_proc_bind_master = 2,
   omp_proc_bind_close = 3,
   omp_proc_bind_spread = 4
 };

openmp/libomptarget/DeviceRTL/src/CMakeLists.txt

 target_sources(omptarget.devicertl PRIVATE
   Configuration.cpp
   Debug.cpp
   Kernel.cpp
   Mapping.cpp
+  Memory.cpp
   Misc.cpp
   Parallelism.cpp
   Reduction.cpp
   State.cpp
   Synchronization.cpp
   Tasking.cpp
   Utils.cpp
   Workshare.cpp)

openmp/libomptarget/DeviceRTL/src/Memory.cpp

+//===------- Memory.cpp - OpenMP device runtime memory allocator -- 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
+//
+//===----------------------------------------------------------------------===//
+//
+//
+//===----------------------------------------------------------------------===//
+
+#pragma omp begin declare target device_type(nohost)
+
+#include "Memory.h"
+
+/// AMDGCN Implementation
+///
+///{
+#pragma omp begin declare variant match(device = {arch(amdgcn)})
+
+extern "C" {
+__attribute__((leaf)) unsigned long long
+__ockl_dm_alloc(unsigned long long Size);
+__attribute__((leaf)) void __ockl_dm_dealloc(unsigned long long Addr);
+
+void *malloc(size_t Size) { return (void *)__ockl_dm_alloc(Size); }
+
+void free(void *P) { __ockl_dm_dealloc((unsigned long long)P); }
+}
+
+#pragma omp end declare variant
+
+///}
+
+#pragma omp end declare target

openmp/libomptarget/DeviceRTL/src/State.cpp

 //===------ State.cpp - OpenMP State & ICV interface ------------- 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 "State.h"
 #include "Configuration.h"
 #include "Debug.h"
 #include "Interface.h"
 #include "Mapping.h"
+#include "Memory.h"
 #include "Synchronization.h"
 #include "Types.h"
 #include "Utils.h"
 
 using namespace _OMP;
 
 #pragma omp begin declare target device_type(nohost)
 
 /// Memory implementation
 ///
 ///{
 
 /// Add worst-case padding so that future allocations are properly aligned.
 /// FIXME: The stack shouldn't require worst-case padding. Alignment needs to be
 /// passed in as an argument and the stack rewritten to support it.
 constexpr const uint32_t Alignment = 16;
 
 /// External symbol to access dynamic shared memory.
 extern unsigned char DynamicSharedBuffer[] __attribute__((aligned(Alignment)));
 #pragma omp allocate(DynamicSharedBuffer) allocator(omp_pteam_mem_alloc)
 
 namespace {
 
-/// Fallback implementations are missing to trigger a link time error.
-/// Implementations for new devices, including the host, should go into a
-/// dedicated begin/end declare variant.
-///
-///{
-
-extern "C" {
-__attribute__((leaf)) void *malloc(uint64_t Size);
-__attribute__((leaf)) void free(void *Ptr);
-}
-
-///}
-
-/// AMDGCN implementations of the shuffle sync idiom.
-///
-///{
-#pragma omp begin declare variant match(device = {arch(amdgcn)})
-
-extern "C" {
-void *malloc(uint64_t Size) {
-  // TODO: Use some preallocated space for dynamic malloc.
-  return nullptr;
-}
-
-void free(void *Ptr) {}
-}
-
-#pragma omp end declare variant
-///}
-
 /// A "smart" stack in shared memory.
 ///
 /// The stack exposes a malloc/free interface but works like a stack internally.
 /// In fact, it is a separate stack *per warp*. That means, each warp must push
 /// and pop symmetrically or this breaks, badly. The implementation will (aim
 /// to) detect non-lock-step warps and fallback to malloc/free. The same will
 /// happen if a warp runs out of memory. The master warp in generic memory is
 /// special and is given more memory than the rest.