Fix pool iteration during freelist cleanup (#18406)

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.

Author: kswiecicki

source/adapters/level_zero/queue.cpp

@@ -630,7 +630,11 @@ ur_result_t urQueueRelease(
       UR_CALL(Queue->synchronize());
 
     // Cleanup the allocations from 'AsyncPool' made by this queue.
+
     Queue->Context->AsyncPool.cleanupPoolsForQueue(Queue);
+    for (auto &Pool : Queue->Context->UsmPoolHandles) {
+      Pool->cleanupPoolsForQueue(Queue);
+    }
 
     // Destroy all the fences created associated with this queue.
     for (auto it = Queue->CommandListMap.begin();
@@ -909,6 +913,9 @@ ur_result_t urQueueFinish(
   }
 
   Queue->Context->AsyncPool.cleanupPoolsForQueue(Queue);
+  for (auto &Pool : Queue->Context->UsmPoolHandles) {
+    Pool->cleanupPoolsForQueue(Queue);
+  }
 
   return UR_RESULT_SUCCESS;
 }

source/adapters/level_zero/usm.cpp

@@ -1207,12 +1207,17 @@ ur_usm_pool_handle_t_::getPoolByHandle(const umf_memory_pool_handle_t UmfPool) {
 }
 
 void ur_usm_pool_handle_t_::cleanupPools() {
-  PoolManager.forEachPool([&](UsmPool *p) { return p->AsyncPool.cleanup(); });
+  PoolManager.forEachPool([&](UsmPool *p) {
+    p->AsyncPool.cleanup();
+    return true;
+  });
 }
 
 void ur_usm_pool_handle_t_::cleanupPoolsForQueue(ur_queue_handle_t Queue) {
-  PoolManager.forEachPool(
-      [&](UsmPool *p) { return p->AsyncPool.cleanupForQueue(Queue); });
+  PoolManager.forEachPool([&](UsmPool *p) {
+    p->AsyncPool.cleanupForQueue(Queue);
+    return true;
+  });
 }
 
 bool ur_usm_pool_handle_t_::hasPool(const umf_memory_pool_handle_t Pool) {

source/adapters/level_zero/usm.hpp

@@ -22,6 +22,8 @@ usm::DisjointPoolAllConfigs InitializeDisjointPoolConfig();
 struct UsmPool {
   UsmPool(umf::pool_unique_handle_t Pool) : UmfPool(std::move(Pool)) {}
   umf::pool_unique_handle_t UmfPool;
+  // 'AsyncPool' needs to be declared after 'UmfPool' so its destructor is
+  // invoked first.
   EnqueuedPool AsyncPool;
 };