[Libmoptarget] Enable the shared allocator for AMDGPU

[jhuber6]

Summary

Currently, the AMDGPU plugin did not support the TARGET_ALLOC_SHARED
allocation kind. We used the fine-grained memory allocator for the
"host" alloc when this is most likely not what is intended. Fine-grained
memory can be accessed by all agents, so it should be considered shared.
This patch removes the use of fine-grained memory for the host
allocator. A later patch will add support for this via the
hsa_amd_memory_lock method.

[kevinsala]

LGTM

openmp/libomptarget/plugins-nextgen/amdgpu/src/rtl.cpp

 #include "hsa/hsa.h"
 #include "hsa/hsa_ext_amd.h"
 #elif __has_include("hsa.h")
 #include "hsa.h"
 #include "hsa_ext_amd.h"
 #endif
 #else
 #include "hsa/hsa.h"
-#include "hsa_ext_amd.h"
+#include "hsa/hsa_ext_amd.h"
 #endif
 
 namespace llvm {
 namespace omp {
 namespace target {
 namespace plugin {
 
 /// Forward declarations for all specialized data structures.
struct AMDHostDeviceTy : public AMDGenericDeviceTy {
   /// kernel args and host pinned memory allocations.
   Error init() {
     if (auto Err = initMemoryPools())
       return Err;
 
     if (auto Err = ArgsMemoryManager.init(getArgsMemoryPool()))
       return Err;
 
-    if (auto Err = PinnedMemoryManager.init(getHostMemoryPool()))
+    if (auto Err = PinnedMemoryManager.init(getFineGrainedMemoryPool()))
       return Err;
 
     return Plugin::success();
   }
 
   /// Deinitialize memory pools and managers.
   Error deinit() {
     if (auto Err = deinitMemoryPools())
for (hsa_agent_t Agent : Agents) {
         return Err;
     }
     return Plugin::success();
   }
 
   /// Get one of the host agents. Return always the first agent.
   hsa_agent_t getAgent() const override { return Agents[0]; }
 
-  /// Get a memory pool for host pinned allocations.
-  AMDGPUMemoryPoolTy &getHostMemoryPool() {
+  /// Get a memory pool for fine-grained allocations.
+  AMDGPUMemoryPoolTy &getFineGrainedMemoryPool() {
     assert(!FineGrainedMemoryPools.empty() && "No fine-grained mempool");
     // Retrive any memory pool.
     return *FineGrainedMemoryPools[0];
   }
 
   /// Get a memory pool for kernel args allocations.
   AMDGPUMemoryPoolTy &getArgsMemoryPool() {
     assert(!ArgsMemoryPools.empty() && "No kernelargs mempool");
int free(void *TgtPtr, TargetAllocTy Kind) override {
 
     AMDGPUMemoryPoolTy *MemoryPool = nullptr;
     switch (Kind) {
     case TARGET_ALLOC_DEFAULT:
     case TARGET_ALLOC_DEVICE:
       MemoryPool = CoarseGrainedMemoryPools[0];
       break;
     case TARGET_ALLOC_HOST:
-      MemoryPool = &HostDevice.getHostMemoryPool();
       break;
     case TARGET_ALLOC_SHARED:
-      // TODO: Not supported yet. We could look at fine-grained host memory
-      // pools that are accessible by this device. The allocation should be made
-      // explicitly accessible if it is not yet.
+      MemoryPool = &HostDevice.getFineGrainedMemoryPool();
       break;
     }
 
     if (!MemoryPool) {
       REPORT("No memory pool for the specified allocation kind\n");
       return OFFLOAD_FAIL;
     }
 
void *AMDGPUDeviceTy::allocate(size_t Size, void *, TargetAllocTy Kind) {
   // Find the correct memory pool.
   AMDGPUMemoryPoolTy *MemoryPool = nullptr;
   switch (Kind) {
   case TARGET_ALLOC_DEFAULT:
   case TARGET_ALLOC_DEVICE:
     MemoryPool = CoarseGrainedMemoryPools[0];
     break;
   case TARGET_ALLOC_HOST:
-    MemoryPool = &HostDevice.getHostMemoryPool();
     break;
   case TARGET_ALLOC_SHARED:
-    // TODO: Not supported yet. We could look at fine-grained host memory
-    // pools that are accessible by this device. The allocation should be made
-    // explicitly accessible if it is not yet.
+    MemoryPool = &HostDevice.getFineGrainedMemoryPool();
     break;
   }
 
   if (!MemoryPool) {
     REPORT("No memory pool for the specified allocation kind\n");
     return nullptr;
   }
 
   // Allocate from the corresponding memory pool.
   void *Alloc = nullptr;
   if (Error Err = MemoryPool->allocate(Size, &Alloc)) {
     REPORT("%s\n", toString(std::move(Err)).data());
     return nullptr;
   }
 
-  if (Kind == TARGET_ALLOC_HOST && Alloc) {
+  if (Alloc && (Kind == TARGET_ALLOC_HOST || Kind == TARGET_ALLOC_SHARED)) {
     auto &KernelAgents = Plugin::get<AMDGPUPluginTy>().getKernelAgents();
 
-    // Enable all kernel agents to access the host pinned buffer.
+    // Enable all kernel agents to access the buffer.
     if (auto Err = MemoryPool->enableAccess(Alloc, Size, KernelAgents)) {
       REPORT("%s\n", toString(std::move(Err)).data());
       return nullptr;
     }
   }
 
   return Alloc;
 }
 
 } // namespace plugin
 } // namespace target
 } // namespace omp
 } // namespace llvm

openmp/libomptarget/test/api/omp_device_managed_memory.c

-// RUN: %libomptarget-compile-run-and-check-nvptx64-nvidia-cuda
-// REQUIRES: nvptx64-nvidia-cuda
+// RUN: %libomptarget-compile-run-and-check-generic
 
 #include <omp.h>
 #include <stdio.h>
 
 void *llvm_omp_target_alloc_shared(size_t, int);
 void llvm_omp_target_free_shared(void *, int);
 
 int main() {

openmp/libomptarget/test/api/omp_host_pinned_memory.c

 // RUN: %libomptarget-compile-run-and-check-generic
+// UNSUPPORTED: amdgcn-amd-amdhsa
 
 #include <omp.h>
 #include <stdio.h>
 
 // Allocate pinned memory on the host
 void *llvm_omp_target_alloc_host(size_t, int);
 void llvm_omp_target_free_host(void *, int);
 

openmp/libomptarget/test/api/omp_host_pinned_memory_alloc.c

 // RUN: %libomptarget-compile-run-and-check-generic
+// UNSUPPORTED: amdgcn-amd-amdhsa
 
 #include <omp.h>
 #include <stdio.h>
 
 int main() {
   const int N = 64;
 
   int *hst_ptr = omp_alloc(N * sizeof(int), llvm_omp_target_host_mem_alloc);