[Backend] Enable default device sharing across handles for multi threaded handle+device creation (#114)

### Issue

The `Multi-threaded Handle creation` test was failing flakily on my
recent local runs:

```
/home/srajeshk/fusilli/tests/test_handle.cpp:83: FAILED:
  REQUIRE( !creationFailed.load() )
with expansion:
  false
```

The test spawns 32 threads that simultaneously attempt to create CPU
handles, and at least one handle creation was failing.

### Root Cause

IREE's HAL drivers (both `local-task` for CPU and HIP for GPU) provide a
**single default device** per configuration. When multiple threads
concurrently called `iree_runtime_instance_try_create_default_device()`,
each attempting to create its own device:

1. The first thread succeeded in creating the device
2. Subsequent threads failed because the default device already existed

IREE's default device model seems to support one "default" device per
driver, not one per handle.

### Fix

Implement **device sharing** across handles using a `weak_ptr` caching
pattern (similar to the existing `createSharedInstance()` pattern for
the runtime instance):

**For CPU (`createCPUDevice`):**
- Single shared device across all CPU handles
- Uses `static std::weak_ptr<iree_hal_device_t>` with mutex protection
- First handle creates the device; subsequent handles reuse it

**For AMDGPU (`createAMDGPUDevice`):**
- Devices cached by `(deviceId, stream)` configuration
- Uses `static std::map<key, std::weak_ptr<iree_hal_device_t>>` with
mutex protection
- Handles with the same configuration share the device; different
configurations get separate devices

### Changes

| File | Change |
|------|--------|
| `backend.h` | Added `IreeHalDeviceSharedPtrType`, removed unused
`IreeHalDeviceUniquePtrType` |
| `handle.h` | Changed `device_` from `unique_ptr` to `shared_ptr` |
| `runtime.h` | Implemented device sharing in `createCPUDevice()` and
`createAMDGPUDevice()` |
| `test_handle.cpp` | Added multi-threaded test for AMDGPU backend |

### Test Plan

- [x] `Multi-threaded Handle creation` test passes (CPU backend)
- [x] `Multi-threaded Handle creation AMDGPU` test passes (GPU backend)

### Disclaimer

- PR description and code fixes generated with assistance from Claude
4.5 Opus under my supervision.

---------

Signed-off-by: Sambhav Jain <[email protected]>

Author: sjain-stanford

include/fusilli/backend/backend.h

@@ -356,8 +356,7 @@ struct IreeHalBufferViewDeleter {
 // Aliases for IREE runtime types with custom deleters.
 using IreeRuntimeInstanceSharedPtrType =
     std::shared_ptr<iree_runtime_instance_t>;
-using IreeHalDeviceUniquePtrType =
-    std::unique_ptr<iree_hal_device_t, IreeHalDeviceDeleter>;
+using IreeHalDeviceSharedPtrType = std::shared_ptr<iree_hal_device_t>;
 using IreeRuntimeSessionUniquePtrType =
     std::unique_ptr<iree_runtime_session_t, IreeRuntimeSessionDeleter>;
 using IreeHalBufferViewUniquePtrType =

include/fusilli/backend/handle.h

@@ -160,7 +160,10 @@ class Handle {
   // `device_` depends on `backend_` and `instance_`.
   Backend backend_;
   IreeRuntimeInstanceSharedPtrType instance_;
-  IreeHalDeviceUniquePtrType device_;
+  // Use shared_ptr to allow sharing devices across handles. For CPU, there's
+  // only one logical device shared by all handles. For GPU, devices with the
+  // same (deviceId, stream) configuration are shared.
+  IreeHalDeviceSharedPtrType device_;
 };
 
 } // namespace fusilli

include/fusilli/backend/runtime.h

@@ -44,6 +44,7 @@
 #include <iree/runtime/api.h>
 
 #include <cstdint>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <string>
@@ -61,7 +62,7 @@ namespace fusilli {
 // Create static singleton IREE runtime instance shared across handles/threads.
 inline ErrorOr<IreeRuntimeInstanceSharedPtrType>
 Handle::createSharedInstance() {
-  // Mutex for thread-safe initialization of weakInstance.
+  // Mutex for thread-safe initialization of the shared instance.
   static std::mutex instanceMutex;
 
   // Static weak_ptr to the IREE runtime instance ensures that the
@@ -72,9 +73,7 @@ Handle::createSharedInstance() {
   // static variable goes out of scope upon program termination.
   static std::weak_ptr<iree_runtime_instance_t> weakInstance;
 
-  // If multiple threads simultaneously request a handle, they will
-  // race into `createSharedInstance()` but only one will succeed in
-  // creating the instance, and others will use it.
+  // Serialize access to the weak_ptr check-then-create logic.
   std::lock_guard<std::mutex> lock(instanceMutex);
 
   // Try to get the shared_ptr from the weak_ptr (if it exists).
@@ -106,15 +105,38 @@ Handle::createSharedInstance() {
 inline ErrorObject Handle::createCPUDevice() {
   FUSILLI_LOG_LABEL_ENDL("INFO: Creating per-handle IREE HAL device");
 
-  iree_hal_device_t *rawDevice = nullptr;
-  FUSILLI_CHECK_ERROR(iree_runtime_instance_try_create_default_device(
-      instance_.get(), iree_make_cstring_view(kHalDriver.at(backend_)),
-      &rawDevice));
+  // Mutex for thread-safe access to the shared CPU device.
+  static std::mutex cpuDeviceMutex;
 
-  // Wrap the raw device ptr with a unique_ptr and custom deleter
-  // for lifetime management.
-  device_ = IreeHalDeviceUniquePtrType(rawDevice);
+  // Static weak_ptr to the CPU device ensures that the device is only
+  // created once and shared across all CPU handles. This is necessary
+  // because IREE's local-task driver typically provides a single default
+  // device. The device is released when the last handle using it goes
+  // out of scope.
+  static std::weak_ptr<iree_hal_device_t> weakCpuDevice;
+
+  // Serialize access to the weak_ptr check-then-create logic.
+  std::lock_guard<std::mutex> lock(cpuDeviceMutex);
+
+  // Try to get the shared_ptr from the weak_ptr (if it exists).
+  IreeHalDeviceSharedPtrType sharedDevice = weakCpuDevice.lock();
+
+  // If weak_ptr expired, create a new CPU device.
+  if (sharedDevice == nullptr) {
+    FUSILLI_LOG_LABEL_ENDL("INFO: Creating shared CPU device");
+    iree_hal_device_t *rawDevice = nullptr;
+    FUSILLI_CHECK_ERROR(iree_runtime_instance_try_create_default_device(
+        instance_.get(), iree_make_cstring_view(kHalDriver.at(backend_)),
+        &rawDevice));
+
+    // Wrap the raw device ptr with a shared_ptr and custom deleter
+    // for lifetime management.
+    sharedDevice =
+        IreeHalDeviceSharedPtrType(rawDevice, IreeHalDeviceDeleter());
+    weakCpuDevice = sharedDevice;
+  }
 
+  device_ = sharedDevice;
   return ok();
 }
 
@@ -130,6 +152,36 @@ inline ErrorObject Handle::createAMDGPUDevice(int deviceId, uintptr_t stream) {
 // when building Fusilli without AMDGPU support (which disables IREE HAL
 // HIP driver from being built).
 #ifdef FUSILLI_ENABLE_AMDGPU
+  // Mutex for thread-safe access to the GPU device cache.
+  static std::mutex gpuDeviceMutex;
+
+  // Cache key for AMDGPU devices: (deviceId, stream) pair.
+  // Devices with the same configuration are shared across handles.
+  using GpuDeviceKey = std::pair<int, uintptr_t>;
+  static std::map<GpuDeviceKey, std::weak_ptr<iree_hal_device_t>>
+      gpuDeviceCache;
+
+  // Serialize access to the cache check-then-create logic.
+  std::lock_guard<std::mutex> lock(gpuDeviceMutex);
+
+  // Clean up all expired entries while we hold the lock.
+  std::erase_if(gpuDeviceCache, // C++20
+                [](const auto &entry) { return entry.second.expired(); });
+
+  GpuDeviceKey key{deviceId, stream};
+
+  // Try to get an existing device from the cache.
+  if (auto it = gpuDeviceCache.find(key); it != gpuDeviceCache.end()) {
+    // Entry exists and is valid (we just cleaned expired ones).
+    FUSILLI_LOG_LABEL_ENDL("INFO: Reusing cached AMDGPU device");
+    // Lock the weak_ptr to get a shared_ptr and assign to device_.
+    device_ = it->second.lock();
+    return ok();
+  }
+
+  // Create a new device since none exists in cache for this configuration.
+  FUSILLI_LOG_LABEL_ENDL("INFO: Creating new AMDGPU device");
+
   // Device parms.
   iree_hal_hip_device_params_t params;
   setDefaultIreeHalHipDeviceParams(&params);
@@ -149,9 +201,14 @@ inline ErrorObject Handle::createAMDGPUDevice(int deviceId, uintptr_t stream) {
       driver, HIP_DEVICE_ID_TO_IREE_DEVICE_ID(deviceId), /*param_count=*/0,
       /*params=*/nullptr, iree_allocator_system(), &rawDevice));
 
-  // Wrap the raw device ptr with a unique_ptr and custom deleter
+  // Wrap the raw device ptr with a shared_ptr and custom deleter
   // for lifetime management.
-  device_ = IreeHalDeviceUniquePtrType(rawDevice);
+  IreeHalDeviceSharedPtrType sharedDevice(rawDevice, IreeHalDeviceDeleter());
+
+  // Cache the device for future handles with the same configuration.
+  gpuDeviceCache[key] = sharedDevice;
+
+  device_ = sharedDevice;
   return ok();
 #else
   return ErrorObject(ErrorCode::InternalError,

tests/test_handle.cpp

@@ -84,6 +84,48 @@ TEST_CASE("Multi-threaded Handle creation", "[handle][thread]") {
   REQUIRE(handles.size() == kNumThreads);
 }
 
+#ifdef FUSILLI_ENABLE_AMDGPU
+TEST_CASE("Multi-threaded Handle creation AMDGPU", "[handle][thread][amdgpu]") {
+  constexpr int kNumThreads = 32;
+
+  std::vector<std::thread> threads;
+  threads.reserve(kNumThreads);
+
+  std::vector<Handle> handles;
+  handles.reserve(kNumThreads);
+
+  // Create a barrier to force threads to start simultaneously.
+  std::barrier startBarrier(kNumThreads);
+
+  // Atomic flag to track failures during handle creation.
+  std::atomic<bool> creationFailed{false};
+
+  // Mutex for pushing to handles in a thread-safe manner.
+  std::mutex handlesMutex;
+
+  for (size_t i = 0; i < kNumThreads; ++i) {
+    threads.emplace_back([&]() {
+      // Wait at the barrier until all threads reach this point.
+      startBarrier.arrive_and_wait();
+      // Create the handle for AMDGPU backend.
+      auto handleOrError = Handle::create(Backend::AMDGPU);
+      if (isError(handleOrError)) {
+        creationFailed.store(true);
+        return;
+      }
+      std::lock_guard<std::mutex> lock(handlesMutex);
+      handles.push_back(std::move(*handleOrError));
+    });
+  }
+  // Wait for all threads to finish.
+  for (auto &t : threads)
+    t.join();
+
+  REQUIRE(!creationFailed.load());
+  REQUIRE(handles.size() == kNumThreads);
+}
+#endif
+
 TEST_CASE("Handle creation with deviceId and stream, CPU backend should fail",
           "[handle]") {
   // Attempting to create CPU handle with with a specific deviceId and stream