Improve kernel launch

Author: kevinsala

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

@@ -439,6 +439,20 @@ 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 arguments before launching the kernel.
   KernelLaunchParamsTy
   prepareArgs(GenericDeviceTy &GenericDevice, void **ArgPtrs,

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

@@ -518,66 +518,77 @@ Error GenericKernelTy::printLaunchInfoDetails(GenericDeviceTy &GenericDevice,
   return Plugin::success();
 }
 
-Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
-                              ptrdiff_t *ArgOffsets, KernelArgsTy &KernelArgs,
-                              AsyncInfoWrapperTy &AsyncInfoWrapper) const {
-  llvm::SmallVector<void *, 16> Args;
-  llvm::SmallVector<void *, 16> Ptrs;
-
-  uint32_t NumThreads[3] = {KernelArgs.ThreadLimit[0],
-                            KernelArgs.ThreadLimit[1],
-                            KernelArgs.ThreadLimit[2]};
-  uint32_t NumBlocks[3] = {KernelArgs.NumTeams[0], KernelArgs.NumTeams[1],
-                           KernelArgs.NumTeams[2]};
-  if (!isBareMode()) {
-    NumThreads[0] = getNumThreads(GenericDevice, NumThreads);
-    NumBlocks[0] = getNumBlocks(GenericDevice, NumBlocks, KernelArgs.Tripcount,
-                                NumThreads[0], KernelArgs.ThreadLimit[0] > 0);
-  }
-
-  uint32_t MaxBlockMemSize = GenericDevice.getMaxBlockSharedMemSize();
-  uint32_t DynBlockMemSize = KernelArgs.DynCGroupMem;
-  uint32_t TotalBlockMemSize = StaticBlockMemSize + DynBlockMemSize;
-  if (StaticBlockMemSize > MaxBlockMemSize)
+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;
+  void *DynFallbackPtr = nullptr;
+
+  // No enough block memory to cover the static one. Cannot run the kernel.
+  if (StaticSize > MaxSize)
     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 &&
-           TotalBlockMemSize > MaxBlockMemSize)
+           TotalSize > MaxSize)
     return Plugin::error(
         ErrorCode::INVALID_ARGUMENT,
         "Static and dynamic block memory size exceeds maximum");
 
-  void *DynBlockMemFbPtr = nullptr;
-  uint32_t DynBlockMemLaunchSize = DynBlockMemSize;
-
-  DynCGroupMemFallbackType DynBlockMemFb = DynCGroupMemFallbackType::None;
-  if (DynBlockMemSize && (!GenericDevice.hasNativeBlockSharedMem() ||
-                          TotalBlockMemSize > MaxBlockMemSize)) {
+  DynCGroupMemFallbackType DynFallback = DynCGroupMemFallbackType::None;
+  if (DynSize && (!GenericDevice.hasNativeBlockSharedMem() ||
+                          TotalSize > MaxSize)) {
     // Launch without native dynamic block memory.
-    DynBlockMemLaunchSize = 0;
-    DynBlockMemFb = static_cast<DynCGroupMemFallbackType>(
+    DynNativeSize = 0;
+    DynFallback = static_cast<DynCGroupMemFallbackType>(
         KernelArgs.Flags.DynCGroupMemFallback);
-    if (DynBlockMemFb == DynCGroupMemFallbackType::DefaultMem) {
+    if (DynFallback == DynCGroupMemFallbackType::DefaultMem) {
       // Get global memory as fallback.
       auto AllocOrErr = GenericDevice.dataAlloc(
-          NumBlocks[0] * DynBlockMemSize,
+          NumBlocks[0] * DynSize,
           /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE);
       if (!AllocOrErr)
         return AllocOrErr.takeError();
-
-      DynBlockMemFbPtr = *AllocOrErr;
-      AsyncInfoWrapper.freeAllocationAfterSynchronization(DynBlockMemFbPtr);
+      DynFallbackPtr = *AllocOrErr;
     } else {
       // Do not provide any memory as fallback.
-      DynBlockMemSize = 0;
+      DynSize = 0;
     }
   }
+  return { DynSize, DynNativeSize, DynFallback, DynFallbackPtr };
+}
+
+Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
+                              ptrdiff_t *ArgOffsets, KernelArgsTy &KernelArgs,
+                              AsyncInfoWrapperTy &AsyncInfoWrapper) const {
+  llvm::SmallVector<void *, 16> Args;
+  llvm::SmallVector<void *, 16> Ptrs;
+
+  uint32_t NumThreads[3] = {KernelArgs.ThreadLimit[0],
+                            KernelArgs.ThreadLimit[1],
+                            KernelArgs.ThreadLimit[2]};
+  uint32_t NumBlocks[3] = {KernelArgs.NumTeams[0], KernelArgs.NumTeams[1],
+                           KernelArgs.NumTeams[2]};
+  if (!isBareMode()) {
+    NumThreads[0] = getNumThreads(GenericDevice, NumThreads);
+    NumBlocks[0] = getNumBlocks(GenericDevice, NumBlocks, KernelArgs.Tripcount,
+                                NumThreads[0], KernelArgs.ThreadLimit[0] > 0);
+  }
+
+  auto DynBlockMemInfoOrErr = prepareBlockMemory(GenericDevice, KernelArgs);
+  if (!DynBlockMemInfoOrErr)
+    return DynBlockMemInfoOrErr.takeError();
+
+  DynBlockMemInfoTy &DynBlockMemInfo = *DynBlockMemInfoOrErr;
+  if (DynBlockMemInfo.FallbackPtr)
+    AsyncInfoWrapper.freeAllocationAfterSynchronization(DynBlockMemInfo.FallbackPtr);
 
   auto KernelLaunchEnvOrErr = getKernelLaunchEnvironment(
-      GenericDevice, KernelArgs, DynBlockMemSize, DynBlockMemFb,
-      DynBlockMemFbPtr, AsyncInfoWrapper);
+      GenericDevice, KernelArgs, DynBlockMemInfo.Size, DynBlockMemInfo.Fallback,
+      DynBlockMemInfo.FallbackPtr, AsyncInfoWrapper);
   if (!KernelLaunchEnvOrErr)
     return KernelLaunchEnvOrErr.takeError();
 
@@ -608,7 +619,7 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
           printLaunchInfo(GenericDevice, KernelArgs, NumThreads, NumBlocks))
     return Err;
 
-  return launchImpl(GenericDevice, NumThreads, NumBlocks, DynBlockMemLaunchSize,
+  return launchImpl(GenericDevice, NumThreads, NumBlocks, DynBlockMemInfo.NativeSize,
                     KernelArgs, LaunchParams, AsyncInfoWrapper);
 }