Skip to content

[SYCL] Cache implicit local arg position info #18654

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jun 3, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions sycl/source/detail/kernel_name_based_cache_t.hpp
Original file line number Diff line number Diff line change
@@ -74,6 +74,10 @@ struct FastKernelSubcacheT {
struct KernelNameBasedCacheT {
FastKernelSubcacheT FastKernelSubcache;
std::optional<bool> UsesAssert;
// Implicit local argument position is represented by an optional int, this
// uses another optional on top of that to represent lazy initialization of
// the cached value.
std::optional<std::optional<int>> ImplicitLocalArgPos;
};

} // namespace detail
24 changes: 18 additions & 6 deletions sycl/source/detail/program_manager/program_manager.cpp
Original file line number Diff line number Diff line change
@@ -1841,12 +1841,24 @@ void ProgramManager::cacheKernelImplicitLocalArg(RTDeviceBinaryImage &Img) {
}
}

std::optional<int>
ProgramManager::kernelImplicitLocalArgPos(KernelNameStrRefT KernelName) const {
auto it = m_KernelImplicitLocalArgPos.find(KernelName);
if (it != m_KernelImplicitLocalArgPos.end())
return it->second;
return {};
std::optional<int> ProgramManager::kernelImplicitLocalArgPos(
KernelNameStrRefT KernelName,
KernelNameBasedCacheT *KernelNameBasedCachePtr) const {
auto getLocalArgPos = [&]() -> std::optional<int> {
auto it = m_KernelImplicitLocalArgPos.find(KernelName);
if (it != m_KernelImplicitLocalArgPos.end())
return it->second;
return {};
};

if (!KernelNameBasedCachePtr)
return getLocalArgPos();
std::optional<std::optional<int>> &ImplicitLocalArgPos =
KernelNameBasedCachePtr->ImplicitLocalArgPos;
if (!ImplicitLocalArgPos.has_value()) {
ImplicitLocalArgPos = getLocalArgPos();
}
return ImplicitLocalArgPos.value();
}

static bool isBfloat16DeviceLibImage(sycl_device_binary RawImg,
5 changes: 3 additions & 2 deletions sycl/source/detail/program_manager/program_manager.hpp
Original file line number Diff line number Diff line change
@@ -373,8 +373,9 @@ class ProgramManager {

SanitizerType kernelUsesSanitizer() const { return m_SanitizerFoundInImage; }

std::optional<int>
kernelImplicitLocalArgPos(KernelNameStrRefT KernelName) const;
std::optional<int> kernelImplicitLocalArgPos(
KernelNameStrRefT KernelName,
KernelNameBasedCacheT *KernelNameBasedCachePtr) const;

std::set<RTDeviceBinaryImage *>
getRawDeviceImages(const std::vector<kernel_id> &KernelIDs);
12 changes: 7 additions & 5 deletions sycl/source/detail/scheduler/commands.cpp
Original file line number Diff line number Diff line change
@@ -2389,8 +2389,9 @@ static ur_result_t SetKernelParamsAndLaunch(
const std::function<void *(Requirement *Req)> &getMemAllocationFunc,
bool IsCooperative, bool KernelUsesClusterLaunch,
uint32_t WorkGroupMemorySize, const RTDeviceBinaryImage *BinImage,
KernelNameStrRefT KernelName, void *KernelFuncPtr = nullptr,
int KernelNumArgs = 0,
KernelNameStrRefT KernelName,
KernelNameBasedCacheT *KernelNameBasedCachePtr,
void *KernelFuncPtr = nullptr, int KernelNumArgs = 0,
detail::kernel_param_desc_t (*KernelParamDescGetter)(int) = nullptr,
bool KernelHasSpecialCaptures = true) {
const AdapterPtr &Adapter = Queue.getAdapter();
@@ -2437,7 +2438,8 @@ static ur_result_t SetKernelParamsAndLaunch(
}

std::optional<int> ImplicitLocalArg =
ProgramManager::getInstance().kernelImplicitLocalArgPos(KernelName);
ProgramManager::getInstance().kernelImplicitLocalArgPos(
KernelName, KernelNameBasedCachePtr);
// Set the implicit local memory buffer to support
// get_work_group_scratch_memory. This is for backend not supporting
// CUDA-style local memory setting. Note that we may have -1 as a position,
@@ -2752,8 +2754,8 @@ void enqueueImpKernel(
*Queue, Args, DeviceImageImpl, Kernel, NDRDesc, EventsWaitList,
OutEventImpl, EliminatedArgMask, getMemAllocationFunc,
KernelIsCooperative, KernelUsesClusterLaunch, WorkGroupMemorySize,
BinImage, KernelName, KernelFuncPtr, KernelNumArgs,
KernelParamDescGetter, KernelHasSpecialCaptures);
BinImage, KernelName, KernelNameBasedCachePtr, KernelFuncPtr,
KernelNumArgs, KernelParamDescGetter, KernelHasSpecialCaptures);
}
if (UR_RESULT_SUCCESS != Error) {
// If we have got non-success error code, let's analyze it to emit nice