https://github.com/adurang created https://github.com/llvm/llvm-project/pull/223761
Move preparations related to OpenMP KLE and dynamicCGroupMem fallback out of the plugins and into libomptarget. Resructure Device::launch as it grew too large. --- <sub>Stack created with <a href="https://github.com/github/gh-stack";>GitHub Stacks CLI</a> • <a href="https://gh.io/stacks-feedback";>Give Feedback 💬</a></sub> >From b19fbd2b24f6ee54da9e6ffe57fbd71de7965b50 Mon Sep 17 00:00:00 2001 From: "Duran, Alex" <[email protected]> Date: Tue, 15 Sep 2026 09:40:11 -0700 Subject: [PATCH] [offload][omp] Move OpenMP KLE to libomptarget Move preparations related to OpenMP KLE and dynamicCGroupMem fallback out of the plugins and into libomptarget. Resructure Device::launch as it grew too large. --- offload/include/device.h | 1 + offload/libompaccsupport/PluginManager.cpp | 1 + offload/libompaccsupport/device.cpp | 352 +++++++++++++++--- .../common/include/PluginInterface.h | 55 +-- .../common/src/PluginInterface.cpp | 146 +------- .../common/src/RecordReplay.cpp | 2 +- 6 files changed, 301 insertions(+), 256 deletions(-) diff --git a/offload/include/device.h b/offload/include/device.h index 9c5fa061ec608..c4bcec5a71805 100644 --- a/offload/include/device.h +++ b/offload/include/device.h @@ -49,6 +49,7 @@ struct KernelLaunchInfoTy { uint32_t MaxNumThreads = 0; uint32_t PreferredNumThreads = 0; uint32_t ReductionDataSize = 0; + uint32_t StaticBlockMemSize = 0; llvm::omp::OMPTgtExecModeFlags Mode = llvm::omp::OMP_TGT_EXEC_MODE_BARE; bool isBareMode() const { return Mode == llvm::omp::OMP_TGT_EXEC_MODE_BARE; } diff --git a/offload/libompaccsupport/PluginManager.cpp b/offload/libompaccsupport/PluginManager.cpp index babd9347655d2..9c9f2edb3fca9 100644 --- a/offload/libompaccsupport/PluginManager.cpp +++ b/offload/libompaccsupport/PluginManager.cpp @@ -529,6 +529,7 @@ static int loadImagesOntoDevice(DeviceTy &Device) { ? std::max(Cfg.MinThreads, int32_t(GenericDevice.getDefaultNumThreads())) : GenericDevice.getDefaultNumThreads(); + LaunchInfo.StaticBlockMemSize = Kernel->getStaticBlockMemSize(); Device.setKernelLaunchInfo(DeviceEntry.Address, LaunchInfo); } diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp index 14ab45709d0d5..a198a4a900835 100644 --- a/offload/libompaccsupport/device.cpp +++ b/offload/libompaccsupport/device.cpp @@ -389,6 +389,152 @@ static void resolveKernelLaunchParams(void **const TgtArgs, LaunchArgs.Args = &Ptrs[0]; } +namespace { +/// 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; +}; +} // namespace + +/// Prepare the block memory buffer requested for the kernel and execute the +/// specified fallback if necessary. +static llvm::Expected<DynBlockMemConfTy> +prepareBlockMemory(GenericDeviceTy &GenericDevice, + const KernelLaunchInfoTy &KernelEnv, uint32_t DynCGroupMem, + DynCGroupMemFallbackType DynCGroupMemFallback, + uint32_t NumBlocks) { + uint32_t MaxBlockMemSize = GenericDevice.getMaxBlockSharedMemSize(); + uint32_t DynBlockMemSize = DynCGroupMem; + uint32_t TotalBlockMemSize = KernelEnv.StaticBlockMemSize + DynBlockMemSize; + uint32_t DynNativeBlockMemSize = DynBlockMemSize; + void *DynFallbackPtr = nullptr; + + // No enough block memory to cover the static one. Cannot run the kernel. + if (KernelEnv.StaticBlockMemSize > MaxBlockMemSize) + return error::createOffloadError( + error::ErrorCode::INVALID_ARGUMENT, + "Static block memory size exceeds maximum"); + // No enough block memory to cover dynamic one, and the fallback is aborting. + if (DynCGroupMemFallback == DynCGroupMemFallbackType::Abort && + TotalBlockMemSize > MaxBlockMemSize) + return error::createOffloadError( + error::ErrorCode::INVALID_ARGUMENT, + "Requested block memory size (static + dynamic) exceeds maximum"); + + DynCGroupMemFallbackType DynFallback = DynCGroupMemFallbackType::None; + if (DynBlockMemSize && TotalBlockMemSize > MaxBlockMemSize) { + // Launch without native dynamic block memory. + DynNativeBlockMemSize = 0; + DynFallback = DynCGroupMemFallback; + if (DynFallback != DynCGroupMemFallbackType::DefaultMem) { + // Do not provide any memory as fallback. + DynBlockMemSize = 0; + } else { + // Get global memory as fallback. + auto AllocOrErr = GenericDevice.dataAlloc( + NumBlocks * DynBlockMemSize, + /*HostPtr=*/nullptr, TARGET_ALLOC_DEVICE, /*Alignment=*/0); + if (!AllocOrErr) + return AllocOrErr.takeError(); + DynFallbackPtr = *AllocOrErr; + } + } + return DynBlockMemConfTy{DynBlockMemSize, DynNativeBlockMemSize, DynFallback, + DynFallbackPtr}; +} + +static void freeAfterSynchronization(GenericDeviceTy &GenericDevice, + AsyncInfoTy &AsyncInfo, void *Ptr, + TargetAllocTy Kind) { + AsyncInfo.addPostProcessingFunction([&GenericDevice, Ptr, Kind]() -> int { + if (auto Err = GenericDevice.dataDelete(Ptr, Kind)) { + REPORT() << "Failure to free device memory " << Ptr << ": " + << toString(std::move(Err)); + return OFFLOAD_FAIL; + } + return OFFLOAD_SUCCESS; + }); +} + +/// Return a device pointer to a new kernel launch environment, or null if +/// this launch has no reserved dyn_ptr slot to store one in. \p NumBlocks0 is +/// the number of blocks for this launch and is used to size the reduction +/// buffer. +static llvm::Expected<KernelLaunchEnvironmentTy *> getKernelLaunchEnvironment( + GenericDeviceTy &GenericDevice, const KernelLaunchArgsTy &LaunchArgs, + const KernelLaunchInfoTy &KernelEnv, + const DynBlockMemConfTy &DynBlockMemConf, uint32_t DynCGroupMem, + void **DynPtrSlot, AsyncInfoTy &AsyncInfo, uint32_t NumBlocks0) { + // Ctor/Dtor have no arguments, replaying uses the original kernel launch + // environment, and launches with no reserved dyn_ptr slot (e.g. older + // compiler versions, or non-OpenMP launches) have nowhere to store one. + if ((GenericDevice.getRecordReplay() && + GenericDevice.getRecordReplay()->isReplaying()) || + !DynPtrSlot) + return nullptr; + + const bool NeedsReductionBuffer = KernelEnv.ReductionDataSize != 0; + if (NeedsReductionBuffer && LaunchArgs.OmpABIVersion < OMP_KERNEL_ARG_VERSION) + return error::createOffloadError( + error::ErrorCode::INVALID_BINARY, + "kernel was built against an older OpenMP kernel-launch-environment " + "ABI (v%u); current runtime requires v%u for cross-team reductions", + LaunchArgs.OmpABIVersion, OMP_KERNEL_ARG_VERSION); + if (!NeedsReductionBuffer && !DynCGroupMem) + return reinterpret_cast<KernelLaunchEnvironmentTy *>(~0); + + auto AllocOrErr = GenericDevice.dataAlloc( + sizeof(KernelLaunchEnvironmentTy), + /*HostPtr=*/nullptr, TARGET_ALLOC_DEVICE, /*Alignment=*/0); + if (!AllocOrErr) + return AllocOrErr.takeError(); + + // Remember to free the memory later. + freeAfterSynchronization(GenericDevice, AsyncInfo, *AllocOrErr, + TARGET_ALLOC_DEVICE); + + // Use the KLE in the __tgt_async_info to ensure a stable address for the + // async data transfer. + auto &LocalKLE = + static_cast<__tgt_async_info *>(AsyncInfo)->KernelLaunchEnvironment; + LocalKLE = KernelLaunchEnvironmentTy{}; + LocalKLE.DynCGroupMemSize = DynBlockMemConf.Size; + LocalKLE.DynCGroupMemFbPtr = DynBlockMemConf.FallbackPtr; + LocalKLE.DynCGroupMemFb = DynBlockMemConf.Fallback; + LocalKLE.ReductionBuffer = nullptr; + + if (NeedsReductionBuffer) { + // Use number of teams many buffer elements. + auto ReductionAllocOrErr = GenericDevice.dataAlloc( + uint64_t(KernelEnv.ReductionDataSize) * NumBlocks0, + /*HostPtr=*/nullptr, TARGET_ALLOC_DEVICE, /*Alignment=*/0); + if (!ReductionAllocOrErr) + return ReductionAllocOrErr.takeError(); + LocalKLE.ReductionBuffer = *ReductionAllocOrErr; + // Remember to free the memory later. + freeAfterSynchronization(GenericDevice, AsyncInfo, *ReductionAllocOrErr, + TARGET_ALLOC_DEVICE); + } + + INFO(OMP_INFOTYPE_DATA_TRANSFER, GenericDevice.getDeviceId(), + "Copying data from host to device, HstPtr=" DPxMOD ", TgtPtr=" DPxMOD + ", Size=%" PRId64 ", Name=KernelLaunchEnv\n", + DPxPTR(&LocalKLE), DPxPTR(*AllocOrErr), + sizeof(KernelLaunchEnvironmentTy)); + + if (auto Err = GenericDevice.dataSubmit( + *AllocOrErr, &LocalKLE, sizeof(KernelLaunchEnvironmentTy), AsyncInfo)) + return Err; + return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr); +} + /// Get the effective number of threads for the kernel based on the /// user-defined number of threads. static uint32_t getEffectiveNumThreads(GenericDeviceTy &GenericDevice, @@ -505,32 +651,27 @@ getEffectiveNumBlocks(GenericDeviceTy &GenericDevice, uint32_t UserNumBlocks, GenericDevice.getBlockLimit(EffectiveNumThreads)); } -// Run region on device -int32_t DeviceTy::launchKernel(void *TgtEntryPtr, void **TgtVarsPtr, - ptrdiff_t *TgtOffsets, KernelArgsTy &KernelArgs, - KernelReplayOutcomeTy *ReplayOutcome, - AsyncInfoTy &AsyncInfo) { - llvm::SmallVector<void *> Args, Ptrs; - llvm::SmallVector<int64_t> ArgSizes; - +/// Build the base KernelLaunchArgsTy for a launch from the public +/// KernelArgsTy and the kernel's cached launch-geometry properties. +static KernelLaunchArgsTy buildLaunchArgs(const KernelArgsTy &KernelArgs, + KernelReplayOutcomeTy *ReplayOutcome, + const KernelLaunchInfoTy &KernelEnv) { KernelLaunchArgsTy LaunchArgs; LaunchArgs.OmpABIVersion = KernelArgs.Version; LaunchArgs.ReplayOutcome = ReplayOutcome; LaunchArgs.ArgSizes = KernelArgs.ArgSizes; LaunchArgs.Tripcount = KernelArgs.Tripcount; - LaunchArgs.DynCGroupMem = KernelArgs.DynCGroupMem; llvm::copy(KernelArgs.UserNumBlocks, LaunchArgs.UserNumBlocks); llvm::copy(KernelArgs.UserThreadLimit, LaunchArgs.UserThreadLimit); LaunchArgs.Flags.Cooperative = KernelArgs.Flags.Cooperative; - LaunchArgs.Flags.DynCGroupMemFallback = KernelArgs.Flags.DynCGroupMemFallback; - - KernelLaunchInfoTy KernelEnv = getKernelLaunchInfo(TgtEntryPtr); - LaunchArgs.KernelEnvironment.ReductionDataSize = KernelEnv.ReductionDataSize; - LaunchArgs.KernelEnvironment.MaxNumThreads = KernelEnv.MaxNumThreads; - - const bool StrictBlocks = KernelArgs.Flags.StrictBlocks; - const bool StrictThreads = KernelArgs.Flags.StrictThreads; + LaunchArgs.MaxNumThreads = KernelEnv.MaxNumThreads; + return LaunchArgs; +} +/// Assert the launch geometry invariants expected by the plugin layer. +static void checkLaunchInvariants(const KernelLaunchArgsTy &LaunchArgs, + const KernelArgsTy &KernelArgs, + const KernelLaunchInfoTy &KernelEnv) { // Multidimensional is only supported with bare mode for now. assert(KernelEnv.isBareMode() || LaunchArgs.UserThreadLimit[1] == 1 && @@ -540,41 +681,58 @@ int32_t DeviceTy::launchKernel(void *TgtEntryPtr, void **TgtVarsPtr, "Non-bare mode should only use the first thread and block " "dimensions"); - assert(!StrictBlocks || + assert(!KernelArgs.Flags.StrictBlocks || LaunchArgs.UserNumBlocks[0] > 0 && LaunchArgs.UserNumBlocks[1] > 0 && LaunchArgs.UserNumBlocks[2] > 0 && "Strict requires number of blocks greater than zero"); - assert(!StrictThreads || + assert(!KernelArgs.Flags.StrictThreads || LaunchArgs.UserThreadLimit[0] > 0 && LaunchArgs.UserThreadLimit[1] > 0 && LaunchArgs.UserThreadLimit[2] > 0 && "Strict requires number of threads greater than zero"); +} + +/// Calculate or adjust, in place, the effective number of threads and blocks +/// for the first dimension, unless the caller requested strict counts. +static void adjustEffectiveGeometry(GenericDeviceTy &GenericDevice, + KernelLaunchArgsTy &LaunchArgs, + const KernelArgsTy &KernelArgs, + const KernelLaunchInfoTy &KernelEnv) { + const bool StrictBlocks = KernelArgs.Flags.StrictBlocks; + const bool StrictThreads = KernelArgs.Flags.StrictThreads; + if (StrictThreads && StrictBlocks) + return; + + assert(!KernelEnv.isBareMode() && + "bare kernel launches must request strict thread/block counts"); // Record whether the user actually requested a thread limit (thread_limit // clause) before possibly overwriting UserThreadLimit[0] below with the // computed effective value. const bool ThreadLimitFromUser = LaunchArgs.UserThreadLimit[0] > 0; - // Calculate or adjust the effective number of threads and blocks for the - // first dimension, if the caller didn't request strict counts. - if (!StrictThreads || !StrictBlocks) { - assert(!KernelEnv.isBareMode() && - "bare kernel launches must request strict thread/block counts"); - - GenericDeviceTy &GenericDevice = RTL->getDevice(RTLDeviceID); - uint32_t EffectiveNumThreads = LaunchArgs.UserThreadLimit[0]; - if (!StrictThreads) - EffectiveNumThreads = - getEffectiveNumThreads(GenericDevice, EffectiveNumThreads, KernelEnv); + uint32_t EffectiveNumThreads = LaunchArgs.UserThreadLimit[0]; + if (!StrictThreads) + EffectiveNumThreads = + getEffectiveNumThreads(GenericDevice, EffectiveNumThreads, KernelEnv); - if (!StrictBlocks) - LaunchArgs.UserNumBlocks[0] = getEffectiveNumBlocks( - GenericDevice, LaunchArgs.UserNumBlocks[0], LaunchArgs.Tripcount, - EffectiveNumThreads, StrictThreads, ThreadLimitFromUser, KernelEnv); + if (!StrictBlocks) + LaunchArgs.UserNumBlocks[0] = getEffectiveNumBlocks( + GenericDevice, LaunchArgs.UserNumBlocks[0], LaunchArgs.Tripcount, + EffectiveNumThreads, StrictThreads, ThreadLimitFromUser, KernelEnv); - LaunchArgs.UserThreadLimit[0] = EffectiveNumThreads; - } + LaunchArgs.UserThreadLimit[0] = EffectiveNumThreads; +} +/// Flatten the kernel arguments into \p LaunchArgs.Args. Returns the address +/// of the element reserved for the kernel launch environment (dyn_ptr), or +/// null if this launch has no such slot. +static void **resolveArgsAndDynPtrSlot(KernelArgsTy &KernelArgs, + void **TgtVarsPtr, ptrdiff_t *TgtOffsets, + llvm::SmallVector<void *> &Args, + llvm::SmallVector<void *> &Ptrs, + llvm::SmallVector<int64_t> &ArgSizes, + KernelLaunchArgsTy &LaunchArgs) { if (KernelArgs.Flags.IsCUDA) { // Kernel languages (CUDA/HIP) pass an already-flattened argument-pointer // array through KernelArgs.ArgPtrs instead of using the OpenMP @@ -583,34 +741,106 @@ int32_t DeviceTy::launchKernel(void *TgtEntryPtr, void **TgtVarsPtr, reinterpret_cast<KernelLaunchParamsTy *>(KernelArgs.ArgPtrs); LaunchArgs.NumArgs = LaunchParams->NumArgs; LaunchArgs.Args = LaunchParams->Args; - } else { - resolveKernelLaunchParams(TgtVarsPtr, TgtOffsets, KernelArgs.NumArgs, Args, - Ptrs, LaunchArgs); - // The dyn_ptr slot is reserved by the host (version >= 4) or by - // upgradeKernelArgs (version 3) as the last element of the argument - // array. Version 3 device kernels expect it first instead, so rotate it - // to the front to match that ABI. - if (KernelArgs.NumArgs > 0 && - KernelArgs.Version >= OMP_KERNEL_ARG_MIN_VERSION_WITH_DYN_PTR) { - if (KernelArgs.Version == OMP_KERNEL_ARG_MIN_VERSION_WITH_DYN_PTR) { - std::rotate(Args.begin(), Args.end() - 1, Args.end()); - LaunchArgs.DynPtrSlot = &Args[0]; - - // Keep ArgSizes in sync with the rotated Args, if present. - if (LaunchArgs.ArgSizes) { - ArgSizes.assign(LaunchArgs.ArgSizes, - LaunchArgs.ArgSizes + KernelArgs.NumArgs); - std::rotate(ArgSizes.begin(), ArgSizes.end() - 1, ArgSizes.end()); - LaunchArgs.ArgSizes = ArgSizes.data(); - } - } else { - LaunchArgs.DynPtrSlot = &Args[KernelArgs.NumArgs - 1]; - } - } + return nullptr; + } + + resolveKernelLaunchParams(TgtVarsPtr, TgtOffsets, KernelArgs.NumArgs, Args, + Ptrs, LaunchArgs); + + if (KernelArgs.NumArgs == 0 || + KernelArgs.Version < OMP_KERNEL_ARG_MIN_VERSION_WITH_DYN_PTR) + return nullptr; + + // The dyn_ptr slot is reserved by the host (version >= 4) or by + // upgradeKernelArgs (version 3) as the last element of the argument array. + // Version 3 device kernels expect it first instead, so rotate it to the + // front to match that ABI. + if (KernelArgs.Version != OMP_KERNEL_ARG_MIN_VERSION_WITH_DYN_PTR) + return &Args[KernelArgs.NumArgs - 1]; + + std::rotate(Args.begin(), Args.end() - 1, Args.end()); + + // Keep ArgSizes in sync with the rotated Args, if present. + if (LaunchArgs.ArgSizes) { + ArgSizes.assign(LaunchArgs.ArgSizes, + LaunchArgs.ArgSizes + KernelArgs.NumArgs); + std::rotate(ArgSizes.begin(), ArgSizes.end() - 1, ArgSizes.end()); + LaunchArgs.ArgSizes = ArgSizes.data(); + } + return &Args[0]; +} + +/// Compute the dynamic block-memory configuration for this launch, filling in +/// \p LaunchArgs.DynCGroupMem with the native size to request, and, if this +/// launch has a reserved dyn_ptr slot (\p DynPtrSlot), the device-side kernel +/// launch environment. +static llvm::Error +prepareDynamicLaunchState(GenericDeviceTy &GenericDevice, + const KernelLaunchInfoTy &KernelEnv, + KernelLaunchArgsTy &LaunchArgs, uint32_t DynCGroupMem, + DynCGroupMemFallbackType DynCGroupMemFallback, + void **DynPtrSlot, AsyncInfoTy &AsyncInfo) { + uint32_t NumBlocksTotal = LaunchArgs.UserNumBlocks[0] * + LaunchArgs.UserNumBlocks[1] * + LaunchArgs.UserNumBlocks[2]; + auto DynBlockMemConfOrErr = + prepareBlockMemory(GenericDevice, KernelEnv, DynCGroupMem, + DynCGroupMemFallback, NumBlocksTotal); + if (!DynBlockMemConfOrErr) + return DynBlockMemConfOrErr.takeError(); + + DynBlockMemConfTy &DynBlockMemConf = *DynBlockMemConfOrErr; + LaunchArgs.DynCGroupMem = DynBlockMemConf.NativeSize; + if (DynBlockMemConf.FallbackPtr) + freeAfterSynchronization(GenericDevice, AsyncInfo, + DynBlockMemConf.FallbackPtr, TARGET_ALLOC_DEVICE); + + auto KernelLaunchEnvOrErr = getKernelLaunchEnvironment( + GenericDevice, LaunchArgs, KernelEnv, DynBlockMemConf, DynCGroupMem, + DynPtrSlot, AsyncInfo, LaunchArgs.UserNumBlocks[0]); + if (!KernelLaunchEnvOrErr) + return KernelLaunchEnvOrErr.takeError(); + + // Fill in the kernel launch environment (dyn_ptr) if this launch has a + // reserved slot for it. When replaying, getKernelLaunchEnvironment() + // returns null so the recorded value already in the slot is preserved. + if (DynPtrSlot && *KernelLaunchEnvOrErr) + *DynPtrSlot = *KernelLaunchEnvOrErr; + + return llvm::Error::success(); +} + +// Run region on device +int32_t DeviceTy::launchKernel(void *TgtEntryPtr, void **TgtVarsPtr, + ptrdiff_t *TgtOffsets, KernelArgsTy &KernelArgs, + KernelReplayOutcomeTy *ReplayOutcome, + AsyncInfoTy &AsyncInfo) { + llvm::SmallVector<void *> Args, Ptrs; + llvm::SmallVector<int64_t> ArgSizes; + + GenericDeviceTy &GenericDevice = RTL->getDevice(RTLDeviceID); + KernelLaunchInfoTy KernelEnv = getKernelLaunchInfo(TgtEntryPtr); + KernelLaunchArgsTy LaunchArgs = + buildLaunchArgs(KernelArgs, ReplayOutcome, KernelEnv); + + checkLaunchInvariants(LaunchArgs, KernelArgs, KernelEnv); + adjustEffectiveGeometry(GenericDevice, LaunchArgs, KernelArgs, KernelEnv); + + void **DynPtrSlot = resolveArgsAndDynPtrSlot( + KernelArgs, TgtVarsPtr, TgtOffsets, Args, Ptrs, ArgSizes, LaunchArgs); + + auto DynCGroupMemFallback = static_cast<DynCGroupMemFallbackType>( + KernelArgs.Flags.DynCGroupMemFallback); + if (auto Err = prepareDynamicLaunchState( + GenericDevice, KernelEnv, LaunchArgs, KernelArgs.DynCGroupMem, + DynCGroupMemFallback, DynPtrSlot, AsyncInfo)) { + REPORT() << "Failure to prepare launch state for kernel " << TgtEntryPtr + << ": " << toString(std::move(Err)); + return OFFLOAD_FAIL; } auto *Kernel = reinterpret_cast<GenericKernelTy *>(TgtEntryPtr); - INFO(OMP_INFOTYPE_PLUGIN_KERNEL, RTL->getDevice(RTLDeviceID).getDeviceId(), + INFO(OMP_INFOTYPE_PLUGIN_KERNEL, GenericDevice.getDeviceId(), "Launching kernel %s with [%u,%u,%u] blocks and [%u,%u,%u] threads in " "%s mode\n", Kernel->getName(), LaunchArgs.UserNumBlocks[0], diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h index da2b09cfbc7ea..e2a67bbb7804c 100644 --- a/offload/plugins-nextgen/common/include/PluginInterface.h +++ b/offload/plugins-nextgen/common/include/PluginInterface.h @@ -316,18 +316,6 @@ 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; -}; - /// Tracker of virtual memory address reservations. template <typename HandleTy> class VMemTrackerTy { struct EntryTy { @@ -442,11 +430,6 @@ struct KernelLaunchArgsTy { /// Size of the argument data in bytes, one entry per \p Args element, /// possibly null. int64_t *ArgSizes = nullptr; - /// Address of the element of \p Args reserved for the kernel launch - /// environment (dyn_ptr), or null if this launch has no such slot. The - /// caller owns the storage it points into; the plugin fills it in once it - /// has computed the actual (device-side) value. - void **DynPtrSlot = nullptr; /// Tripcount for the teams / distribute loop, 0 otherwise. uint64_t Tripcount = 0; /// Amount of dynamic cgroup memory requested. @@ -455,18 +438,12 @@ struct KernelLaunchArgsTy { uint32_t UserNumBlocks[3] = {0, 0, 0}; /// User-requested number of threads (for x,y,z dimension). uint32_t UserThreadLimit[3] = {0, 0, 0}; - struct { - /// Size in bytes of a single cross-team reduction buffer element for - /// this kernel, or 0 if the kernel does not need a reduction buffer. - uint32_t ReductionDataSize = 0; - /// Maximum number of threads per block that this kernel may use. - uint32_t MaxNumThreads = 0; - } KernelEnvironment; + /// Maximum number of threads per block that this kernel may use. + uint32_t MaxNumThreads = 0; struct { uint64_t Cooperative : 1; // Was this kernel spawned as cooperative. - uint64_t DynCGroupMemFallback : 2; // The fallback for dynamic cgroup mem. - uint64_t Unused : 61; - } Flags = {0, 0, 0}; + uint64_t Unused : 63; + } Flags = {0, 0}; /// Set by the caller when replaying a previously recorded kernel launch, so /// the plugin can report the outcome back; null for a normal launch. KernelReplayOutcomeTy *ReplayOutcome = nullptr; @@ -488,10 +465,7 @@ struct GenericKernelTy { /// Launch the kernel on the specific device. The device must be the same /// one used to initialize the kernel. \p LaunchArgs.Args is the flattened - /// argument-pointer array to pass to the kernel, with any offsets already - /// resolved. \p LaunchArgs.DynPtrSlot, if non-null, points at the element - /// of it reserved for the kernel launch environment (dyn_ptr); the caller - /// owns the storage it points into. + /// argument-pointer array to pass to the kernel, with any offsets. Error launch(GenericDeviceTy &GenericDevice, KernelLaunchArgsTy &LaunchArgs, AsyncInfoWrapperTy &AsyncInfoWrapper) const; virtual Error launchImpl(GenericDeviceTy &GenericDevice, @@ -530,15 +504,6 @@ struct GenericKernelTy { return *ImagePtr; } - /// Return a device pointer to a new kernel launch environment. - /// - /// \p NumBlocks0 is the number of blocks for this launch and is used to size - /// the reduction buffer. - Expected<KernelLaunchEnvironmentTy *> getKernelLaunchEnvironment( - GenericDeviceTy &GenericDevice, const KernelLaunchArgsTy &LaunchArgs, - const DynBlockMemConfTy &DynBlockMemConf, - AsyncInfoWrapperTy &AsyncInfoWrapper, uint32_t NumBlocks0) const; - /// Indicate whether an execution mode is valid. static bool isValidExecutionMode(OMPTgtExecModeFlags ExecutionMode) { switch (ExecutionMode) { @@ -566,13 +531,6 @@ struct GenericKernelTy { uint32_t NumBlocks[3]) const; private: - /// Prepare the block memory buffer requested for the kernel and execute the - /// specified fallback if necessary. - Expected<DynBlockMemConfTy> - prepareBlockMemory(GenericDeviceTy &GenericDevice, - const KernelLaunchArgsTy &LaunchArgs, - uint32_t NumBlocks) const; - /// The kernel name. std::string Name; @@ -582,9 +540,6 @@ struct GenericKernelTy { protected: /// The static memory sized per block. uint32_t StaticBlockMemSize = 0; - - /// The prototype kernel launch environment. - KernelLaunchEnvironmentTy KernelLaunchEnvironment; }; /// Information about an allocation, when it has been allocated, and when/if it diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp index cb8dcdc857cdc..d8d3ce0fcb343 100644 --- a/offload/plugins-nextgen/common/src/PluginInterface.cpp +++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp @@ -75,78 +75,6 @@ Error GenericKernelTy::init(GenericDeviceTy &GenericDevice, return initImpl(GenericDevice, Image); } -Expected<KernelLaunchEnvironmentTy *> -GenericKernelTy::getKernelLaunchEnvironment( - GenericDeviceTy &GenericDevice, const KernelLaunchArgsTy &LaunchArgs, - const DynBlockMemConfTy &DynBlockMemConf, - AsyncInfoWrapperTy &AsyncInfoWrapper, uint32_t NumBlocks0) const { - // Ctor/Dtor have no arguments, replaying uses the original kernel launch - // environment, and launches with no reserved dyn_ptr slot (e.g. older - // compiler versions, or non-OpenMP launches) have nowhere to store one. - if ((GenericDevice.getRecordReplay() && - GenericDevice.getRecordReplay()->isReplaying()) || - !LaunchArgs.DynPtrSlot) - return nullptr; - - const bool NeedsReductionBuffer = - LaunchArgs.KernelEnvironment.ReductionDataSize != 0; - if (NeedsReductionBuffer && LaunchArgs.OmpABIVersion < OMP_KERNEL_ARG_VERSION) - return Plugin::error(ErrorCode::INVALID_BINARY, - "kernel was built against an older OpenMP " - "kernel-launch-environment ABI (v%u); current " - "runtime requires v%u for cross-team reductions", - LaunchArgs.OmpABIVersion, OMP_KERNEL_ARG_VERSION); - if (!NeedsReductionBuffer && !LaunchArgs.DynCGroupMem) - return reinterpret_cast<KernelLaunchEnvironmentTy *>(~0); - - auto AllocOrErr = GenericDevice.dataAlloc( - sizeof(KernelLaunchEnvironmentTy), - /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE, /*Alignment=*/0); - if (!AllocOrErr) - return AllocOrErr.takeError(); - - // Remember to free the memory later. - AsyncInfoWrapper.freeAllocationAfterSynchronization( - *AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE); - - /// Use the KLE in the __tgt_async_info to ensure a stable address for the - /// async data transfer. - auto &LocalKLE = (*AsyncInfoWrapper).KernelLaunchEnvironment; - LocalKLE = KernelLaunchEnvironment; - - LocalKLE.DynCGroupMemSize = DynBlockMemConf.Size; - LocalKLE.DynCGroupMemFbPtr = DynBlockMemConf.FallbackPtr; - LocalKLE.DynCGroupMemFb = DynBlockMemConf.Fallback; - LocalKLE.ReductionBuffer = nullptr; - - if (NeedsReductionBuffer) { - // Use number of teams many buffer elements. - auto AllocOrErr = GenericDevice.dataAlloc( - uint64_t(LaunchArgs.KernelEnvironment.ReductionDataSize) * NumBlocks0, - /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE, - /*Alignment=*/0); - if (!AllocOrErr) - return AllocOrErr.takeError(); - LocalKLE.ReductionBuffer = *AllocOrErr; - // Remember to free the memory later. - AsyncInfoWrapper.freeAllocationAfterSynchronization( - *AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE); - } - - INFO(OMP_INFOTYPE_DATA_TRANSFER, GenericDevice.getDeviceId(), - "Copying data from host to device, HstPtr=" DPxMOD ", TgtPtr=" DPxMOD - ", Size=%" PRId64 ", Name=KernelLaunchEnv\n", - DPxPTR(&LocalKLE), DPxPTR(*AllocOrErr), - sizeof(KernelLaunchEnvironmentTy)); - - auto Err = GenericDevice.dataSubmit(*AllocOrErr, &LocalKLE, - sizeof(KernelLaunchEnvironmentTy), - AsyncInfoWrapper); - if (Err) - return Err; - return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr); -} - Error GenericKernelTy::printLaunchInfo(GenericDeviceTy &GenericDevice, const KernelLaunchArgsTy &LaunchArgs, uint32_t NumThreads[3], @@ -161,53 +89,6 @@ Error GenericKernelTy::printLaunchInfoDetails( return Plugin::success(); } -Expected<DynBlockMemConfTy> -GenericKernelTy::prepareBlockMemory(GenericDeviceTy &GenericDevice, - const KernelLaunchArgsTy &LaunchArgs, - uint32_t NumBlocks) const { - uint32_t MaxBlockMemSize = GenericDevice.getMaxBlockSharedMemSize(); - uint32_t DynBlockMemSize = LaunchArgs.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 (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. - if (static_cast<DynCGroupMemFallbackType>( - LaunchArgs.Flags.DynCGroupMemFallback) == - DynCGroupMemFallbackType::Abort && - TotalBlockMemSize > MaxBlockMemSize) - return Plugin::error( - ErrorCode::INVALID_ARGUMENT, - "Requested block memory size (static + dynamic) exceeds maximum"); - - DynCGroupMemFallbackType DynFallback = DynCGroupMemFallbackType::None; - if (DynBlockMemSize && TotalBlockMemSize > MaxBlockMemSize) { - // Launch without native dynamic block memory. - DynNativeBlockMemSize = 0; - DynFallback = static_cast<DynCGroupMemFallbackType>( - LaunchArgs.Flags.DynCGroupMemFallback); - if (DynFallback != DynCGroupMemFallbackType::DefaultMem) { - // Do not provide any memory as fallback. - DynBlockMemSize = 0; - } else { - // Get global memory as fallback. - auto AllocOrErr = GenericDevice.dataAlloc( - NumBlocks * DynBlockMemSize, - /*HostPtr=*/nullptr, TargetAllocTy::TARGET_ALLOC_DEVICE, - /*Alignment=*/0); - if (!AllocOrErr) - return AllocOrErr.takeError(); - DynFallbackPtr = *AllocOrErr; - } - } - return DynBlockMemConfTy{DynBlockMemSize, DynNativeBlockMemSize, DynFallback, - DynFallbackPtr}; -} - Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, KernelLaunchArgsTy &LaunchArgs, AsyncInfoWrapperTy &AsyncInfoWrapper) const { @@ -218,29 +99,6 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, LaunchArgs.UserNumBlocks[1], LaunchArgs.UserNumBlocks[2]}; - auto DynBlockMemConfOrErr = prepareBlockMemory( - GenericDevice, LaunchArgs, - EffectiveNumBlocks[0] * EffectiveNumBlocks[1] * EffectiveNumBlocks[2]); - if (!DynBlockMemConfOrErr) - return DynBlockMemConfOrErr.takeError(); - - DynBlockMemConfTy &DynBlockMemConf = *DynBlockMemConfOrErr; - if (DynBlockMemConf.FallbackPtr) - AsyncInfoWrapper.freeAllocationAfterSynchronization( - DynBlockMemConf.FallbackPtr, TargetAllocTy::TARGET_ALLOC_DEVICE); - - auto KernelLaunchEnvOrErr = - getKernelLaunchEnvironment(GenericDevice, LaunchArgs, DynBlockMemConf, - AsyncInfoWrapper, EffectiveNumBlocks[0]); - if (!KernelLaunchEnvOrErr) - return KernelLaunchEnvOrErr.takeError(); - - // Fill in the kernel launch environment (dyn_ptr) if this launch has a - // reserved slot for it. When replaying, getKernelLaunchEnvironment() - // returns null so the recorded value already in the slot is preserved. - if (LaunchArgs.DynPtrSlot && *KernelLaunchEnvOrErr) - *LaunchArgs.DynPtrSlot = *KernelLaunchEnvOrErr; - if (auto Err = printLaunchInfo(GenericDevice, LaunchArgs, EffectiveNumThreads, EffectiveNumBlocks)) return Err; @@ -255,7 +113,7 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, // Record the kernel prologue data before kernel launch. auto RRHandleOrErr = RecordReplay->recordPrologue( *this, LaunchArgs, EffectiveNumBlocks, EffectiveNumThreads, - DynBlockMemConf.NativeSize); + LaunchArgs.DynCGroupMem); if (!RRHandleOrErr) return RRHandleOrErr.takeError(); RRHandle = *RRHandleOrErr; @@ -263,7 +121,7 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, if (auto Err = launchImpl(GenericDevice, EffectiveNumThreads, EffectiveNumBlocks, - DynBlockMemConf.NativeSize, LaunchArgs, AsyncInfoWrapper)) + LaunchArgs.DynCGroupMem, LaunchArgs, AsyncInfoWrapper)) return Err; if (RecordReplay) { diff --git a/offload/plugins-nextgen/common/src/RecordReplay.cpp b/offload/plugins-nextgen/common/src/RecordReplay.cpp index 768e85338a140..f97c26ceb2c52 100644 --- a/offload/plugins-nextgen/common/src/RecordReplay.cpp +++ b/offload/plugins-nextgen/common/src/RecordReplay.cpp @@ -279,7 +279,7 @@ Error NativeRecordReplayTy::recordDescImpl( // Export minimum and maximum for allowed number of threads. If zero, it means // there was no restriction provided by the program. - uint32_t MaxThreads = LaunchArgs.KernelEnvironment.MaxNumThreads; + uint32_t MaxThreads = LaunchArgs.MaxNumThreads; json::Array JsonThreadsLimits; JsonThreadsLimits.push_back(1); JsonThreadsLimits.push_back(MaxThreads); _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
