[llvm] [Offload] Update allocations to include device (PR #154733)

[Piotr Balcer via llvm-commits]

================
@@ -535,30 +536,37 @@ Error olMemAlloc_impl(ol_device_handle_t Device, ol_alloc_type_t Type,
 
   *AllocationOut = *Alloc;
   {
-    std::lock_guard<std::mutex> Lock(OffloadContext::get().AllocInfoMapMutex);
-    OffloadContext::get().AllocInfoMap.insert_or_assign(
-        *Alloc, AllocInfo{Device, Type});
+    std::lock_guard<std::mutex> Lock(OffloadContext::get().AllocInfoListMutex);
+    OffloadContext::get().AllocInfoList.emplace_back(
+        AllocInfo{*AllocationOut, Device, Type});
   }
   return Error::success();
 }
 
-Error olMemFree_impl(void *Address) {
-  ol_device_handle_t Device;
-  ol_alloc_type_t Type;
+Error olMemFree_impl(ol_device_handle_t Device, void *Address) {
+  AllocInfo Removed;
   {
-    std::lock_guard<std::mutex> Lock(OffloadContext::get().AllocInfoMapMutex);
-    if (!OffloadContext::get().AllocInfoMap.contains(Address))
-      return createOffloadError(ErrorCode::INVALID_ARGUMENT,
-                                "address is not a known allocation");
+    std::lock_guard<std::mutex> Lock(OffloadContext::get().AllocInfoListMutex);
+
+    auto &List = OffloadContext::get().AllocInfoList;
+    auto Entry = std::find_if(List.begin(), List.end(), [&](AllocInfo &Entry) {
----------------
pbalcer wrote:

It makes liboffload's `free` (and, really, based on your other comment, any function that accepts an USM pointer) performance to be `O(N)` where N is the number of allocations. Unless we can reasonably say that N is going to be very small, and I'm not sure we can, because e.g., SYCL doesn't restrict users from arbitrarily allocating memory, there will be a noticeable performance difference at least for some applications.

https://github.com/llvm/llvm-project/pull/154733