honeykrisp: validate the existing Aurora capture/replay heap

src/asahi/lib/agx_device.c

@@ -1319,16 +1319,14 @@ agx_open_device(void *memctx, struct agx_device *dev)
       goto fail;
    }
 
-   /* Keep captured addresses unavailable to ordinary allocations between
-    * capture and replay. Reserve after all fixed firmware/alias holes, and
-    * for native AGX3 and M5 using the common allocator; other winsyses own their VA. */
-   if (agx_chip_native_agx3(dev->chip) || agx_chip_g17_m5(dev->chip)) {
-      dev->replay_heap_initialized = agx_capture_replay_heap_init(
-         &dev->main_heap, &dev->replay_heap, user_size);
-      if (!dev->replay_heap_initialized) {
-         fprintf(stderr, "Failed to reserve capture/replay GPU VA\n");
-         goto fail;
-      }
+   /* Reserve captured addresses after fixed firmware/alias holes so ordinary
+    * allocations cannot take them between capture and replay. */
+   dev->replay_heap_initialized = agx_capture_replay_heap_init(
+      &dev->main_heap, &dev->replay_heap, user_size, &dev->replay_start,
+      &dev->replay_size);
+   if (!dev->replay_heap_initialized) {
+      fprintf(stderr, "Failed to reserve capture/replay GPU VA\n");
+      goto fail;
    }
 
    dev->vm_id = vm_create.vm_id;

src/asahi/lib/agx_device.h

@@ -190,6 +190,9 @@ struct agx_device {
    struct util_vma_heap replay_heap;
    /* Heap selection must stay stable when replay_heap.free_size reaches 0. */
    bool replay_heap_initialized;
+   /* The capture/replay range, for validating a replayed address. */
+   uint64_t replay_start;
+   uint64_t replay_size;
    /* Native command-pool BO aliases with representable 32-bit context offsets.
     * Protected by vma_lock like the other heaps. */
    struct util_vma_heap g17p_vdm_alias_heap;
@@ -534,6 +537,10 @@ struct agx_va *agx_va_alloc(struct agx_device *dev, uint64_t size_B,
                             uint64_t fixed_va);
 void agx_va_free(struct agx_device *dev, struct agx_va *va, bool unbind);
 
+/* Whether a fixed range, with its guard, lies inside the capture/replay range. */
+bool agx_va_capture_valid(const struct agx_device *dev, uint64_t addr,
+                          uint64_t size_B, uint64_t align_B);
+
 static inline struct drm_asahi_cmd_header
 agx_cmd_header(bool compute, uint16_t barrier_vdm, uint16_t barrier_cdm)
 {

src/asahi/lib/agx_va.c

@@ -18,13 +18,42 @@ agx_vma_heap(struct agx_device *dev, enum agx_va_flags flags)
    return &dev->main_heap;
 }
 
+bool
+agx_va_capture_valid(const struct agx_device *dev, uint64_t addr,
+                     uint64_t size_B, uint64_t align_B)
+{
+   if (!dev->replay_heap_initialized)
+      return false;
+
+   const uint64_t end = dev->replay_start + dev->replay_size;
+
+   /* Written to never wrap: size_B comes from the application. */
+   if (size_B > dev->replay_size || dev->guard_size > dev->replay_size - size_B)
+      return false;
+   if (addr < dev->replay_start || addr >= end)
+      return false;
+   if (align_B == 0 || addr % align_B != 0)
+      return false;
+
+   return addr - dev->replay_start <=
+          dev->replay_size - (size_B + dev->guard_size);
+}
+
 struct agx_va *
 agx_va_alloc(struct agx_device *dev, uint64_t size_B, uint64_t align_B,
              enum agx_va_flags flags, uint64_t fixed_va)
 {
    assert((fixed_va != 0) == !!(flags & AGX_VA_FIXED));
    assert((fixed_va % align_B) == 0);
 
+   /* A fixed capture/replay request outside the range is not a request the
+    * heap can answer: util_vma_heap_alloc_addr asserts on a range that wraps.
+    * The caller checks agx_va_capture_valid to tell it from a range in use.
+    */
+   if ((flags & AGX_VA_FIXED) && (flags & AGX_VA_CAPTURE_REPLAY) &&
+       !agx_va_capture_valid(dev, fixed_va, size_B, align_B))
+      return NULL;
+
    /* All allocations need a guard at the end to prevent overreads.
     *
     * TODO: Even with soft fault?

src/asahi/lib/agx_va_runtime_hole.c

@@ -43,7 +43,8 @@ agx_g17p_vdm_alias_heap_init(struct util_vma_heap *heap)
 bool
 agx_capture_replay_heap_init(struct util_vma_heap *main_heap,
                              struct util_vma_heap *replay_heap,
-                             uint64_t user_size)
+                             uint64_t user_size, uint64_t *out_base,
+                             uint64_t *out_size)
 {
    const uint64_t alignment = UINT64_C(1) << 30;
    const uint64_t size = MIN2(user_size / 8, UINT64_C(64) << 30) &
@@ -55,6 +56,8 @@ agx_capture_replay_heap_init(struct util_vma_heap *main_heap,
    if (!base)
       return false;
    util_vma_heap_init(replay_heap, base, size);
+   *out_base = base;
+   *out_size = size;
    return true;
 }
 

src/asahi/lib/agx_va_runtime_hole.h

@@ -136,7 +136,8 @@ bool agx_va_overlaps_runtime_graph_hole(enum agx_chip chip, uint64_t addr,
 bool agx_g17p_vdm_alias_heap_init(struct util_vma_heap *heap);
 bool agx_capture_replay_heap_init(struct util_vma_heap *main_heap,
                                   struct util_vma_heap *replay_heap,
-                                  uint64_t user_size);
+                                  uint64_t user_size, uint64_t *out_base,
+                                  uint64_t *out_size);
 bool agx_g17p_vdm_alias_range_valid(uint64_t addr, uint64_t size);
 bool agx_g17p_vdm_alias_compact(uint64_t addr, uint32_t *offset);
 

src/asahi/vulkan/hk_buffer.c

@@ -5,6 +5,8 @@
  * SPDX-License-Identifier: MIT
  */
 #include "hk_buffer.h"
+
+#include <inttypes.h>
 #include "hk_acceleration_structure.h"
 
 #include <errno.h>
@@ -209,15 +211,38 @@ hk_CreateBuffer(VkDevice device, const VkBufferCreateInfo *pCreateInfo,
          if (dev->dev.replay_heap_initialized)
             flags |= AGX_VA_CAPTURE_REPLAY;
          bda_fixed_addr = hk_get_bda_replay_addr(pCreateInfo);
-         if (bda_fixed_addr != 0)
+         if (bda_fixed_addr != 0) {
+            /* An address that cannot be in the range is not available, which
+             * is not the same as running out of address space.
+             */
+            if (dev->dev.replay_heap_initialized &&
+                !agx_va_capture_valid(&dev->dev, bda_fixed_addr, vma_size_B,
+                                      alignment)) {
+               vk_buffer_destroy(&dev->vk, pAllocator, &buffer->vk);
+               return vk_errorf(dev, VK_ERROR_INVALID_OPAQUE_CAPTURE_ADDRESS,
+                                "Capture address 0x%" PRIx64
+                                " is outside the capture range",
+                                bda_fixed_addr);
+            }
             flags |= AGX_VA_FIXED;
+         }
       }
 
       buffer->va =
          agx_va_alloc(&dev->dev, vma_size_B, alignment, flags, bda_fixed_addr);
 
       if (!buffer->va) {
          vk_buffer_destroy(&dev->vk, pAllocator, &buffer->vk);
+
+         /* A fixed request has one place to go. If it fails the address is in
+          * use. A request without an address fails when the range is full.
+          */
+         if (bda_fixed_addr != 0) {
+            return vk_errorf(dev, VK_ERROR_INVALID_OPAQUE_CAPTURE_ADDRESS,
+                             "Capture address 0x%" PRIx64 " is in use",
+                             bda_fixed_addr);
+         }
+
          return vk_errorf(dev, VK_ERROR_OUT_OF_DEVICE_MEMORY,
                           "Sparse VMA allocation failed");
       }

src/asahi/vulkan/meson.build

@@ -504,6 +504,12 @@ if with_tests
       '-Wl,--wrap=os_get_available_system_memory',
       '-Wl,--wrap=os_get_total_physical_memory']),
     suite : ['asahi'])
