Add fixes and improvements

Author: kevinsala

offload/libomptarget/device.cpp

@@ -373,9 +373,9 @@ bool DeviceTy::isAccessiblePtr(const void *Ptr, size_t Size) {
 }
 
 uint64_t DeviceTy::getMaxSharedTeamMemory() {
-  InfoTreeNode Info = RTL->query_device_info(RTLDeviceID);
+  InfoTreeNode Info = RTL->obtain_device_info(RTLDeviceID);
 
-  auto EntryOpt = Info.get(DeviceInfo::WORK_GROUP_SHARED_MEM_SIZE);
+  auto EntryOpt = Info.get(DeviceInfo::WORK_GROUP_LOCAL_MEM_SIZE);
   if (!EntryOpt)
     return 0;
 

offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h

@@ -52,7 +52,6 @@ typedef enum {
   HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_GRANULE = 6,
   HSA_AMD_MEMORY_POOL_INFO_RUNTIME_ALLOC_ALIGNMENT = 7,
   HSA_AMD_MEMORY_POOL_INFO_ACCESSIBLE_BY_ALL = 15,
-  HSA_AMD_MEMORY_POOL_INFO_ALLOC_MAX_SIZE = 16,
 } hsa_amd_memory_pool_info_t;
 
 typedef enum {

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

@@ -299,6 +299,18 @@ struct InfoTreeNode {
   }
 };
 
+/// Configuration of dynamic block memory needed for launching a kernel.
+struct DynBlockMemConfTy {
+  /// The size of the dynamic block memory buffer.
+  uint32_t Size = 0;
+  /// The size of dynamic shared memory natively provided by the device.
+  uint32_t NativeSize = 0;
+  /// The fallback that was triggered (if any).
+  DynCGroupMemFallbackType Fallback = DynCGroupMemFallbackType::None;
+  /// The fallback pointer if global memory was used as alternative.
+  void *FallbackPtr = nullptr;
+};
+
 /// Class wrapping a __tgt_device_image and its offload entry table on a
 /// specific device. This class is responsible for storing and managing
 /// the offload entries for an image on a device.
@@ -386,10 +398,11 @@ struct GenericKernelTy {
   }
 
   /// Return a device pointer to a new kernel launch environment.
-  Expected<KernelLaunchEnvironmentTy *> getKernelLaunchEnvironment(
-      GenericDeviceTy &GenericDevice, const KernelArgsTy &KernelArgs,
-      uint32_t BlockMemSize, DynCGroupMemFallbackType DynBlockMemFb,
-      void *DynBlockMemFbPtr, AsyncInfoWrapperTy &AsyncInfoWrapper) const;
+  Expected<KernelLaunchEnvironmentTy *>
+  getKernelLaunchEnvironment(GenericDeviceTy &GenericDevice,
+                             const KernelArgsTy &KernelArgs,
+                             const DynBlockMemConfTy &DynBlockMemConf,
+                             AsyncInfoWrapperTy &AsyncInfoWrapper) const;
 
   /// Indicate whether an execution mode is valid.
   static bool isValidExecutionMode(OMPTgtExecModeFlags ExecutionMode) {
@@ -435,19 +448,11 @@ struct GenericKernelTy {
                                        uint32_t NumBlocks[3]) const;
 
 private:
-  /// Information about the dynamic block memory needed for launching a kernel.
-  struct DynBlockMemInfoTy {
-    /// The size of the dynamic block memory buffer.
-    uint32_t Size = 0;
-    /// The size of dynamic shared memory natively provided by the device.
-    uint32_t NativeSize = 0;
-    /// The fallback that was triggered (if any).
-    DynCGroupMemFallbackType DynBlockMemFb = DynCGroupMemFallbackType::None;
-    /// The fallback pointer if global memory was used as alternative.
-    void *FallbackPtr = nullptr;
-  };
-
-  Expected<DynBlockMemInfoTy> prepareBlockMemory(GenericDeviceTy &GenericDevice, KernelArgsTy &KernelArgs);
+  /// Prepare the block memory buffer requested for the kernel and execute the
+  /// specified fallback if necessary.
+  Expected<DynBlockMemConfTy> prepareBlockMemory(GenericDeviceTy &GenericDevice,
+                                                 KernelArgsTy &KernelArgs,
+                                                 uint32_t NumBlocks) const;
 
   /// Prepare the arguments before launching the kernel.
   KernelLaunchParamsTy

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

@@ -437,8 +437,8 @@ Error GenericKernelTy::init(GenericDeviceTy &GenericDevice,
 Expected<KernelLaunchEnvironmentTy *>
 GenericKernelTy::getKernelLaunchEnvironment(
     GenericDeviceTy &GenericDevice, const KernelArgsTy &KernelArgs,
-    uint32_t BlockMemSize, DynCGroupMemFallbackType DynBlockMemFb,
-    void *DynBlockMemFbPtr, AsyncInfoWrapperTy &AsyncInfoWrapper) const {
+    const DynBlockMemConfTy &DynBlockMemConf,
+    AsyncInfoWrapperTy &AsyncInfoWrapper) const {
   // Ctor/Dtor have no arguments, replaying uses the original kernel launch
   // environment. Older versions of the compiler do not generate a kernel
   // launch environment.
@@ -480,9 +480,9 @@ GenericKernelTy::getKernelLaunchEnvironment(
     LocalKLE.ReductionBuffer = nullptr;
   }
 
-  LocalKLE.DynCGroupMemSize = BlockMemSize;
-  LocalKLE.DynCGroupMemFbPtr = DynBlockMemFbPtr;
-  LocalKLE.DynCGroupMemFb = DynBlockMemFb;
+  LocalKLE.DynCGroupMemSize = DynBlockMemConf.Size;
+  LocalKLE.DynCGroupMemFbPtr = DynBlockMemConf.FallbackPtr;
+  LocalKLE.DynCGroupMemFb = DynBlockMemConf.Fallback;
 
   INFO(OMP_INFOTYPE_DATA_TRANSFER, GenericDevice.getDeviceId(),
        "Copying data from host to device, HstPtr=" DPxMOD ", TgtPtr=" DPxMOD
@@ -518,47 +518,51 @@ Error GenericKernelTy::printLaunchInfoDetails(GenericDeviceTy &GenericDevice,
   return Plugin::success();
 }
 
-Expected<DynBlockMemInfoTy> prepareBlockMemory(GenericDeviceTy &GenericDevice, KernelArgsTy &KernelArgs) {
-  uint32_t MaxSize = GenericDevice.getMaxBlockSharedMemSize();
-  uint32_t DynSize = KernelArgs.DynCGroupMem;
-  uint32_t TotalSize = StaticSize + DynSize;
-  uint32_t DynNativeSize = DynSize;
+Expected<DynBlockMemConfTy>
+GenericKernelTy::prepareBlockMemory(GenericDeviceTy &GenericDevice,
+                                    KernelArgsTy &KernelArgs,
+                                    uint32_t NumBlocks) const {
+  uint32_t MaxBlockMemSize = GenericDevice.getMaxBlockSharedMemSize();
+  uint32_t DynBlockMemSize = KernelArgs.DynCGroupMem;
+  uint32_t TotalBlockMemSize = StaticBlockMemSize + DynBlockMemSize;
+  uint32_t DynNativeBlockMemSize = DynBlockMemSize;
   void *DynFallbackPtr = nullptr;
 
   // No enough block memory to cover the static one. Cannot run the kernel.
-  if (StaticSize > MaxSize)
+  if (StaticBlockMemSize > MaxBlockMemSize)
     return Plugin::error(ErrorCode::INVALID_ARGUMENT,
                          "Static block memory size exceeds maximum");
   // No enough block memory to cover dynamic one, and the fallback is aborting.
   else if (static_cast<DynCGroupMemFallbackType>(
                KernelArgs.Flags.DynCGroupMemFallback) ==
                DynCGroupMemFallbackType::Abort &&
-           TotalSize > MaxSize)
+           TotalBlockMemSize > MaxBlockMemSize)
     return Plugin::error(
         ErrorCode::INVALID_ARGUMENT,
         "Static and dynamic block memory size exceeds maximum");
 
   DynCGroupMemFallbackType DynFallback = DynCGroupMemFallbackType::None;
-  if (DynSize && (!GenericDevice.hasNativeBlockSharedMem() ||
-                          TotalSize > MaxSize)) {
+  if (DynBlockMemSize && (!GenericDevice.hasNativeBlockSharedMem() ||
+                          TotalBlockMemSize > MaxBlockMemSize)) {
     // Launch without native dynamic block memory.
-    DynNativeSize = 0;
+    DynNativeBlockMemSize = 0;
     DynFallback = static_cast<DynCGroupMemFallbackType>(
         KernelArgs.Flags.DynCGroupMemFallback);
     if (DynFallback == DynCGroupMemFallbackType::DefaultMem) {
       // Get global memory as fallback.
       auto AllocOrErr = GenericDevice.dataAlloc(
-          NumBlocks[0] * DynSize,
+          NumBlocks * DynBlockMemSize,
           /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE);
       if (!AllocOrErr)
         return AllocOrErr.takeError();
       DynFallbackPtr = *AllocOrErr;
     } else {
       // Do not provide any memory as fallback.
-      DynSize = 0;
+      DynBlockMemSize = 0;
     }
   }
-  return { DynSize, DynNativeSize, DynFallback, DynFallbackPtr };
+  return DynBlockMemConfTy{DynBlockMemSize, DynNativeBlockMemSize, DynFallback,
+                           DynFallbackPtr};
 }
 
 Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
@@ -578,17 +582,18 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
                                 NumThreads[0], KernelArgs.ThreadLimit[0] > 0);
   }
 
-  auto DynBlockMemInfoOrErr = prepareBlockMemory(GenericDevice, KernelArgs);
-  if (!DynBlockMemInfoOrErr)
-    return DynBlockMemInfoOrErr.takeError();
+  auto DynBlockMemConfOrErr =
+      prepareBlockMemory(GenericDevice, KernelArgs, NumBlocks[0]);
+  if (!DynBlockMemConfOrErr)
+    return DynBlockMemConfOrErr.takeError();
 
-  DynBlockMemInfoTy &DynBlockMemInfo = *DynBlockMemInfoOrErr;
-  if (DynBlockMemInfo.FallbackPtr)
-    AsyncInfoWrapper.freeAllocationAfterSynchronization(DynBlockMemInfo.FallbackPtr);
+  DynBlockMemConfTy &DynBlockMemConf = *DynBlockMemConfOrErr;
+  if (DynBlockMemConf.FallbackPtr)
+    AsyncInfoWrapper.freeAllocationAfterSynchronization(
+        DynBlockMemConf.FallbackPtr);
 
   auto KernelLaunchEnvOrErr = getKernelLaunchEnvironment(
-      GenericDevice, KernelArgs, DynBlockMemInfo.Size, DynBlockMemInfo.Fallback,
-      DynBlockMemInfo.FallbackPtr, AsyncInfoWrapper);
+      GenericDevice, KernelArgs, DynBlockMemConf, AsyncInfoWrapper);
   if (!KernelLaunchEnvOrErr)
     return KernelLaunchEnvOrErr.takeError();
 
@@ -619,8 +624,9 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
           printLaunchInfo(GenericDevice, KernelArgs, NumThreads, NumBlocks))
     return Err;
 
-  return launchImpl(GenericDevice, NumThreads, NumBlocks, DynBlockMemInfo.NativeSize,
-                    KernelArgs, LaunchParams, AsyncInfoWrapper);
+  return launchImpl(GenericDevice, NumThreads, NumBlocks,
+                    DynBlockMemConf.NativeSize, KernelArgs, LaunchParams,
+                    AsyncInfoWrapper);
 }
 
 KernelLaunchParamsTy GenericKernelTy::prepareArgs(
@@ -2044,7 +2050,7 @@ InfoTreeNode GenericPluginTy::obtain_device_info(int32_t DeviceId) {
            toString(std::move(Err)).data());
     return InfoTreeNode{};
   }
-  return *InfoOrErr;
+  return std::move(*InfoOrErr);
 }
 
 void GenericPluginTy::print_device_info(int32_t DeviceId) {

offload/tools/deviceinfo/llvm-offload-device-info.cpp

@@ -205,8 +205,6 @@ ol_result_t printDevice(std::ostream &S, ol_device_handle_t D) {
       S, D, OL_DEVICE_INFO_MAX_MEM_ALLOC_SIZE, "Max Mem Allocation Size", "B"));
   OFFLOAD_ERR(printDeviceValue<uint64_t>(S, D, OL_DEVICE_INFO_GLOBAL_MEM_SIZE,
                                          "Global Mem Size", "B"));
-  OFFLOAD_ERR(printDeviceValue<uint64_t>(S, D, OL_DEVICE_INFO_WORK_GROUP_SHARED_MEM_SIZE,
-                                         "Work Group Shared Mem Size", "B"));
   OFFLOAD_ERR(
       printDeviceValue<uint64_t>(S, D, OL_DEVICE_INFO_WORK_GROUP_LOCAL_MEM_SIZE,
                                  "Work Group Shared Mem Size", "B"));