kk: Add grid struct for dispatches for convenience

Signed-off-by: Aitor Camacho <aitor@lunarg.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41849>

src/kosmickrisp/vulkan/kk_cmd_buffer.c

@@ -495,7 +495,7 @@ kk_cmd_buffer_flush_push_descriptors(struct kk_cmd_buffer *cmd,
 }
 
 void
-kk_dispatch_precomp(struct kk_cmd_buffer *cmd, struct mtl_size grid,
+kk_dispatch_precomp(struct kk_cmd_buffer *cmd, struct kk_grid grid,
                     bool pre_gfx, enum libkk_program idx, void *data,
                     size_t data_size)
 {
@@ -517,7 +517,12 @@ kk_dispatch_precomp(struct kk_cmd_buffer *cmd, struct mtl_size grid,
       .y = prog->info.workgroup_size[1],
       .z = prog->info.workgroup_size[2],
    };
-   mtl_dispatch_threads(encoder, grid, local_size);
+
+   if (grid.mode == KK_GRID_DIRECT)
+      mtl_dispatch_threads(encoder, grid.size, local_size);
+   else
+      mtl_dispatch_threadgroups_with_indirect_buffer(encoder, grid.indirect,
+                                                     grid.offset, local_size);
 }
 
 void

src/kosmickrisp/vulkan/kk_cmd_buffer.h

@@ -274,7 +274,53 @@ uint64_t kk_upload_descriptor_root(struct kk_cmd_buffer *cmd,
 void kk_cmd_buffer_flush_push_descriptors(struct kk_cmd_buffer *cmd,
                                           struct kk_descriptor_state *desc);
 
-void kk_dispatch_precomp(struct kk_cmd_buffer *cmd, struct mtl_size grid,
+enum kk_grid_mode {
+   KK_GRID_DIRECT = 0u,
+   KK_GRID_INDIRECT,
+};
+struct kk_grid {
+   enum kk_grid_mode mode;
+   union {
+      struct {
+         uint32_t offset;
+         mtl_buffer *indirect;
+      };
+      struct mtl_size size;
+   };
+};
+
+static struct kk_grid
+kk_grid_3d(uint32_t x, uint32_t y, uint32_t z)
+{
+   return (struct kk_grid){
+      .mode = KK_GRID_DIRECT,
+      .size = {x, y, z},
+   };
+}
+
+static struct kk_grid
+kk_grid_2d(uint32_t x, uint32_t y)
+{
+   return kk_grid_3d(x, y, 1u);
+}
+
+static struct kk_grid
+kk_grid_1d(uint32_t x)
+{
+   return kk_grid_3d(x, 1u, 1u);
+}
+
+static struct kk_grid
+kk_grid_indirect(mtl_buffer *indirect, uint32_t offset)
+{
+   return (struct kk_grid){
+      .mode = KK_GRID_INDIRECT,
+      .indirect = indirect,
+      .offset = offset,
+   };
+}
+
+void kk_dispatch_precomp(struct kk_cmd_buffer *cmd, struct kk_grid grid,
                          bool pre_gfx, enum libkk_program idx, void *data,
                          size_t data_size);
 

src/kosmickrisp/vulkan/kk_cmd_dispatch.c

@@ -49,7 +49,7 @@ kk_predicate_compute(struct kk_cmd_buffer *cmd, uint64_t indirect_addr_out,
     * generated commands, constructing an indirect command buffer on the GPU
     * which only contains the commands to run if the condition is true. For the
     * time being, we apply predicates by zeroing out disabled indirect data */
-   struct mtl_size grid = {1u, 1u, 1u};
+   struct kk_grid grid = kk_grid_1d(1u);
    if (cmd->state.cond_render.inverted) {
       libkk_predicate_indirect_eq_zero(cmd, grid, false, indirect_addr_out,
                                        indirect_addr_in, cond_addr, 3u, 3u);

src/kosmickrisp/vulkan/kk_cmd_draw.c

@@ -1008,7 +1008,7 @@ kk_predicate_draws(struct kk_cmd_buffer *cmd, struct kk_draw_data data)
     * generated commands, constructing an indirect command buffer on the GPU
     * which only contains the commands to run if the condition is true. For the
     * time being, we apply predicates by zeroing out disabled indirect data */
-   struct mtl_size grid = {data.draw_count, 1u, 1u};
+   struct kk_grid grid = kk_grid_1d(data.draw_count);
    for (uint32_t i = 0; i < data.predicate_count; i++) {
       uint64_t addr = data.predicates[i].gpu_addr;
       switch (data.predicates[i].op) {
@@ -1083,8 +1083,8 @@ kk_unroll_geometry(struct kk_cmd_buffer *cmd, struct kk_draw_data data)
       .mode = data.prim,
    };
 
-   struct mtl_size grid = {1024 * data.draw_count, 1, 1};
-   libkk_unroll_geometry_struct(cmd, grid, true, info);
+   libkk_unroll_geometry_struct(cmd, kk_grid_1d(1024 * data.draw_count), true,
+                                info);
 
    data.indirect_draws.buffer = out_draws.buffer;
    data.indirect_draws.offset = out_draws.offset;

src/kosmickrisp/vulkan/kk_encoder.c

@@ -147,8 +147,7 @@ upload_queue_writes(struct kk_cmd_buffer *cmd)
       /* kk_cmd_allocate_buffer sets the cmd buffer error so we can just exit */
       if (unlikely(!data_gpu.gpu))
          return;
-      struct mtl_size grid = {count, 1, 1};
-      libkk_write_u32_array(cmd, grid, false, data_gpu.gpu);
+      libkk_write_u32_array(cmd, kk_grid_1d(count), false, data_gpu.gpu);
       enc->imm_writes.size = 0u;
    }
 
@@ -160,10 +159,10 @@ upload_queue_writes(struct kk_cmd_buffer *cmd)
             util_dynarray_element(&enc->copy_query_pool_result_infos,
                                   struct kk_copy_query_pool_results_info, i);
 
-         struct mtl_size grid = {push_data->query_count, 1, 1};
          const struct libkk_copy_queries_args *data =
             (const struct libkk_copy_queries_args *)push_data;
-         libkk_copy_queries_struct(cmd, grid, false, *data);
+         libkk_copy_queries_struct(cmd, kk_grid_1d(push_data->query_count),
+                                   false, *data);
       }
       enc->copy_query_pool_result_infos.size = 0u;
    }

src/kosmickrisp/vulkan/kk_query_pool.c

@@ -244,8 +244,7 @@ emit_zero_queries(struct kk_cmd_buffer *cmd, struct kk_query_pool *pool,
       .reports_per_query = kk_reports_per_query(pool),
       .set_available = set_available,
    };
-   struct mtl_size grid = {.x = num_queries, .y = 1u, .z = 1u};
-   libkk_reset_query_struct(cmd, grid, false, info);
+   libkk_reset_query_struct(cmd, kk_grid_1d(num_queries), false, info);
 }
 
 VKAPI_ATTR void VKAPI_CALL