|
10 | 10 |
|
11 | 11 | #include "context.hpp"
|
12 | 12 |
|
| 13 | +#include <mutex> |
| 14 | +#include <set> |
| 15 | +#include <unordered_map> |
| 16 | + |
13 | 17 | ur_result_t cl_adapter::getDevicesFromContext(
|
14 | 18 | ur_context_handle_t hContext,
|
15 | 19 | std::unique_ptr<std::vector<cl_device_id>> &DevicesInCtx) {
|
@@ -89,10 +93,17 @@ urContextGetInfo(ur_context_handle_t hContext, ur_context_info_t propName,
|
89 | 93 | case UR_CONTEXT_INFO_NUM_DEVICES:
|
90 | 94 | case UR_CONTEXT_INFO_DEVICES:
|
91 | 95 | case UR_CONTEXT_INFO_REFERENCE_COUNT: {
|
92 |
| - |
93 |
| - CL_RETURN_ON_FAILURE( |
| 96 | + size_t CheckPropSize = 0; |
| 97 | + auto ClResult = |
94 | 98 | clGetContextInfo(cl_adapter::cast<cl_context>(hContext), CLPropName,
|
95 |
| - propSize, pPropValue, pPropSizeRet)); |
| 99 | + propSize, pPropValue, &CheckPropSize); |
| 100 | + if (pPropValue && CheckPropSize != propSize) { |
| 101 | + return UR_RESULT_ERROR_INVALID_SIZE; |
| 102 | + } |
| 103 | + CL_RETURN_ON_FAILURE(ClResult); |
| 104 | + if (pPropSizeRet) { |
| 105 | + *pPropSizeRet = CheckPropSize; |
| 106 | + } |
96 | 107 | return UR_RESULT_SUCCESS;
|
97 | 108 | }
|
98 | 109 | default:
|
@@ -130,8 +141,53 @@ UR_APIEXPORT ur_result_t UR_APICALL urContextCreateWithNativeHandle(
|
130 | 141 | }
|
131 | 142 |
|
132 | 143 | UR_APIEXPORT ur_result_t UR_APICALL urContextSetExtendedDeleter(
|
133 |
| - [[maybe_unused]] ur_context_handle_t hContext, |
134 |
| - [[maybe_unused]] ur_context_extended_deleter_t pfnDeleter, |
135 |
| - [[maybe_unused]] void *pUserData) { |
136 |
| - return UR_RESULT_ERROR_UNSUPPORTED_FEATURE; |
| 144 | + ur_context_handle_t hContext, ur_context_extended_deleter_t pfnDeleter, |
| 145 | + void *pUserData) { |
| 146 | + static std::unordered_map<ur_context_handle_t, |
| 147 | + std::set<ur_context_extended_deleter_t>> |
| 148 | + ContextCallbackMap; |
| 149 | + static std::mutex ContextCallbackMutex; |
| 150 | + |
| 151 | + { |
| 152 | + std::lock_guard<std::mutex> Lock(ContextCallbackMutex); |
| 153 | + // Callbacks can only be registered once and we need to avoid double |
| 154 | + // allocating. |
| 155 | + if (ContextCallbackMap.count(hContext) && |
| 156 | + ContextCallbackMap[hContext].count(pfnDeleter)) { |
| 157 | + return UR_RESULT_SUCCESS; |
| 158 | + } |
| 159 | + |
| 160 | + ContextCallbackMap[hContext].insert(pfnDeleter); |
| 161 | + } |
| 162 | + |
| 163 | + struct ContextCallback { |
| 164 | + void execute() { |
| 165 | + pfnDeleter(pUserData); |
| 166 | + { |
| 167 | + std::lock_guard<std::mutex> Lock(*CallbackMutex); |
| 168 | + (*CallbackMap)[hContext].erase(pfnDeleter); |
| 169 | + if ((*CallbackMap)[hContext].empty()) { |
| 170 | + CallbackMap->erase(hContext); |
| 171 | + } |
| 172 | + } |
| 173 | + delete this; |
| 174 | + } |
| 175 | + ur_context_handle_t hContext; |
| 176 | + ur_context_extended_deleter_t pfnDeleter; |
| 177 | + void *pUserData; |
| 178 | + std::unordered_map<ur_context_handle_t, |
| 179 | + std::set<ur_context_extended_deleter_t>> *CallbackMap; |
| 180 | + std::mutex *CallbackMutex; |
| 181 | + }; |
| 182 | + auto Callback = |
| 183 | + new ContextCallback({hContext, pfnDeleter, pUserData, &ContextCallbackMap, |
| 184 | + &ContextCallbackMutex}); |
| 185 | + auto ClCallback = [](cl_context, void *pUserData) { |
| 186 | + auto *C = static_cast<ContextCallback *>(pUserData); |
| 187 | + C->execute(); |
| 188 | + }; |
| 189 | + CL_RETURN_ON_FAILURE(clSetContextDestructorCallback( |
| 190 | + cl_adapter::cast<cl_context>(hContext), ClCallback, Callback)); |
| 191 | + |
| 192 | + return UR_RESULT_SUCCESS; |
137 | 193 | }
|
0 commit comments