nvk: Add a new arena data structure

This is about 80% of what nvk_heap does today but with the explicit goal
of just pulling out the growable memory arena part out and separating
all of the heap logic.

Reviewed-by: Mel Henning <mhenning@darkrefraction.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/35327>

src/nouveau/vulkan/meson.build

@@ -45,6 +45,8 @@ nvk_files = files(
   'nvk_indirect_execution_set.h',
   'nvk_instance.c',
   'nvk_instance.h',
+  'nvk_mem_arena.c',
+  'nvk_mem_arena.h',
   'nvk_mme.c',
   'nvk_mme.h',
   'nvk_nir_lower_descriptors.c',

src/nouveau/vulkan/nvk_mem_arena.c

@@ -0,0 +1,185 @@
+/*
+ * Copyright © 2025 Collabora Ltd. and Red Hat Inc.
+ * SPDX-License-Identifier: MIT
+ */
+#include "nvk_mem_arena.h"
+
+#include "nvk_device.h"
+
+#include "util/u_atomic.h"
+
+VkResult
+nvk_mem_arena_init(struct nvk_device *dev, struct nvk_mem_arena *arena,
+                   enum nvkmd_mem_flags mem_flags,
+                   enum nvkmd_mem_map_flags map_flags,
+                   bool contiguous, uint64_t max_size_B)
+{
+   VkResult result;
+
+   memset(arena, 0, sizeof(*arena));
+
+   arena->mem_flags = mem_flags;
+   if (map_flags)
+      arena->mem_flags |= NVKMD_MEM_CAN_MAP;
+   arena->map_flags = map_flags;
+
+   assert(util_is_power_of_two_nonzero64(max_size_B));
+   assert(max_size_B >= NVK_MEM_ARENA_MIN_SIZE);
+   assert(max_size_B <= NVK_MEM_ARENA_MAX_SIZE);
+   arena->max_mem_count =
+      util_logbase2(max_size_B) - NVK_MEM_ARENA_MIN_SIZE_LOG2 + 1;
+   arena->mem_count = 0;
+
+   if (contiguous) {
+      result = nvkmd_dev_alloc_va(dev->nvkmd, &dev->vk.base,
+                                  0 /* va_flags */, 0 /* pte_kind */,
+                                  max_size_B, 0 /* align_B */,
+                                  0 /* fixed_addr */,
+                                  &arena->contig_va);
+      if (result != VK_SUCCESS)
+         return result;
+   }
+
+   simple_mtx_init(&arena->mutex, mtx_plain);
+
+   return VK_SUCCESS;
+}
+
+void
+nvk_mem_arena_finish(struct nvk_device *dev, struct nvk_mem_arena *arena)
+{
+   /* Freeing the VA will unbind all the memory */
+   if (nvk_mem_arena_is_contiguous(arena))
+      nvkmd_va_free(arena->contig_va);
+
+   for (uint32_t mem_idx = 0; mem_idx < arena->mem_count; mem_idx++)
+      nvkmd_mem_unref(arena->mem[mem_idx].mem);
+
+   simple_mtx_destroy(&arena->mutex);
+}
+
+VkResult
+nvk_mem_arena_grow_locked(struct nvk_device *dev, struct nvk_mem_arena *arena,
+                          uint64_t *addr_out, uint64_t *new_mem_size_B_out)
+{
+   const uint32_t mem_count = nvk_mem_arena_mem_count(arena);
+   VkResult result;
+
+   if (mem_count >= arena->max_mem_count) {
+      return vk_errorf(dev, VK_ERROR_OUT_OF_DEVICE_MEMORY,
+                       "Arena has already hit its maximum size");
+   }
+
+   const uint64_t mem_size_B = nvk_mem_arena_mem_size_B(mem_count);
+
+   struct nvkmd_mem *mem;
+   if (nvk_mem_arena_is_mapped(arena)) {
+      result = nvkmd_dev_alloc_mapped_mem(dev->nvkmd, &dev->vk.base,
+                                          mem_size_B, 0, arena->mem_flags,
+                                          arena->map_flags, &mem);
+   } else {
+      result = nvkmd_dev_alloc_mem(dev->nvkmd, &dev->vk.base,
+                                   mem_size_B, 0, arena->mem_flags, &mem);
+   }
+   if (result != VK_SUCCESS)
+      return result;
+
+   uint64_t addr;
+   if (nvk_mem_arena_is_contiguous(arena)) {
+      const uint64_t mem_offset_B =
+         nvk_contiguous_mem_arena_mem_offset_B(mem_count);
+      result = nvkmd_va_bind_mem(arena->contig_va, &dev->vk.base,
+                                 mem_offset_B, mem, 0, mem_size_B);
+      if (result != VK_SUCCESS) {
+         nvkmd_mem_unref(mem);
+         return result;
+      }
+      addr = arena->contig_va->addr + mem_offset_B;
+   } else {
+      addr = mem->va->addr;
+   }
+
+   arena->mem[mem_count] = (struct nvk_arena_mem) {
+      .mem = mem,
+      .addr = addr,
+   };
+   if (p_atomic_xchg(&arena->mem_count, mem_count + 1) != mem_count) {
+      return vk_errorf(dev, VK_ERROR_UNKNOWN,
+                       "Raced in arena grow.  This is an internal driver bug "
+                       "and things are now in an unknown state.");
+   }
+
+   if (addr_out != NULL)
+      *addr_out = addr;
+   if (new_mem_size_B_out != NULL)
+      *new_mem_size_B_out = mem_size_B;
+
+   return VK_SUCCESS;
+}
+
+static uint32_t
+nvk_mem_arena_find_mem_by_addr(const struct nvk_mem_arena *arena, uint64_t addr)
+{
+   if (nvk_mem_arena_is_contiguous(arena)) {
+      assert(addr >= arena->contig_va->addr);
+      assert(addr < arena->contig_va->addr + nvk_mem_arena_size_B(arena));
+      const uint64_t arena_offset_B = addr - arena->contig_va->addr;
+      return nvk_contiguous_mem_arena_find_mem_by_offset(arena, arena_offset_B);
+   } else {
+      const uint32_t mem_count = nvk_mem_arena_mem_count(arena);
+
+      /* Start at the top because, given a random address, there's a 50%
+       * liklihood that it's in the largest mem.
+       */
+      for (int32_t mem_idx = mem_count - 1; mem_idx >= 0; mem_idx--) {
+         const struct nvk_arena_mem *mem = &arena->mem[mem_idx];
+         const uint64_t mem_size_B = nvk_mem_arena_mem_size_B(mem_idx);
+         if (addr >= mem->addr && addr < mem->addr + mem_size_B)
+            return mem_idx;
+      }
+
+      unreachable("Not an arena address");
+   }
+}
+
+void *
+nvk_mem_arena_map(const struct nvk_mem_arena *arena,
+                  uint64_t addr, size_t map_range_B)
+{
+   uint32_t mem_idx = nvk_mem_arena_find_mem_by_addr(arena, addr);
+   const struct nvk_arena_mem *mem = &arena->mem[mem_idx];
+
+   assert(addr >= mem->addr);
+   const uint64_t mem_offset_B = addr - mem->addr;
+
+   ASSERTED const uint64_t mem_size_B = nvk_mem_arena_mem_size_B(mem_idx);
+   assert(mem_offset_B + map_range_B <= mem_size_B);
+
+   return mem->mem->map + mem_offset_B;
+}
+
+void
+nvk_mem_arena_copy_to_gpu(const struct nvk_mem_arena *arena,
+                          uint64_t dst_addr, const void *src, size_t size_B)
+{
+   assert(nvk_mem_arena_is_mapped(arena));
+
+   while (size_B) {
+      uint32_t mem_idx = nvk_mem_arena_find_mem_by_addr(arena, dst_addr);
+      const struct nvk_arena_mem *mem = &arena->mem[mem_idx];
+      const uint64_t mem_size_B = nvk_mem_arena_mem_size_B(mem_idx);
+
+      assert(dst_addr >= mem->addr);
+      const uint64_t mem_offset_B = dst_addr - mem->addr;
+      assert(mem_offset_B < mem_size_B);
+
+      /* We can't copy past the end of the mem */
+      const size_t copy_size_B = MIN2(size_B, mem_size_B - mem_offset_B);
+
+      memcpy(mem->mem->map + mem_offset_B, src, copy_size_B);
+
+      dst_addr += copy_size_B;
+      src += copy_size_B;
+      size_B -= copy_size_B;
+   }
+}

src/nouveau/vulkan/nvk_mem_arena.h

@@ -0,0 +1,195 @@
+/*
+ * Copyright © 2025 Collabora Ltd. and Red Hat Inc.
+ * SPDX-License-Identifier: MIT
+ */
+#ifndef NVK_MEM_ARENA_H
+#define NVK_MEM_ARENA_H
+
+#include "nvk_private.h"
+
+#include "util/simple_mtx.h"
+#include "util/u_atomic.h"
+#include "nvkmd/nvkmd.h"
+
+struct nvk_device;
+
+#define NVK_MEM_ARENA_MIN_SIZE_LOG2 16
+#define NVK_MEM_ARENA_MAX_SIZE_LOG2 32
+#define NVK_MEM_ARENA_MIN_SIZE (1ull << NVK_MEM_ARENA_MIN_SIZE_LOG2)
+#define NVK_MEM_ARENA_MAX_SIZE (1ull << NVK_MEM_ARENA_MAX_SIZE_LOG2)
+#define NVK_MEM_ARENA_MAX_MEM_COUNT (NVK_MEM_ARENA_MAX_SIZE_LOG2 - \
+                                     NVK_MEM_ARENA_MIN_SIZE_LOG2 + 1)
+
+static inline uint64_t
+nvk_mem_arena_mem_size_B(uint32_t mem_idx)
+{
+   /* The first two are both NVK_MEM_ARENA_MIN_SIZE and then we double after
+    * that.
+    */
+   return mem_idx == 0 ? NVK_MEM_ARENA_MIN_SIZE
+                       : ((NVK_MEM_ARENA_MIN_SIZE >> 1) << mem_idx);
+}
+
+static inline uint64_t
+nvk_contiguous_mem_arena_mem_offset_B(uint32_t mem_idx)
+{
+   /* The first one is at offset 0 and offset_B == size_B after that. */
+   return mem_idx == 0 ? 0 : ((NVK_MEM_ARENA_MIN_SIZE >> 1) << mem_idx);
+}
+
+struct nvk_arena_mem {
+   struct nvkmd_mem *mem;
+   uint64_t addr;
+};
+
+/** A growable pool of GPU memory
+ *
+ * This data structure does not provide any special allocation or address
+ * management.  It just provides the growable memory area.  Users of this
+ * struct are expected to wrap it in something which provides the desired
+ * allocation structure on top of it.
+ */
+struct nvk_mem_arena {
+   enum nvkmd_mem_flags mem_flags;
+   enum nvkmd_mem_map_flags map_flags;
+
+   /** Used to lock this arena
+    *
+    * This lock MUST be held when calling nvk_mem_arena_grow_locked().
+    */
+   simple_mtx_t mutex;
+
+   /* VA for contiguous heaps, NULL otherwise */
+   struct nvkmd_va *contig_va;
+
+   /* Maximum mem_count for this arena */
+   uint32_t max_mem_count;
+
+   /* Number of nvk_arena_mem.  This value is an atomic which is only ever
+    * increased and only after the new nvk_arena_mem has been populated so
+    * it's always safe to fetch it and then look at mem[i] for i < mem_count
+    * without taking the lock.
+    */
+   uint32_t mem_count;
+
+   struct nvk_arena_mem mem[NVK_MEM_ARENA_MAX_MEM_COUNT];
+};
+
+VkResult nvk_mem_arena_init(struct nvk_device *dev,
+                            struct nvk_mem_arena *arena,
+                            enum nvkmd_mem_flags mem_flags,
+                            enum nvkmd_mem_map_flags map_flags,
+                            bool contiguous, uint64_t max_size_B);
+
+void nvk_mem_arena_finish(struct nvk_device *dev,
+                          struct nvk_mem_arena *arena);
+
+static inline uint64_t
+nvk_mem_arena_is_contiguous(const struct nvk_mem_arena *arena)
+{
+   return arena->contig_va != NULL;
+}
+
+static inline uint64_t
+nvk_mem_arena_is_mapped(const struct nvk_mem_arena *arena)
+{
+   return arena->map_flags != 0;
+}
+
+/* After calling this function, it is safe to look at any arena->mem[i]
+ * where i is less than the returned count.
+ */
+static inline uint32_t
+nvk_mem_arena_mem_count(const struct nvk_mem_arena *arena)
+{
+   return p_atomic_read(&arena->mem_count);
+}
+
+static inline uint64_t
+nvk_mem_arena_size_B(const struct nvk_mem_arena *arena)
+{
+   uint32_t mem_count = nvk_mem_arena_mem_count(arena);
+   return nvk_contiguous_mem_arena_mem_offset_B(mem_count);
+}
+
+static inline uint64_t
+nvk_contiguous_mem_arena_base_address(const struct nvk_mem_arena *arena)
+{
+   assert(nvk_mem_arena_is_contiguous(arena));
+   return arena->contig_va->addr;
+}
+
+/** Grows the arena by doubling its size
+ *
+ * arena->mutex MUST be held when calling this function
+ */
+VkResult nvk_mem_arena_grow_locked(struct nvk_device *dev,
+                                   struct nvk_mem_arena *arena,
+                                   uint64_t *addr_out,
+                                   uint64_t *new_mem_size_B_out);
+
+static inline uint32_t
+nvk_contiguous_mem_arena_find_mem_by_offset(const struct nvk_mem_arena *arena,
+                                            uint64_t arena_offset_B)
+{
+   assert(nvk_mem_arena_is_contiguous(arena));
+
+   assert((arena_offset_B >> (NVK_MEM_ARENA_MIN_SIZE_LOG2 - 1)) <= UINT32_MAX);
+   const uint32_t mem_idx =
+      util_logbase2(arena_offset_B >> (NVK_MEM_ARENA_MIN_SIZE_LOG2 - 1));
+
+   assert(mem_idx < nvk_mem_arena_mem_count(arena));
+
+   assert(arena_offset_B >= nvk_contiguous_mem_arena_mem_offset_B(mem_idx));
+   assert(arena_offset_B < nvk_contiguous_mem_arena_mem_offset_B(mem_idx + 1));
+   assert(nvk_contiguous_mem_arena_mem_offset_B(mem_idx + 1) ==
+            nvk_contiguous_mem_arena_mem_offset_B(mem_idx) +
+            nvk_mem_arena_mem_size_B(mem_idx));
+
+   ASSERTED const struct nvk_arena_mem *mem = &arena->mem[mem_idx];
+   ASSERTED uint64_t addr = arena->contig_va->addr + arena_offset_B;
+   assert(addr >= mem->addr);
+   assert(addr < mem->addr + nvk_mem_arena_mem_size_B(mem_idx));
+
+   return mem_idx;
+}
+
+/** An optimized version of `nvk_mem_arena_map()` for contiguous arenas
+ *
+ * See `nvk_mem_arena_map()` for restrictions on the mapped pointer.  Unlike
+ * `nvk_mem_arena_map()`, this takes an offset instead of an address.
+ */
+static inline void *
+nvk_contiguous_mem_arena_map_offset(const struct nvk_mem_arena *arena,
+                                    uint64_t arena_offset_B,
+                                    ASSERTED size_t map_range_B)
+{
+   assert(nvk_mem_arena_is_mapped(arena));
+   const uint32_t mem_idx =
+      nvk_contiguous_mem_arena_find_mem_by_offset(arena, arena_offset_B);
+   const uint64_t mem_offset_B =
+      arena_offset_B - nvk_contiguous_mem_arena_mem_offset_B(mem_idx);
+
+   ASSERTED const uint64_t mem_size_B = nvk_mem_arena_mem_size_B(mem_idx);
+   assert(mem_offset_B + map_range_B <= mem_size_B);
+
+   return arena->mem[mem_idx].mem->map + mem_offset_B;
+}
+
+/** Returns a pointer to the CPU map of the given GPU address
+ *
+ * While nvk_mem_arena can ensure contiguous GPU addresses if requested (see
+ * nvk_mem_arena_init()), CPU addresses may not be contiguous.  However, if
+ * `dst_addr` is aligned to some power of two alignment align_B and
+ * `align_B <= NVK_MEM_ARENA_MIN_SIZE`, then the returned pointer will be
+ * valid for at least `align_B` bytes.  For larger or unaligned allocations,
+ * use nvk_mem_arena_copy_to_gpu() instead.
+ */
+void *nvk_mem_arena_map(const struct nvk_mem_arena *arena,
+                        uint64_t addr, size_t map_range_B);
+
+void nvk_mem_arena_copy_to_gpu(const struct nvk_mem_arena *arena,
+                               uint64_t dst_addr,
+                               const void *src, size_t size_B);
+
+#endif /* NVK_MEM_ARENA_H */