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

Author: hidekisaito

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

@@ -2924,6 +2924,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"),
@@ -4339,6 +4341,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 {
@@ -4460,6 +4464,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;
 
@@ -4549,6 +4559,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;
 
@@ -5246,6 +5260,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/include/PluginInterface.h

@@ -1051,9 +1051,6 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
     return OMPX_MinThreadsForLowTripCount;
   }
 
-  virtual uint32_t getUseCoarseGrain() {
-    return OMPX_EnableCoarseAllocs;
-  }
   /// Whether or not to reuse blocks for high trip count loops.
   /// @see OMPX_ReuseBlocksForHighTripCount
   bool getReuseBlocksForHighTripCount() {
@@ -1212,9 +1209,6 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   BoolEnvar OMPX_ReuseBlocksForHighTripCount =
       BoolEnvar("LIBOMPTARGET_REUSE_BLOCKS_FOR_HIGH_TRIP_COUNT", true);
 
-  /// Envar to enable coasre alocs.
-  BoolEnvar OMPX_EnableCoarseAllocs;
-
 protected:
   /// Environment variables defined by the LLVM OpenMP implementation
   /// regarding the initial number of streams and events.

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

@@ -927,7 +927,6 @@ GenericDeviceTy::GenericDeviceTy(GenericPluginTy &Plugin, int32_t DeviceId,
       OMPX_InitialNumEvents("LIBOMPTARGET_NUM_INITIAL_EVENTS", 1),
       OMPX_NumMultiDevices("LIBOMPTARGET_NUM_MULTI_DEVICES", 0),
       OMPX_EnableRuntimeAutotuning("OMPX_ENABLE_RUNTIME_AUTOTUNING", false),
-      OMPX_EnableCoarseAllocs("OMPX_ENABLE_COARSE_ALLOCS", false),
       DeviceId(DeviceId), GridValues(OMPGridValues),
       PeerAccesses(NumDevices, PeerAccessState::PENDING), PeerAccessesLock(),
       PinnedAllocs(*this), RPCServer(nullptr) {
@@ -2207,11 +2206,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 && getDevice(DeviceId).getUseCoarseGrain() )
-      set_coarse_grain_mem_region(DeviceId, HostPtr, Size);
-
     return *AllocOrErr;
   }();
   T.res(R);