hrx-queue-affinity.patch

# libhrx: stream-carried queue affinity
#
# SUPERSEDED: this work was committed to hrx-system as 2082d042 ("Stream queue
# affinity and a narrower dispatch barrier") and is contained in the ref
# scripts/bootstrap-hrx.sh pins (5927b0e0). The file is kept as a record only;
# do not apply it. The one change it does NOT carry — the allocator.c
# device-allocation import — lives in patches/0003 and is applied by bootstrap.
#
# Backup of uncommitted work in reference/hrx-system, which is gitignored in
# this repo. Apply from the root of an hrx-system checkout:
#     git apply /path/to/hrx-queue-affinity.patch
#
# Base commit: f6ccdeb9ed470bfbf76e132194564bb9afdfde06
#     [Loom] Model checked device execution with kernel launches (#267)
#
# Adds HRX_DEVICE_PROPERTY_QUEUE_COUNT, hrx_stream_create_on_queue() and
# hrx_stream_get_queue_affinity(). Affinity is a bitmask (bit i = queue i);
# zero means any queue, which is what hrx_stream_create() passes, so existing
# callers are unchanged. Out-of-range bits return HRX_STATUS_OUT_OF_RANGE.
#
# CAVEAT: this checkout already carried one unrelated local change before this
# work - a narrowed dispatch-to-dispatch barrier in libhrx/src/libhrx/stream.c
# around lines 468-481. Separate it before upstreaming.

