This patch implements a runtime check whether the system's USM setting

matches the one required by the target image. If necessary, the zero-copy
mode will be enabled.

Change-Id: Ic9950e15d0ee4664910c1a575e192634f3227898

Author: ThorBl

openmp/libomptarget/include/rtl.h

@@ -19,8 +19,6 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/DynamicLibrary.h"
 
-#include "omptarget.h"
-
 #include <list>
 #include <map>
 #include <mutex>

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

@@ -30,6 +30,8 @@
 #include "UtilitiesRTL.h"
 #include "omptarget.h"
 
+#include "hsakmt/hsakmt.h"
+
 #include "print_tracing.h"
 
 #include "memtype.h"
@@ -46,8 +48,6 @@
 #include "llvm/Support/Program.h"
 #include "llvm/Support/raw_ostream.h"
 
-#include "hsakmt/hsakmt.h"
-
 #if defined(__has_include)
 #if __has_include("hsa/hsa.h")
 #include "hsa/hsa.h"
@@ -80,12 +80,12 @@
 #define OMPT_IF_TRACING_ENABLED(stmts)
 #endif
 
-#define CHECK_KMT_ERROR(val) kmtCheck((val), #val, __FILE__, __LINE__)
+#define KMT_EXPECT_SUCCESS(val) kmtExpectSucc((val), #val, __FILE__, __LINE__)
 template <typename T>
