Add ray query feature support

src/script.cc

@@ -142,6 +142,7 @@ bool Script::IsKnownFeature(const std::string& name) const {
              ".shaderSubgroupExtendedTypes" ||
          name == "CooperativeMatrixFeaturesKHR.cooperativeMatrix" ||
          name == "RayTracingPipelineFeaturesKHR.rayTracingPipeline" ||
+         name == "RayQueryFeaturesKHR.rayQuery" ||
          name == "AccelerationStructureFeaturesKHR.accelerationStructure" ||
          name == "BufferDeviceAddressFeatures.bufferDeviceAddress" ||
          name == "ShaderLongVectorFeaturesEXT.longVector";

src/vulkan/compute_pipeline.cc

@@ -23,14 +23,18 @@ namespace vulkan {
 
 ComputePipeline::ComputePipeline(
     Device* device,
+    BlasesMap* blases,
+    TlasesMap* tlases,
     uint32_t fence_timeout_ms,
     bool pipeline_runtime_layer_enabled,
     const std::vector<VkPipelineShaderStageCreateInfo>& shader_stage_info)
     : Pipeline(PipelineType::kCompute,
                device,
                fence_timeout_ms,
                pipeline_runtime_layer_enabled,
-               shader_stage_info) {}
+               shader_stage_info),
+      blases_(blases),
+      tlases_(tlases) {}
 
 ComputePipeline::~ComputePipeline() = default;
 

src/vulkan/compute_pipeline.h

@@ -29,6 +29,8 @@ class ComputePipeline : public Pipeline {
  public:
   ComputePipeline(
       Device* device,
+      BlasesMap* blases,
+      TlasesMap* tlases,
       uint32_t fence_timeout_ms,
       bool pipeline_runtime_layer_enabled,
       const std::vector<VkPipelineShaderStageCreateInfo>& shader_stage_info);
@@ -38,9 +40,15 @@ class ComputePipeline : public Pipeline {
 
   Result Compute(uint32_t x, uint32_t y, uint32_t z, bool is_timed_execution);
 
+  BlasesMap* GetBlases() override { return blases_; }
+  TlasesMap* GetTlases() override { return tlases_; }
+
  private:
   Result CreateVkComputePipeline(const VkPipelineLayout& pipeline_layout,
                                  VkPipeline* pipeline);
+
+  BlasesMap* blases_;
+  TlasesMap* tlases_;
 };
 
 }  // namespace vulkan

src/vulkan/device.cc

@@ -109,6 +109,7 @@ const char kBufferDeviceAddress[] =
     "BufferDeviceAddressFeatures.bufferDeviceAddress";
 const char kRayTracingPipeline[] =
     "RayTracingPipelineFeaturesKHR.rayTracingPipeline";
+const char kRayQuery[] = "RayQueryFeaturesKHR.rayQuery";
 
 struct BaseOutStructure {
   VkStructureType sType;
@@ -583,6 +584,7 @@ Result Device::Initialize(
   VkPhysicalDeviceBufferDeviceAddressFeatures* bda_ptrs = nullptr;
   VkPhysicalDeviceRayTracingPipelineFeaturesKHR* ray_tracing_pipeline_ptrs =
       nullptr;
+  VkPhysicalDeviceRayQueryFeaturesKHR* ray_query_ptrs = nullptr;
 
   void* ptr = available_features2.pNext;
   while (ptr != nullptr) {
@@ -639,6 +641,9 @@ Result Device::Initialize(
         ray_tracing_pipeline_ptrs =
             static_cast<VkPhysicalDeviceRayTracingPipelineFeaturesKHR*>(ptr);
         break;
+      case VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_QUERY_FEATURES_KHR:
+        ray_query_ptrs = static_cast<VkPhysicalDeviceRayQueryFeaturesKHR*>(ptr);
+        break;
       case VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES:
         vulkan11_ptrs = static_cast<VkPhysicalDeviceVulkan11Features*>(ptr);
         break;
@@ -772,6 +777,9 @@ Result Device::Initialize(
             "provided");
       }
     }
+    if (feature == kRayQuery && ray_query_ptrs == nullptr) {
+      return amber::Result("Ray query requested but feature not returned");
+    }
 
     // Next check the fields of the feature structures.
 
@@ -937,6 +945,9 @@ Result Device::Initialize(
         cooperative_matrix_ptrs->cooperativeMatrix != VK_TRUE) {
       return amber::Result("Missing cooperative matrix feature");
     }
+    if (feature == kRayQuery && ray_query_ptrs->rayQuery != VK_TRUE) {
+      return amber::Result("Missing ray query feature");
+    }
   }
 
   if (!AreAllExtensionsSupported(available_extensions,

src/vulkan/engine_vulkan.cc

@@ -219,12 +219,12 @@ Result EngineVulkan::CreatePipeline(amber::Pipeline* pipeline) {
         pool_.get(), shader_group_create_info);
   } else if (pipeline->GetType() == PipelineType::kCompute) {
     vk_pipeline = std::make_unique<ComputePipeline>(
-        device_.get(), engine_data.fence_timeout_ms,
+        device_.get(), &blases_, &tlases_, engine_data.fence_timeout_ms,
         engine_data.pipeline_runtime_layer_enabled, stage_create_info);
     r = vk_pipeline->AsCompute()->Initialize(pool_.get());
   } else {
     vk_pipeline = std::make_unique<GraphicsPipeline>(
-        device_.get(), pipeline->GetColorAttachments(),
+        device_.get(), &blases_, &tlases_, pipeline->GetColorAttachments(),
         pipeline->GetDepthStencilBuffer(), pipeline->GetResolveTargets(),
         engine_data.fence_timeout_ms,
         engine_data.pipeline_runtime_layer_enabled, stage_create_info);
@@ -341,17 +341,15 @@ Result EngineVulkan::CreatePipeline(amber::Pipeline* pipeline) {
     }
   }
 
-  if (info.vk_pipeline->IsRayTracing()) {
-    for (const auto& tlas_info : pipeline->GetTLASes()) {
-      auto cmd = std::make_unique<TLASCommand>(pipeline);
-      cmd->SetDescriptorSet(tlas_info.descriptor_set);
-      cmd->SetBinding(tlas_info.binding);
-      cmd->SetTLAS(tlas_info.tlas);
+  for (const auto& tlas_info : pipeline->GetTLASes()) {
+    auto cmd = std::make_unique<TLASCommand>(pipeline);
+    cmd->SetDescriptorSet(tlas_info.descriptor_set);
+    cmd->SetBinding(tlas_info.binding);
+    cmd->SetTLAS(tlas_info.tlas);
 
-      r = info.vk_pipeline->AddTLASDescriptor(cmd.get());
-      if (!r.IsSuccess()) {
-        return r;
-      }
+    r = info.vk_pipeline->AddTLASDescriptor(cmd.get());
+    if (!r.IsSuccess()) {
+      return r;
     }
   }
 

src/vulkan/graphics_pipeline.cc

@@ -387,6 +387,8 @@ class RenderPassGuard {
 
 GraphicsPipeline::GraphicsPipeline(
     Device* device,
+    BlasesMap* blases,
+    TlasesMap* tlases,
     const std::vector<amber::Pipeline::BufferInfo>& color_buffers,
     amber::Pipeline::BufferInfo depth_stencil_buffer,
     const std::vector<amber::Pipeline::BufferInfo>& resolve_targets,
@@ -398,7 +400,9 @@ GraphicsPipeline::GraphicsPipeline(
                fence_timeout_ms,
                pipeline_runtime_layer_enabled,
                shader_stage_info),
-      depth_stencil_buffer_(depth_stencil_buffer) {
+      depth_stencil_buffer_(depth_stencil_buffer),
+      blases_(blases),
+      tlases_(tlases) {
   for (const auto& info : color_buffers) {
     color_buffers_.push_back(&info);
   }

src/vulkan/graphics_pipeline.h

@@ -41,6 +41,8 @@ class GraphicsPipeline : public Pipeline {
  public:
   GraphicsPipeline(
       Device* device,
+      BlasesMap* blases,
+      TlasesMap* tlases,
       const std::vector<amber::Pipeline::BufferInfo>& color_buffers,
       amber::Pipeline::BufferInfo depth_stencil_buffer,
       const std::vector<amber::Pipeline::BufferInfo>& resolve_targets,
@@ -69,6 +71,9 @@ class GraphicsPipeline : public Pipeline {
   uint32_t GetWidth() const { return frame_width_; }
   uint32_t GetHeight() const { return frame_height_; }
 
+  BlasesMap* GetBlases() override { return blases_; }
+  TlasesMap* GetTlases() override { return tlases_; }
+
   void SetPatchControlPoints(uint32_t points) {
     patch_control_points_ = points;
   }
@@ -106,6 +111,8 @@ class GraphicsPipeline : public Pipeline {
   uint32_t clear_stencil_ = 0;
   float clear_depth_ = 1.0f;
   uint32_t patch_control_points_ = 3;
+  BlasesMap* blases_;
+  TlasesMap* tlases_;
 };
 
 }  // namespace vulkan

src/vulkan/pipeline_test.cc

@@ -26,7 +26,7 @@ using VulkanPipelineTest = testing::Test;
 TEST_F(VulkanPipelineTest, AddBufferDescriptorAddPushConstant) {
   amber::Pipeline amber_pipeline(PipelineType::kCompute);
   std::vector<VkPipelineShaderStageCreateInfo> create_infos;
-  ComputePipeline pipeline(nullptr, 0, false, create_infos);
+  ComputePipeline pipeline(nullptr, nullptr, nullptr, 0, false, create_infos);
 
   auto cmd = std::make_unique<BufferCommand>(
       BufferCommand::BufferType::kPushConstant, &amber_pipeline);
@@ -38,7 +38,7 @@ TEST_F(VulkanPipelineTest, AddBufferDescriptorAddPushConstant) {
 TEST_F(VulkanPipelineTest, AddBufferDescriptorAddBufferTwice) {
   amber::Pipeline amber_pipeline(PipelineType::kCompute);
   std::vector<VkPipelineShaderStageCreateInfo> create_infos;
-  ComputePipeline pipeline(nullptr, 0, false, create_infos);
+  ComputePipeline pipeline(nullptr, nullptr, nullptr, 0, false, create_infos);
 
   auto cmd = std::make_unique<BufferCommand>(
       BufferCommand::BufferType::kUniform, &amber_pipeline);