diff --git a/libhrx/cts/core/hrx_loader.cpp b/libhrx/cts/core/hrx_loader.cpp
index 78d0b4e..1af7d37 100644
--- a/libhrx/cts/core/hrx_loader.cpp
+++ b/libhrx/cts/core/hrx_loader.cpp
@@ -133,6 +133,8 @@ void HrxLoader::load(const std::string& path) {
   LOAD(semaphore_signal);
 
   LOAD(stream_create);
+  LOAD(stream_create_on_queue);
+  LOAD(stream_get_queue_affinity);
   LOAD(stream_retain);
   LOAD(stream_release);
   LOAD(stream_synchronize);
diff --git a/libhrx/cts/core/hrx_loader.hpp b/libhrx/cts/core/hrx_loader.hpp
index 3fbde44..fb90830 100644
--- a/libhrx/cts/core/hrx_loader.hpp
+++ b/libhrx/cts/core/hrx_loader.hpp
@@ -91,6 +91,8 @@ class HrxLoader {
 
   // Streams.
   decltype(&hrx_stream_create) stream_create;
+  decltype(&hrx_stream_create_on_queue) stream_create_on_queue;
+  decltype(&hrx_stream_get_queue_affinity) stream_get_queue_affinity;
   decltype(&hrx_stream_retain) stream_retain;
   decltype(&hrx_stream_release) stream_release;
   decltype(&hrx_stream_synchronize) stream_synchronize;
diff --git a/libhrx/cts/tests/stream/stream_test.cpp b/libhrx/cts/tests/stream/stream_test.cpp
index cbbb357..7c0f687 100644
--- a/libhrx/cts/tests/stream/stream_test.cpp
+++ b/libhrx/cts/tests/stream/stream_test.cpp
@@ -2,9 +2,46 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include <catch2/catch_test_macros.hpp>
+#include <cstdint>
 
 #include "hrx_test_fixture.hpp"
 
+namespace {
+
+uint32_t queue_count(hrx_device_t device) {
+  uint32_t count = 0;
+  REQUIRE_OK(hrx().device_get_property(device, HRX_DEVICE_PROPERTY_QUEUE_COUNT,
+                                       &count, sizeof(count)));
+  return count;
+}
+
+// Fills a buffer through |stream| and reads it back, so the queue the stream
+// named is proven to be a queue that actually runs work.
+void fill_and_verify(hrx_stream_t stream, uint32_t pattern) {
+  constexpr size_t kByteCount = 1024;
+  hrx_buffer_t buf = nullptr;
+  REQUIRE_OK(hrx().buffer_allocate(
+      stream, kByteCount,
+      HRX_MEMORY_TYPE_HOST_LOCAL | HRX_MEMORY_TYPE_DEVICE_VISIBLE,
+      HRX_BUFFER_USAGE_DEFAULT | HRX_BUFFER_USAGE_MAPPING_SCOPED, &buf));
+
+  REQUIRE_OK(hrx().stream_fill_buffer(stream, buf, 0, kByteCount, &pattern,
+                                      sizeof(pattern)));
+  REQUIRE_OK(hrx().stream_synchronize(stream));
+
+  void* ptr = nullptr;
+  REQUIRE_OK(hrx().buffer_map(buf, HRX_MAP_READ, 0, kByteCount, &ptr));
+  const uint32_t* data = static_cast<const uint32_t*>(ptr);
+  for (size_t i = 0; i < kByteCount / sizeof(uint32_t); ++i) {
+    REQUIRE(data[i] == pattern);
+  }
+  REQUIRE_OK(hrx().buffer_unmap(buf));
+
+  hrx().buffer_release(buf);
+}
+
+}  // namespace
+
 TEST_CASE_METHOD(HrxTestFixture, "Stream create and release", "[stream]") {
   hrx_stream_t stream = nullptr;
   REQUIRE_OK(hrx().stream_create(device_, 0, &stream));
@@ -76,3 +113,74 @@ TEST_CASE_METHOD(HrxTestFixture, "Stream get_timeline_position",
 
   hrx().stream_release(stream);
 }
+
+TEST_CASE_METHOD(HrxTestFixture, "Device reports at least one queue",
+                 "[stream][queue]") {
+  REQUIRE(queue_count(device_) >= 1);
+}
+
+TEST_CASE_METHOD(HrxTestFixture, "Stream create leaves the queue unbound",
+                 "[stream][queue]") {
+  hrx_stream_t stream = nullptr;
+  REQUIRE_OK(hrx().stream_create(device_, 0, &stream));
+
+  hrx_queue_affinity_t affinity = ~hrx_queue_affinity_t{0};
+  REQUIRE_OK(hrx().stream_get_queue_affinity(stream, &affinity));
+  REQUIRE(affinity == 0);
+
+  fill_and_verify(stream, 0x5EED5EEDu);
+  hrx().stream_release(stream);
+}
+
+TEST_CASE_METHOD(HrxTestFixture, "Stream keeps the queue it was created on",
+                 "[stream][queue]") {
+  const uint32_t queues = queue_count(device_);
+  for (uint32_t i = 0; i < queues; ++i) {
+    const hrx_queue_affinity_t requested = hrx_queue_affinity_t{1} << i;
+    hrx_stream_t stream = nullptr;
+    REQUIRE_OK(hrx().stream_create_on_queue(device_, 0, requested, &stream));
+
+    hrx_queue_affinity_t affinity = 0;
+    REQUIRE_OK(hrx().stream_get_queue_affinity(stream, &affinity));
+    REQUIRE(affinity == requested);
+
+    fill_and_verify(stream, 0xA5000000u | i);
+    hrx().stream_release(stream);
+  }
+}
+
+TEST_CASE_METHOD(HrxTestFixture, "Streams on distinct queues both complete",
+                 "[stream][queue]") {
+  if (queue_count(device_) < 2) {
+    SUCCEED("device has a single queue");
+    return;
+  }
+
+  hrx_stream_t first = nullptr;
+  hrx_stream_t second = nullptr;
+  REQUIRE_OK(hrx().stream_create_on_queue(device_, 0, 1, &first));
+  REQUIRE_OK(hrx().stream_create_on_queue(device_, 0, 2, &second));
+
+  fill_and_verify(first, 0x11111111u);
+  fill_and_verify(second, 0x22222222u);
+
+  hrx().stream_release(second);
+  hrx().stream_release(first);
+}
+
+TEST_CASE_METHOD(HrxTestFixture, "Stream rejects a queue the device lacks",
+                 "[stream][queue]") {
+  const uint32_t queues = queue_count(device_);
+  if (queues >= 64) {
+    SUCCEED("every affinity bit names a real queue");
+    return;
+  }
+
+  hrx_stream_t stream = nullptr;
+  hrx_status_t status = hrx().stream_create_on_queue(
+      device_, 0, hrx_queue_affinity_t{1} << queues, &stream);
+  REQUIRE(!hrx_status_is_ok(status));
+  REQUIRE(hrx().status_code(status) == HRX_STATUS_OUT_OF_RANGE);
+  REQUIRE(stream == nullptr);
+  hrx().status_ignore(status);
+}
diff --git a/libhrx/include/hrx_runtime.h b/libhrx/include/hrx_runtime.h
index bcb29a5..2fada92 100644
--- a/libhrx/include/hrx_runtime.h
+++ b/libhrx/include/hrx_runtime.h
@@ -190,6 +190,9 @@ typedef enum hrx_device_property_t {
   HRX_DEVICE_PROPERTY_MAX_SHARED_MEMORY,
   HRX_DEVICE_PROPERTY_CLOCK_RATE,
   HRX_DEVICE_PROPERTY_PCI_BUS_ID,
+  // uint32_t. Hardware queues that can be named individually; defines the
+  // valid bit range of a hrx_queue_affinity_t on this device.
+  HRX_DEVICE_PROPERTY_QUEUE_COUNT,
 } hrx_device_property_t;
 
 // Memory type bitfield. Values match iree_hal_memory_type_t.
@@ -252,6 +255,9 @@ typedef struct hrx_semaphore_list_t {
   size_t count;
 } hrx_semaphore_list_t;
 
+// Selects hardware queues: bit i is queue i of HRX_DEVICE_PROPERTY_QUEUE_COUNT.
+// Zero means any queue. Multiple bits mean any one of them, chosen by the
+// implementation.
 typedef uint64_t hrx_queue_affinity_t;
 
 // Borrowed string storage valid only for the duration documented by the
@@ -570,6 +576,21 @@ HRX_API hrx_status_t hrx_semaphore_signal(hrx_semaphore_t semaphore,
 HRX_API hrx_status_t hrx_stream_create(hrx_device_t device, uint32_t flags,
                                        hrx_stream_t* stream);
 
+// Creates a stream whose every submission is issued on |queue_affinity|: its
+// command buffers, flushes, timeline barriers and stream-ordered allocations.
+// Zero is what hrx_stream_create() passes and leaves the queue to the
+// implementation.
+//
+// The affinity is fixed at creation because a command buffer is recorded for
+// the queue it will later run on; a stream that recorded on one queue and
+// submitted on another would be recorded against the wrong physical device.
+//
+// Returns HRX_STATUS_OUT_OF_RANGE if a bit names a queue the device does not
+// have.
+HRX_API hrx_status_t hrx_stream_create_on_queue(
+    hrx_device_t device, uint32_t flags, hrx_queue_affinity_t queue_affinity,
+    hrx_stream_t* stream);
+
 HRX_API void hrx_stream_retain(hrx_stream_t stream);
 HRX_API void hrx_stream_release(hrx_stream_t stream);
 
@@ -592,6 +613,10 @@ HRX_API hrx_status_t hrx_stream_get_semaphore(hrx_stream_t stream,
 HRX_API hrx_status_t hrx_stream_get_device(hrx_stream_t stream,
                                            hrx_device_t* device);
 
+// Returns the affinity the stream was created with; zero means any queue.
+HRX_API hrx_status_t hrx_stream_get_queue_affinity(
+    hrx_stream_t stream, hrx_queue_affinity_t* queue_affinity);
+
 HRX_API hrx_status_t hrx_stream_get_timeline_position(
     hrx_stream_t stream, hrx_timeline_point_t* position);
 
diff --git a/libhrx/src/libhrx/buffer.c b/libhrx/src/libhrx/buffer.c
index 68b63c3..a8848b0 100644
--- a/libhrx/src/libhrx/buffer.c
+++ b/libhrx/src/libhrx/buffer.c
@@ -82,7 +82,8 @@ hrx_status_t hrx_buffer_allocate(hrx_stream_t stream, size_t size,
                                       params, &buf->hal_pool);
   if (iree_status_is_ok(status)) {
     status = iree_hal_device_queue_alloca(
-        stream->device->hal_device, IREE_HAL_QUEUE_AFFINITY_ANY, wait_list,
+        stream->device->hal_device,
+        hrx_normalize_queue_affinity(stream->queue_affinity), wait_list,
         signal_list, buf->hal_pool, params, (iree_device_size_t)size,
         IREE_HAL_ALLOCA_FLAG_NONE, &buf->hal_buffer);
   }
diff --git a/libhrx/src/libhrx/device.c b/libhrx/src/libhrx/device.c
index 4d6fa98..4c58ca4 100644
--- a/libhrx/src/libhrx/device.c
+++ b/libhrx/src/libhrx/device.c
@@ -31,6 +31,23 @@ hrx_status_t hrx_device_query_total_memory_from_spec(
   return hrx_ok_status();
 }
 
+hrx_status_t hrx_device_query_queue_count(hrx_device_t device,
+                                          uint32_t* out_count) {
+  if (!device || !out_count) {
+    return hrx_make_status(HRX_STATUS_INVALID_ARGUMENT, "NULL argument");
+  }
+  const iree_hal_device_queue_spec_t* queues =
+      iree_hal_device_spec_queues(iree_hal_device_spec(device->hal_device));
+  uint64_t total = 0;
+  for (iree_host_size_t i = 0; queues && i < queues->family_count; ++i) {
+    total += queues->families[i].queue_count;
+  }
+  // A queue past the width of an affinity mask cannot be named at all.
+  if (total > IREE_HAL_MAX_QUEUES) total = IREE_HAL_MAX_QUEUES;
+  *out_count = total > 0 ? (uint32_t)total : 1u;
+  return hrx_ok_status();
+}
+
 static hrx_status_t hrx_device_sample_memory(
     hrx_device_t device, iree_device_size_t* out_total,
     iree_device_size_t* out_available) {
@@ -107,6 +124,13 @@ hrx_status_t hrx_device_get_property(hrx_device_t device,
       *(uint64_t*)value = (uint64_t)total_bytes;
       return status;
     }
+    case HRX_DEVICE_PROPERTY_QUEUE_COUNT: {
+      if (value_size < sizeof(uint32_t)) {
+        return hrx_make_status(HRX_STATUS_OUT_OF_RANGE,
+                               "buffer too small for uint32_t");
+      }
+      return hrx_device_query_queue_count(device, (uint32_t*)value);
+    }
     case HRX_DEVICE_PROPERTY_COMPUTE_UNITS:
     case HRX_DEVICE_PROPERTY_MAX_WORKGROUP_SIZE: {
       if (value_size < sizeof(uint32_t)) {
diff --git a/libhrx/src/libhrx/event.c b/libhrx/src/libhrx/event.c
index ff0c4bd..c149bc3 100644
--- a/libhrx/src/libhrx/event.c
+++ b/libhrx/src/libhrx/event.c
@@ -110,7 +110,8 @@ hrx_status_t hrx_event_record(hrx_event_t event, hrx_stream_t stream) {
   };
 
   iree_status_t iree_status = iree_hal_device_queue_barrier(
-      stream->device->hal_device, IREE_HAL_QUEUE_AFFINITY_ANY, wait_list,
+      stream->device->hal_device,
+      hrx_normalize_queue_affinity(stream->queue_affinity), wait_list,
       signal_list, IREE_HAL_EXECUTE_FLAG_NONE);
   if (!iree_status_is_ok(iree_status)) {
     return hrx_status_from_iree(iree_status);
@@ -222,7 +223,8 @@ hrx_status_t hrx_stream_wait_event(hrx_stream_t stream, hrx_event_t event) {
   };
 
   iree_status_t iree_status = iree_hal_device_queue_barrier(
-      stream->device->hal_device, IREE_HAL_QUEUE_AFFINITY_ANY, wait_list,
+      stream->device->hal_device,
+      hrx_normalize_queue_affinity(stream->queue_affinity), wait_list,
       signal_list, IREE_HAL_EXECUTE_FLAG_NONE);
   if (!iree_status_is_ok(iree_status)) {
     return hrx_status_from_iree(iree_status);
diff --git a/libhrx/src/libhrx/graph_exec.c b/libhrx/src/libhrx/graph_exec.c
index 52a464d..58f4823 100644
--- a/libhrx/src/libhrx/graph_exec.c
+++ b/libhrx/src/libhrx/graph_exec.c
@@ -553,7 +553,8 @@ hrx_status_t hrx_graph_exec_launch(hrx_graph_exec_t exec, hrx_stream_t stream) {
     HRX_RETURN_AND_END_ZONE_IF_IREE_ERROR(
         z0,
         iree_hal_device_queue_execute(
-            stream->device->hal_device, IREE_HAL_QUEUE_AFFINITY_ANY, wait_list,
+            stream->device->hal_device,
+            hrx_normalize_queue_affinity(stream->queue_affinity), wait_list,
             signal_list, stream->pending_cb,
             iree_hal_buffer_binding_table_empty(), IREE_HAL_EXECUTE_FLAG_NONE));
     stream->timepoint = next_value;
@@ -645,6 +646,10 @@ hrx_status_t hrx_graph_exec_launch(hrx_graph_exec_t exec, hrx_stream_t stream) {
         .payload_values = signal_vals,
     };
 
+    // Graph blocks stay on any queue even when |stream| names one: the block
+    // command buffers were recorded at instantiate time, when no stream was
+    // known. Ordering still holds through the stream's timeline semaphore.
+    // Honouring the stream's affinity here needs the affinity at instantiate.
     switch (block->type) {
       case HRX_GRAPH_BLOCK_TYPE_QUEUE_BARRIER:
         status = iree_hal_device_queue_barrier(
diff --git a/libhrx/src/libhrx/hrx_internal.h b/libhrx/src/libhrx/hrx_internal.h
index a9059ab..88e4e79 100644
--- a/libhrx/src/libhrx/hrx_internal.h
+++ b/libhrx/src/libhrx/hrx_internal.h
@@ -166,6 +166,13 @@ static inline hrx_status_t hrx_iree_dispatch_flags_from_hrx(
   return hrx_ok_status();
 }
 
+// The hrx ABI spells "any queue" as zero; the HAL spells it as all bits set.
+static inline iree_hal_queue_affinity_t hrx_normalize_queue_affinity(
+    hrx_queue_affinity_t affinity) {
+  return affinity == 0 ? IREE_HAL_QUEUE_AFFINITY_ANY
+                       : (iree_hal_queue_affinity_t)affinity;
+}
+
 // Buffer view metadata.
 static_assert(HRX_ELEMENT_TYPE_NONE == IREE_HAL_ELEMENT_TYPE_NONE,
               "element type mismatch");
@@ -280,6 +287,9 @@ typedef struct hrx_stream_s {
   iree_hal_command_buffer_t* pending_cb;
   bool has_pending_work;
   uint32_t flags;
+  // Fixed at creation and used by every submission the stream makes; recording
+  // and submitting must name the same queue. Zero means any.
+  hrx_queue_affinity_t queue_affinity;
 } hrx_stream_s;
 
 //===----------------------------------------------------------------------===//
@@ -635,6 +645,13 @@ hrx_status_t hrx_ensure_shared_state(void);
 hrx_status_t hrx_device_query_total_memory_from_spec(
     hrx_device_t device, bool* out_known, iree_device_size_t* out_total);
 
+// Queues the device declares, flattened into the queue-affinity bit space.
+//
+// Returns 1 when the HAL spec declares no queue families: every device services
+// at least one queue, and hrx will not hand out a bit it cannot vouch for.
+hrx_status_t hrx_device_query_queue_count(hrx_device_t device,
+                                          uint32_t* out_count);
+
 // Convert iree_status_t to hrx_status_t.
 hrx_status_t hrx_status_from_iree(iree_status_t iree_status);
 
diff --git a/libhrx/src/libhrx/queue_ops.c b/libhrx/src/libhrx/queue_ops.c
index 9c59624..cfe00a0 100644
--- a/libhrx/src/libhrx/queue_ops.c
+++ b/libhrx/src/libhrx/queue_ops.c
@@ -28,12 +28,6 @@ static iree_hal_semaphore_list_t hrx_to_iree_semaphore_list(
 // Max semaphores per direct queue op (stack-allocated arrays).
 #define HRX_MAX_QUEUE_SEMAPHORES 16
 
-static iree_hal_queue_affinity_t hrx_normalize_queue_affinity(
-    hrx_queue_affinity_t affinity) {
-  return affinity == 0 ? IREE_HAL_QUEUE_AFFINITY_ANY
-                       : (iree_hal_queue_affinity_t)affinity;
-}
-
 typedef struct hrx_host_call_thunk_t {
   hrx_host_call_fn_t callback;
   void* user_data;
diff --git a/libhrx/src/libhrx/stream.c b/libhrx/src/libhrx/stream.c
index 9765896..7d5319a 100644
--- a/libhrx/src/libhrx/stream.c
+++ b/libhrx/src/libhrx/stream.c
@@ -13,10 +13,13 @@
 static hrx_status_t hrx_stream_begin_cb(hrx_stream_t stream) {
   if (stream->pending_cb) return hrx_ok_status();
 
+  // The affinity given here also picks the physical device the buffer is
+  // recorded against, so it must be the one hrx_stream_flush() submits with.
   iree_status_t status = iree_hal_command_buffer_create(
       stream->device->hal_device, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
       IREE_HAL_COMMAND_CATEGORY_TRANSFER | IREE_HAL_COMMAND_CATEGORY_DISPATCH,
-      IREE_HAL_QUEUE_AFFINITY_ANY, /*binding_capacity=*/0, &stream->pending_cb);
+      hrx_normalize_queue_affinity(stream->queue_affinity),
+      /*binding_capacity=*/0, &stream->pending_cb);
   if (!iree_status_is_ok(status)) {
     return hrx_status_from_iree(status);
   }
@@ -42,13 +45,43 @@ static iree_status_t hrx_stream_record_ordering_barrier(hrx_stream_t stream) {
       IREE_HAL_EXECUTION_BARRIER_FLAG_NONE, 1, &memory_barrier, 0, NULL);
 }
 
+// Rejects affinity bits naming queues the device does not have. Zero (any) is
+// always valid.
+static hrx_status_t hrx_stream_validate_queue_affinity(
+    hrx_device_t device, hrx_queue_affinity_t affinity) {
+  if (affinity == 0) return hrx_ok_status();
+  uint32_t queue_count = 0;
+  hrx_status_t status = hrx_device_query_queue_count(device, &queue_count);
+  if (!hrx_status_is_ok(status)) return status;
+  const hrx_queue_affinity_t supported =
+      queue_count >= IREE_HAL_MAX_QUEUES
+          ? ~(hrx_queue_affinity_t)0
+          : (((hrx_queue_affinity_t)1 << queue_count) - 1);
+  if ((affinity & ~supported) != 0) {
+    return hrx_make_status(HRX_STATUS_OUT_OF_RANGE,
+                           "queue affinity names a queue the device lacks");
+  }
+  return hrx_ok_status();
+}
+
 hrx_status_t hrx_stream_create(hrx_device_t device, uint32_t flags,
                                hrx_stream_t* stream) {
+  return hrx_stream_create_on_queue(device, flags, /*queue_affinity=*/0,
+                                    stream);
+}
+
+hrx_status_t hrx_stream_create_on_queue(hrx_device_t device, uint32_t flags,
+                                        hrx_queue_affinity_t queue_affinity,
+                                        hrx_stream_t* stream) {
   if (!device || !stream) {
     return hrx_make_status(HRX_STATUS_INVALID_ARGUMENT,
                            "device or stream is NULL");
   }
 
+  hrx_status_t affinity_status =
+      hrx_stream_validate_queue_affinity(device, queue_affinity);
+  if (!hrx_status_is_ok(affinity_status)) return affinity_status;
+
   hrx_stream_s* s = (hrx_stream_s*)calloc(1, sizeof(hrx_stream_s));
   if (!s) {
     return hrx_make_status(HRX_STATUS_OUT_OF_MEMORY,
@@ -59,6 +92,7 @@ hrx_status_t hrx_stream_create(hrx_device_t device, uint32_t flags,
   s->device = device;
   hrx_device_retain(s->device);
   s->flags = flags;
+  s->queue_affinity = queue_affinity;
   s->timepoint = 0;
   s->has_pending_work = false;
   s->pending_cb = NULL;
@@ -139,7 +173,8 @@ hrx_status_t hrx_stream_flush(hrx_stream_t stream) {
   iree_hal_buffer_binding_table_t binding_table =
       iree_hal_buffer_binding_table_empty();
   status = iree_hal_device_queue_execute(
-      stream->device->hal_device, IREE_HAL_QUEUE_AFFINITY_ANY, wait_list,
+      stream->device->hal_device,
+      hrx_normalize_queue_affinity(stream->queue_affinity), wait_list,
       signal_list, stream->pending_cb, binding_table, /*flags=*/0);
   if (!iree_status_is_ok(status)) {
     HRX_RETURN_AND_END_ZONE(z0, hrx_status_from_iree(status));
@@ -213,6 +248,16 @@ hrx_status_t hrx_stream_get_device(hrx_stream_t stream, hrx_device_t* device) {
   return hrx_ok_status();
 }
 
+hrx_status_t hrx_stream_get_queue_affinity(
+    hrx_stream_t stream, hrx_queue_affinity_t* queue_affinity) {
+  if (!stream || !queue_affinity) {
+    return hrx_make_status(HRX_STATUS_INVALID_ARGUMENT,
+                           "stream or queue_affinity is NULL");
+  }
+  *queue_affinity = stream->queue_affinity;
+  return hrx_ok_status();
+}
+
 hrx_status_t hrx_stream_get_timeline_position(hrx_stream_t stream,
                                               hrx_timeline_point_t* position) {
   if (!stream || !position) {
@@ -269,7 +314,8 @@ hrx_status_t hrx_stream_wait_on(hrx_stream_t stream,
   };
 
   iree_status_t iree_status = iree_hal_device_queue_barrier(
-      stream->device->hal_device, IREE_HAL_QUEUE_AFFINITY_ANY, wait_list,
+      stream->device->hal_device,
+      hrx_normalize_queue_affinity(stream->queue_affinity), wait_list,
       signal_list, /*flags=*/0);
   if (!iree_status_is_ok(iree_status)) {
     HRX_RETURN_AND_END_ZONE(z0, hrx_status_from_iree(iree_status));
@@ -461,9 +507,23 @@ hrx_status_t hrx_stream_dispatch(hrx_stream_t stream,
     HRX_RETURN_AND_END_ZONE(z0, hrx_status_from_iree(iree_status));
   }
 
-  iree_status = hrx_stream_record_ordering_barrier(stream);
-  if (!iree_status_is_ok(iree_status)) {
-    HRX_RETURN_AND_END_ZONE(z0, hrx_status_from_iree(iree_status));
+  // Dispatch->dispatch + write->read, not retire->issue / ALL. The old
+  // barrier drained the pipe between every kernel. Decode is hundreds of
+  // launches; that drain was the wall. This is the visibility the next
+  // kernel actually needs.
+  {
+    iree_hal_memory_barrier_t memory_barrier = {
+        .source_scope = IREE_HAL_ACCESS_SCOPE_DISPATCH_WRITE,
+        .target_scope = IREE_HAL_ACCESS_SCOPE_DISPATCH_READ |
+                        IREE_HAL_ACCESS_SCOPE_CONSTANT_READ,
+    };
+    iree_status = iree_hal_command_buffer_execution_barrier(
+        stream->pending_cb, IREE_HAL_EXECUTION_STAGE_DISPATCH,
+        IREE_HAL_EXECUTION_STAGE_DISPATCH, IREE_HAL_EXECUTION_BARRIER_FLAG_NONE,
+        1, &memory_barrier, 0, NULL);
+    if (!iree_status_is_ok(iree_status)) {
+      HRX_RETURN_AND_END_ZONE(z0, hrx_status_from_iree(iree_status));
+    }
   }
 
   stream->has_pending_work = true;
diff --git a/runtime/src/iree/hal/drivers/amdgpu/logical_device.c b/runtime/src/iree/hal/drivers/amdgpu/logical_device.c
index 5355e97..be7c4cf 100644
--- a/runtime/src/iree/hal/drivers/amdgpu/logical_device.c
+++ b/runtime/src/iree/hal/drivers/amdgpu/logical_device.c
@@ -1915,12 +1915,13 @@ static iree_status_t iree_hal_amdgpu_logical_device_create_device_spec(
     const iree_hal_amdgpu_physical_device_t* physical_device =
         logical_device->physical_devices[i];
     if (IREE_UNLIKELY(physical_device->device_ordinal > UINT32_MAX ||
-                      physical_device->host_queue_count > UINT32_MAX)) {
+                      physical_device->host_queue_capacity > UINT32_MAX)) {
       status = iree_make_status(
           IREE_STATUS_OUT_OF_RANGE,
           "AMDGPU device spec physical row out of range: "
           "device_ordinal=%" PRIhsz ", queue_count=%" PRIhsz,
-          physical_device->device_ordinal, physical_device->host_queue_count);
+          physical_device->device_ordinal,
+          physical_device->host_queue_capacity);
       break;
     }
 
@@ -1947,7 +1948,11 @@ static iree_status_t iree_hal_amdgpu_logical_device_create_device_spec(
     physical_params->numa.node_id = physical_device->host_numa_node;
     physical_params->physical_ordinal =
         (uint32_t)physical_device->device_ordinal;
-    physical_params->queue_count = (uint32_t)physical_device->host_queue_count;
+    // Capacity, not count: the spec is built before the frontier is assigned,
+    // and the AQL queues are only created there, so host_queue_count is still
+    // zero here. Capacity is how many queues the device will have.
+    physical_params->queue_count =
+        (uint32_t)physical_device->host_queue_capacity;
     physical_params->compute_unit_count = physical_device->compute_unit_count;
     physical_params->wavefront_size = physical_device->wavefront_size;
     physical_params->maximum_waves_per_compute_unit =