[AMDGPU][Offload] GFX90A coarse grain shared alloc option

Author: hidekisaito

offload/plugins-nextgen/amdgpu/src/rtl.cpp

@@ -2680,6 +2680,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
         OMPX_ApuMaps("OMPX_APU_MAPS", false),
         OMPX_EnableGFX90ACoarseGrainUsmMaps(
             "OMPX_ENABLE_GFX90A_COARSE_GRAIN_USM_MAPS", false),
+        OMPX_EnableGFX90ACoarseGrainSharedAlloc(
+            "OMPX_ENABLE_GFX90A_COARSE_GRAIN_SHARED_ALLOC", false),
         OMPX_StrictSanityChecks("OMPX_STRICT_SANITY_CHECKS", false),
         OMPX_SyncCopyBack("LIBOMPTARGET_SYNC_COPY_BACK", true),
         OMPX_APUPrefaultMemcopy("LIBOMPTARGET_APU_PREFAULT_MEMCOPY", "true"),
@@ -4093,6 +4095,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   /// - Coarse graining upon USM map on MI200 needs to be enabled.
   void specialBehaviorHandling() {
     EnableGFX90ACoarseGrainUsmMaps = OMPX_EnableGFX90ACoarseGrainUsmMaps;
+    EnableGFX90ACoarseGrainSharedAlloc =
+        OMPX_EnableGFX90ACoarseGrainSharedAlloc;
   }
 
   bool IsGfx90aCoarseGrainUsmMapEnabledImpl() override final {
@@ -4210,6 +4214,12 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   /// OMPX_DISABLE_USM_MAPS
   BoolEnvar OMPX_EnableGFX90ACoarseGrainUsmMaps;
 
+  /// Value of OMPX_ENABLE_GFX90A_COARSE_GRAIN_SHARED_ALLOC.
+  /// Use on MI200 systems to enable coarse grain
+  /// allocation of TARGET_ALLOC_SHARED memory.
+  /// Default is fine grain allocation.
+  BoolEnvar OMPX_EnableGFX90ACoarseGrainSharedAlloc;
+
   /// Makes warnings turn into fatal errors
   BoolEnvar OMPX_StrictSanityChecks;
 
@@ -4295,6 +4305,10 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   // residing on the page goes through implicit/explicit OpenMP map.
   bool EnableGFX90ACoarseGrainUsmMaps = false;
 
+  // Set by OMPX_ENABLE_GFX90A_COARSE_GRAIN_SHARED_ALLOC environment variable.
+  // If set, TARGET_ALLOC_SHARED is allocated on coarse grain memory on MI200
+  bool EnableGFX90ACoarseGrainSharedAlloc = false;
+
   /// True if in multi-device mode.
   bool IsMultiDeviceEnabled = false;
 
@@ -4954,6 +4968,11 @@ void *AMDGPUDeviceTy::allocate(size_t Size, void *, TargetAllocTy Kind) {
     break;
   }
 
+  if (Kind == TARGET_ALLOC_SHARED && IsEquippedWithGFX90A &&
+      EnableGFX90ACoarseGrainSharedAlloc) {
+    MemoryPool = CoarseGrainedMemoryPools[0];
+  }
+
   if (!MemoryPool) {
     REPORT("No memory pool for the specified allocation kind\n");
     return nullptr;

offload/plugins-nextgen/common/src/PluginInterface.cpp

@@ -2030,11 +2030,6 @@ void *GenericPluginTy::data_alloc(int32_t DeviceId, int64_t Size, void *HostPtr,
     }
     assert(*AllocOrErr && "Null pointer upon successful allocation");
 
-    // Method has no effect when the CUDA Plugin is used.
-    // This method can only be called if HostPtr is not null.
-    if (HostPtr && Kind == TARGET_ALLOC_SHARED)
-      set_coarse_grain_mem_region(DeviceId, HostPtr, Size);
-
     return *AllocOrErr;
   }();
   T.res(R);