spirv: handle ray query intrinsics

v2: Fixup comment (Caio)
    Use generated builders (Caio)

v3: Update spirv2dxil CI expectations

Signed-off-by: Lionel Landwerlin <lionel.g.landwerlin@intel.com>
Reviewed-by: Caio Oliveira <caio.oliveira@intel.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/13718>

src/compiler/spirv/spirv_to_nir.c

@@ -805,6 +805,7 @@ vtn_types_compatible(struct vtn_builder *b,
       return true;
 
    case vtn_base_type_accel_struct:
+   case vtn_base_type_ray_query:
       return true;
 
    case vtn_base_type_function:
@@ -845,6 +846,7 @@ vtn_type_copy(struct vtn_builder *b, struct vtn_type *src)
    case vtn_base_type_sampled_image:
    case vtn_base_type_event:
    case vtn_base_type_accel_struct:
+   case vtn_base_type_ray_query:
       /* Nothing more to do */
       break;
 
@@ -1810,11 +1812,30 @@ vtn_handle_type(struct vtn_builder *b, SpvOp opcode,
       val->type->type = glsl_uint64_t_type();
       break;
 
-   case SpvOpTypeOpaque:
+
+   case SpvOpTypeOpaque: {
       val->type->base_type = vtn_base_type_struct;
       const char *name = vtn_string_literal(b, &w[2], count - 2, NULL);
       val->type->type = glsl_struct_type(NULL, 0, name, false);
       break;
+   }
+
+   case SpvOpTypeRayQueryKHR: {
+      val->type->base_type = vtn_base_type_ray_query;
+      const char *name = "RayQueryKHR";
+      val->type->type = glsl_struct_type(NULL, 0, name, false);
+      /* We may need to run queries on helper invocations. Here the parser
+       * doesn't go through a deeper analysis on whether the result of a query
+       * will be used in derivative instructions.
+       *
+       * An implementation willing to optimize this would look through the IR
+       * and check if any derivative instruction uses the result of a query
+       * and drop this flag if not.
+       */
+      if (b->shader->info.stage == MESA_SHADER_FRAGMENT)
+         val->type->access = ACCESS_INCLUDE_HELPERS;
+      break;
+   }
 
    case SpvOpTypeEvent:
       val->type->base_type = vtn_base_type_event;
@@ -5290,6 +5311,7 @@ vtn_handle_variable_or_type_instruction(struct vtn_builder *b, SpvOp opcode,
    case SpvOpTypeQueue:
    case SpvOpTypePipe:
    case SpvOpTypeAccelerationStructureKHR:
+   case SpvOpTypeRayQueryKHR:
       vtn_handle_type(b, opcode, w, count);
       break;
 
@@ -5577,6 +5599,146 @@ vtn_handle_write_packed_primitive_indices(struct vtn_builder *b, SpvOp opcode,
    }
 }
 
+struct ray_query_value {
+   nir_ray_query_value     nir_value;
+   const struct glsl_type *glsl_type;
+};
+
+static struct ray_query_value
+spirv_to_nir_type_ray_query_intrinsic(struct vtn_builder *b,
+                                      SpvOp opcode)
+{
+   switch (opcode) {
+#define CASE(_spv, _nir, _type) case SpvOpRayQueryGet##_spv:            \
+      return (struct ray_query_value) { .nir_value = nir_ray_query_value_##_nir, .glsl_type = _type }
+      CASE(RayTMinKHR,                                            tmin,                               glsl_floatN_t_type(32));
+      CASE(RayFlagsKHR,                                           flags,                              glsl_uint_type());
+      CASE(WorldRayDirectionKHR,                                  world_ray_direction,                glsl_vec_type(3));
+      CASE(WorldRayOriginKHR,                                     world_ray_origin,                   glsl_vec_type(3));
+      CASE(IntersectionTypeKHR,                                   intersection_type,                  glsl_uint_type());
+      CASE(IntersectionTKHR,                                      intersection_t,                     glsl_floatN_t_type(32));
+      CASE(IntersectionInstanceCustomIndexKHR,                    intersection_instance_custom_index, glsl_int_type());
+      CASE(IntersectionInstanceIdKHR,                             intersection_instance_id,           glsl_int_type());
+      CASE(IntersectionInstanceShaderBindingTableRecordOffsetKHR, intersection_instance_sbt_index,    glsl_uint_type());
+      CASE(IntersectionGeometryIndexKHR,                          intersection_geometry_index,        glsl_int_type());
+      CASE(IntersectionPrimitiveIndexKHR,                         intersection_primitive_index,       glsl_int_type());
+      CASE(IntersectionBarycentricsKHR,                           intersection_barycentrics,          glsl_vec_type(2));
+      CASE(IntersectionFrontFaceKHR,                              intersection_front_face,            glsl_bool_type());
+      CASE(IntersectionCandidateAABBOpaqueKHR,                    intersection_candidate_aabb_opaque, glsl_bool_type());
+      CASE(IntersectionObjectToWorldKHR,                          intersection_object_to_world,       glsl_matrix_type(glsl_get_base_type(glsl_float_type()), 3, 4));
+      CASE(IntersectionWorldToObjectKHR,                          intersection_world_to_object,       glsl_matrix_type(glsl_get_base_type(glsl_float_type()), 3, 4));
+      CASE(IntersectionObjectRayOriginKHR,                        intersection_object_ray_origin,     glsl_vec_type(3));
+      CASE(IntersectionObjectRayDirectionKHR,                     intersection_object_ray_direction,  glsl_vec_type(3));
+#undef CASE
+   default:
+      vtn_fail_with_opcode("Unhandled opcode", opcode);
+   }
+}
+
+static void
+ray_query_load_intrinsic_create(struct vtn_builder *b, SpvOp opcode,
+                                const uint32_t *w, nir_ssa_def *src0,
+                                nir_ssa_def *src1)
+{
+   struct ray_query_value value =
+      spirv_to_nir_type_ray_query_intrinsic(b, opcode);
+
+   if (glsl_type_is_matrix(value.glsl_type)) {
+      const struct glsl_type *elem_type = glsl_get_array_element(value.glsl_type);
+      const unsigned elems = glsl_get_length(value.glsl_type);
+
+      struct vtn_ssa_value *ssa = vtn_create_ssa_value(b, value.glsl_type);
+      for (unsigned i = 0; i < elems; i++) {
+         ssa->elems[i]->def =
+            nir_build_rq_load(&b->nb,
+                              glsl_get_vector_elements(elem_type),
+                              glsl_get_bit_size(elem_type),
+                              src0, src1,
+                              .base = value.nir_value,
+                              .column = i);
+      }
+
+      vtn_push_ssa_value(b, w[2], ssa);
+   } else {
+      assert(glsl_type_is_vector_or_scalar(value.glsl_type));
+
+      vtn_push_nir_ssa(b, w[2],
+                       nir_rq_load(&b->nb,
+                                   glsl_get_vector_elements(value.glsl_type),
+                                   glsl_get_bit_size(value.glsl_type),
+                                   src0, src1));
+   }
+}
+
+static void
+vtn_handle_ray_query_intrinsic(struct vtn_builder *b, SpvOp opcode,
+                               const uint32_t *w, unsigned count)
+{
+   switch (opcode) {
+   case SpvOpRayQueryInitializeKHR: {
+      nir_intrinsic_instr *intrin =
+         nir_intrinsic_instr_create(b->nb.shader,
+                                    nir_intrinsic_rq_initialize);
+      /* The sources are in the same order in the NIR intrinsic */
+      for (unsigned i = 0; i < 8; i++)
+         intrin->src[i] = nir_src_for_ssa(vtn_ssa_value(b, w[i + 1])->def);
+      nir_builder_instr_insert(&b->nb, &intrin->instr);
+      break;
+   }
+
+   case SpvOpRayQueryTerminateKHR:
+      nir_rq_terminate(&b->nb, vtn_ssa_value(b, w[1])->def);
+      break;
+
+   case SpvOpRayQueryProceedKHR:
+      vtn_push_nir_ssa(b, w[2],
+                       nir_rq_proceed(&b->nb, 1, vtn_ssa_value(b, w[3])->def));
+      break;
+
+   case SpvOpRayQueryGenerateIntersectionKHR:
+      nir_rq_generate_intersection(&b->nb,
+                                   vtn_ssa_value(b, w[1])->def,
+                                   vtn_ssa_value(b, w[2])->def);
+      break;
+
+   case SpvOpRayQueryConfirmIntersectionKHR:
+      nir_rq_confirm_intersection(&b->nb, vtn_ssa_value(b, w[1])->def);
+      break;
+
+   case SpvOpRayQueryGetIntersectionTKHR:
+   case SpvOpRayQueryGetIntersectionTypeKHR:
+   case SpvOpRayQueryGetIntersectionInstanceCustomIndexKHR:
+   case SpvOpRayQueryGetIntersectionInstanceIdKHR:
+   case SpvOpRayQueryGetIntersectionInstanceShaderBindingTableRecordOffsetKHR:
+   case SpvOpRayQueryGetIntersectionGeometryIndexKHR:
+   case SpvOpRayQueryGetIntersectionPrimitiveIndexKHR:
+   case SpvOpRayQueryGetIntersectionBarycentricsKHR:
+   case SpvOpRayQueryGetIntersectionFrontFaceKHR:
+   case SpvOpRayQueryGetIntersectionObjectRayDirectionKHR:
+   case SpvOpRayQueryGetIntersectionObjectRayOriginKHR:
+   case SpvOpRayQueryGetIntersectionObjectToWorldKHR:
+   case SpvOpRayQueryGetIntersectionWorldToObjectKHR:
+      ray_query_load_intrinsic_create(b, opcode, w,
+                                      vtn_ssa_value(b, w[3])->def,
+                                      nir_i2b1(&b->nb, vtn_ssa_value(b, w[4])->def));
+      break;
+
+   case SpvOpRayQueryGetRayTMinKHR:
+   case SpvOpRayQueryGetRayFlagsKHR:
+   case SpvOpRayQueryGetWorldRayDirectionKHR:
+   case SpvOpRayQueryGetWorldRayOriginKHR:
+   case SpvOpRayQueryGetIntersectionCandidateAABBOpaqueKHR:
+      ray_query_load_intrinsic_create(b, opcode, w,
+                                      vtn_ssa_value(b, w[3])->def,
+                                      /* Committed value is ignored for these */
+                                      nir_imm_bool(&b->nb, false));
+      break;
+
+   default:
+      vtn_fail_with_opcode("Unhandled opcode", opcode);
+   }
+}
+
 static bool
 vtn_handle_body_instruction(struct vtn_builder *b, SpvOp opcode,
                             const uint32_t *w, unsigned count)
@@ -6006,6 +6168,32 @@ vtn_handle_body_instruction(struct vtn_builder *b, SpvOp opcode,
       vtn_handle_ray_intrinsic(b, opcode, w, count);
       break;
 
+   case SpvOpRayQueryInitializeKHR:
+   case SpvOpRayQueryTerminateKHR:
+   case SpvOpRayQueryGenerateIntersectionKHR:
+   case SpvOpRayQueryConfirmIntersectionKHR:
+   case SpvOpRayQueryProceedKHR:
+   case SpvOpRayQueryGetIntersectionTypeKHR:
+   case SpvOpRayQueryGetRayTMinKHR:
+   case SpvOpRayQueryGetRayFlagsKHR:
+   case SpvOpRayQueryGetIntersectionTKHR:
+   case SpvOpRayQueryGetIntersectionInstanceCustomIndexKHR:
+   case SpvOpRayQueryGetIntersectionInstanceIdKHR:
+   case SpvOpRayQueryGetIntersectionInstanceShaderBindingTableRecordOffsetKHR:
+   case SpvOpRayQueryGetIntersectionGeometryIndexKHR:
+   case SpvOpRayQueryGetIntersectionPrimitiveIndexKHR:
+   case SpvOpRayQueryGetIntersectionBarycentricsKHR:
+   case SpvOpRayQueryGetIntersectionFrontFaceKHR:
+   case SpvOpRayQueryGetIntersectionCandidateAABBOpaqueKHR:
+   case SpvOpRayQueryGetIntersectionObjectRayDirectionKHR:
+   case SpvOpRayQueryGetIntersectionObjectRayOriginKHR:
+   case SpvOpRayQueryGetWorldRayDirectionKHR:
+   case SpvOpRayQueryGetWorldRayOriginKHR:
+   case SpvOpRayQueryGetIntersectionObjectToWorldKHR:
+   case SpvOpRayQueryGetIntersectionWorldToObjectKHR:
+      vtn_handle_ray_query_intrinsic(b, opcode, w, count);
+      break;
+
    case SpvOpLifetimeStart:
    case SpvOpLifetimeStop:
       break;

src/compiler/spirv/vtn_private.h

@@ -336,6 +336,7 @@ enum vtn_base_type {
    vtn_base_type_sampler,
    vtn_base_type_sampled_image,
    vtn_base_type_accel_struct,
+   vtn_base_type_ray_query,
    vtn_base_type_function,
    vtn_base_type_event,
 };

src/compiler/spirv/vtn_variables.c

@@ -1791,6 +1791,12 @@ vtn_get_call_payload_for_location(struct vtn_builder *b, uint32_t location_id)
             "or RayPayloadKHR and location %d", location);
 }
 
+static bool
+vtn_type_is_ray_query(struct vtn_type *type)
+{
+   return vtn_type_without_array(type)->base_type == vtn_base_type_ray_query;
+}
+
 static void
 vtn_create_variable(struct vtn_builder *b, struct vtn_value *val,
                     struct vtn_type *ptr_type, SpvStorageClass storage_class,
@@ -1887,6 +1893,7 @@ vtn_create_variable(struct vtn_builder *b, struct vtn_value *val,
 
       var->var->data.mode = nir_mode;
       var->var->data.location = -1;
+      var->var->data.ray_query = vtn_type_is_ray_query(var->type);
       var->var->interface_type = NULL;
       break;
 

src/microsoft/ci/spirv2dxil_reference.txt

@@ -1272,11 +1272,11 @@ Test:SpvParserTest_Impl_GenericVulkanShader_GLSL450MemoryModel.spvasm:main|GLCom
 Test:SpvParserTest_Impl_GenericVulkanShader_SimpleMemoryModel.spvasm:main|GLCompute: Pass
 Test:SpvParserTest_Impl_GenericVulkanShader_VulkanMemoryModel.spvasm:main|GLCompute: Fail
 SPIR-V WARNING:
-    In file ../src/compiler/spirv/spirv_to_nir.c:4687
+    In file ../src/compiler/spirv/spirv_to_nir.c:4708
     Unsupported SPIR-V capability: SpvCapabilityVulkanMemoryModel (5345)
     28 bytes into the SPIR-V binary
 SPIR-V parsing FAILED:
-    In file ../src/compiler/spirv/spirv_to_nir.c:4841
+    In file ../src/compiler/spirv/spirv_to_nir.c:4862
     Vulkan memory model is unsupported by this driver
     68 bytes into the SPIR-V binary
 Compilation failed