kk: Use residency sets for user allocations

Memory allocated through vkAllocateMemory is now tracked in a
VkDevice level MTLResidencySet that will be committed before
queue submission to ensure residency.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/38505>

src/kosmickrisp/bridge/meson.build

@@ -20,6 +20,7 @@ if host_machine.system() == 'darwin'
     'mtl_heap.m',
     'mtl_library.m',
     'mtl_render_state.m',
+    'mtl_residency_set.m',
     'mtl_sampler.m',
     'mtl_sync.m',
     'mtl_texture.m',
@@ -36,6 +37,7 @@ else
     'stubs/mtl_heap.c',
     'stubs/mtl_library.c',
     'stubs/mtl_render_state.c',
+    'stubs/mtl_residency_set.c',
     'stubs/mtl_sampler.c',
     'stubs/mtl_sync.c',
     'stubs/mtl_texture.c',

src/kosmickrisp/bridge/mtl_bridge.h

@@ -31,6 +31,7 @@
 #include "mtl_heap.h"
 #include "mtl_library.h"
 #include "mtl_render_state.h"
+#include "mtl_residency_set.h"
 #include "mtl_sampler.h"
 #include "mtl_sync.h"
 #include "mtl_texture.h"

src/kosmickrisp/bridge/mtl_residency_set.h

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2025 LunarG, Inc.
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: MIT
+ */
+
+#ifndef MTL_RESIDENCY_SET_H
+#define MTL_RESIDENCY_SET_H 1
+
+#include "mtl_types.h"
+
+mtl_residency_set *mtl_new_residency_set(mtl_device *device);
+void mtl_residency_set_add_allocation(mtl_residency_set *residency_set,
+                                      mtl_allocation *allocation);
+void mtl_residency_set_remove_allocation(mtl_residency_set *residency_set,
+                                         mtl_allocation *allocation);
+void mtl_residency_set_commit(mtl_residency_set *residency_set);
+void mtl_residency_set_request_residency(mtl_residency_set *residency_set);
+void mtl_residency_set_end_residency(mtl_residency_set *residency_set);
+
+#endif /* MTL_RESIDENCY_SET_H */

src/kosmickrisp/bridge/mtl_residency_set.m

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2025 LunarG, Inc.
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "mtl_residency_set.h"
+
+#include <Metal/MTLDevice.h>
+#include <Metal/MTLResidencySet.h>
+
+mtl_residency_set *
+mtl_new_residency_set(mtl_device *device)
+{
+   @autoreleasepool {
+      id<MTLDevice> dev = (id<MTLDevice>)device;
+      MTLResidencySetDescriptor *setDescriptor;
+      setDescriptor = [[MTLResidencySetDescriptor alloc] init];
+      setDescriptor.initialCapacity = 100;
+      NSError *error;
+      id<MTLResidencySet> set = [dev newResidencySetWithDescriptor:setDescriptor
+                                                             error:&error];
+
+      if (error != nil) {
+         fprintf(stderr, "Failed to create MTLResidencySet: %s\n", [error.localizedDescription UTF8String]);
+      }
+
+      return set;
+   }
+}
+
+void
+mtl_residency_set_add_allocation(mtl_residency_set *residency_set,
+                                 mtl_allocation *allocation)
+{
+   @autoreleasepool {
+      id<MTLResidencySet> set = (id<MTLResidencySet>)residency_set;
+      id<MTLAllocation> alloc = (id<MTLAllocation>)allocation;
+      [set addAllocation:alloc];
+   }
+}
+
+void
+mtl_residency_set_remove_allocation(mtl_residency_set *residency_set,
+                                    mtl_allocation *allocation)
+{
+   @autoreleasepool {
+      id<MTLResidencySet> set = (id<MTLResidencySet>)residency_set;
+      id<MTLAllocation> alloc = (id<MTLAllocation>)allocation;
+      [set removeAllocation:alloc];
+   }
+}
+
+void
+mtl_residency_set_commit(mtl_residency_set *residency_set)
+{
+   @autoreleasepool {
+      id<MTLResidencySet> set = (id<MTLResidencySet>)residency_set;
+      [set commit];
+   }
+}
+
+void
+mtl_residency_set_request_residency(mtl_residency_set *residency_set)
+{
+   @autoreleasepool {
+      id<MTLResidencySet> set = (id<MTLResidencySet>)residency_set;
+      [set requestResidency];
+   }
+}
+
+void
+mtl_residency_set_end_residency(mtl_residency_set *residency_set)
+{
+   @autoreleasepool {
+      id<MTLResidencySet> set = (id<MTLResidencySet>)residency_set;
+      [set endResidency];
+   }
+}

