From d1ce6c15ac341c92359456fd36c0ef06102a0b09 Mon Sep 17 00:00:00 2001
From: Jakob Bornecrantz <jakob@collabora.com>
Date: Wed, 4 Oct 2023 15:32:46 +0100
Subject: [PATCH] c/render: Use sub-allocation for mesh shader UBOs

---
 src/xrt/compositor/main/comp_renderer.c      | 41 ++++++++++++++++++--
 src/xrt/compositor/render/render_gfx.c       | 26 +++----------
 src/xrt/compositor/render/render_interface.h |  5 +--
 3 files changed, 46 insertions(+), 26 deletions(-)

diff --git a/src/xrt/compositor/main/comp_renderer.c b/src/xrt/compositor/main/comp_renderer.c
index e5f42ba9d..6fdab768a 100644
--- a/src/xrt/compositor/main/comp_renderer.c
+++ b/src/xrt/compositor/main/comp_renderer.c
@@ -806,6 +806,36 @@ do_gfx_mesh(struct comp_renderer *r,
             VkImageView src_image_views[2],
             const struct xrt_normalized_rect src_norm_rects[2])
 {
+	struct vk_bundle *vk = &r->c->base.vk;
+	VkResult ret;
+
+	/*
+	 * Reserve UBOs and fill in any data a head of time, if we ever want to
+	 * copy UBO data this lets us do that easily write a copy command before
+	 * the other gfx commands.
+	 */
+	struct render_sub_alloc ubos[2];
+	for (uint32_t i = 0; i < 2; i++) {
+
+		struct render_gfx_mesh_ubo_data data = {
+		    .vertex_rot = vertex_rots[i],
+		    .post_transform = src_norm_rects[i],
+		};
+
+		ret = render_sub_alloc_ubo_alloc_and_write( //
+		    vk,                                     // vk_bundle
+		    &rr->ubo_tracker,                       // rsat
+		    &data,                                  // ptr
+		    sizeof(data),                           // size
+		    &ubos[i]);                              // out_rsa
+		VK_CHK_WITH_GOTO(ret, "render_sub_alloc_ubo_alloc_and_write", err_no_memory);
+	}
+
+
+	/*
+	 * Do command writing here.
+	 */
+
 	render_gfx_begin_target( //
 	    rr,                  //
 	    rtr);                //
@@ -819,15 +849,20 @@ do_gfx_mesh(struct comp_renderer *r,
 		render_gfx_distortion(   //
 		    rr,                  //
 		    i,                   //
-		    &vertex_rots[i],     //
+		    &ubos[i],            //
 		    src_samplers[i],     //
-		    src_image_views[i],  //
-		    &src_norm_rects[i]); //
+		    src_image_views[i]); //
 
 		render_gfx_end_view(rr);
 	}
 
 	render_gfx_end_target(rr);
+
+	return;
+
+err_no_memory:
+	// Allocator reset at end of frame, nothing to clean up.
+	VK_ERROR(vk, "Could not allocate all UBOs for frame, that's really strange and shouldn't happen!");
 }
 
 /*!
diff --git a/src/xrt/compositor/render/render_gfx.c b/src/xrt/compositor/render/render_gfx.c
index eebeea531..6bc505eec 100644
--- a/src/xrt/compositor/render/render_gfx.c
+++ b/src/xrt/compositor/render/render_gfx.c
@@ -356,6 +356,7 @@ update_mesh_discriptor_set(struct vk_bundle *vk,
                            VkImageView image_view,
                            uint32_t ubo_binding,
                            VkBuffer buffer,
+                           VkDeviceSize offset,
                            VkDeviceSize size,
                            VkDescriptorSet descriptor_set)
 {
@@ -367,7 +368,7 @@ update_mesh_discriptor_set(struct vk_bundle *vk,
 
 	VkDescriptorBufferInfo buffer_info = {
 	    .buffer = buffer,
-	    .offset = 0,
+	    .offset = offset,
 	    .range = size,
 	};
 
@@ -714,12 +715,8 @@ render_gfx_end_view(struct render_gfx *rr)
 }
 
 void
-render_gfx_distortion(struct render_gfx *rr,
-                      uint32_t view_index,
-                      const struct xrt_matrix_2x2 *vertex_rot,
-                      VkSampler sampler,
-                      VkImageView image_view,
-                      const struct xrt_normalized_rect *src_rect)
+render_gfx_distortion(
+    struct render_gfx *rr, uint32_t view_index, struct render_sub_alloc *ubo, VkSampler sampler, VkImageView image_view)
 {
 	struct vk_bundle *vk = vk_from_rr(rr);
 	struct render_resources *r = rr->r;
@@ -728,20 +725,8 @@ render_gfx_distortion(struct render_gfx *rr,
 
 	struct render_gfx_view *v = &rr->views[view_index];
 
-	struct render_buffer *ubo = &r->mesh.ubos[view_index];
 	VkDescriptorSet descriptor_set = v->mesh.descriptor_set;
 
-	/*
-	 * UBO data.
-	 */
-
-	struct render_gfx_mesh_ubo_data data = {
-	    .vertex_rot = *vertex_rot,
-	    .post_transform = *src_rect,
-	};
-
-	render_buffer_write(vk, ubo, &data, sizeof(struct render_gfx_mesh_ubo_data));
-
 
 	/*
 	 * Descriptors and pipeline.
@@ -754,7 +739,8 @@ render_gfx_distortion(struct render_gfx *rr,
 	    image_view,             // image_view
 	    r->mesh.ubo_binding,    // ubo_binding
 	    ubo->buffer,            // buffer
-	    VK_WHOLE_SIZE,          // size
+	    ubo->offset,            // offset
+	    ubo->size,              // size
 	    descriptor_set);        // descriptor_set
 
 	VkDescriptorSet descriptor_sets[1] = {descriptor_set};
diff --git a/src/xrt/compositor/render/render_interface.h b/src/xrt/compositor/render/render_interface.h
index dbc96d2d2..65a660e0c 100644
--- a/src/xrt/compositor/render/render_interface.h
+++ b/src/xrt/compositor/render/render_interface.h
@@ -898,10 +898,9 @@ render_gfx_end_view(struct render_gfx *rr);
 void
 render_gfx_distortion(struct render_gfx *rr,
                       uint32_t view_index,
-                      const struct xrt_matrix_2x2 *vertex_rot,
+                      struct render_sub_alloc *ubo,
                       VkSampler sampler,
-                      VkImageView image_view,
-                      const struct xrt_normalized_rect *src_rect);
+                      VkImageView image_view);
 
 /*!
  * @}