Skip to content

Commit fc3b83a

Browse files
committed
[UR][L0] Fix pool iteration during freelist cleanup
The forEachPool method stops iterating if the callback function returns false. As a result, during EnqueuedPools cleanup, pools later in the chain could be skipped if ealier pools didn't require cleanup. Additionally, ensure that user-created pools associated with the given queue handle are properly cleaned up during urQueueFinish and urQueueRelease calls.
1 parent 51e8d95 commit fc3b83a

File tree

3 files changed

+17
-3
lines changed

3 files changed

+17
-3
lines changed

unified-runtime/source/adapters/level_zero/queue.cpp

+7
Original file line numberDiff line numberDiff line change
@@ -630,7 +630,11 @@ ur_result_t urQueueRelease(
630630
UR_CALL(Queue->synchronize());
631631

632632
// Cleanup the allocations from 'AsyncPool' made by this queue.
633+
633634
Queue->Context->AsyncPool.cleanupPoolsForQueue(Queue);
635+
for (auto &Pool : Queue->Context->UsmPoolHandles) {
636+
Pool->cleanupPoolsForQueue(Queue);
637+
}
634638

635639
// Destroy all the fences created associated with this queue.
636640
for (auto it = Queue->CommandListMap.begin();
@@ -909,6 +913,9 @@ ur_result_t urQueueFinish(
909913
}
910914

911915
Queue->Context->AsyncPool.cleanupPoolsForQueue(Queue);
916+
for (auto &Pool : Queue->Context->UsmPoolHandles) {
917+
Pool->cleanupPoolsForQueue(Queue);
918+
}
912919

913920
return UR_RESULT_SUCCESS;
914921
}

unified-runtime/source/adapters/level_zero/usm.cpp

+8-3
Original file line numberDiff line numberDiff line change
@@ -1207,12 +1207,17 @@ ur_usm_pool_handle_t_::getPoolByHandle(const umf_memory_pool_handle_t UmfPool) {
12071207
}
12081208

12091209
void ur_usm_pool_handle_t_::cleanupPools() {
1210-
PoolManager.forEachPool([&](UsmPool *p) { return p->AsyncPool.cleanup(); });
1210+
PoolManager.forEachPool([&](UsmPool *p) {
1211+
p->AsyncPool.cleanup();
1212+
return true;
1213+
});
12111214
}
12121215

12131216
void ur_usm_pool_handle_t_::cleanupPoolsForQueue(ur_queue_handle_t Queue) {
1214-
PoolManager.forEachPool(
1215-
[&](UsmPool *p) { return p->AsyncPool.cleanupForQueue(Queue); });
1217+
PoolManager.forEachPool([&](UsmPool *p) {
1218+
p->AsyncPool.cleanupForQueue(Queue);
1219+
return true;
1220+
});
12161221
}
12171222

12181223
bool ur_usm_pool_handle_t_::hasPool(const umf_memory_pool_handle_t Pool) {

unified-runtime/source/adapters/level_zero/usm.hpp

+2
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ usm::DisjointPoolAllConfigs InitializeDisjointPoolConfig();
2222
struct UsmPool {
2323
UsmPool(umf::pool_unique_handle_t Pool) : UmfPool(std::move(Pool)) {}
2424
umf::pool_unique_handle_t UmfPool;
25+
// 'AsyncPool' needs to be declared after 'UmfPool' so its destructor is
26+
// invoked first.
2527
EnqueuedPool AsyncPool;
2628
};
2729

0 commit comments

Comments
 (0)