+  test('hk_va', executable('hk_va_test',
+    files('tests/test-hk-va.c'),
+    include_directories : [inc_gallium, inc_gallium_aux, inc_include, inc_src,
+      inc_asahi, inc_virtio_gpu, inc_virtio_vdrm, include_directories('.')],
+    dependencies : [hk_deps], link_with : [libhk],
+    link_args : [ld_args_gc_sections]), suite : ['asahi'])
   test('hk_tes_query', executable('hk_tes_query_test', files('tests/test-hk-tes-query.c'),
     include_directories : [inc_gallium, inc_gallium_aux, inc_include, inc_src,
       inc_asahi, inc_virtio_gpu, inc_virtio_vdrm, include_directories('.')],

src/asahi/vulkan/tests/test-hk-va.c

@@ -0,0 +1,141 @@
+/*
+ * Copyright 2026 Joshua Warren
+ * SPDX-License-Identifier: MIT
+ */
+
+#include <stdio.h>
+#include <stdlib.h>
+#include "agx_bo.h"
+#include "agx_device.h"
+#include "agx_va_runtime_hole.h"
+
+/* A user window like the one agx_open_device builds: 1 TiB of user VA. */
+#define USER_START 0x2000000000ull
+#define USER_SIZE  0x10000000000ull
+#define PAGE       16384ull
+
+static unsigned failures;
+
+#define CHECK(cond)                                                            \
+   do {                                                                        \
+      if (!(cond)) {                                                           \
+         fprintf(stderr, "%s:%d: CHECK failed: %s\n", __FILE__, __LINE__,      \
+                 #cond);                                                       \
+         failures++;                                                           \
+      }                                                                        \
+   } while (0)
+
+static struct agx_device *
+open_va(void)
+{
+   struct agx_device *dev = calloc(1, sizeof(*dev));
+   simple_mtx_init(&dev->vma_lock, mtx_plain);
+   dev->guard_size = PAGE;
+   util_vma_heap_init(&dev->main_heap, USER_START, USER_SIZE);
+   dev->replay_heap_initialized = agx_capture_replay_heap_init(
+      &dev->main_heap, &dev->replay_heap, USER_SIZE, &dev->replay_start,
+      &dev->replay_size);
+   return dev;
+}
+
+static void
+close_va(struct agx_device *dev)
+{
+   util_vma_heap_finish(&dev->replay_heap);
+   util_vma_heap_finish(&dev->main_heap);
+   simple_mtx_destroy(&dev->vma_lock);
+   free(dev);
+}
+
+static bool
+inside(const struct agx_device *dev, const struct agx_va *va)
+{
+   return va->addr < dev->replay_start + dev->replay_size &&
+          dev->replay_start < va->addr + va->size_B;
+}
+
+int
+main(void)
+{
+   struct agx_device *dev = open_va();
+   CHECK(dev->replay_heap_initialized);
+   const uint64_t start = dev->replay_start, size = dev->replay_size;
+   const uint64_t end = start + size, sz = 4 * PAGE;
+
+   /* The range is carved out of the user window and is aligned. */
+   CHECK(size == 64ull << 30);
+   CHECK(start >= USER_START && end <= USER_START + USER_SIZE);
+   CHECK(start % (1ull << 30) == 0);
+
+   /* Valid and invalid replay addresses. The guard counts as part of the
+    * allocation, and nothing may wrap. */
+   CHECK(agx_va_capture_valid(dev, start, sz, PAGE));
+   CHECK(agx_va_capture_valid(dev, end - sz - dev->guard_size, sz, PAGE));
+   CHECK(!agx_va_capture_valid(dev, end - sz, sz, PAGE));
+   CHECK(!agx_va_capture_valid(dev, end, sz, PAGE));
+   CHECK(!agx_va_capture_valid(dev, start - PAGE, sz, PAGE));
+   CHECK(!agx_va_capture_valid(dev, 0, sz, PAGE));
+   CHECK(!agx_va_capture_valid(dev, UINT64_MAX - PAGE + 1, sz, PAGE));
+   CHECK(!agx_va_capture_valid(dev, start + 4096, sz, PAGE));
+   CHECK(!agx_va_capture_valid(dev, start, UINT64_MAX - 100, PAGE));
+   CHECK(!agx_va_capture_valid(dev, start, size, PAGE));
+   CHECK(!agx_va_capture_valid(dev, start, sz, 0));
+
+   /* A fixed request inside the range works once, then the address is in use.
+    * One outside the range is refused without asserting. */
+   struct agx_va *a = agx_va_alloc(dev, sz, PAGE,
+                                   AGX_VA_CAPTURE_REPLAY | AGX_VA_FIXED,
+                                   start + 8 * PAGE);
+   CHECK(a && a->addr == start + 8 * PAGE);
+   CHECK(!agx_va_alloc(dev, sz, PAGE, AGX_VA_CAPTURE_REPLAY | AGX_VA_FIXED,
+                       start + 8 * PAGE));
+   CHECK(!agx_va_alloc(dev, sz, PAGE, AGX_VA_CAPTURE_REPLAY | AGX_VA_FIXED,
+                       end - PAGE));
+   CHECK(!agx_va_alloc(dev, sz, PAGE, AGX_VA_CAPTURE_REPLAY | AGX_VA_FIXED,
+                       UINT64_MAX - PAGE + 1));
+   agx_va_free(dev, a, false);
+   a = agx_va_alloc(dev, sz, PAGE, AGX_VA_CAPTURE_REPLAY | AGX_VA_FIXED,
+                    start + 8 * PAGE);
+   CHECK(a && a->addr == start + 8 * PAGE);
+   agx_va_free(dev, a, false);
+
+   /* Ordinary allocations never land in the range, even when the ordinary
+    * heap runs dry. The ordinary heap holds 960 GiB, so 4 GiB steps run it
+    * dry after about 240 successes; the cap only stops a runaway loop. */
+   enum { MAX_FILL = 512 };
+   struct agx_va *filled[MAX_FILL];
+   unsigned n_filled = 0;
+   while (n_filled < MAX_FILL) {
+      struct agx_va *va = agx_va_alloc(dev, 4ull << 30, PAGE, 0, 0);
+      if (!va)
+         break;
+      CHECK(!inside(dev, va));
+      filled[n_filled++] = va;
+   }
+   CHECK(n_filled >= 200 && n_filled < MAX_FILL);
+
+   /* With the ordinary heap dry, a small ordinary request either fails or
+    * lands outside the range, and an oversized one fails. */
+   struct agx_va *small = agx_va_alloc(dev, sz, PAGE, 0, 0);
+   CHECK(!small || !inside(dev, small));
+   agx_va_free(dev, small, false);
+   struct agx_va *big = agx_va_alloc(dev, 4 * USER_SIZE, PAGE, 0, 0);
+   CHECK(!big);
+
+   /* Replay buffers without an address come from the range, dry heap or not. */
+   struct agx_va *r = agx_va_alloc(dev, sz, PAGE, AGX_VA_CAPTURE_REPLAY, 0);
+   CHECK(r && inside(dev, r));
+   agx_va_free(dev, r, false);
+
+   for (unsigned i = 0; i < n_filled; ++i)
+      agx_va_free(dev, filled[i], false);
+
+   close_va(dev);
+
+   if (failures) {
+      fprintf(stderr, "%u checks failed\n", failures);
+      return 1;
+   }
+   puts("PASS: capture/replay address range");
+   return 0;
+}