Merge pull request #110 from AMD-Lightning-Internal/amd/dev/rlieberm/CG-rename-OMPX_DISABLE_USM_MAPS-1

[Offload][AMDGPU] Rename OMPX_DISABLE_USM_MAPS env var and associated…

Author: ronlieb

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

@@ -2678,8 +2678,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
             // setting default to true here appears to solve random sdma problem
             "LIBOMPTARGET_AMDGPU_USE_MULTIPLE_SDMA_ENGINES", false),
         OMPX_ApuMaps("OMPX_APU_MAPS", false),
-        OMPX_DisableUsmMaps("OMPX_DISABLE_USM_MAPS", true),
-        OMPX_NoMapChecks("OMPX_DISABLE_MAPS", true),
+        OMPX_EnableGFX90ACoarseGrainUsmMaps(
+            "OMPX_ENABLE_GFX90A_COARSE_GRAIN_USM_MAPS", false),
         OMPX_StrictSanityChecks("OMPX_STRICT_SANITY_CHECKS", false),
         OMPX_SyncCopyBack("LIBOMPTARGET_SYNC_COPY_BACK", true),
         OMPX_APUPrefaultMemcopy("LIBOMPTARGET_APU_PREFAULT_MEMCOPY", "true"),
@@ -2944,7 +2944,7 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     if (auto Err = checkIfMI300x())
       return Err;
 
-    // detect special cases for MI200 and MI300A
+    // detect special cases for MI200
     specialBehaviorHandling();
 
     // detect ROCm-specific environment variables
