| //===--- Level Zero Target RTL Implementation -----------------------------===// | |
| // | |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. | |
| // See https://llvm.org/LICENSE.txt for license information. | |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| // | |
| //===----------------------------------------------------------------------===// | |
| // | |
| // GenericKernel implementation for SPIR-V/Xe machine. | |
| // | |
| //===----------------------------------------------------------------------===// | |
| #include "L0Kernel.h" | |
| #include "L0Device.h" | |
| #include "L0Plugin.h" | |
| #include "L0Program.h" | |
| namespace llvm::omp::target::plugin { | |
| Error L0KernelTy::readKernelProperties(L0ProgramTy &Program) { | |
| const auto &L0Device = L0DeviceTy::makeL0Device(Program.getDevice()); | |
| auto &KernelPR = getProperties(); | |
| ze_kernel_properties_t KP = {}; | |
| KP.stype = ZE_STRUCTURE_TYPE_KERNEL_PROPERTIES; | |
| KP.pNext = nullptr; | |
| ze_kernel_preferred_group_size_properties_t KPrefGRPSize = {}; | |
| KPrefGRPSize.stype = ZE_STRUCTURE_TYPE_KERNEL_PREFERRED_GROUP_SIZE_PROPERTIES; | |
| KPrefGRPSize.pNext = nullptr; | |
| if (L0Device.getDriverAPIVersion() >= ZE_API_VERSION_1_2) | |
| KP.pNext = &KPrefGRPSize; | |
| CALL_ZE_RET_ERROR(zeKernelGetProperties, zeKernel, &KP); | |
| KernelPR.SIMDWidth = KP.maxSubgroupSize; | |
| KernelPR.Width = KP.maxSubgroupSize; | |
| KernelPR.NumKernelArgs = KP.numKernelArgs; | |
| if (KP.pNext) | |
| KernelPR.Width = KPrefGRPSize.preferredMultiple; | |
| if (!L0Device.isDeviceArch(DeviceArchTy::DeviceArch_Gen)) { | |
| KernelPR.Width = (std::max)(KernelPR.Width, 2 * KernelPR.SIMDWidth); | |
| } | |
| KernelPR.MaxThreadGroupSize = KP.maxSubgroupSize * KP.maxNumSubgroups; | |
| return Plugin::success(); | |
| } | |
| Error L0KernelTy::buildKernel(L0ProgramTy &Program) { | |
| const auto *KernelName = getName(); | |
| auto Module = Program.findModuleFromKernelName(KernelName); | |
| if (!Module) | |
| return Plugin::error(ErrorCode::NOT_FOUND, | |
| "kernel '%s' not found in the program", KernelName); | |
| ze_kernel_desc_t KernelDesc = {ZE_STRUCTURE_TYPE_KERNEL_DESC, nullptr, 0, | |
| KernelName}; | |
| CALL_ZE_RET_ERROR(zeKernelCreate, Module, &KernelDesc, &zeKernel); | |
| if (auto Err = readKernelProperties(Program)) | |
| return Err; | |
| return Plugin::success(); | |
| } | |
| Error L0KernelTy::initImpl(GenericDeviceTy &GenericDevice, | |
| DeviceImageTy &Image) { | |
| auto &Program = L0ProgramTy::makeL0Program(Image); | |
| if (auto Err = buildKernel(Program)) | |
| return Err; | |
| Program.addKernel(this); | |
| return Plugin::success(); | |
| } | |
| ze_group_size_t L0KernelTy::createKernelGroups(L0DeviceTy &L0Device, | |
| L0LaunchEnvTy &KEnv, | |
| uint32_t NumThreads[3], | |
| uint32_t NumBlocks[3]) const { | |
| assert(NumThreads[0] > 0 && NumThreads[1] > 0 && NumThreads[2] > 0 && | |
| "Pre-computed ThreadLimit values must be non-zero"); | |
| assert(NumBlocks[0] > 0 && NumBlocks[1] > 0 && NumBlocks[2] > 0 && | |
| "Pre-computed NumTeams values must be non-zero"); | |
| KEnv.GroupCounts = {NumBlocks[0], NumBlocks[1], NumBlocks[2]}; | |
| // Respect max group size attribute in the kernel. | |
| uint32_t MaxGroupSize = KEnv.KernelPR.MaxThreadGroupSize; | |
| KEnv.GroupSizes.groupSizeX = std::min<uint32_t>(MaxGroupSize, NumThreads[0]); | |
| KEnv.GroupSizes.groupSizeY = std::min<uint32_t>(MaxGroupSize, NumThreads[1]); | |
| KEnv.GroupSizes.groupSizeZ = std::min<uint32_t>(MaxGroupSize, NumThreads[2]); | |
| auto DeviceId = L0Device.getDeviceId(); | |
| ODBG(OLDT_Kernel) << "Device " << DeviceId << ": Team sizes = {" | |
| << KEnv.GroupSizes.groupSizeX << ", " | |
| << KEnv.GroupSizes.groupSizeY << ", " | |
| << KEnv.GroupSizes.groupSizeZ << "}"; | |
| ODBG(OLDT_Kernel) << "Device " << DeviceId << ": Number of teams = {" | |
| << KEnv.GroupCounts.groupCountX << ", " | |
| << KEnv.GroupCounts.groupCountY << ", " | |
| << KEnv.GroupCounts.groupCountZ << "}"; | |
| return KEnv.GroupSizes; | |
| } | |
| Error L0KernelTy::setIndirectFlags(L0DeviceTy &L0Device, | |
| LevelZeroPluginContextTy &UserCtx, | |
| L0LaunchEnvTy &KEnv) const { | |
| // Set Kernel Indirect flags. | |
| ze_kernel_indirect_access_flags_t Flags = 0; | |
| Flags |= L0Device.getMemAllocator(TARGET_ALLOC_HOST).getIndirectFlags(); | |
| Flags |= L0Device.getMemAllocator(TARGET_ALLOC_DEVICE).getIndirectFlags(); | |
| // Allocations made through a liboffload context live in its own allocators. | |
| Flags |= UserCtx.getIndirectFlags(); | |
| if (KEnv.KernelPR.IndirectAccessFlags != Flags) { | |
| // Combine with common access flags. | |
| const auto FinalFlags = L0Device.getIndirectFlags() | Flags; | |
| CALL_ZE_RET_ERROR(zeKernelSetIndirectAccess, zeKernel, FinalFlags); | |
| ODBG(OLDT_Kernel) << "Setting indirect access flags " | |
| << reinterpret_cast<void *>(FinalFlags); | |
| KEnv.KernelPR.IndirectAccessFlags = Flags; | |
| } | |
| return Plugin::success(); | |
| } | |
| Error L0KernelTy::launchImpl(GenericDeviceTy &GenericDevice, | |
| uint32_t NumThreads[3], uint32_t NumBlocks[3], | |
| uint32_t DynBlockMemSize, | |
| KernelLaunchArgsTy &LaunchArgs, | |
| AsyncInfoWrapperTy &AsyncInfoWrapper) const { | |
| if (DynBlockMemSize > 0) | |
| return Plugin::error(ErrorCode::UNSUPPORTED, | |
| "dynamic shared memory is unsupported in L0 plugin"); | |
| auto &L0Device = L0DeviceTy::makeL0Device(GenericDevice); | |
| __tgt_async_info *AsyncInfo = AsyncInfoWrapper; | |
| assert(AsyncInfo && "AsyncInfo must be provided for L0 kernel launch"); | |
| auto zeKernel = getZeKernel(); | |
| auto DeviceId = L0Device.getDeviceId(); | |
| ODBG(OLDT_Kernel) << "Device " << DeviceId << ": Launching kernel " | |
| << zeKernel << "..."; | |
| auto *IdStr = L0Device.getZeIdCStr(); | |
| bool IsCooperative = LaunchArgs.Flags.Cooperative; | |
| if (IsCooperative && !L0Device.supportsCooperativeKernels()) { | |
| return Plugin::error( | |
| ErrorCode::UNSUPPORTED, | |
| "cooperative kernel launch is not supported by the device"); | |
| } | |
| auto QueueOrErr = L0Device.getOrCreateQueue(AsyncInfo); | |
| if (!QueueOrErr) | |
| return QueueOrErr.takeError(); | |
| auto *Queue = *QueueOrErr; | |
| auto &KernelPR = getProperties(); | |
| L0LaunchEnvTy KEnv(KernelPR, LaunchArgs); | |
| // Protect from kernel preparation to submission as kernels are shared. | |
| KEnv.Lock.lock(); | |
| createKernelGroups(L0Device, KEnv, NumThreads, NumBlocks); | |
| // Validate cooperative kernel launch constraints | |
| if (IsCooperative) { | |
| uint32_t MaxCooperativeGroupCount = 0; | |
| CALL_ZE_RET_ERROR(zeKernelSuggestMaxCooperativeGroupCount, zeKernel, | |
| &MaxCooperativeGroupCount); | |
| uint32_t TotalGroupCount = KEnv.GroupCounts.groupCountX * | |
| KEnv.GroupCounts.groupCountY * | |
| KEnv.GroupCounts.groupCountZ; | |
| if (TotalGroupCount > MaxCooperativeGroupCount) { | |
| KernelPR.Mtx.unlock(); | |
| return Plugin::error( | |
| ErrorCode::INVALID_ARGUMENT, | |
| "cooperative kernel launch failed: requested %u groups exceeds " | |
| "maximum %u cooperative groups supported by device", | |
| TotalGroupCount, MaxCooperativeGroupCount); | |
| } | |
| ODBG(OLDT_Kernel) << "Device " << DeviceId | |
| << ": Cooperative kernel validated: using " | |
| << TotalGroupCount | |
| << " groups (max: " << MaxCooperativeGroupCount << ")"; | |
| } | |
| // With pointer-array arguments, zeCommandListAppendLaunchKernelWithArguments | |
| // folds group-size, per-argument set, and launch into a single call. | |
| if (LaunchArgs.NumArgs != KernelPR.NumKernelArgs) | |
| return Plugin::error( | |
| ErrorCode::INVALID_ARGUMENT, | |
| "Number of arguments (%u) does not match the number of arguments " | |
| "expected by the kernel (%u)", | |
| LaunchArgs.NumArgs, KernelPR.NumKernelArgs); | |
| assert(Queue->getUserCtx() && "Queue must belong to a plugin context"); | |
| if (auto Err = setIndirectFlags(L0Device, *Queue->getUserCtx(), KEnv)) | |
| return Err; | |
| // The next call should unlock the KernelLock internally. | |
| if (auto Err = Queue->launchKernel(zeKernel, KEnv)) | |
| return Err; | |
| ODBG(OLDT_Kernel) << "Device " << DeviceId << ": Submitted kernel " | |
| << zeKernel << " to device " << IdStr; | |
| return Plugin::success(); | |
| } | |
| Expected<uint32_t> | |
| L0KernelTy::getMaxCooperativeGroupCount(GenericDeviceTy &GenericDevice, | |
| const uint32_t NumThreads[3], | |
| uint32_t DynBlockMemSize) const { | |
| ze_result_t Res = zeKernelSetGroupSize(zeKernel, NumThreads[0], NumThreads[1], | |
| NumThreads[2]); | |
| if (Res != ZE_RESULT_SUCCESS) | |
| return Plugin::error(ErrorCode::UNSUPPORTED, | |
| "failed to set group size for cooperative launch"); | |
| uint32_t MaxCooperativeGroupCount = 0; | |
| Res = zeKernelSuggestMaxCooperativeGroupCount(zeKernel, | |
| &MaxCooperativeGroupCount); | |
| if (Res != ZE_RESULT_SUCCESS) | |
| return Plugin::error(ErrorCode::UNSUPPORTED, | |
| "failed to query max cooperative group count"); | |
| return MaxCooperativeGroupCount; | |
| } | |
| } // namespace llvm::omp::target::plugin |