[OpenMP] Make coarse graining not dependent on USM info in the plugin.

This patch cleans up coarse graining upon mapping for MI200 and it removes
a dependency on the requirement flags in the plugin, to ease downstream merging of
llvm#80345.

Change-Id: I0bad9113c106c238c2089bf8097789a27f6899ea

Author: carlobertolli

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

@@ -2704,17 +2704,6 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     if (auto Err = AMDGPUSignalManager.init(OMPX_InitialNumSignals))
       return Err;
 
-    // Initialize memspace table to keep track of coarse grain memory regions
-    // in USM mode
-    if (Plugin::get().getRequiresFlags() & OMP_REQ_UNIFIED_SHARED_MEMORY) {
-      // TODO: add framework for multiple systems supporting
-      // unified_shared_memory
-      coarse_grain_mem_tab = new AMDGPUMemTypeBitFieldTable(
-          AMDGPU_X86_64_SystemConfiguration::max_addressable_byte +
-              1, // memory size
-          AMDGPU_X86_64_SystemConfiguration::page_size);
-    }
-
     // Take the second timepoints and compute the required metadata.
     OMPT_IF_ENABLED(completeH2DTimeRate(HostRef1, DeviceRef1););
 
@@ -2782,6 +2771,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     // Invalidate agent reference.
     Agent = {0};
 
+    delete CoarseGrainMemoryTable;
+
     return Plugin::success();
   }
 