@@ -3499,8 +3499,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
                                  bool set_attr = true) override final {
     // If the table has not yet been created, check if the gpu arch is
     // MI200 and create it, but only if USM Map is enabled.
-    if (!IsEquippedWithGFX90A || OMPX_DisableUsmMaps)
-      return Plugin::success();
+    if (!IsEquippedWithGFX90A || !EnableGFX90ACoarseGrainUsmMaps)
+      return Plugin::error("Invalid request to set coarse grain mode");
     if (!CoarseGrainMemoryTable)
       CoarseGrainMemoryTable = new AMDGPUMemTypeBitFieldTable(
           AMDGPU_X86_64_SystemConfiguration::max_addressable_byte +
@@ -4090,21 +4090,13 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   }
 
   /// Determines if
-  /// - Map checks should be disabled
-  /// - Coarse graining upon map on MI200 needs to be disabled.
-  /// - Prefaulting GPU page tables on MI300A needs to be enabled.
+  /// - Coarse graining upon USM map on MI200 needs to be enabled.
   void specialBehaviorHandling() {
-    if (OMPX_NoMapChecks.get() == false) {
-      NoUSMMapChecks = false;
-    }
-
-    if (OMPX_DisableUsmMaps.get() == true) {
-      EnableFineGrainedMemory = true;
-    }
+    EnableGFX90ACoarseGrainUsmMaps = OMPX_EnableGFX90ACoarseGrainUsmMaps;
   }
 
-  bool IsFineGrainedMemoryEnabledImpl() override final {
-    return EnableFineGrainedMemory;
+  bool IsGfx90aCoarseGrainUsmMapEnabledImpl() override final {
+    return !EnableGFX90ACoarseGrainUsmMaps;
   }
 
   bool hasAPUDeviceImpl() override final { return IsAPU; }
@@ -4207,17 +4199,16 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   /// automatic zero-copy behavior on non-APU GPUs.
   BoolEnvar OMPX_ApuMaps;
 
-  /// Value of OMPX_DISABLE_USM_MAPS. Use on MI200
-  /// systems to disable both device memory
-  /// allocations and host-device memory copies upon
-  /// map, and coarse graining of mapped variables.
-  BoolEnvar OMPX_DisableUsmMaps;
-
-  /// Value of OMPX_DISABLE_MAPS. Turns off map table checks
-  /// in libomptarget in unified_shared_memory mode. Legacy:
-  /// never turned to false (unified_shared_memory mode is
-  /// currently always without map checks.
-  BoolEnvar OMPX_NoMapChecks;
+  /// Value of OMPX_ENABLE_GFX90A_COARSE_GRAIN_USM_MAPS.
+  /// Use on MI200 systems to enable coarse graining
+  /// of mapped variables (and other variables partially
+  /// or fully on the same memory page) under unified
+  /// shared memory.
+  ///
+  /// It was enabled by default up to Rocm6.3
+  /// and env var spelling for controlling it was
+  /// OMPX_DISABLE_USM_MAPS
+  BoolEnvar OMPX_EnableGFX90ACoarseGrainUsmMaps;
 
   /// Makes warnings turn into fatal errors
   BoolEnvar OMPX_StrictSanityChecks;
@@ -4298,14 +4289,24 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
   /// False otherwise.
   bool IsXnackEnabled = false;
 
-  // Set by OMPX_DISABLE_USM_MAPS environment variable.
-  // If set, fine graned memory is used for maps instead of coarse grained.
-  bool EnableFineGrainedMemory = false;
+  // Set by OMPX_ENABLE_GFX90A_COARSE_GRAIN_USM_MAPS environment variable.
+  // If set, under unified shared memory on MI200, fine grained memory page
+  // is switched to coarse grain (and stay coarse grain) if a variable
+  // residing on the page goes through implicit/explicit OpenMP map.
+  bool EnableGFX90ACoarseGrainUsmMaps = false;
 
-  /// Set by OMPX_DISABLE_MAPS environment variable.
-  // If false, map checks are performed also in unified_shared_memory mode.
-  // TODO: this feature is non functional.
-  bool NoUSMMapChecks = true;
+  /// True if in multi-device mode.
+  bool IsMultiDeviceEnabled = false;
+
+public:
+  /// Return if it is an MI300 series device.
+  bool checkIfMI300Device() {
+    // Include MI300, MI300X, MI308.
+    llvm::StringRef StrGfxName(ComputeUnitKind);
+    return llvm::StringSwitch<bool>(StrGfxName)
+        .Case("gfx942", true)
+        .Default(false);
+  }
 };
 
 Error AMDGPUDeviceImageTy::loadExecutable(const AMDGPUDeviceTy &Device) {
@@ -4964,10 +4965,10 @@ void *AMDGPUDeviceTy::allocate(size_t Size, void *, TargetAllocTy Kind) {
     REPORT("%s\n", toString(std::move(Err)).data());
     return nullptr;
   }
-  // FIXME: Maybe this should be guarded by hasgfx90a
-  if (MemoryPool == CoarseGrainedMemoryPools[0]) {
-    // printf(" Device::allocate calling setCoarseGrainMemoryImpl(Alloc, Size,
-    // false)\n");
+  if (MemoryPool == CoarseGrainedMemoryPools[0] && IsEquippedWithGFX90A &&
+      EnableGFX90ACoarseGrainUsmMaps) {
+    // Need to register in the coarse grain usm map table
+    // if not already registered.
     if (auto Err = setCoarseGrainMemoryImpl(Alloc, Size, /*set_attr=*/false)) {
       REPORT("%s\n", toString(std::move(Err)).data());
       return nullptr;

offload/plugins-nextgen/common/include/PluginInterface.h

@@ -836,10 +836,10 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
   virtual bool supportsUnifiedMemoryImpl() { return false; }
 
   // Returns true if coarse graining of mapped variables is
-  // disabled on MI200 GPUs.
-  // virtual bool IsFineGrainedMemoryEnabled() { return false; }
-  bool IsFineGrainedMemoryEnabled();
-  virtual bool IsFineGrainedMemoryEnabledImpl() { return false; }
+  // enabled on MI200 GPUs.
+  // virtual bool IsGfx90aCoarseGrainUsmMapEnabled() { return false; }
+  bool IsGfx90aCoarseGrainUsmMapEnabled();
+  virtual bool IsGfx90aCoarseGrainUsmMapEnabledImpl() { return false; }
 
   /// Create an event.
   Error createEvent(void **EventPtrStorage);
@@ -1254,8 +1254,9 @@ struct GenericPluginTy {
   /// Returns if this device supports USM.
   bool supports_unified_memory(int32_t DeviceId);
 
-  /// Returns if fine grained memory is supported.
-  bool is_fine_grained_memory_enabled(int32_t DeviceId);
+  /// Returns if GFX90A coarse graining of OpenMP mapped
+  /// variables is enabled under unified shared memory.
+  bool is_gfx90a_coarse_grain_usm_map_enabled(int32_t DeviceId);
 
   /// Returns if managed memory is supported.
   bool is_system_supporting_managed_memory(int32_t DeviceId);

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

@@ -1606,8 +1606,8 @@ bool GenericDeviceTy::supportsUnifiedMemory() {
   return supportsUnifiedMemoryImpl();
 }
 
-bool GenericDeviceTy::IsFineGrainedMemoryEnabled() {
-  return IsFineGrainedMemoryEnabledImpl();
+bool GenericDeviceTy::IsGfx90aCoarseGrainUsmMapEnabled() {
+  return IsGfx90aCoarseGrainUsmMapEnabledImpl();
 }
 
 Error GenericDeviceTy::prepopulatePageTable(void *ptr, int64_t size) {
@@ -1932,9 +1932,11 @@ bool GenericPluginTy::supports_unified_memory(int32_t DeviceId) {
   return R;
 }
 
-bool GenericPluginTy::is_fine_grained_memory_enabled(int32_t DeviceId) {
+bool GenericPluginTy::is_gfx90a_coarse_grain_usm_map_enabled(int32_t DeviceId) {
   auto T = logger::log<bool>(__func__, DeviceId);
-  auto R = [&]() { return getDevice(DeviceId).IsFineGrainedMemoryEnabled(); }();
+  auto R = [&]() {
+    return getDevice(DeviceId).IsGfx90aCoarseGrainUsmMapEnabled();
+  }();
   T.res(R);
   return R;
 }

offload/src/OpenMP/API.cpp

@@ -696,6 +696,11 @@ EXTERN void omp_register_coarse_grain_mem(void *ptr, size_t size, int setattr) {
     FATAL_MESSAGE(omp_get_default_device(), "%s",
                   toString(DeviceOrErr.takeError()).c_str());
 
+  if (!(DeviceOrErr->RTL->is_gfx90a(omp_get_default_device()) &&
+        DeviceOrErr->RTL->is_gfx90a_coarse_grain_usm_map_enabled(
+            omp_get_default_device())))
+    return;
+
   bool set_attr = (setattr == 1) ? true : false;
   DeviceOrErr->RTL->set_coarse_grain_mem(omp_get_default_device(), ptr, size,
                                          set_attr);

offload/src/OpenMP/Mapping.cpp

@@ -271,7 +271,8 @@ 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.
       if (Device.RTL->is_gfx90a(Device.DeviceID) && HstPtrBegin &&
-          (!Device.RTL->is_fine_grained_memory_enabled(Device.DeviceID))) {
+          (!Device.RTL->is_gfx90a_coarse_grain_usm_map_enabled(
+              Device.DeviceID))) {
         Device.RTL->set_coarse_grain_mem_region(Device.DeviceID, HstPtrBegin,
                                                 Size);
         INFO(OMP_INFOTYPE_MAPPING_CHANGED, Device.DeviceID,

offload/test/mapping/coarse_grain.cpp

@@ -4,11 +4,13 @@
 // RUN: | %fcheck-generic -check-prefix=CHECK_FINE
 
 // RUN: %libomptarget-compilexx-generic
-// RUN: env OMPX_DISABLE_USM_MAPS=1 HSA_XNACK=1 LIBOMPTARGET_INFO=30 %libomptarget-run-generic 2>&1 \
+// RUN: env OMPX_ENABLE_GFX90A_COARSE_GRAIN_USM_MAPS=0 OMPX_DISABLE_USM_MAPS=1 HSA_XNACK=1 LIBOMPTARGET_INFO=30 \
+// RUN: %libomptarget-run-generic 2>&1 \
 // RUN: | %fcheck-generic -check-prefix=CHECK_FINE
 
 // RUN: %libomptarget-compilexx-generic
-// RUN: env OMPX_DISABLE_USM_MAPS=0 HSA_XNACK=1 LIBOMPTARGET_INFO=30 %libomptarget-run-generic 2>&1 \
+// RUN: env OMPX_ENABLE_GFX90A_COARSE_GRAIN_USM_MAPS=1 OMPX_DISABLE_USM_MAPS=0 HSA_XNACK=1 LIBOMPTARGET_INFO=30 \
+// RUN: %libomptarget-run-generic 2>&1 \
 // RUN: | %fcheck-generic -check-prefix=CHECK
 
 // UNSUPPORTED: aarch64-unknown-linux-gnu