lavapipe: Implement KHR_ray_tracing_maintenance1

Reviewed-By: Mike Blumenkrantz <michael.blumenkrantz@gmail.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/28187>

src/gallium/frontends/lavapipe/lvp_device.c

@@ -140,6 +140,7 @@ static const struct vk_device_extension_table lvp_device_extensions_supported =
    .KHR_push_descriptor                   = true,
    .KHR_pipeline_library                  = true,
    .KHR_ray_query                         = true,
+   .KHR_ray_tracing_maintenance1          = true,
    .KHR_ray_tracing_pipeline              = true,
    .KHR_relaxed_block_layout              = true,
    .KHR_sampler_mirror_clamp_to_edge      = true,
@@ -490,6 +491,10 @@ lvp_get_features(const struct lvp_physical_device *pdevice,
       /* VK_KHR_ray_query */
       .rayQuery = true,
 
+      /* VK_KHR_ray_tracing_maintenance1 */
+      .rayTracingMaintenance1 = true,
+      .rayTracingPipelineTraceRaysIndirect2 = true,
+
       /* VK_KHR_ray_tracing_pipeline */
       .rayTracingPipeline = true,
       .rayTracingPipelineShaderGroupHandleCaptureReplay = false,

src/gallium/frontends/lavapipe/lvp_execute.c

@@ -4565,6 +4565,35 @@ handle_trace_rays_indirect(struct vk_cmd_queue_entry *cmd, struct rendering_stat
    state->pctx->launch_grid(state->pctx, &state->trace_rays_info);
 }
 
+static void
+handle_trace_rays_indirect2(struct vk_cmd_queue_entry *cmd, struct rendering_state *state)
+{
+   struct vk_cmd_trace_rays_indirect2_khr *trace = &cmd->u.trace_rays_indirect2_khr;
+
+   emit_ray_tracing_state(state);
+
+   size_t indirect_offset;
+   VkBuffer _indirect = get_buffer(state, (void *)(uintptr_t)trace->indirect_device_address, &indirect_offset);
+   VK_FROM_HANDLE(lvp_buffer, indirect, _indirect);
+
+   struct pipe_transfer *transfer;
+   const uint8_t *map = pipe_buffer_map(state->pctx, indirect->bo, PIPE_MAP_READ, &transfer);
+   map += indirect_offset;
+   const VkTraceRaysIndirectCommand2KHR *src = (const void *)map;
+
+   VkTraceRaysIndirectCommand2KHR *command = lvp_push_internal_buffer(
+      state, MESA_SHADER_COMPUTE, sizeof(VkTraceRaysIndirectCommand2KHR));
+   *command = *src;
+
+   state->trace_rays_info.grid[0] = DIV_ROUND_UP(src->width, state->trace_rays_info.block[0]);
+   state->trace_rays_info.grid[1] = DIV_ROUND_UP(src->height, state->trace_rays_info.block[1]);
+   state->trace_rays_info.grid[2] = DIV_ROUND_UP(src->depth, state->trace_rays_info.block[2]);
+
+   state->pctx->buffer_unmap(state->pctx, transfer);
+
+   state->pctx->launch_grid(state->pctx, &state->trace_rays_info);
+}
+
 void lvp_add_enqueue_cmd_entrypoints(struct vk_device_dispatch_table *disp)
 {
    struct vk_device_dispatch_table cmd_enqueue_dispatch;
@@ -4713,6 +4742,7 @@ void lvp_add_enqueue_cmd_entrypoints(struct vk_device_dispatch_table *disp)
    ENQUEUE_CMD(CmdWriteAccelerationStructuresPropertiesKHR)
 
    ENQUEUE_CMD(CmdSetRayTracingPipelineStackSizeKHR)
+   ENQUEUE_CMD(CmdTraceRaysIndirect2KHR)
    ENQUEUE_CMD(CmdTraceRaysIndirectKHR)
    ENQUEUE_CMD(CmdTraceRaysKHR)
 
@@ -5101,6 +5131,9 @@ static void lvp_execute_cmd_buffer(struct list_head *cmds,
          break;
       case VK_CMD_SET_RAY_TRACING_PIPELINE_STACK_SIZE_KHR:
          break;
+      case VK_CMD_TRACE_RAYS_INDIRECT2_KHR:
+         handle_trace_rays_indirect2(cmd, state);
+         break;
       case VK_CMD_TRACE_RAYS_INDIRECT_KHR:
          handle_trace_rays_indirect(cmd, state);
          break;

src/gallium/frontends/lavapipe/lvp_pipeline.c

@@ -350,6 +350,7 @@ compile_spirv(struct lvp_device *pdevice, const VkPipelineShaderStageCreateInfo
          .runtime_descriptor_array = true,
          .shader_enqueue = true,
          .ray_query = true,
+         .ray_cull_mask = true,
          .ray_tracing = true,
       },
       .ubo_addr_format = nir_address_format_vec2_index_32bit_offset,

src/gallium/frontends/lavapipe/lvp_ray_tracing_pipeline.c

@@ -260,6 +260,7 @@ struct lvp_ray_traversal_state {
 struct lvp_ray_tracing_state {
    nir_variable *bvh_base;
    nir_variable *flags;
+   nir_variable *cull_mask;
    nir_variable *sbt_offset;
    nir_variable *sbt_stride;
    nir_variable *miss_index;
@@ -319,6 +320,7 @@ lvp_ray_tracing_state_init(nir_shader *nir, struct lvp_ray_tracing_state *state)
 {
    state->bvh_base = nir_variable_create(nir, nir_var_shader_temp, glsl_uint64_t_type(), "bvh_base");
    state->flags = nir_variable_create(nir, nir_var_shader_temp, glsl_uint_type(), "flags");
+   state->cull_mask = nir_variable_create(nir, nir_var_shader_temp, glsl_uint_type(), "cull_mask");
    state->sbt_offset = nir_variable_create(nir, nir_var_shader_temp, glsl_uint_type(), "sbt_offset");
    state->sbt_stride = nir_variable_create(nir, nir_var_shader_temp, glsl_uint_type(), "sbt_stride");
    state->miss_index = nir_variable_create(nir, nir_var_shader_temp, glsl_uint_type(), "miss_index");
@@ -736,6 +738,7 @@ lvp_trace_ray(nir_builder *b, struct lvp_ray_tracing_pipeline_compiler *compiler
 
    nir_store_var(b, state->bvh_base, bvh_base, 0x1);
    nir_store_var(b, state->flags, flags, 0x1);
+   nir_store_var(b, state->cull_mask, cull_mask, 0x1);
    nir_store_var(b, state->sbt_offset, sbt_offset, 0x1);
    nir_store_var(b, state->sbt_stride, sbt_stride, 0x1);
    nir_store_var(b, state->miss_index, miss_index, 0x1);
@@ -979,6 +982,9 @@ lvp_lower_ray_tracing_instr(nir_builder *b, nir_instr *instr, void *data)
       def = lvp_mul_vec3_mat(b, nir_load_var(b, state->dir), wto_matrix, false);
       break;
    }
+   case nir_intrinsic_load_cull_mask:
+      def = nir_iand_imm(b, nir_load_var(b, state->cull_mask), 0xFF);
+      break;
    /* Ray tracing stack lowering */
    case nir_intrinsic_load_scratch: {
       nir_src_rewrite(&intr->src[0], nir_iadd(b, nir_load_var(b, state->stack_ptr), intr->src[0].ssa));