@@ -3226,11 +3217,21 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   }
 
   Error setCoarseGrainMemoryImpl(void *ptr, int64_t size) override final {
-    // track coarse grain memory pages in local table
-    coarse_grain_mem_tab->insert((const uintptr_t)ptr, size);
+    // If the table has not yet been created, check if the gpu arch is
+    // MI200 and create it.
+    if (!IsEquippedWithGFX90A)
+      return Plugin::success();
+    if (!CoarseGrainMemoryTable)
+      CoarseGrainMemoryTable = new AMDGPUMemTypeBitFieldTable(
+          AMDGPU_X86_64_SystemConfiguration::max_addressable_byte +
+              1, // memory size
+          AMDGPU_X86_64_SystemConfiguration::page_size);
 
-    // Instruct ROCr that the [ptr, ptr+size-1] pages are
-    // coarse grain
+    // track coarse grain memory pages in local table for user queries.
+    CoarseGrainMemoryTable->insert((const uintptr_t)ptr, size);
+
+    // Ask ROCr to turn [ptr, ptr+size-1] pages to
+    // coarse grain.
     hsa_amd_svm_attribute_pair_t tt;
     tt.attribute = HSA_AMD_SVM_ATTRIB_GLOBAL_FLAG;
     tt.value = HSA_AMD_SVM_GLOBAL_FLAG_COARSE_GRAINED;
@@ -3244,13 +3245,12 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
 
   uint32_t queryCoarseGrainMemoryImpl(const void *ptr,
                                       int64_t size) override final {
-
-    // if the table is not yet allocated, it means we have not yet gone through
-    // an OpenMP pragma or API that would provoke intialization of the RTL
-    if (!coarse_grain_mem_tab)
+    // If the table has not yet been created it means that
+    // no memory has yet been set to coarse grain.
+    if (!CoarseGrainMemoryTable)
       return 0;
 
-    return coarse_grain_mem_tab->contains((const uintptr_t)ptr, size);
+    return CoarseGrainMemoryTable->contains((const uintptr_t)ptr, size);
   }
 
   Error prepopulatePageTableImpl(void *ptr, int64_t size) override final {
@@ -3936,7 +3936,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   AMDHostDeviceTy &HostDevice;
 
   // Data structure used to keep track of coarse grain memory regions
-  AMDGPUMemTypeBitFieldTable *coarse_grain_mem_tab = nullptr;
+  // on MI200 in unified_shared_memory programs only.
+  AMDGPUMemTypeBitFieldTable *CoarseGrainMemoryTable = nullptr;
 
   /// Pointer to the preallocated device memory pool
   void *PreAllocatedDeviceMemoryPool;

openmp/libomptarget/plugins-nextgen/amdgpu/utils/memtype.h

@@ -47,6 +47,11 @@ class AMDGPUMemTypeBitFieldTable {
     tab = (uint64_t *)calloc(tab_size, sizeof(uint64_t));
   }
 
+  ~AMDGPUMemTypeBitFieldTable() {
+    if (tab)
+      free(tab);
+  }
+
   // Set all pages touched by address in the range [base, base+size-1]
   // \arg base : pointer to first byte of the memory area whose
   // type should become of the tracked type
@@ -102,7 +107,7 @@ class AMDGPUMemTypeBitFieldTable {
   // memory type. For any bit:
   // 0 = page is *not* of tracked type
   // 1 = page is of tracked type
-  uint64_t *tab;
+  uint64_t *tab = nullptr;
 };
 
 #endif //__MEMTYPE_H__

openmp/libomptarget/src/OpenMP/API.cpp

@@ -654,10 +654,17 @@ EXTERN int omp_target_disassociate_ptr(const void *HostPtr, int DeviceNum) {
 }
 
 EXTERN int omp_is_coarse_grain_mem_region(void *ptr, size_t size) {
-  DeviceTy &Device = *PM->getDevice(omp_get_default_device());
-  if (!Device.RTL->query_coarse_grain_mem_region)
+  if (!(PM->getRequirements() & OMP_REQ_UNIFIED_SHARED_MEMORY))
     return 0;
-  return Device.RTL->query_coarse_grain_mem_region(Device.DeviceID, ptr, size);
+  auto DeviceOrErr = PM->getDevice(omp_get_default_device());
+  if (!DeviceOrErr)
+    FATAL_MESSAGE(omp_get_default_device(), "%s",
+                  toString(DeviceOrErr.takeError()).c_str());
+
+  if (!DeviceOrErr->RTL->query_coarse_grain_mem_region)
+    return 0;
+  return DeviceOrErr->RTL->query_coarse_grain_mem_region(
+      omp_get_default_device(), ptr, size);
 }
 
 EXTERN void *omp_get_mapped_ptr(const void *Ptr, int DeviceNum) {

openmp/libomptarget/src/OpenMP/Mapping.cpp

@@ -266,16 +266,20 @@ TargetPointerResultTy MappingInfoTy::getTargetPointer(
       // memory as coarse-grained. The usage of coarse-grained memory can be
       // overriden by setting the env-var OMPX_DISABLE_USM_MAPS=1.
       // This is not done for APUs.
-      if (!(Device.RTL->has_apu_device(Device.DeviceID) ||
-            Device.RTL->has_USM_capable_dGPU(Device.DeviceID)) &&
-          Device.RTL->is_fine_grained_memory_enabled(Device.DeviceID) &&
-          HstPtrBegin && Device.RTL->set_coarse_grain_mem_region) {
+      if (Device.RTL->has_USM_capable_dGPU(Device.DeviceID) && HstPtrBegin &&
+          (!Device.RTL->is_fine_grained_memory_enabled(Device.DeviceID)) &&
+          Device.RTL->set_coarse_grain_mem_region) {
         Device.RTL->set_coarse_grain_mem_region(Device.DeviceID, HstPtrBegin,
                                                 Size);
+        INFO(OMP_INFOTYPE_MAPPING_CHANGED, Device.DeviceID,
+             "Memory pages for HstPtrBegin " DPxMOD " Size=%" PRId64
+             " switched to coarse grain\n",
+             DPxPTR((uintptr_t)HstPtrBegin), Size);
       }
 
       // If we are here, it means that we are either in auto zero-copy or USM.
-      // Enable GPU page table prefaulting if selected by the user.
+      // Enable GPU page table prefaulting if selected by the user. This feature
+      // is only enabled for APUs.
       if (Device.EagerZeroCopyMaps) {
         Device.RTL->prepopulate_page_table(Device.DeviceID, HstPtrBegin, Size);
         INFO(OMP_INFOTYPE_MAPPING_CHANGED, Device.DeviceID,

openmp/libomptarget/test/lit.cfg

@@ -117,6 +117,7 @@ if config.libomptarget_has_libc:
 # For all other targets, we currently assume it is.
 supports_unified_shared_memory = True
 supports_apu = False
+is_mi200 = False
 if config.libomptarget_current_target.startswith('nvptx'):
   try:
     cuda_arch = int(config.cuda_test_arch[:3])
@@ -137,10 +138,14 @@ elif config.libomptarget_current_target.startswith('amdgcn'):
         (config.amdgpu_test_arch.startswith("gfx942") and
          evaluate_bool_env(config.environment['IS_APU']))):
        supports_apu = True
-if supports_unified_shared_memory:
-   config.available_features.add('unified_shared_memory')
-if supports_apu:
-   config.available_features.add('apu')
+    if (config.amdgpu_test_arch.startswith("gfx90a")):
+       is_mi200 = True
+    if supports_unified_shared_memory:
+        config.available_features.add('unified_shared_memory')
+    if supports_apu:
+        config.available_features.add('apu')
+    if is_mi200:
+        config.available_features.add('mi200')
 
 # Setup environment to find dynamic library at runtime
 if config.operating_system == 'Windows':

openmp/libomptarget/test/mapping/coarse_grain.cpp

@@ -0,0 +1,65 @@
+// clang-format off
+// RUN: %libomptarget-compilexx-generic
+// RUN: env HSA_XNACK=1 LIBOMPTARGET_INFO=30 %libomptarget-run-generic 2>&1 \
+// RUN: | %fcheck-generic -check-prefix=CHECK
+
+// RUN: %libomptarget-compilexx-generic
+// RUN: env OMPX_DISABLE_USM_MAPS=1 HSA_XNACK=1 LIBOMPTARGET_INFO=30 %libomptarget-run-generic 2>&1 \
+// RUN: | %fcheck-generic -check-prefix=CHECK_FINE
+
+// UNSUPPORTED: aarch64-unknown-linux-gnu
+// UNSUPPORTED: aarch64-unknown-linux-gnu-LTO
+// UNSUPPORTED: nvptx64-nvidia-cuda
+// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
+// UNSUPPORTED: x86_64-pc-linux-gnu
+// UNSUPPORTED: x86_64-pc-linux-gnu-LTO
+
+// REQUIRES: unified_shared_memory
+// REQUIRES: mi200
+
+// clang-format on
+
+#include <cstdio>
+#include <omp.h>
+
+#pragma omp requires unified_shared_memory
+
+int main() {
+  const size_t n = 1024;
+
+  double *a = new double[n];
+  // clang-format off
+  // CHECK: Memory pages for HstPtrBegin 0x{{.*}} Size=8192 switched to coarse grain
+  // CHECK: Before mapping, memory is fine grain.
+  // CHECK_FINE: Before mapping, memory is fine grain.
+  // clang-format on
+  if (omp_is_coarse_grain_mem_region(a, n * sizeof(double)))
+    printf("Before mapping, memory is coarse grain.\n");
+  else
+    printf("Before mapping, memory is fine grain.\n");
+
+#pragma omp target enter data map(to : a[:n])
+
+  // CHECK: After mapping, memory is coarse grain.
+  // CHECK_FINE: After mapping, memory is fine grain.
+  if (omp_is_coarse_grain_mem_region(a, n * sizeof(double)))
+    printf("After mapping, memory is coarse grain.\n");
+  else
+    printf("After mapping, memory is fine grain.\n");
+
+#pragma omp target exit data map(from : a[:n])
+
+  // CHECK: After removing map, memory is still coarse grain.
+  // CHECK_FINE: After removing map, memory is back to fine grain.
+  if (omp_is_coarse_grain_mem_region(a, n * sizeof(double)))
+    printf("After removing map, memory is still coarse grain.\n");
+  else
+    printf("After removing map, memory is back to fine grain.\n");
+
+// Plugins must be initialized for unified_shared_memory requirement
+// to be added. An empty target region is enough for that initialization.
+#pragma omp target
+  {}
+
+  return 0;
+}