-int kmtCheck(T err, const char *const func, const char *const file,
-              const int line) {
+int kmtExpectSucc(T err, const char *const func, const char *const file,
+                  const int line) {
   if (err != HSAKMT_STATUS_SUCCESS) {
-    DP("HsaKmt Error at: %s : %u \n", file, line);
+    FAILURE_MESSAGE("HsaKmt Error at: %s : %u \n", file, line);
     return -1;
   }
   return 0;
@@ -95,8 +95,6 @@ int kmtCheck(T err, const char *const func, const char *const file,
 extern void setOmptTimestamp(uint64_t Start, uint64_t End);
 extern void setOmptHostToDeviceRate(double Slope, double Offset);
 
-
-
 /// HSA system clock frequency
 double TicksToTime = 1.0;
 
@@ -2504,7 +2502,7 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
 
   /// AMDGPU returns the product of the number of compute units and the waves
   /// per compute unit.
-  uint64_t requestedRPCPortCount() const override  {
+  uint64_t requestedRPCPortCount() const override {
     return HardwareParallelism;
   }
 
@@ -3581,8 +3579,8 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
 
     NoMapChecks = BoolEnvar("OMPX_DISABLE_MAPS", true);
     DisableUsmMaps = BoolEnvar("OMPX_DISABLE_USM_MAPS", false);
-    APUMaps = BoolEnvar("OMPX_APU_MAPS", false);
-    HSAXnack = BoolEnvar("HSA_XNACK", false);
+    HsaXnack = BoolEnvar("HSA_XNACK", false);
+    IsHsaXnackDefined = HsaXnack.isPresent();
   }
 
   void setUpEnv() override final {
@@ -3594,36 +3592,13 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     if (DisableUsmMaps.get() == true) {
       EnableFineGrainedMemory = true;
     }
-
-    if (hasAPUDevice() || hasGfx90aDevice()) {
-      // OMPX_APU_MAPS is a temporary env variable
-      // that should always be used with HSA_XNACK=1:
-      // error if it is not. Once this is made default behavior
-      // for USM=OFF, HSA_XNACK=1, then we can remove the error
-      // as the behavior is only triggered by HSA_XNACK value
-      if ((APUMaps.get() == true) && (HSAXnack.get() == false)) {
-        FATAL_MESSAGE0(1, "OMPX_APU_MAPS behavior requires HSA_XNACK=1");
-      }
-
-      assert(!(Plugin::get().getRequiresFlags() & OMP_REQ_UNDEFINED));
-      // Last condition of the following if statement is a workaround due to
-      // the presence of OMPX_APU_MAPS to enable unified shared memory mode
-      // for default programs run with HSA_XNACK=1. Remove once the
-      // OMPX_APU_MAPS mode is made default. Formerly implemented in
-      // RTLsTy::disableAPUMapsForUSM
-
-      if ((APUMaps.get() == true) && (HSAXnack.get() == true) &&
-          !(Plugin::get().getRequiresFlags() & OMP_REQ_UNIFIED_SHARED_MEMORY)) {
-        DisableAllocationsForMapsOnApus = true;
-      }
-    } else if (APUMaps.get() == true && HSAXnack == true) {
-      FATAL_MESSAGE0(1,
-                     "OMPX_APU_MAPS and HSA_XNACK enabled on non-APU system");
-    }
   }
 
   /// Check whether the image is compatible with an AMDGPU device.
-  Expected<bool> isImageCompatible(__tgt_image_info *Info) const override {
+  Expected<bool>
+  isImageCompatible(__tgt_image_info *Info,
+                    __tgt_device_image *TgtImage) const override {
+
     for (hsa_agent_t Agent : KernelAgents) {
       std::string Target;
       auto Err = utils::iterateAgentISAs(Agent, [&](hsa_isa_t ISA) {
@@ -3650,14 +3625,47 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
       if (utils::isImageCompatibleWithEnv(Info, Target))
         return true;
     }
+
+    // Check if the system's XNACK mode matches the one required by the
+    // image. Print a warning if not.
+    utils::checkImageCompatibilityWithSystemXnackMode(TgtImage,
+                                                      IsXnackEnabled());
+
     return false;
   }
 
-  void
-  checkAndAdjustXnackStatus(__tgt_device_image *TgtImage) const override final {
-    if (utils::wasBinaryBuiltWithXnackEnabled(TgtImage)) {
-      // TODO: Implement me. This is part of another ticket. It requries methods
-      // to manipulate the XNACK status.
+  void checkAndAdjustUsmModeForTargetImage(
+      __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()))
+      return;
+
+    bool IsXnackRequired =
+        Plugin::get().getRequiresFlags() & OMP_REQ_UNIFIED_SHARED_MEMORY;
+
+    utils::XnackBuildMode BinaryXnackMode =
+        utils::extractXnackModeFromBinary(TgtImage);
+
+    DisableAllocationsForMapsOnApus = false;
+
+    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");
+      }
     }
   }
 
@@ -3727,6 +3735,15 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
     return HSA_STATUS_ERROR;
   }
 
+  bool IsXnackEnabled() const {
+    // -1 instructs the runtime to query the XNACK status without modifying it.
+    int32_t enable = -1;
+    KMT_EXPECT_SUCCESS(hsaKmtOpenKFD());
+    KMT_EXPECT_SUCCESS(hsaKmtGetXNACKMode(&enable));
+    KMT_EXPECT_SUCCESS(hsaKmtCloseKFD());
+    return (enable > 0);
+  }
+
   bool checkForDeviceByGFXName(const llvm::StringRef GfxLookUpName) {
     bool CheckForMI300A =
         (GfxLookUpName.find_insensitive("gfx940") != llvm::StringRef::npos);
@@ -3776,8 +3793,8 @@ struct AMDGPUPluginTy final : public GenericPluginTy {
 
   BoolEnvar NoMapChecks;
   BoolEnvar DisableUsmMaps;
-  BoolEnvar APUMaps;
-  BoolEnvar HSAXnack;
+  BoolEnvar HsaXnack;
+  bool IsHsaXnackDefined{false};
 
   // Set by OMPX_APU_MAPS environment variable.
   // If set, maps cause no copy operations. USM is used instead. Allocated

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

@@ -81,6 +81,13 @@ enum COV_OFFSETS : uint32_t {
   PER_DEVICE_PREALLOC_SIZE = 131072
 };
 
+enum XnackBuildMode : short {
+  XNACK_UNSUPPORTED = -1,
+  XNACK_MINUS = 0,
+  XNACK_PLUS = 1,
+  XNACK_ANY = 2
+};
+
 /// Parse a TargetID to get processor arch and feature map.
 /// Returns processor subarch.
 /// Returns TargetID features in \p FeatureMap argument.
@@ -169,25 +176,41 @@ bool isImageCompatibleWithEnv(const __tgt_image_info *Info,
   return true;
 }
 
-// Check target image for XNACK option (XNACK+, XNACK-ANY, XNACK-)
-[[nodiscard]] bool
-wasBinaryBuiltWithXnackEnabled(__tgt_device_image *TgtImage) {
+// Check target image for XNACK mode (XNACK+, XNACK-ANY, XNACK-)
+[[nodiscard]] XnackBuildMode
+extractXnackModeFromBinary(__tgt_device_image *TgtImage) {
   assert((TgtImage != nullptr) && "TgtImage is nullptr.");
   u_int16_t EFlags = elf_get_eflags(TgtImage);
 
   unsigned XnackFlags = EFlags & ELF::EF_AMDGPU_FEATURE_XNACK_V4;
 
   switch (XnackFlags) {
-  case ELF::EF_AMDGPU_FEATURE_XNACK_ANY_V4:
   case ELF::EF_AMDGPU_FEATURE_XNACK_ON_V4:
-    return true;
+    return XnackBuildMode::XNACK_PLUS;
+  case ELF::EF_AMDGPU_FEATURE_XNACK_ANY_V4:
+    return XnackBuildMode::XNACK_ANY;
   case ELF::EF_AMDGPU_FEATURE_XNACK_OFF_V4:
+    return XnackBuildMode::XNACK_MINUS;
   case ELF::EF_AMDGPU_FEATURE_XNACK_UNSUPPORTED_V4:
-    return false;
+    return XnackBuildMode::XNACK_UNSUPPORTED;
   default:
     FAILURE_MESSAGE("Unknown XNACK flag!\n");
   }
-  return false;
+  return XNACK_MINUS;
+}
+
+void checkImageCompatibilityWithSystemXnackMode(__tgt_device_image *TgtImage,
+                                                bool IsXnackEnabled) {
+  XnackBuildMode ImageXnackMode = utils::extractXnackModeFromBinary(TgtImage);
+  if ((IsXnackEnabled && !ImageXnackMode)) {
+    FAILURE_MESSAGE(
+        "Image is not compatible with current XNACK mode! XNACK is enabled "
+        "on the system but image was compiled with xnack-.\n");
+  } else if (!IsXnackEnabled && (ImageXnackMode == 1)) {
+    FAILURE_MESSAGE("Image is not compatible with current XNACK mode! "
+                    "XNACK is disabled on the system. However, the image "
+                    "requires xnack+.\n");
+  }
 }
 
 struct KernelMetaDataTy {

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

@@ -476,8 +476,7 @@ uint32_t GenericKernelTy::getNumThreads(GenericDeviceTy &GenericDevice,
 uint64_t GenericKernelTy::getNumBlocks(GenericDeviceTy &GenericDevice,
                                        uint32_t NumTeamsClause[3],
                                        uint64_t LoopTripCount,
-                                       uint32_t &NumThreads
-				       ) const {
+                                       uint32_t &NumThreads) const {
   assert(NumTeamsClause[1] == 0 && NumTeamsClause[2] == 0 &&
          "Multi dimensional launch not supported yet.");
 
@@ -1501,21 +1500,15 @@ int32_t __tgt_rtl_is_valid_binary_info(__tgt_device_image *TgtImage,
   if (!Plugin::isActive())
     return false;
 
-
   if (!__tgt_rtl_is_valid_binary(TgtImage))
     return false;
-
-  // Need to call this method before 'isImageCompatibleCheck' in order to adjust
-  // settings.
-  Plugin::get().checkAndAdjustXnackStatus(TgtImage);
-
   // A subarchitecture was not specified. Assume it is compatible.
   if (!Info->Arch)
     return true;
 
   // Check the compatibility with all the available devices. Notice the
   // devices may not be initialized yet.
-  auto CompatibleOrErr = Plugin::get().isImageCompatible(Info);
+  auto CompatibleOrErr = Plugin::get().isImageCompatible(Info, TgtImage);
   if (!CompatibleOrErr) {
     // This error should not abort the execution, so we just inform the user
     // through the debug system.
@@ -1625,6 +1618,8 @@ __tgt_target_table *__tgt_rtl_load_binary(int32_t DeviceId,
     return nullptr;
   }
 
+  Plugin::get().checkAndAdjustUsmModeForTargetImage(TgtImage);
+
   __tgt_target_table *Table = *TableOrErr;
   assert(Table != nullptr && "Invalid table");
 

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.h

@@ -1008,7 +1008,7 @@ struct GenericPluginTy {
 
   virtual void setUpEnv() {}
   virtual void
-  checkAndAdjustXnackStatus(__tgt_device_image *TgtImage) const {}
+  checkAndAdjustUsmModeForTargetImage(__tgt_device_image *TgtImage) {}
 
   /// Get the ELF code to recognize the binary image of this plugin.
   virtual uint16_t getMagicElfBits() const = 0;
@@ -1059,7 +1059,9 @@ struct GenericPluginTy {
   /// Indicate if an image is compatible with the plugin devices. Notice that
   /// this function may be called before actually initializing the devices. So
   /// we could not move this function into GenericDeviceTy.
-  virtual Expected<bool> isImageCompatible(__tgt_image_info *Info) const = 0;
+  virtual Expected<bool>
+  isImageCompatible(__tgt_image_info *Info,
+                    __tgt_device_image *TgtImage) const = 0;
 
   /// Indicate whether the plugin supports empty images.
   virtual bool supportsEmptyImages() const { return false; }

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

@@ -1024,7 +1024,9 @@ struct CUDAPluginTy final : public GenericPluginTy {
   }
 
   /// Check whether the image is compatible with the available CUDA devices.
-  Expected<bool> isImageCompatible(__tgt_image_info *Info) const override {
+  Expected<bool>
+  isImageCompatible(__tgt_image_info *Info,
+                    __tgt_device_image *TgtImage) const override {
     for (int32_t DevId = 0; DevId < getNumDevices(); ++DevId) {
       CUdevice Device;
       CUresult Res = cuDeviceGet(&Device, DevId);

openmp/libomptarget/plugins-nextgen/generic-elf-64bit/src/rtl.cpp

@@ -414,7 +414,9 @@ struct GenELF64PluginTy final : public GenericPluginTy {
   }
 
   /// All images (ELF-compatible) should be compatible with this plugin.
-  Expected<bool> isImageCompatible(__tgt_image_info *Info) const override {
+  Expected<bool>
+  isImageCompatible(__tgt_image_info *Info,
+                    __tgt_device_image *TgtImage) const override {
     return true;
   }
 

openmp/libomptarget/src/CMakeLists.txt

@@ -9,8 +9,16 @@
 # Build offloading library libomptarget.so.
 #
 ##===----------------------------------------------------------------------===##
+# libhsakmt.a
+find_library ( HSAKMT_LIB libhsakmt.a REQURIED HINTS ${CMAKE_INSTALL_PREFIX} PATHS /opt/rocm)
+get_filename_component ( HSAKMT_LIB_PATH ${HSAKMT_LIB} DIRECTORY )
+link_directories (${HSAKMT_LIB_PATH})
 
-libomptarget_say("Building offloading runtime library libomptarget.")
+# lib_amdgpu
+pkg_check_modules(drm_amdgpu REQUIRED IMPORTED_TARGET libdrm_amdgpu)
+
+# libnuma
+find_library(numa libnuma.so REQUIRED HINTS /usr/lib/x86_64-linux-gnu/)
 
 add_llvm_library(omptarget
   SHARED
@@ -61,7 +69,14 @@ endif()
 
 find_library(LLVM_OFFLOAD_ARCH LLVMOffloadArch HINTS ${LLVM_LIBRARY_DIR} ${LLVM_BINARY_DIR} PATH_SUFFIXES lib REQUIRED)
 
-target_link_libraries(omptarget PRIVATE
+target_include_directories(omptarget PRIVATE ${LIBOMPTARGET_INCLUDE_DIR})
+
+target_link_libraries(
+  omptarget
+  PRIVATE hsakmt
+  PRIVATE drm_amdgpu
+  PRIVATE numa
+  PRIVATE
   ${CMAKE_DL_LIBS}
   ${LLVM_OFFLOAD_ARCH}
   "-Wl,--no-allow-shlib-undefined")