src/kosmickrisp/bridge/mtl_types.h

@@ -39,6 +39,8 @@ typedef void mtl_stencil_descriptor;
 typedef void mtl_depth_stencil_descriptor;
 typedef void mtl_depth_stencil_state;
 typedef void mtl_render_pass_attachment_descriptor;
+typedef void mtl_residency_set;
+typedef void mtl_allocation;
 
 /** ENUMS */
 enum mtl_cpu_cache_mode {

src/kosmickrisp/bridge/stubs/mtl_residency_set.c

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2025 LunarG, Inc.
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: MIT
+ */
+
+#include "mtl_residency_set.h"
+
+mtl_residency_set *
+mtl_new_residency_set(mtl_device *device)
+{
+   return NULL;
+}
+
+void
+mtl_residency_set_add_allocation(mtl_residency_set *residency_set,
+                                 mtl_allocation *allocation)
+{
+}
+
+void
+mtl_residency_set_remove_allocation(mtl_residency_set *residency_set,
+                                    mtl_allocation *allocation)
+{
+}
+
+void
+mtl_residency_set_commit(mtl_residency_set *residency_set)
+{
+}
+
+void
+mtl_residency_set_request_residency(mtl_residency_set *residency_set)
+{
+}
+
+void
+mtl_residency_set_end_residency(mtl_residency_set *residency_set)
+{
+}

src/kosmickrisp/vulkan/kk_cmd_dispatch.c

@@ -87,17 +87,6 @@ kk_flush_compute_state(struct kk_cmd_buffer *cmd)
    if (desc->root_dirty)
       kk_upload_descriptor_root(cmd, VK_PIPELINE_BIND_POINT_COMPUTE);
 
-   /* Make user allocated heaps resident */
-   simple_mtx_lock(&dev->user_heap_cache.mutex);
-   if (cmd->encoder->main.user_heap_hash != dev->user_heap_cache.hash) {
-      cmd->encoder->main.user_heap_hash = dev->user_heap_cache.hash;
-      mtl_heap **heaps = util_dynarray_begin(&dev->user_heap_cache.handles);
-      uint32_t count =
-         util_dynarray_num_elements(&dev->user_heap_cache.handles, mtl_heap *);
-      mtl_compute_use_heaps(enc, heaps, count);
-   }
-   simple_mtx_unlock(&dev->user_heap_cache.mutex);
-
    struct kk_bo *root_buffer = desc->root.root_buffer;
    if (root_buffer)
       mtl_compute_set_buffer(enc, root_buffer->map, 0, 0);

src/kosmickrisp/vulkan/kk_cmd_draw.c

@@ -824,18 +824,6 @@ kk_flush_draw_state(struct kk_cmd_buffer *cmd)
    if (desc->root_dirty)
       kk_upload_descriptor_root(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS);
 
-   /* Make user allocated heaps resident */
-   struct kk_device *dev = kk_cmd_buffer_device(cmd);
-   simple_mtx_lock(&dev->user_heap_cache.mutex);
-   if (cmd->encoder->main.user_heap_hash != dev->user_heap_cache.hash) {
-      cmd->encoder->main.user_heap_hash = dev->user_heap_cache.hash;
-      mtl_heap **heaps = util_dynarray_begin(&dev->user_heap_cache.handles);
-      uint32_t count =
-         util_dynarray_num_elements(&dev->user_heap_cache.handles, mtl_heap *);
-      mtl_render_use_heaps(enc, heaps, count);
-   }
-   simple_mtx_unlock(&dev->user_heap_cache.mutex);
-
    struct kk_bo *root_buffer = desc->root.root_buffer;
    if (root_buffer) {
       mtl_set_vertex_buffer(enc, root_buffer->map, 0, 0);

src/kosmickrisp/vulkan/kk_device.c

@@ -233,8 +233,9 @@ kk_CreateDevice(VkPhysicalDevice physicalDevice,
    if (result != VK_SUCCESS)
       goto fail_sampler_heap;
 
-   simple_mtx_init(&dev->user_heap_cache.mutex, mtx_plain);
-   dev->user_heap_cache.handles = UTIL_DYNARRAY_INIT;
+   simple_mtx_init(&dev->user_residency_set.mutex, mtx_plain);
+   dev->user_residency_set.residency_set =
+      mtl_new_residency_set(dev->mtl_handle);
 
    kk_parse_device_environment_options(dev);
 
@@ -270,8 +271,11 @@ kk_DestroyDevice(VkDevice _device, const VkAllocationCallbacks *pAllocator)
    /* Meta first since it may destroy Vulkan objects */
    kk_device_finish_meta(dev);
 
-   util_dynarray_fini(&dev->user_heap_cache.handles);
-   simple_mtx_destroy(&dev->user_heap_cache.mutex);
+   /* Need to end the residency otherwise stopping a capture crashes the
+    * program... */
+   mtl_residency_set_end_residency(dev->user_residency_set.residency_set);
+   mtl_release(dev->user_residency_set.residency_set);
+   simple_mtx_destroy(&dev->user_residency_set.mutex);
    kk_device_finish_lib(dev);
    kk_query_table_finish(dev, &dev->occlusion_queries);
    kk_destroy_sampler_heap(dev, &dev->samplers);
@@ -356,17 +360,26 @@ kk_GetDeviceProcAddr(VkDevice _device, const char *pName)
 void
 kk_device_add_user_heap(struct kk_device *dev, mtl_heap *heap)
 {
-   simple_mtx_lock(&dev->user_heap_cache.mutex);
-   util_dynarray_append(&dev->user_heap_cache.handles, heap);
-   dev->user_heap_cache.hash += 1u;
-   simple_mtx_unlock(&dev->user_heap_cache.mutex);
+   simple_mtx_lock(&dev->user_residency_set.mutex);
+   mtl_residency_set_add_allocation(dev->user_residency_set.residency_set,
+                                    heap);
+   simple_mtx_unlock(&dev->user_residency_set.mutex);
 }
 
 void
 kk_device_remove_user_heap(struct kk_device *dev, mtl_heap *heap)
 {
-   simple_mtx_lock(&dev->user_heap_cache.mutex);
-   util_dynarray_delete_unordered(&dev->user_heap_cache.handles, mtl_heap *,
-                                  heap);
-   simple_mtx_unlock(&dev->user_heap_cache.mutex);
+   simple_mtx_lock(&dev->user_residency_set.mutex);
+   mtl_residency_set_remove_allocation(dev->user_residency_set.residency_set,
+                                       heap);
+   simple_mtx_unlock(&dev->user_residency_set.mutex);
+}
+
+void
+kk_device_make_resources_resident(struct kk_device *dev)
+{
+   simple_mtx_lock(&dev->user_residency_set.mutex);
+   mtl_residency_set_commit(dev->user_residency_set.residency_set);
+   mtl_residency_set_request_residency(dev->user_residency_set.residency_set);
+   simple_mtx_unlock(&dev->user_residency_set.mutex);
 }

src/kosmickrisp/vulkan/kk_device.h

@@ -32,10 +32,9 @@ enum kk_device_lib_pipeline {
    KK_LIB_COUNT,
 };
 
-struct kk_user_heap_cache {
+struct kk_user_residency_set {
    simple_mtx_t mutex;
-   uint32_t hash;
-   struct util_dynarray handles;
+   mtl_residency_set *residency_set;
 };
 
 struct mtl_sampler_packed {
@@ -72,7 +71,7 @@ struct kk_sampler_heap {
 
    struct kk_query_table table;
 
-   /* Map of agx_sampler_packed to hk_rc_sampler */
+   /* Map of mtl_sampler_packed to kk_rc_sampler */
    struct hash_table *ht;
 };
 
@@ -92,7 +91,7 @@ struct kk_device {
 
    /* Track all heaps the user allocated so we can set them all as resident when
     * recording as required by Metal. */
-   struct kk_user_heap_cache user_heap_cache;
+   struct kk_user_residency_set user_residency_set;
 
    mtl_compute_pipeline_state *lib_pipelines[KK_LIB_COUNT];
 
@@ -126,6 +125,7 @@ VkResult kk_device_init_lib(struct kk_device *dev);
 void kk_device_finish_lib(struct kk_device *dev);
 void kk_device_add_user_heap(struct kk_device *dev, mtl_heap *heap);
 void kk_device_remove_user_heap(struct kk_device *dev, mtl_heap *heap);
+void kk_device_make_resources_resident(struct kk_device *dev);
 
 /* Required to create a sampler */
 mtl_sampler *kk_sampler_create(struct kk_device *dev,

src/kosmickrisp/vulkan/kk_queue.c

@@ -31,6 +31,11 @@ kk_queue_submit(struct vk_queue *vk_queue, struct vk_queue_submit *submit)
    if (result != VK_SUCCESS)
       return result;
 
+   /* Ensure any changes to residency are propagated before we submit any work.
+    * All resources should have been allocated before submission. Otherwise,
+    * users are playing with fire. */
+   kk_device_make_resources_resident(dev);
+
    /* Chain with previous sumbission */
    if (queue->wait_fence) {
       util_dynarray_append(&encoder->main.fences, queue->wait_fence);