Add iree_hal_device_group_t to own device topology lifecycle (iree-org#23576)

The topology matrix from iree-org#23573 needs an owner with a clear lifetime
contract. Today devices are passed to the HAL module as a flat array —
there is no object that retains the devices, owns the topology, and
guarantees the topology pointer remains valid for the duration of
execution. Every device holds a raw pointer to the topology it was
assigned, so if the topology is freed while devices are still alive,
those pointers dangle.

`iree_hal_device_group_t` is that owner. It takes already-created
devices, builds the immutable topology matrix from their capabilities,
pushes topology info into each device, and retains all of them. The
group's lifetime brackets the devices': whoever holds the devices
long-term (the HAL module, the CTS harness, a Python session) retains
the group, and the group retains the devices, so the topology pointer in
each device is guaranteed valid.

### Creation API

The builder pattern matches the topology builder it wraps:
stack-allocate, add devices, finalize. Finalize is a consuming operation
— it queries capabilities from all devices, computes edge descriptors,
calls driver-specific refinement for same-driver pairs, builds the
topology matrix, assigns topology info into each device via the new
vtable method, and produces the immutable group. The builder is zeroed
after finalize (whether it succeeds or fails) and cannot be reused.

For the common single-device case (7 of 9 callers),
`iree_hal_device_group_create_from_device` wraps the builder sequence
into a one-liner.

### `assign_topology_info` vtable method

Devices need to receive their topology info after the matrix is built —
the topology doesn't exist yet when the device is created, and the
device's index in the matrix isn't known until group creation. This is a
new vtable method on `iree_hal_device_t` that the group calls during
finalize. All existing driver implementations (local-sync, local-task,
CUDA, HIP, Vulkan, Metal, AMDGPU, null) store the info into their device
struct. The method is called exactly once per device, during group
creation.

### HAL module integration

`iree_hal_module_create` now takes a `iree_hal_device_group_t*` instead
of `(device_count, devices[])`. The module retains the group and
delegates all device access through `iree_hal_device_group_device_count`
/ `iree_hal_device_group_device_at`. This eliminates the flexible array
member from `iree_hal_module_t`, simplifies allocation (fixed-size
struct instead of variable-size), and makes the lifetime contract
explicit: the module holds the group, the group holds the devices and
topology.

All callers are updated — CLI tooling, the high-level runtime session,
TFLite bindings, Python bindings, PJRT, ConstEval, simple_embedding
samples, check_test, and the CTS test harness.

### Testing

A mock device (`hal/testing/mock_device`) provides controllable
capabilities for testing topology construction without requiring real
hardware. The device group tests exercise builder validation (empty
builds, duplicate devices, capacity limits), single-device and
multi-device group creation, topology correctness (self-edges,
cross-device edges with expected interop modes), the convenience
function, and lifetime ordering (group outlives devices). The CTS test
harness creates a device group in `SetUpTestSuite` so every CTS test
runs with topology info assigned.

### Where this is going

The device group is the scheduling domain for the causal execution
system. When the AMDGPU driver gets its frontier-integrated semaphores
and queue operations, the device group's topology matrix is what tells
the scheduler whether a semaphore can be waited on natively or needs
handle import, whether a buffer can be read directly or needs DMA
transfer, and what the relative cost is. The group also becomes the
natural attachment point for collective channel creation and
multi-device resource pools.

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: benvanik

compiler/src/iree/compiler/ConstEval/Runtime.cpp

@@ -126,8 +126,8 @@ CompiledBinary::CompiledBinary() = default;
 CompiledBinary::~CompiledBinary() = default;
 
 void CompiledBinary::deinitialize() {
-  hal_module.reset();
-  main_module.reset();
+  halModule.reset();
+  mainModule.reset();
   context.reset();
   device.reset();
 }
@@ -325,7 +325,7 @@ LogicalResult FunctionCall::invoke(Location loc, StringRef name) {
   // Lookup function.
   iree_vm_function_t function;
   if (auto status = iree_vm_module_lookup_function_by_name(
-          binary.main_module.get(), IREE_VM_FUNCTION_LINKAGE_EXPORT,
+          binary.mainModule.get(), IREE_VM_FUNCTION_LINKAGE_EXPORT,
           iree_string_view_t{name.data(),
                              static_cast<iree_host_size_t>(name.size())},
           &function)) {
@@ -462,29 +462,34 @@ LogicalResult CompiledBinary::initialize(Location loc, void *data,
   }
   iree_hal_driver_release(driver);
 
-  // Create hal module.
+  // Create device group and hal module.
+  iree_hal_device_group_t *deviceGroup = nullptr;
   if (iree_status_is_ok(status)) {
-    std::array<iree_hal_device_t *, 1> devices = {device.get()};
-    status = iree_hal_module_create(
-        runtime.instance.get(), iree_hal_module_device_policy_default(),
-        devices.size(), devices.data(), IREE_HAL_MODULE_FLAG_NONE,
-        iree_hal_module_debug_sink_stdio(stderr), iree_allocator_system(),
-        &hal_module);
+    status = iree_hal_device_group_create_from_device(
+        device.get(), iree_allocator_system(), &deviceGroup);
   }
+  if (iree_status_is_ok(status)) {
+    status = iree_hal_module_create(runtime.instance.get(),
+                                    iree_hal_module_device_policy_default(),
+                                    deviceGroup, IREE_HAL_MODULE_FLAG_NONE,
+                                    iree_hal_module_debug_sink_stdio(stderr),
+                                    iree_allocator_system(), &halModule);
+  }
+  iree_hal_device_group_release(deviceGroup);
 
   // Bytecode module.
   if (iree_status_is_ok(status)) {
     status = iree_vm_bytecode_module_create(
         runtime.instance.get(), IREE_VM_BYTECODE_MODULE_FLAG_NONE,
         iree_make_const_byte_span(data, length), iree_allocator_null(),
-        iree_allocator_system(), &main_module);
+        iree_allocator_system(), &mainModule);
   }
 
   // Create context.
   if (iree_status_is_ok(status)) {
     std::array<iree_vm_module_t *, 2> modules = {
-        hal_module.get(),
-        main_module.get(),
+        halModule.get(),
+        mainModule.get(),
     };
     status = iree_vm_context_create_with_modules(
         runtime.instance.get(), IREE_VM_CONTEXT_FLAG_NONE, modules.size(),

compiler/src/iree/compiler/ConstEval/Runtime.h

@@ -43,8 +43,8 @@ class CompiledBinary {
                                       Type mlirType);
 
   iree::vm::ref<iree_hal_device_t> device;
-  iree::vm::ref<iree_vm_module_t> hal_module;
-  iree::vm::ref<iree_vm_module_t> main_module;
+  iree::vm::ref<iree_vm_module_t> halModule;
+  iree::vm::ref<iree_vm_module_t> mainModule;
   iree::vm::ref<iree_vm_context_t> context;
 
   friend class FunctionCall;

integrations/pjrt/src/iree_pjrt/common/api_impl.cc

@@ -1680,12 +1680,16 @@ iree_status_t ClientInstance::PopulateVMModules(
     iree_hal_device_t* hal_device,
     iree::vm::ref<iree_vm_module_t>& main_module) {
   // HAL module.
+  iree_hal_device_group_t* device_group = nullptr;
+  IREE_RETURN_IF_ERROR(iree_hal_device_group_create_from_device(
+      hal_device, host_allocator(), &device_group));
   modules.push_back({});
-  IREE_RETURN_IF_ERROR(iree_hal_module_create(
-      vm_instance(), iree_hal_module_device_policy_default(),
-      /*device_count=*/1, &hal_device, IREE_HAL_MODULE_FLAG_NONE,
-      iree_hal_module_debug_sink_stdio(stderr), host_allocator(),
-      &modules.back()));
+  iree_status_t status = iree_hal_module_create(
+      vm_instance(), iree_hal_module_device_policy_default(), device_group,
+      IREE_HAL_MODULE_FLAG_NONE, iree_hal_module_debug_sink_stdio(stderr),
+      host_allocator(), &modules.back());
+  iree_hal_device_group_release(device_group);
+  IREE_RETURN_IF_ERROR(status);
 
   // Main module.
   modules.push_back(main_module);

runtime/bindings/python/hal.cc

@@ -1107,38 +1107,40 @@ VmModule CreateHalModule(
         PyExc_ValueError,
         "\"device\" and \"devices\" are mutually exclusive arguments.");
   }
-  std::vector<iree_hal_device_t*> devices_vector;
-  iree_hal_device_t* device_ptr;
-  iree_hal_device_t** devices_ptr;
-  iree_host_size_t device_count;
-  iree_vm_module_t* module = NULL;
+  // Build a device group from the provided device(s).
+  iree_hal_device_group_builder_t group_builder;
+  iree_hal_device_group_builder_initialize(&group_builder);
   if (device) {
-    device_ptr = device.value()->raw_ptr();
-    devices_ptr = &device_ptr;
-    device_count = 1;
+    CheckApiStatus(iree_hal_device_group_builder_add_device(
+                       &group_builder, device.value()->raw_ptr()),
+                   "Error adding device to group builder");
   } else {
-    // Set device related arguments in the case of multiple devices.
-    devices_vector.reserve(devices->size());
     for (auto devicesIt = devices->begin(); devicesIt != devices->end();
          ++devicesIt) {
-      devices_vector.push_back(py::cast<HalDevice*>(*devicesIt)->raw_ptr());
+      CheckApiStatus(
+          iree_hal_device_group_builder_add_device(
+              &group_builder, py::cast<HalDevice*>(*devicesIt)->raw_ptr()),
+          "Error adding device to group builder");
     }
-    devices_ptr = devices_vector.data();
-    device_count = devices_vector.size();
   }
+  iree_hal_device_group_t* device_group = nullptr;
+  CheckApiStatus(iree_hal_device_group_builder_finalize(
+                     &group_builder, iree_allocator_system(), &device_group),
+                 "Error creating device group");
 
   iree_hal_module_debug_sink_t iree_hal_module_debug_sink =
       iree_hal_module_debug_sink_stdio(stderr);
   if (debug_sink) {
     iree_hal_module_debug_sink = (*debug_sink)->AsIreeHalModuleDebugSink();
   }
 
-  CheckApiStatus(
-      iree_hal_module_create(
-          instance->raw_ptr(), iree_hal_module_device_policy_default(),
-          device_count, devices_ptr, IREE_HAL_MODULE_FLAG_NONE,
-          iree_hal_module_debug_sink, iree_allocator_system(), &module),
-      "Error creating hal module");
+  iree_vm_module_t* module = NULL;
+  iree_status_t status = iree_hal_module_create(
+      instance->raw_ptr(), iree_hal_module_device_policy_default(),
+      device_group, IREE_HAL_MODULE_FLAG_NONE, iree_hal_module_debug_sink,
+      iree_allocator_system(), &module);
+  iree_hal_device_group_release(device_group);
+  CheckApiStatus(status, "Error creating hal module");
   VmModule vm_module = VmModule::StealFromRawPtr(module);
   if (debug_sink) {
     // Retain a reference. We want the callback to be valid after

runtime/bindings/tflite/interpreter.c

@@ -60,11 +60,16 @@ static iree_status_t _TfLiteInterpreterPrepareHAL(
       "failed creating the default device for driver '%.*s'",
       (int)driver_name.size, driver_name.data);
 
-  IREE_RETURN_IF_ERROR(iree_hal_module_create(
+  iree_hal_device_group_t* device_group = NULL;
+  IREE_RETURN_IF_ERROR(iree_hal_device_group_create_from_device(
+      interpreter->device, interpreter->allocator, &device_group));
+  status = iree_hal_module_create(
       interpreter->instance, iree_hal_module_device_policy_default(),
-      /*device_count=*/1, &interpreter->device, IREE_HAL_MODULE_FLAG_NONE,
+      device_group, IREE_HAL_MODULE_FLAG_NONE,
       iree_hal_module_debug_sink_stdio(stderr), interpreter->allocator,
-      &interpreter->hal_module));
+      &interpreter->hal_module);
+  iree_hal_device_group_release(device_group);
+  IREE_RETURN_IF_ERROR(status);
 
   return iree_ok_status();
 }

runtime/src/iree/hal/BUILD.bazel

@@ -47,6 +47,8 @@ iree_runtime_cc_library(
         "detail.h",
         "device.c",
         "device.h",
+        "device_group.c",
+        "device_group.h",
         "driver.c",
         "driver.h",
         "driver_registry.c",
@@ -118,3 +120,15 @@ iree_runtime_cc_test(
         "//runtime/src/iree/testing:gtest_main",
     ],
 )
+
+iree_runtime_cc_test(
+    name = "device_group_test",
+    srcs = ["device_group_test.cc"],
+    deps = [
+        ":hal",
+        "//runtime/src/iree/base",
+        "//runtime/src/iree/hal/testing:mock_device",
+        "//runtime/src/iree/testing:gtest",
+        "//runtime/src/iree/testing:gtest_main",
+    ],
+)

runtime/src/iree/hal/CMakeLists.txt

@@ -40,6 +40,8 @@ iree_cc_library(
     "detail.h"
     "device.c"
     "device.h"
+    "device_group.c"
+    "device_group.h"
     "driver.c"
     "driver.h"
     "driver_registry.c"
@@ -111,4 +113,17 @@ iree_cc_test(
     iree::testing::gtest_main
 )
 
+iree_cc_test(
+  NAME
+    device_group_test
+  SRCS
+    "device_group_test.cc"
+  DEPS
+    ::hal
+    iree::base
+    iree::hal::testing::mock_device
+    iree::testing::gtest
+    iree::testing::gtest_main
+)
+
 ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###

runtime/src/iree/hal/api.h

@@ -18,6 +18,7 @@
 #include "iree/hal/channel_provider.h"  // IWYU pragma: export
 #include "iree/hal/command_buffer.h"    // IWYU pragma: export
 #include "iree/hal/device.h"            // IWYU pragma: export
+#include "iree/hal/device_group.h"      // IWYU pragma: export
 #include "iree/hal/driver.h"            // IWYU pragma: export
 #include "iree/hal/driver_registry.h"   // IWYU pragma: export
 #include "iree/hal/event.h"             // IWYU pragma: export

runtime/src/iree/hal/cts/cts_test_base.h

@@ -89,10 +89,12 @@ static iree_status_t TryGetDriver(const std::string& driver_name,
 class CTSTestResources {
  public:
   static iree_hal_driver_t* driver_;
+  static iree_hal_device_group_t* device_group_;
   static iree_hal_device_t* device_;
   static iree_hal_allocator_t* device_allocator_;
 };
 /*static*/ iree_hal_driver_t* CTSTestResources::driver_ = NULL;
+/*static*/ iree_hal_device_group_t* CTSTestResources::device_group_ = NULL;
 /*static*/ iree_hal_device_t* CTSTestResources::device_ = NULL;
 /*static*/ iree_hal_allocator_t* CTSTestResources::device_allocator_ = NULL;
 
@@ -129,6 +131,13 @@ class CTSTestBase : public BaseType, public CTSTestResources {
     IREE_CHECK_OK(status);
     device_ = device;
 
+    // Create a device group so the device gets topology info assigned.
+    // The group must outlive the device (it holds a raw topology pointer).
+    iree_hal_device_group_t* device_group = NULL;
+    IREE_CHECK_OK(iree_hal_device_group_create_from_device(
+        device_, iree_allocator_system(), &device_group));
+    device_group_ = device_group;
+
     device_allocator_ = iree_hal_device_allocator(device_);
     iree_hal_allocator_retain(device_allocator_);
   }
@@ -142,6 +151,12 @@ class CTSTestBase : public BaseType, public CTSTestResources {
       iree_hal_device_release(device_);
       device_ = NULL;
     }
+    // Release the device group after the device — the device holds a raw
+    // pointer to the group's embedded topology.
+    if (device_group_) {
+      iree_hal_device_group_release(device_group_);
+      device_group_ = NULL;
+    }
     if (driver_) {
       iree_hal_driver_release(driver_);
       driver_ = NULL;

runtime/src/iree/hal/device.c

@@ -96,6 +96,14 @@ IREE_API_EXPORT iree_status_t iree_hal_device_refine_topology_edge(
                                                             dst_device, edge);
 }
 
+IREE_API_EXPORT iree_status_t iree_hal_device_assign_topology_info(
+    iree_hal_device_t* device,
+    const iree_hal_device_topology_info_t* topology_info) {
+  IREE_ASSERT_ARGUMENT(device);
+  IREE_ASSERT_ARGUMENT(topology_info);
+  return _VTABLE_DISPATCH(device, assign_topology_info)(device, topology_info);
+}
+
 IREE_API_EXPORT iree_hal_semaphore_compatibility_t
 iree_hal_device_query_semaphore_compatibility(iree_hal_device_t* device,
                                               iree_hal_semaphore_t* semaphore) {