runtime/src/iree/hal/device_group.c - 3p/openxla/iree - Git at Google

 // Copyright 2026 The IREE Authors
 //
 // Licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

 #include "iree/hal/device_group.h"

 #include <string.h>

 #include "iree/async/frontier_tracker.h"
 #include "iree/base/internal/atomics.h"
 #include "iree/hal/topology_builder.h"

 //===----------------------------------------------------------------------===//
 // iree_hal_device_group_t
 //===----------------------------------------------------------------------===//

 struct iree_hal_device_group_t {
   // Retain/release counter for the group object.
   iree_atomic_ref_count_t ref_count;

   // Host allocator used to free the group object.
   iree_allocator_t host_allocator;

   // Number of devices in the group.
   iree_host_size_t device_count;

   // Retained frontier tracker shared by all devices in the group.
   iree_async_frontier_tracker_t* frontier_tracker;

   // Retained devices. Order defines topology indices (device i = devices[i]).
   iree_hal_device_t* devices[IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT];

   // Cached driver name for each device (backed by the device's own storage,
   // valid for the device's lifetime).
   iree_string_view_t driver_names[IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT];

   // Immutable topology matrix built during creation.
   // Embedded (not heap-allocated) so devices can hold a stable pointer to it
   // for the group's lifetime.
   iree_hal_topology_t topology;
 };

 IREE_API_EXPORT iree_status_t iree_hal_device_group_create_from_device(
     iree_hal_device_t* device, iree_async_frontier_tracker_t* frontier_tracker,
     iree_allocator_t host_allocator, iree_hal_device_group_t** out_group) {
   IREE_ASSERT_ARGUMENT(device);
   IREE_ASSERT_ARGUMENT(frontier_tracker);
   IREE_ASSERT_ARGUMENT(out_group);
   *out_group = NULL;
   iree_hal_device_group_builder_t builder;
   iree_hal_device_group_builder_initialize(&builder, frontier_tracker);
   iree_status_t status =
       iree_hal_device_group_builder_add_device(&builder, device);
   if (iree_status_is_ok(status)) {
     status = iree_hal_device_group_builder_finalize(&builder, host_allocator,
                                                     out_group);
   } else {
     iree_hal_device_group_builder_deinitialize(&builder);
   }
   return status;
 }

 static void iree_hal_device_group_destroy(iree_hal_device_group_t* group) {
   IREE_ASSERT_ARGUMENT(group);
   IREE_TRACE_ZONE_BEGIN(z0);

   for (iree_host_size_t i = 0; i < group->device_count; ++i) {
     iree_hal_device_release(group->devices[i]);
   }
   iree_async_frontier_tracker_release(group->frontier_tracker);

   iree_allocator_t host_allocator = group->host_allocator;
   iree_allocator_free(host_allocator, group);

   IREE_TRACE_ZONE_END(z0);
 }

 IREE_API_EXPORT void iree_hal_device_group_retain(
     iree_hal_device_group_t* group) {
   if (group) {
     iree_atomic_ref_count_inc(&group->ref_count);
   }
 }

 IREE_API_EXPORT void iree_hal_device_group_release(
     iree_hal_device_group_t* group) {
   if (group && iree_atomic_ref_count_dec(&group->ref_count) == 1) {
     iree_hal_device_group_destroy(group);
   }
 }

 IREE_API_EXPORT iree_host_size_t
 iree_hal_device_group_device_count(const iree_hal_device_group_t* group) {
   IREE_ASSERT_ARGUMENT(group);
   return group->device_count;
 }

 IREE_API_EXPORT iree_hal_device_t* iree_hal_device_group_device_at(
     const iree_hal_device_group_t* group, iree_host_size_t index) {
   IREE_ASSERT_ARGUMENT(group);
   if (index >= group->device_count) return NULL;
   return group->devices[index];
 }

 IREE_API_EXPORT const iree_hal_topology_t* iree_hal_device_group_topology(
     const iree_hal_device_group_t* group) {
   IREE_ASSERT_ARGUMENT(group);
   return &group->topology;
 }

 //===----------------------------------------------------------------------===//
 // iree_hal_device_group_builder_t
 //===----------------------------------------------------------------------===//

 IREE_API_EXPORT void iree_hal_device_group_builder_initialize(
     iree_hal_device_group_builder_t* builder,
     iree_async_frontier_tracker_t* frontier_tracker) {
   IREE_ASSERT_ARGUMENT(builder);
   IREE_ASSERT_ARGUMENT(frontier_tracker);
   memset(builder, 0, sizeof(*builder));
   builder->frontier_tracker = frontier_tracker;
   iree_async_frontier_tracker_retain(frontier_tracker);
 }

 IREE_API_EXPORT void iree_hal_device_group_builder_deinitialize(
     iree_hal_device_group_builder_t* builder) {
   IREE_ASSERT_ARGUMENT(builder);
   iree_async_frontier_tracker_release(builder->frontier_tracker);
   memset(builder, 0, sizeof(*builder));
 }

 IREE_API_EXPORT iree_status_t iree_hal_device_group_builder_add_device(
     iree_hal_device_group_builder_t* builder, iree_hal_device_t* device) {
   IREE_ASSERT_ARGUMENT(builder);
   IREE_ASSERT_ARGUMENT(device);
   if (builder->count >= IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT) {
     return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED,
                             "device group builder at max capacity (%d)",
                             IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT);
   }
   builder->devices[builder->count++] = device;
   return iree_ok_status();
 }

 // Computes pre-populated bitmaps for a device's topology info by scanning
 // its edges in the topology matrix.
 static void iree_hal_device_group_compute_bitmaps(
     const iree_hal_topology_t* topology, uint32_t device_index,
     iree_hal_device_topology_info_t* out_info) {
   uint32_t device_count = topology->device_count;
   out_info->can_wait_from = 0;
   out_info->can_signal_to = 0;
   out_info->can_import_from = 0;
   out_info->can_p2p_with = 0;

   for (uint32_t j = 0; j < device_count; ++j) {
     if (j == device_index) continue;

     // can_wait_from: can device_index wait on device j's semaphores?
     // Edge[j][device_index] describes how device_index interacts with j's
     // resources.
     iree_hal_topology_edge_t edge_from_j =
         iree_hal_topology_query_edge(topology, j, device_index);
     if (iree_hal_topology_edge_wait_mode(edge_from_j.lo) !=
         IREE_HAL_TOPOLOGY_INTEROP_MODE_NONE) {
       out_info->can_wait_from |= (iree_hal_topology_device_bitmap_t)1 << j;
     }

     // can_import_from: can device_index import buffers from device j?
     // Set if either non-coherent or coherent buffers can be shared without
     // a host-staged copy.
     iree_hal_topology_interop_mode_t nc_read =
         iree_hal_topology_edge_buffer_read_mode_noncoherent(edge_from_j.lo);
     iree_hal_topology_interop_mode_t c_read =
         iree_hal_topology_edge_buffer_read_mode_coherent(edge_from_j.lo);
     if (nc_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_NATIVE ||
         nc_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_IMPORT ||
         c_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_NATIVE ||
         c_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_IMPORT) {
       out_info->can_import_from |= (iree_hal_topology_device_bitmap_t)1 << j;
     }

     // can_p2p_with: P2P access between device_index and device j?
     if (nc_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_NATIVE ||
         c_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_NATIVE ||
         (iree_hal_topology_edge_capability_flags(edge_from_j.lo) &
          IREE_HAL_TOPOLOGY_CAPABILITY_P2P_COPY)) {
       out_info->can_p2p_with |= (iree_hal_topology_device_bitmap_t)1 << j;
     }

     // can_signal_to: can device_index signal to device j?
     // Edge[device_index][j] describes how j interacts with device_index's
     // resources — the signal_mode tells us if j can observe our signals.
     iree_hal_topology_edge_t edge_to_j =
         iree_hal_topology_query_edge(topology, device_index, j);
     if (iree_hal_topology_edge_signal_mode(edge_to_j.lo) !=
         IREE_HAL_TOPOLOGY_INTEROP_MODE_NONE) {
       out_info->can_signal_to |= (iree_hal_topology_device_bitmap_t)1 << j;
     }
   }
 }

 IREE_API_EXPORT iree_status_t iree_hal_device_group_create_with_replacements(
     iree_hal_device_group_t* source_group,
     iree_hal_device_group_replacement_callback_t replacement_callback,
     iree_allocator_t host_allocator, iree_hal_device_group_t** out_group) {
   IREE_ASSERT_ARGUMENT(source_group);
   IREE_ASSERT_ARGUMENT(out_group);
   *out_group = NULL;
   IREE_TRACE_ZONE_BEGIN(z0);

   if (IREE_UNLIKELY(!replacement_callback.fn)) {
     IREE_TRACE_ZONE_END(z0);
     return iree_make_status(
         IREE_STATUS_INVALID_ARGUMENT,
         "device group replacement callback function is required");
   }

   const iree_host_size_t device_count = source_group->device_count;
   iree_hal_device_group_t* group = NULL;
   iree_status_t status =
       iree_allocator_malloc(host_allocator, sizeof(*group), (void**)&group);
   if (!iree_status_is_ok(status)) {
     IREE_TRACE_ZONE_END(z0);
     return status;
   }
   memset(group, 0, sizeof(*group));
   iree_atomic_ref_count_init(&group->ref_count);
   group->host_allocator = host_allocator;
   group->device_count = device_count;
   group->frontier_tracker = source_group->frontier_tracker;
   iree_async_frontier_tracker_retain(group->frontier_tracker);
   group->topology = source_group->topology;

   for (iree_host_size_t i = 0; i < device_count && iree_status_is_ok(status);
        ++i) {
     iree_hal_device_t* replacement_device = NULL;
     status = replacement_callback.fn(source_group, i, source_group->devices[i],
                                      replacement_callback.user_data,
                                      &replacement_device);
     if (iree_status_is_ok(status) && IREE_UNLIKELY(!replacement_device)) {
       status = iree_make_status(
           IREE_STATUS_INVALID_ARGUMENT,
           "replacement device %" PRIhsz " must not be NULL", i);
     }
     if (iree_status_is_ok(status)) {
       group->devices[i] = replacement_device;
     } else {
       iree_hal_device_release(replacement_device);
     }
     group->driver_names[i] = source_group->driver_names[i];
   }

   iree_host_size_t assigned_device_count = 0;
   for (iree_host_size_t i = 0; i < device_count && iree_status_is_ok(status);
        ++i) {
     iree_hal_device_topology_info_t topology_info;
     memset(&topology_info, 0, sizeof(topology_info));
     topology_info.topology_index = (uint32_t)i;
     topology_info.topology = &group->topology;
     topology_info.self_edge =
         iree_hal_topology_query_edge(&group->topology, (uint32_t)i, (uint32_t)i)
             .lo;
     topology_info.frontier.tracker = group->frontier_tracker;
     topology_info.frontier.base_axis = iree_async_axis_make_queue(
         iree_async_frontier_tracker_session_epoch(group->frontier_tracker),
         iree_async_frontier_tracker_machine_index(group->frontier_tracker),
         (uint8_t)i, /*queue_index=*/0);

     iree_hal_device_group_compute_bitmaps(&group->topology, (uint32_t)i,
                                           &topology_info);

     status =
         iree_hal_device_assign_topology_info(group->devices[i], &topology_info);
     if (iree_status_is_ok(status)) {
       assigned_device_count = i + 1;
     }
   }

   if (iree_status_is_ok(status)) {
     *out_group = group;
   } else {
     for (iree_host_size_t i = 0; i < assigned_device_count; ++i) {
       status = iree_status_join(status, iree_hal_device_assign_topology_info(
                                             group->devices[i], NULL));
     }
     iree_hal_device_group_release(group);
   }

   IREE_TRACE_ZONE_END(z0);
   return status;
 }

 IREE_API_EXPORT iree_status_t iree_hal_device_group_builder_finalize(
     iree_hal_device_group_builder_t* builder, iree_allocator_t host_allocator,
     iree_hal_device_group_t** out_group) {
   IREE_ASSERT_ARGUMENT(builder);
   IREE_ASSERT_ARGUMENT(out_group);
   *out_group = NULL;
   IREE_TRACE_ZONE_BEGIN(z0);

   iree_host_size_t device_count = builder->count;
   if (device_count == 0) {
     iree_hal_device_group_builder_deinitialize(builder);
     IREE_TRACE_ZONE_END(z0);
     return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                             "device group builder has no devices");
   }

   // Allocate the group.
   iree_hal_device_group_t* group = NULL;
   iree_status_t status =
       iree_allocator_malloc(host_allocator, sizeof(*group), (void**)&group);
   if (!iree_status_is_ok(status)) {
     iree_hal_device_group_builder_deinitialize(builder);
     IREE_TRACE_ZONE_END(z0);
     return status;
   }
   memset(group, 0, sizeof(*group));
   iree_atomic_ref_count_init(&group->ref_count);
   group->host_allocator = host_allocator;
   group->device_count = device_count;
   group->frontier_tracker = builder->frontier_tracker;
   builder->frontier_tracker = NULL;

   // Retain all devices and cache driver names.
   for (iree_host_size_t i = 0; i < device_count; ++i) {
     group->devices[i] = builder->devices[i];
     iree_hal_device_retain(group->devices[i]);
     group->driver_names[i] = iree_hal_device_id(group->devices[i]);
   }

   // Invalidate the builder now — we've taken everything we need.
   memset(builder, 0, sizeof(*builder));

   // Query capabilities from all devices.
   iree_hal_device_capabilities_t
       capabilities[IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT];
   memset(capabilities, 0, sizeof(capabilities));
   for (iree_host_size_t i = 0; i < device_count && iree_status_is_ok(status);
        ++i) {
     status =
         iree_hal_device_query_capabilities(group->devices[i], &capabilities[i]);
   }

   // Build topology.
   iree_hal_topology_builder_t topology_builder;
   iree_hal_topology_builder_initialize(&topology_builder,
                                        (uint32_t)device_count);

   // Compute and set edges for all device pairs.
   for (uint32_t i = 0; i < (uint32_t)device_count && iree_status_is_ok(status);
        ++i) {
     for (uint32_t j = 0;
          j < (uint32_t)device_count && iree_status_is_ok(status); ++j) {
       if (i == j) continue;  // Self-edges are pre-initialized.

       // Compute base edge from capabilities.
       iree_hal_topology_edge_t edge = iree_hal_topology_edge_from_capabilities(
           &capabilities[i], &capabilities[j], group->driver_names[i],
           group->driver_names[j]);

       // Allow same-driver devices to refine the edge with hardware-specific
       // knowledge (e.g., NVLink topology, Infinity Fabric link widths).
       if (iree_string_view_equal(group->driver_names[i],
                                  group->driver_names[j])) {
         status = iree_hal_device_refine_topology_edge(group->devices[i],
                                                       group->devices[j], &edge);
         if (!iree_status_is_ok(status)) break;
       }

       status =
           iree_hal_topology_builder_set_edge(&topology_builder, i, j, edge);
     }
   }

   // Set NUMA nodes from queried capabilities.
   for (uint32_t i = 0; i < (uint32_t)device_count && iree_status_is_ok(status);
        ++i) {
     status = iree_hal_topology_builder_set_numa_node(&topology_builder, i,
                                                      capabilities[i].numa_node);
   }

   // Finalize topology.
   if (iree_status_is_ok(status)) {
     status =
         iree_hal_topology_builder_finalize(&topology_builder, &group->topology);
   }

   // Assign topology info to each device.
   iree_host_size_t assigned_device_count = 0;
   for (iree_host_size_t i = 0; i < device_count && iree_status_is_ok(status);
        ++i) {
     iree_hal_device_topology_info_t topology_info;
     memset(&topology_info, 0, sizeof(topology_info));
     topology_info.topology_index = (uint32_t)i;
     topology_info.topology = &group->topology;
     topology_info.self_edge =
         iree_hal_topology_query_edge(&group->topology, (uint32_t)i, (uint32_t)i)
             .lo;
     topology_info.frontier.tracker = group->frontier_tracker;
     topology_info.frontier.base_axis = iree_async_axis_make_queue(
         iree_async_frontier_tracker_session_epoch(group->frontier_tracker),
         iree_async_frontier_tracker_machine_index(group->frontier_tracker),
         (uint8_t)i, /*queue_index=*/0);

     iree_hal_device_group_compute_bitmaps(&group->topology, (uint32_t)i,
                                           &topology_info);

     status =
         iree_hal_device_assign_topology_info(group->devices[i], &topology_info);
     if (iree_status_is_ok(status)) {
       assigned_device_count = i + 1;
     }
   }

   if (iree_status_is_ok(status)) {
     *out_group = group;
   } else {
     for (iree_host_size_t i = 0; i < assigned_device_count; ++i) {
       status = iree_status_join(status, iree_hal_device_assign_topology_info(
                                             group->devices[i], NULL));
     }
     iree_hal_device_group_release(group);
   }

   IREE_TRACE_ZONE_END(z0);
   return status;
 }
	// Copyright 2026 The IREE Authors
	//
	// Licensed under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

	#include "iree/hal/device_group.h"

	#include <string.h>

	#include "iree/async/frontier_tracker.h"
	#include "iree/base/internal/atomics.h"
	#include "iree/hal/topology_builder.h"

	//===----------------------------------------------------------------------===//
	// iree_hal_device_group_t
	//===----------------------------------------------------------------------===//

	struct iree_hal_device_group_t {
	// Retain/release counter for the group object.
	iree_atomic_ref_count_t ref_count;

	// Host allocator used to free the group object.
	iree_allocator_t host_allocator;

	// Number of devices in the group.
	iree_host_size_t device_count;

	// Retained frontier tracker shared by all devices in the group.
	iree_async_frontier_tracker_t* frontier_tracker;

	// Retained devices. Order defines topology indices (device i = devices[i]).
	iree_hal_device_t* devices[IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT];

	// Cached driver name for each device (backed by the device's own storage,
	// valid for the device's lifetime).
	iree_string_view_t driver_names[IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT];

	// Immutable topology matrix built during creation.
	// Embedded (not heap-allocated) so devices can hold a stable pointer to it
	// for the group's lifetime.
	iree_hal_topology_t topology;
	};

	IREE_API_EXPORT iree_status_t iree_hal_device_group_create_from_device(
	iree_hal_device_t* device, iree_async_frontier_tracker_t* frontier_tracker,
	iree_allocator_t host_allocator, iree_hal_device_group_t** out_group) {
	IREE_ASSERT_ARGUMENT(device);
	IREE_ASSERT_ARGUMENT(frontier_tracker);
	IREE_ASSERT_ARGUMENT(out_group);
	*out_group = NULL;
	iree_hal_device_group_builder_t builder;
	iree_hal_device_group_builder_initialize(&builder, frontier_tracker);
	iree_status_t status =
	iree_hal_device_group_builder_add_device(&builder, device);
	if (iree_status_is_ok(status)) {
	status = iree_hal_device_group_builder_finalize(&builder, host_allocator,
	out_group);
	} else {
	iree_hal_device_group_builder_deinitialize(&builder);
	}
	return status;
	}

	static void iree_hal_device_group_destroy(iree_hal_device_group_t* group) {
	IREE_ASSERT_ARGUMENT(group);
	IREE_TRACE_ZONE_BEGIN(z0);

	for (iree_host_size_t i = 0; i < group->device_count; ++i) {
	iree_hal_device_release(group->devices[i]);
	}
	iree_async_frontier_tracker_release(group->frontier_tracker);

	iree_allocator_t host_allocator = group->host_allocator;
	iree_allocator_free(host_allocator, group);

	IREE_TRACE_ZONE_END(z0);
	}

	IREE_API_EXPORT void iree_hal_device_group_retain(
	iree_hal_device_group_t* group) {
	if (group) {
	iree_atomic_ref_count_inc(&group->ref_count);
	}
	}

	IREE_API_EXPORT void iree_hal_device_group_release(
	iree_hal_device_group_t* group) {
	if (group && iree_atomic_ref_count_dec(&group->ref_count) == 1) {
	iree_hal_device_group_destroy(group);
	}
	}

	IREE_API_EXPORT iree_host_size_t
	iree_hal_device_group_device_count(const iree_hal_device_group_t* group) {
	IREE_ASSERT_ARGUMENT(group);
	return group->device_count;
	}

	IREE_API_EXPORT iree_hal_device_t* iree_hal_device_group_device_at(
	const iree_hal_device_group_t* group, iree_host_size_t index) {
	IREE_ASSERT_ARGUMENT(group);
	if (index >= group->device_count) return NULL;
	return group->devices[index];
	}

	IREE_API_EXPORT const iree_hal_topology_t* iree_hal_device_group_topology(
	const iree_hal_device_group_t* group) {
	IREE_ASSERT_ARGUMENT(group);
	return &group->topology;
	}

	//===----------------------------------------------------------------------===//
	// iree_hal_device_group_builder_t
	//===----------------------------------------------------------------------===//

	IREE_API_EXPORT void iree_hal_device_group_builder_initialize(
	iree_hal_device_group_builder_t* builder,
	iree_async_frontier_tracker_t* frontier_tracker) {
	IREE_ASSERT_ARGUMENT(builder);
	IREE_ASSERT_ARGUMENT(frontier_tracker);
	memset(builder, 0, sizeof(*builder));
	builder->frontier_tracker = frontier_tracker;
	iree_async_frontier_tracker_retain(frontier_tracker);
	}

	IREE_API_EXPORT void iree_hal_device_group_builder_deinitialize(
	iree_hal_device_group_builder_t* builder) {
	IREE_ASSERT_ARGUMENT(builder);
	iree_async_frontier_tracker_release(builder->frontier_tracker);
	memset(builder, 0, sizeof(*builder));
	}

	IREE_API_EXPORT iree_status_t iree_hal_device_group_builder_add_device(
	iree_hal_device_group_builder_t* builder, iree_hal_device_t* device) {
	IREE_ASSERT_ARGUMENT(builder);
	IREE_ASSERT_ARGUMENT(device);
	if (builder->count >= IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT) {
	return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED,
	"device group builder at max capacity (%d)",
	IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT);
	}
	builder->devices[builder->count++] = device;
	return iree_ok_status();
	}

	// Computes pre-populated bitmaps for a device's topology info by scanning
	// its edges in the topology matrix.
	static void iree_hal_device_group_compute_bitmaps(
	const iree_hal_topology_t* topology, uint32_t device_index,
	iree_hal_device_topology_info_t* out_info) {
	uint32_t device_count = topology->device_count;
	out_info->can_wait_from = 0;
	out_info->can_signal_to = 0;
	out_info->can_import_from = 0;
	out_info->can_p2p_with = 0;

	for (uint32_t j = 0; j < device_count; ++j) {
	if (j == device_index) continue;

	// can_wait_from: can device_index wait on device j's semaphores?
	// Edge[j][device_index] describes how device_index interacts with j's
	// resources.
	iree_hal_topology_edge_t edge_from_j =
	iree_hal_topology_query_edge(topology, j, device_index);
	if (iree_hal_topology_edge_wait_mode(edge_from_j.lo) !=
	IREE_HAL_TOPOLOGY_INTEROP_MODE_NONE) {
	out_info->can_wait_from \|= (iree_hal_topology_device_bitmap_t)1 << j;
	}

	// can_import_from: can device_index import buffers from device j?
	// Set if either non-coherent or coherent buffers can be shared without
	// a host-staged copy.
	iree_hal_topology_interop_mode_t nc_read =
	iree_hal_topology_edge_buffer_read_mode_noncoherent(edge_from_j.lo);
	iree_hal_topology_interop_mode_t c_read =
	iree_hal_topology_edge_buffer_read_mode_coherent(edge_from_j.lo);
	if (nc_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_NATIVE \|\|
	nc_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_IMPORT \|\|
	c_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_NATIVE \|\|
	c_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_IMPORT) {
	out_info->can_import_from \|= (iree_hal_topology_device_bitmap_t)1 << j;
	}

	// can_p2p_with: P2P access between device_index and device j?
	if (nc_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_NATIVE \|\|
	c_read == IREE_HAL_TOPOLOGY_INTEROP_MODE_NATIVE \|\|
	(iree_hal_topology_edge_capability_flags(edge_from_j.lo) &
	IREE_HAL_TOPOLOGY_CAPABILITY_P2P_COPY)) {
	out_info->can_p2p_with \|= (iree_hal_topology_device_bitmap_t)1 << j;
	}

	// can_signal_to: can device_index signal to device j?
	// Edge[device_index][j] describes how j interacts with device_index's
	// resources — the signal_mode tells us if j can observe our signals.
	iree_hal_topology_edge_t edge_to_j =
	iree_hal_topology_query_edge(topology, device_index, j);
	if (iree_hal_topology_edge_signal_mode(edge_to_j.lo) !=
	IREE_HAL_TOPOLOGY_INTEROP_MODE_NONE) {
	out_info->can_signal_to \|= (iree_hal_topology_device_bitmap_t)1 << j;
	}
	}
	}

	IREE_API_EXPORT iree_status_t iree_hal_device_group_create_with_replacements(
	iree_hal_device_group_t* source_group,
	iree_hal_device_group_replacement_callback_t replacement_callback,
	iree_allocator_t host_allocator, iree_hal_device_group_t** out_group) {
	IREE_ASSERT_ARGUMENT(source_group);
	IREE_ASSERT_ARGUMENT(out_group);
	*out_group = NULL;
	IREE_TRACE_ZONE_BEGIN(z0);

	if (IREE_UNLIKELY(!replacement_callback.fn)) {
	IREE_TRACE_ZONE_END(z0);
	return iree_make_status(
	IREE_STATUS_INVALID_ARGUMENT,
	"device group replacement callback function is required");
	}

	const iree_host_size_t device_count = source_group->device_count;
	iree_hal_device_group_t* group = NULL;
	iree_status_t status =
	iree_allocator_malloc(host_allocator, sizeof(group), (void*)&group);
	if (!iree_status_is_ok(status)) {
	IREE_TRACE_ZONE_END(z0);
	return status;
	}
	memset(group, 0, sizeof(*group));
	iree_atomic_ref_count_init(&group->ref_count);
	group->host_allocator = host_allocator;
	group->device_count = device_count;
	group->frontier_tracker = source_group->frontier_tracker;
	iree_async_frontier_tracker_retain(group->frontier_tracker);
	group->topology = source_group->topology;

	for (iree_host_size_t i = 0; i < device_count && iree_status_is_ok(status);
	++i) {
	iree_hal_device_t* replacement_device = NULL;
	status = replacement_callback.fn(source_group, i, source_group->devices[i],
	replacement_callback.user_data,
	&replacement_device);
	if (iree_status_is_ok(status) && IREE_UNLIKELY(!replacement_device)) {
	status = iree_make_status(
	IREE_STATUS_INVALID_ARGUMENT,
	"replacement device %" PRIhsz " must not be NULL", i);
	}
	if (iree_status_is_ok(status)) {
	group->devices[i] = replacement_device;
	} else {
	iree_hal_device_release(replacement_device);
	}
	group->driver_names[i] = source_group->driver_names[i];
	}

	iree_host_size_t assigned_device_count = 0;
	for (iree_host_size_t i = 0; i < device_count && iree_status_is_ok(status);
	++i) {
	iree_hal_device_topology_info_t topology_info;
	memset(&topology_info, 0, sizeof(topology_info));
	topology_info.topology_index = (uint32_t)i;
	topology_info.topology = &group->topology;
	topology_info.self_edge =
	iree_hal_topology_query_edge(&group->topology, (uint32_t)i, (uint32_t)i)
	.lo;
	topology_info.frontier.tracker = group->frontier_tracker;
	topology_info.frontier.base_axis = iree_async_axis_make_queue(
	iree_async_frontier_tracker_session_epoch(group->frontier_tracker),
	iree_async_frontier_tracker_machine_index(group->frontier_tracker),
	(uint8_t)i, /queue_index=/0);

	iree_hal_device_group_compute_bitmaps(&group->topology, (uint32_t)i,
	&topology_info);

	status =
	iree_hal_device_assign_topology_info(group->devices[i], &topology_info);
	if (iree_status_is_ok(status)) {
	assigned_device_count = i + 1;
	}
	}

	if (iree_status_is_ok(status)) {
	*out_group = group;
	} else {
	for (iree_host_size_t i = 0; i < assigned_device_count; ++i) {
	status = iree_status_join(status, iree_hal_device_assign_topology_info(
	group->devices[i], NULL));
	}
	iree_hal_device_group_release(group);
	}

	IREE_TRACE_ZONE_END(z0);
	return status;
	}

	IREE_API_EXPORT iree_status_t iree_hal_device_group_builder_finalize(
	iree_hal_device_group_builder_t* builder, iree_allocator_t host_allocator,
	iree_hal_device_group_t** out_group) {
	IREE_ASSERT_ARGUMENT(builder);
	IREE_ASSERT_ARGUMENT(out_group);
	*out_group = NULL;
	IREE_TRACE_ZONE_BEGIN(z0);

	iree_host_size_t device_count = builder->count;
	if (device_count == 0) {
	iree_hal_device_group_builder_deinitialize(builder);
	IREE_TRACE_ZONE_END(z0);
	return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
	"device group builder has no devices");
	}

	// Allocate the group.
	iree_hal_device_group_t* group = NULL;
	iree_status_t status =
	iree_allocator_malloc(host_allocator, sizeof(group), (void*)&group);
	if (!iree_status_is_ok(status)) {
	iree_hal_device_group_builder_deinitialize(builder);
	IREE_TRACE_ZONE_END(z0);
	return status;
	}
	memset(group, 0, sizeof(*group));
	iree_atomic_ref_count_init(&group->ref_count);
	group->host_allocator = host_allocator;
	group->device_count = device_count;
	group->frontier_tracker = builder->frontier_tracker;
	builder->frontier_tracker = NULL;

	// Retain all devices and cache driver names.
	for (iree_host_size_t i = 0; i < device_count; ++i) {
	group->devices[i] = builder->devices[i];
	iree_hal_device_retain(group->devices[i]);
	group->driver_names[i] = iree_hal_device_id(group->devices[i]);
	}

	// Invalidate the builder now — we've taken everything we need.
	memset(builder, 0, sizeof(*builder));

	// Query capabilities from all devices.
	iree_hal_device_capabilities_t
	capabilities[IREE_HAL_TOPOLOGY_MAX_DEVICE_COUNT];
	memset(capabilities, 0, sizeof(capabilities));
	for (iree_host_size_t i = 0; i < device_count && iree_status_is_ok(status);
	++i) {
	status =
	iree_hal_device_query_capabilities(group->devices[i], &capabilities[i]);
	}

	// Build topology.
	iree_hal_topology_builder_t topology_builder;
	iree_hal_topology_builder_initialize(&topology_builder,
	(uint32_t)device_count);

	// Compute and set edges for all device pairs.
	for (uint32_t i = 0; i < (uint32_t)device_count && iree_status_is_ok(status);
	++i) {
	for (uint32_t j = 0;
	j < (uint32_t)device_count && iree_status_is_ok(status); ++j) {
	if (i == j) continue; // Self-edges are pre-initialized.

	// Compute base edge from capabilities.
	iree_hal_topology_edge_t edge = iree_hal_topology_edge_from_capabilities(
	&capabilities[i], &capabilities[j], group->driver_names[i],
	group->driver_names[j]);

	// Allow same-driver devices to refine the edge with hardware-specific
	// knowledge (e.g., NVLink topology, Infinity Fabric link widths).
	if (iree_string_view_equal(group->driver_names[i],
	group->driver_names[j])) {
	status = iree_hal_device_refine_topology_edge(group->devices[i],
	group->devices[j], &edge);
	if (!iree_status_is_ok(status)) break;
	}

	status =
	iree_hal_topology_builder_set_edge(&topology_builder, i, j, edge);
	}
	}

	// Set NUMA nodes from queried capabilities.
	for (uint32_t i = 0; i < (uint32_t)device_count && iree_status_is_ok(status);
	++i) {
	status = iree_hal_topology_builder_set_numa_node(&topology_builder, i,
	capabilities[i].numa_node);
	}

	// Finalize topology.
	if (iree_status_is_ok(status)) {
	status =
	iree_hal_topology_builder_finalize(&topology_builder, &group->topology);
	}

	// Assign topology info to each device.
	iree_host_size_t assigned_device_count = 0;
	for (iree_host_size_t i = 0; i < device_count && iree_status_is_ok(status);
	++i) {
	iree_hal_device_topology_info_t topology_info;
	memset(&topology_info, 0, sizeof(topology_info));
	topology_info.topology_index = (uint32_t)i;
	topology_info.topology = &group->topology;
	topology_info.self_edge =
	iree_hal_topology_query_edge(&group->topology, (uint32_t)i, (uint32_t)i)
	.lo;
	topology_info.frontier.tracker = group->frontier_tracker;
	topology_info.frontier.base_axis = iree_async_axis_make_queue(
	iree_async_frontier_tracker_session_epoch(group->frontier_tracker),
	iree_async_frontier_tracker_machine_index(group->frontier_tracker),
	(uint8_t)i, /queue_index=/0);

	iree_hal_device_group_compute_bitmaps(&group->topology, (uint32_t)i,
	&topology_info);

	status =
	iree_hal_device_assign_topology_info(group->devices[i], &topology_info);
	if (iree_status_is_ok(status)) {
	assigned_device_count = i + 1;
	}
	}

	if (iree_status_is_ok(status)) {
	*out_group = group;
	} else {
	for (iree_host_size_t i = 0; i < assigned_device_count; ++i) {
	status = iree_status_join(status, iree_hal_device_assign_topology_info(
	group->devices[i], NULL));
	}
	iree_hal_device_group_release(group);
	}

	IREE_TRACE_ZONE_END(z0);
	return status;
	}