[OpenMP][MI300A][libomptarget] Add option to make map constructs prefault pages on target GPU.

Preafulting pages allows to execute applications built with xnack- in unified_shared_memory mode.
A translation for each mapped page is added to the GPU page table and there is no need to run the XNACK algorithm when a page is referenced by a GPU thread for the first time.

Change-Id: Ib3939510ecadc05e8a8f9a109d6f9d28eef3811b

Author: carlobertolli

openmp/libomptarget/include/omptargetplugin.h

@@ -37,8 +37,8 @@ int32_t __tgt_rtl_number_of_devices(void);
 // Return if the system is equipped with an APU
 bool __tgt_rtl_has_apu_device(void);
 
-// Returns true, if the system is equipped with an GFX90a
-bool __tgt_rtl_has_gfx90a_device(void);
+// Returns true, if the system is equipped with a dGPU which supports USM.
+bool __tgt_rtl_has_USM_capable_dGPU(void);
 
 bool __tgt_rtl_are_allocations_for_maps_on_apus_disabled(void);
 

openmp/libomptarget/include/rtl.h

@@ -37,8 +37,9 @@ struct RTLInfoTy {
   typedef int32_t(is_data_exchangable_ty)(int32_t, int32_t);
   typedef int32_t(number_of_devices_ty)();
   typedef bool(has_apu_device_ty)();
-  typedef bool(has_gfx90a_device_ty)();
+  typedef bool(has_USM_capable_dGPU_ty)();
   typedef bool(are_allocations_for_maps_on_apus_disabled_ty)();
+  typedef bool(requested_prepopulate_gpu_page_table_ty)();
   typedef bool(is_no_maps_check_ty)();
   typedef bool(is_fine_grained_memory_enabled_ty)();
   typedef int32_t(init_device_ty)(int32_t);
@@ -74,6 +75,7 @@ struct RTLInfoTy {
   typedef int32_t(sync_event_ty)(int32_t, void *);
   typedef int32_t(destroy_event_ty)(int32_t, void *);
   typedef int(set_coarse_grain_mem_region_ty)(int32_t, void *, int64_t);
+  typedef int(prepopulate_page_table_ty)(int32_t, void *, int64_t);
   typedef int32_t(query_coarse_grain_mem_region_ty)(int32_t, void *, int64_t);
   typedef int32_t(enable_access_to_all_agents_ty)(void *, int32_t);
   typedef int32_t(release_async_info_ty)(int32_t, __tgt_async_info *);
@@ -107,9 +109,11 @@ struct RTLInfoTy {
   is_data_exchangable_ty *is_data_exchangable = nullptr;
   number_of_devices_ty *number_of_devices = nullptr;
   has_apu_device_ty *has_apu_device = nullptr;
-  has_gfx90a_device_ty *has_gfx90a_device = nullptr;
+  has_USM_capable_dGPU_ty *has_USM_capable_dGPU = nullptr;
   are_allocations_for_maps_on_apus_disabled_ty
       *are_allocations_for_maps_on_apus_disabled = nullptr;
+  requested_prepopulate_gpu_page_table_ty
+      *requested_prepopulate_gpu_page_table = nullptr;
   is_no_maps_check_ty *is_no_maps_check = nullptr;
   is_fine_grained_memory_enabled_ty *is_fine_grained_memory_enabled = nullptr;
   init_device_ty *init_device = nullptr;
@@ -145,6 +149,7 @@ struct RTLInfoTy {
   data_lock_ty *data_lock = nullptr;
   data_unlock_ty *data_unlock = nullptr;
   set_coarse_grain_mem_region_ty *set_coarse_grain_mem_region = nullptr;
+  prepopulate_page_table_ty *prepopulate_page_table = nullptr;
   query_coarse_grain_mem_region_ty *query_coarse_grain_mem_region = nullptr;
   enable_access_to_all_agents_ty *enable_access_to_all_agents = nullptr;
   data_notify_mapped_ty *data_notify_mapped = nullptr;

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

@@ -2377,6 +2377,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     if (Error Err = AMDImage->loadExecutable(*this))
       return std::move(Err);
 
+    Plugin::get().checkAndAdjustUsmModeForTargetImage(TgtImage);
+
     return AMDImage;
   }
 
@@ -2682,6 +2684,20 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
     return coarse_grain_mem_tab->contains((const uintptr_t)ptr, size);
   }
 
+  Error prepopulatePageTableImpl(void *ptr, int64_t size) override final {
+    // Instruct ROCr that the [ptr, ptr+size-1] pages are
+    // coarse grain
+    hsa_amd_svm_attribute_pair_t tt;
+    tt.attribute = HSA_AMD_SVM_ATTRIB_AGENT_ACCESSIBLE_IN_PLACE;
+    tt.value = Agent.handle;
+    hsa_status_t err = hsa_amd_svm_attributes_set(ptr, size, &tt, 1);
+    if (err != HSA_STATUS_SUCCESS) {
+      return Plugin::error("Failed to prepopulate GPU page table.");
+    }
+
+    return Plugin::success();
+  }
+
   /// Create an event.
   Error createEventImpl(void **EventPtrStorage) override {
     AMDGPUEventTy **Event = reinterpret_cast<AMDGPUEventTy **>(EventPtrStorage);
@@ -3419,8 +3435,10 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
 #endif
 
     // Initialize flags for device type:
-    hasGfx90aDevice();
     hasAPUDevice();
+    // check for dGPUs with USM support
+    hasGfx90aDevice();
+    hasMI300xDevice();
 
     readEnvVars();
 
@@ -3463,7 +3481,19 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
 #define ALDEBARAN_MAJOR 9
 #define ALDEBARAN_STEPPING 10
 
-  bool hasGfx90aDevice() override final {
+  bool hasMI300xDevice() {
+    if (HasMi300xDevice != -1)
+      return HasMi300xDevice;
+
+    if (!Initialized)
+      FATAL_MESSAGE(1, "%s", "hasMI300xDevice called on uninitialized plugin");
+    // On splinter the MI300X identifies itself as a GFX941. Use GFX name to
+    // distinguish for testing.
+    HasMi300xDevice = checkForDeviceByGFXName("gfx941");
+    return HasMi300xDevice;
+  }
+
+  bool hasGfx90aDevice() {
     if (HasGFX90ADevice != -1)
       return HasGFX90ADevice;
 
@@ -3474,10 +3504,18 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     return HasGFX90ADevice;
   }
 
+  bool hasDGpuWithUsmSupport() override final {
+    return hasGfx90aDevice() || hasMI300xDevice();
+  }
+
   bool AreAllocationsForMapsOnApusDisabled() override final {
     return DisableAllocationsForMapsOnApus;
   }
 
+  bool requestedPrepopulateGPUPageTable() override final {
+    return PrepopulateGPUPageTable;
+  }
+
   bool IsNoMapsCheck() override final { return NoUSMMapChecks; }
 
   bool IsFineGrainedMemoryEnabled() override final {
@@ -3491,6 +3529,8 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     NoMapChecks = BoolEnvar("OMPX_DISABLE_MAPS", true);
     DisableUsmMaps = BoolEnvar("OMPX_DISABLE_USM_MAPS", false);
     HsaXnack = BoolEnvar("HSA_XNACK", false);
+    APUPrefault = BoolEnvar("OMPX_EAGER_ZERO_COPY_MAPS", false);
+    ZeroCopyForMapsOnUsm = BoolEnvar("OMPX_APU_MAPS", false);
   }
 
   void setUpEnv() override final {
@@ -3502,6 +3542,15 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     if (DisableUsmMaps.get() == true) {
       EnableFineGrainedMemory = true;
     }
+
+    if (hasAPUDevice()) {
+      // OMPX_EAGER_ZERO_COPY_MAPS=1 && HSA_XNACK=0 (XNACK-disabled)
+      // && default (non-USM) program
+      if ((APUPrefault.get() == true) && !IsXnackEnabled() &&
+          !(Plugin::get().getRequiresFlags() & OMP_REQ_UNIFIED_SHARED_MEMORY)) {
+        PrepopulateGPUPageTable = true;
+      }
+    }
   }
 
   /// Check whether the image is compatible with an AMDGPU device.
@@ -3545,38 +3594,74 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
   }
 
   void checkAndAdjustUsmModeForTargetImage(
-      __tgt_device_image *TgtImage) override final {
+      const __tgt_device_image *TgtImage) override final {
     assert((TgtImage != nullptr) && "TgtImage is nullptr");
     assert(!(Plugin::get().getRequiresFlags() & OMP_REQ_UNDEFINED) &&
            "Requires flags are not set.");
 
-    if (!(hasAPUDevice() || hasGfx90aDevice()))
+    if (!(hasAPUDevice() || hasDGpuWithUsmSupport()))
       return;
 
     bool IsXnackRequired =
         Plugin::get().getRequiresFlags() & OMP_REQ_UNIFIED_SHARED_MEMORY;
-
     utils::XnackBuildMode BinaryXnackMode =
         utils::extractXnackModeFromBinary(TgtImage);
 
+    if (IsXnackRequired) {
+      handleImageRequiresUsmMode(BinaryXnackMode);
+    } else {
+      handleDefaultMode(BinaryXnackMode);
+    }
+  }
+
+  void handleImageRequiresUsmMode(utils::XnackBuildMode xnackImageMode) {
+    bool IsXnackActiveOnSystem = IsXnackEnabled();
+
+    if ((xnackImageMode == utils::XnackBuildMode::XNACK_ANY) ||
+        (xnackImageMode == utils::XnackBuildMode::XNACK_PLUS &&
+         IsXnackActiveOnSystem) ||
+        (xnackImageMode == utils::XnackBuildMode::XNACK_MINUS &&
+         !IsXnackActiveOnSystem)) {
+      DisableAllocationsForMapsOnApus = true; // Zero-copy
+
+      if (APUPrefault.get() && hasAPUDevice())
+        PrepopulateGPUPageTable = true; // Pre-faulting
+    }
+
+    if (!IsXnackActiveOnSystem &&
+        (xnackImageMode != utils::XnackBuildMode::XNACK_PLUS)) {
+      FAILURE_MESSAGE(
+          "Running a program that requries XNACK on a system where XNACK is "
+          "disabled! This may potentially cause memory errors! Just saying.\n");
+    }
+  }
+
+  void handleDefaultMode(utils::XnackBuildMode xnackImageMode) {
+    // assuming that copying is required
     DisableAllocationsForMapsOnApus = false;
+    bool IsXnackActiveOnSystem = IsXnackEnabled();
 
-    if (IsXnackEnabled()) {
-      if (!IsXnackRequired) {
-        switch (BinaryXnackMode) {
-        case utils::XnackBuildMode::XNACK_PLUS:
-        case utils::XnackBuildMode::XNACK_ANY:
-          DisableAllocationsForMapsOnApus = true; // Zero-copy
-        }
-        return;
-      }
-    } else {
-      if (IsXnackRequired) {
-        FAILURE_MESSAGE(
-            "XNACK is disabled. However, the program requires XNACK "
-            "support. Enable XNACK and re-run the program.\n");
+    if (IsXnackActiveOnSystem &&
+        (hasAPUDevice() || ZeroCopyForMapsOnUsm.get()) &&
+        ((xnackImageMode == utils::XnackBuildMode::XNACK_ANY) ||
+         (xnackImageMode == utils::XnackBuildMode::XNACK_PLUS))) {
+      DisableAllocationsForMapsOnApus = true; // Zero-copy
+
+      if (hasAPUDevice() && APUPrefault.get()) {
+        PrepopulateGPUPageTable = true; // Pre-faulting
       }
+      return;
+    }
+
+    if (!IsXnackActiveOnSystem && hasAPUDevice() && APUPrefault.get() &&
+        ((xnackImageMode == utils::XnackBuildMode::XNACK_ANY) ||
+         (xnackImageMode == utils::XnackBuildMode::XNACK_MINUS))) {
+      DisableAllocationsForMapsOnApus = true; // Zero-copy
+      PrepopulateGPUPageTable = true;         // Pre-faulting
+      return;
     }
+
+    return;
   }
 
   /// This plugin does not support exchanging data between two devices.
@@ -3649,11 +3734,10 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     return ((HsaXnack.get()) || (utils::IsXnackEnabledViaKernelParam()));
   }
 
-  bool checkForDeviceByGFXName(const llvm::StringRef GfxLookUpName) {
-    bool CheckForMI300A =
-        (GfxLookUpName.find_insensitive("gfx940") != llvm::StringRef::npos);
+  bool checkForDeviceByGFXName(const llvm::StringRef GfxLookUpName,
+                               char mi300Specifier = ' ') {
+
     char GfxName[64];
-    llvm::StringRef GfxNameRef = llvm::StringRef(GfxName);
 
     for (hsa_agent_t GPUAgent : KernelAgents) {
       std::memset((void *)&GfxName, 0, sizeof(char) * 64);
@@ -3664,22 +3748,37 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
       if (Status != HSA_STATUS_SUCCESS)
         continue;
 
-      if (GfxLookUpName.find_insensitive(GfxNameRef) != llvm::StringRef::npos) {
-        // Special handling for MI300. We will have to distinguish between an
-        // MI300A and X
-        if (CheckForMI300A) {
-          uint32_t ChipID = 0;
-          Status = hsa_agent_get_info(
-              GPUAgent, (hsa_agent_info_t)HSA_AMD_AGENT_INFO_CHIP_ID, &ChipID);
+      llvm::StringRef GfxNameRef = llvm::StringRef(GfxName);
+
+      if (GfxLookUpName.equals_insensitive(GfxNameRef)) {
+        if (mi300Specifier == ' ')
+          return true;
 
-          if (Status != HSA_STATUS_SUCCESS) {
-            continue;
-          }
+        // Special handling for MI300. We will have to distinguish between
+        // an MI300A and X
+        uint32_t ChipID = 0;
+        Status = hsa_agent_get_info(
+            GPUAgent, (hsa_agent_info_t)HSA_AMD_AGENT_INFO_CHIP_ID, &ChipID);
 
-          if ((ChipID & 0x1))
-            continue;
+        if (Status != HSA_STATUS_SUCCESS) {
+          continue;
+        }
+
+        bool IsMi300X = ChipID & 0x1;
+
+        switch (mi300Specifier) {
+        case 'A':
+        case 'a':
+          if (!IsMi300X)
+            return true;
+          break;
+        case 'x':
+          if (IsMi300X) // We are looking for a MI300X
+            return true;
+          break;
+        default:
+          FAILURE_MESSAGE("Unknown MI300 specifier!\n");
         }
-        return true;
       }
     }
     return false;
@@ -3693,18 +3792,31 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
   /// Flag that shows if device is a GFX90A AMD GPU
   int16_t HasGFX90ADevice{-1};
 
+  int16_t HasMi300xDevice{-1};
+
   /// Flag that shows if device is an APU device
   int16_t HasAPUDevice{-1};
 
   BoolEnvar NoMapChecks;
   BoolEnvar DisableUsmMaps;
   BoolEnvar HsaXnack;
+  BoolEnvar APUPrefault;
+
+  // Set by OMPX_APU_MAPS
+  // Enables code that detect if zero copying is possible. If so, the variable
+  // DisableAllocationsForMapsOnApus is set to 'true'.
+  BoolEnvar ZeroCopyForMapsOnUsm;
 
-  // Set by OMPX_APU_MAPS environment variable.
   // If set, maps cause no copy operations. USM is used instead. Allocated
-  // memory remains coarse grained.
+  // memory remains coarse grained. The variable is only considered to be set if
+  // ZeroCopyForMapsOnUsm (OMPX_APU_MAPS) is set.
   bool DisableAllocationsForMapsOnApus{false};
 
+  // Set by OMPX_EAGER_ZERO_COPY_MAPS environment variable.
+  // If set, map clauses provoke prefaulting of the GPU
+  // page table.
+  bool PrepopulateGPUPageTable{false};
+
   // Set by OMPX_DISABLE_MAPS environment variable.
   // When active (default value), maps are ignored by the runtime
   bool NoUSMMapChecks{true};

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

@@ -178,7 +178,7 @@ bool isImageCompatibleWithEnv(const __tgt_image_info *Info,
 
 // Check target image for XNACK mode (XNACK+, XNACK-ANY, XNACK-)
 [[nodiscard]] XnackBuildMode
-extractXnackModeFromBinary(__tgt_device_image *TgtImage) {
+extractXnackModeFromBinary(const __tgt_device_image *TgtImage) {
   assert((TgtImage != nullptr) && "TgtImage is nullptr.");
   u_int16_t EFlags = elf_get_eflags(TgtImage);