Keep IREE device groups alive per handle

Moves HAL device group ownership to Handle so IREE topology and frontier state stay valid for the lifetime of the device.

Co-Authored-By: GPT 5 <codex@openai.com>
Signed-off-by: Sambhav Jain <sambhav@alumni.stanford.edu>

Author: sjain-stanford

include/fusilli/backend/backend.h

@@ -377,6 +377,14 @@ struct IreeHalDeviceDeleter {
   }
 };
 
+// Custom deleter for IREE HAL device group.
+struct IreeHalDeviceGroupDeleter {
+  void operator()(iree_hal_device_group_t *deviceGroup) const {
+    if (deviceGroup)
+      iree_hal_device_group_release(deviceGroup);
+  }
+};
+
 // Custom deleter for IREE VM context.
 struct IreeVmContextDeleter {
   void operator()(iree_vm_context_t *context) const {
@@ -413,6 +421,8 @@ struct IreeHalBufferViewDeleter {
 using IreeVmInstanceSharedPtrType = std::shared_ptr<iree_vm_instance_t>;
 using IreeHalDeviceUniquePtrType =
     std::unique_ptr<iree_hal_device_t, IreeHalDeviceDeleter>;
+using IreeHalDeviceGroupUniquePtrType =
+    std::unique_ptr<iree_hal_device_group_t, IreeHalDeviceGroupDeleter>;
 using IreeVmContextUniquePtrType =
     std::unique_ptr<iree_vm_context_t, IreeVmContextDeleter>;
 using IreeAsyncFrontierTrackerUniquePtrType =

include/fusilli/backend/handle.h

@@ -60,6 +60,8 @@ class Handle {
                          "Handle::create got an unknown backend");
     }
 
+    FUSILLI_CHECK_ERROR(handle.createDeviceGroup());
+
     return ok(std::move(handle));
   }
 
@@ -104,6 +106,7 @@ class Handle {
     // Lazy create handle-specific IREE HAL device and populate the handle.
     FUSILLI_CHECK_ERROR(
         handle.createAMDGPUDevice(/*deviceId=*/deviceId, /*stream=*/stream));
+    FUSILLI_CHECK_ERROR(handle.createDeviceGroup());
 
     return ok(std::move(handle));
   }
@@ -121,8 +124,7 @@ class Handle {
 
   Backend getBackend() const { return backend_; }
 
-  // Allow Graph and Buffer to access private Handle methods `getDevice()`
-  // and `getInstance()`.
+  // Allow Graph and Buffer to access private Handle methods.
   friend class Graph;
   friend class Buffer;
 
@@ -139,6 +141,10 @@ class Handle {
   // Definition in `fusilli/backend/runtime.h`.
   ErrorObject createAMDGPUDevice(int deviceId, uintptr_t stream);
 
+  // Creates IREE HAL device group for this handle's device. Definition in
+  // `fusilli/backend/runtime.h`.
+  ErrorObject createDeviceGroup();
+
   // Private constructor (use factory `create` method for handle creation).
   Handle(Backend backend, IreeVmInstanceSharedPtrType instance)
       : backend_(backend), instance_(std::move(instance)) {}
@@ -149,6 +155,12 @@ class Handle {
   // valid as long as this handle exists.
   iree_hal_device_t *getDevice() const { return device_.get(); }
 
+  // Returns a raw pointer to the underlying IREE HAL device group.
+  // WARNING: The returned raw pointer is not safe to store since
+  // its lifetime is tied to the `Handle` object and only
+  // valid as long as this handle exists.
+  iree_hal_device_group_t *getDeviceGroup() const { return deviceGroup_.get(); }
+
   // Returns a raw pointer to the underlying IREE VM instance.
   // WARNING: The returned raw pointer is not safe to store since
   // its lifetime is tied to the `Handle` objects and only
@@ -160,6 +172,7 @@ class Handle {
   Backend backend_;
   IreeVmInstanceSharedPtrType instance_;
   IreeHalDeviceUniquePtrType device_;
+  IreeHalDeviceGroupUniquePtrType deviceGroup_;
 };
 
 } // namespace fusilli

include/fusilli/backend/runtime.h

@@ -159,6 +159,24 @@ inline ErrorObject Handle::createCPUDevice() {
   return ok();
 }
 
+inline ErrorObject Handle::createDeviceGroup() {
+  FUSILLI_LOG_LABEL_ENDL("INFO: Creating per-handle IREE HAL device group");
+  iree_allocator_t allocator = iree_allocator_system();
+
+  iree_async_frontier_tracker_t *rawFrontierTracker = nullptr;
+  FUSILLI_CHECK_ERROR(iree_async_frontier_tracker_create(
+      iree_async_frontier_tracker_options_default(), allocator,
+      &rawFrontierTracker));
+  IreeAsyncFrontierTrackerUniquePtrType frontierTracker(rawFrontierTracker);
+
+  iree_hal_device_group_t *rawDeviceGroup = nullptr;
+  FUSILLI_CHECK_ERROR(iree_hal_device_group_create_from_device(
+      getDevice(), frontierTracker.get(), allocator, &rawDeviceGroup));
+  deviceGroup_ = IreeHalDeviceGroupUniquePtrType(rawDeviceGroup);
+
+  return ok();
+}
+
 // Copied from the IREE runtime code.
 #define HIP_DEVICE_ID_TO_IREE_DEVICE_ID(device)                                \
   (iree_hal_device_id_t)((device) + 1)
@@ -239,21 +257,11 @@ inline ErrorObject Graph::createVmContext(const Handle &handle,
 
   // Create HAL module and register it with the context.
   {
-    iree_async_frontier_tracker_t *rawFrontierTracker = nullptr;
-    FUSILLI_CHECK_ERROR(iree_async_frontier_tracker_create(
-        iree_async_frontier_tracker_options_default(), allocator,
-        &rawFrontierTracker));
-    IreeAsyncFrontierTrackerUniquePtrType frontierTracker(rawFrontierTracker);
-
-    iree_hal_device_group_t *deviceGroup = nullptr;
-    FUSILLI_CHECK_ERROR(iree_hal_device_group_create_from_device(
-        handle.getDevice(), frontierTracker.get(), allocator, &deviceGroup));
     iree_vm_module_t *halModule = nullptr;
     FUSILLI_CHECK_ERROR(iree_hal_module_create(
         handle.getInstance(), iree_hal_module_device_policy_default(),
-        deviceGroup, IREE_HAL_MODULE_FLAG_NONE,
+        handle.getDeviceGroup(), IREE_HAL_MODULE_FLAG_NONE,
         iree_hal_module_debug_sink_null(), allocator, &halModule));
-    iree_hal_device_group_release(deviceGroup);
     iree_status_t status = iree_vm_context_register_modules(
         vmContext_.get(), /*module_count=*/1, &halModule);
     iree_vm_module_release(halModule);