switch device when releasing memory of inactive device

icicle_v3/src/device_api.cpp

@@ -133,10 +133,28 @@ namespace icicle {
 
   eIcicleError DeviceAPI::free_memory(void* ptr, MemoryTracker& tracker) const
   {
-    auto dev = tracker.identify_device(ptr);
-    if (dev == std::nullopt || **dev != get_thread_local_device()) { return eIcicleError::INVALID_DEVICE; }
-    auto err = free_memory(ptr);
-    if (err == eIcicleError::SUCCESS) { tracker.remove_allocation(ptr); }
+    auto ptr_dev = tracker.identify_device(ptr);
+
+    if (ptr_dev == std::nullopt) {
+      ICICLE_LOG_ERROR << "Trying to release host memory from device " << (**ptr_dev).type << " " << (**ptr_dev).id;
+      return eIcicleError::INVALID_DEVICE;
+    }
+    // If releasing memory of non-active device, switch device, release and switch back
+    // Alternatively I would have to consider it an error but it means that user has to switch device before dropping
+    // (when scope is closed)
+    const auto& cur_device = get_thread_local_device();
+    const bool is_active_device = **ptr_dev == cur_device;
+    if (is_active_device) {
+      auto err = free_memory(ptr);
+      if (err == eIcicleError::SUCCESS) { tracker.remove_allocation(ptr); }
+      return err;
+    }
+
+    // Getting here means memory does not belong to active device
+    // not ideal to call runtime.h api from here. TODO: how to better handle that case?
+    auto err = icicle_set_device(**ptr_dev);
+    if (err == eIcicleError::SUCCESS) err = icicle_free(ptr);
+    err = icicle_set_device(cur_device);
     return err;
   }