Improving SBT API

include/daxa/c/command_recorder.h

@@ -205,7 +205,6 @@ typedef struct
     uint32_t width;
     uint32_t height;
     uint32_t depth;
-    uint32_t raygen_handle_offset;
     uint32_t miss_handle_offset;
     uint32_t hit_handle_offset;
     uint32_t callable_handle_offset;
@@ -217,7 +216,6 @@ static daxa_TraceRaysInfo const DAXA_DEFAULT_TRACE_RAYS_INFO = DAXA_ZERO_INIT;
 typedef struct
 {
     uint64_t indirect_device_address;
-    uint32_t raygen_handle_offset;
     uint32_t miss_handle_offset;
     uint32_t hit_handle_offset;
     uint32_t callable_handle_offset;

include/daxa/c/pipeline.h

@@ -38,28 +38,80 @@ static daxa_RayTracingShaderGroupInfo const DAXA_DEFAULT_RAY_TRACING_SHADER_GROU
 
 typedef struct
 {
-    daxa_SpanToConst(daxa_RayTracingShaderInfo) ray_gen_stages;
-    daxa_SpanToConst(daxa_RayTracingShaderInfo) miss_stages;
-    daxa_SpanToConst(daxa_RayTracingShaderInfo) callable_stages;
-    daxa_SpanToConst(daxa_RayTracingShaderInfo) intersection_stages;
-    daxa_SpanToConst(daxa_RayTracingShaderInfo) closest_hit_stages;
-    daxa_SpanToConst(daxa_RayTracingShaderInfo) any_hit_stages;
-    daxa_SpanToConst(daxa_RayTracingShaderGroupInfo) shader_groups;
+    uint64_t device_address;
+    uint64_t stride;
+    uint64_t size;
+} daxa_StridedDeviceAddressRegion;
+
+static daxa_StridedDeviceAddressRegion const DAXA_DEFAULT_STRIDED_DEVICE_ADDRESS_REGION = {
+    .device_address = 0,
+    .stride = 0,
+    .size = 0,
+};
+
+typedef struct
+{
+    daxa_StridedDeviceAddressRegion raygen_region;
+    daxa_StridedDeviceAddressRegion miss_region;
+    daxa_StridedDeviceAddressRegion hit_region;
+    daxa_StridedDeviceAddressRegion callable_region;
+} daxa_RayTracingShaderBindingTable;
+
+typedef struct 
+{
+    daxa_ShaderGroupType type;
+    daxa_StridedDeviceAddressRegion region;
+} daxa_GroupRegionInfo;
+
+static daxa_RayTracingShaderBindingTable const DAXA_DEFAULT_RAY_TRACING_SHADER_BINDING_TABLE = {
+    .raygen_region = DAXA_ZERO_INIT,
+    .miss_region = DAXA_ZERO_INIT,
+    .hit_region = DAXA_ZERO_INIT,
+    .callable_region = DAXA_ZERO_INIT,
+};
+
+typedef struct
+{
+    daxa_BufferId buffer;
+    daxa_SpanToConst(daxa_GroupRegionInfo) group_regions;
+} daxa_RayTracingShaderBindingTableEntries;
+
+typedef struct
+{
+    daxa_SpanToConst(daxa_RayTracingShaderInfo) stages;
+    daxa_SpanToConst(daxa_RayTracingShaderGroupInfo) groups;
     uint32_t max_ray_recursion_depth;
     uint32_t push_constant_size;
     daxa_SmallString name;
 } daxa_RayTracingPipelineInfo;
 
+
+
 DAXA_EXPORT daxa_RayTracingPipelineInfo const *
 daxa_ray_tracing_pipeline_info(daxa_RayTracingPipeline ray_tracing_pipeline);
 
+typedef struct
+{
+    daxa_SpanToConst(uint32_t) group_indices;
+} daxa_BuildShaderBindingTableInfo;
+
+static daxa_BuildShaderBindingTableInfo const DAXA_DEFAULT_BUILD_SHADER_BINDING_TABLE_INFO = DAXA_ZERO_INIT;
+
 DAXA_EXPORT daxa_Result
-daxa_ray_tracing_pipeline_create_default_sbt(daxa_RayTracingPipeline pipeline, daxa_RayTracingShaderBindingTable * out_sbt, daxa_BufferId * out_buffer);
+daxa_ray_tracing_pipeline_create_default_sbt(daxa_RayTracingPipeline pipeline, daxa_RayTracingShaderBindingTableEntries * out_sbts, daxa_BufferId * out_buffer);
+DAXA_EXPORT daxa_Result
+daxa_ray_tracing_pipeline_create_sbt(
+    daxa_RayTracingPipeline pipeline,
+    daxa_RayTracingShaderBindingTableEntries * out_sbts, daxa_BufferId * out_buffer,
+    daxa_BuildShaderBindingTableInfo const * info);
 
 // out_blob must be the size of the group_count * raytracing_properties.shaderGroupHandleSize
 DAXA_EXPORT daxa_Result
 daxa_ray_tracing_pipeline_get_shader_group_handles(daxa_RayTracingPipeline pipeline, void *out_blob);
 
+DAXA_EXPORT uint32_t
+daxa_ray_tracing_pipeline_get_shader_group_count(daxa_RayTracingPipeline pipeline);
+
 DAXA_EXPORT uint64_t
 daxa_ray_tracing_pipeline_inc_refcnt(daxa_RayTracingPipeline pipeline);
 DAXA_EXPORT uint64_t

include/daxa/c/types.h

@@ -248,12 +248,23 @@ typedef enum
     DAXA_QUEUE_FAMILY_MAX_ENUM
 } daxa_QueueFamily;
 
-typedef struct
+typedef enum 
+{
+    DAXA_EXT_SHADER_GROUP__RAYGEN,
+    DAXA_EXT_SHADER_GROUP__MISS,
+    DAXA_EXT_SHADER_GROUP__TRIANGLE,
+    DAXA_EXT_SHADER_GROUP__PROCEDURAL,
+    DAXA_EXT_SHADER_GROUP__CALLABLE,
+    DAXA_EXT_SHADER_GROUP__MAX_ENUM,
+} daxa_ExtendedShaderGroupType;
+
+typedef enum 
 {
-    VkStridedDeviceAddressRegionKHR raygen_region;
-    VkStridedDeviceAddressRegionKHR miss_region;
-    VkStridedDeviceAddressRegionKHR hit_region;
-    VkStridedDeviceAddressRegionKHR callable_region;
-} daxa_RayTracingShaderBindingTable;
+    DAXA_SHADER_GROUP_RAYGEN,
+    DAXA_SHADER_GROUP_MISS,
+    DAXA_SHADER_GROUP_HIT,
+    DAXA_SHADER_GROUP_CALLABLE,
+    DAXA_SHADER_GROUP_MAX_ENUM,
+} daxa_ShaderGroupType;
 
 #endif // #ifndef __DAXA_TYPES_H__

include/daxa/command_recorder.hpp

@@ -136,7 +136,6 @@ namespace daxa
         uint32_t width = 1;
         uint32_t height = 1;
         uint32_t depth = 1;
-        uint32_t raygen_shader_binding_table_offset = {};
         uint32_t miss_shader_binding_table_offset = {};
         uint32_t miss_shader_binding_table_stride = {};
         uint32_t hit_shader_binding_table_offset = {};
@@ -146,7 +145,6 @@ namespace daxa
     struct TraceRaysIndirectInfo
     {
         DeviceAddress indirect_device_address = {};
-        uint32_t raygen_shader_binding_table_offset = {};
         uint32_t miss_shader_binding_table_offset = {};
         uint32_t miss_shader_binding_table_stride = {};
         uint32_t hit_shader_binding_table_offset = {};

include/daxa/pipeline.hpp

@@ -18,19 +18,46 @@ namespace daxa
         static inline constexpr ShaderCreateFlags REQUIRE_FULL_SUBGROUPS  = {0x00000002};
     };
 
+
+
     struct ShaderInfo
     {
         u32 const * byte_code = {};
         u32 byte_code_size = {};
         ShaderCreateFlags create_flags = {};
         Optional<u32> required_subgroup_size = {};
         SmallString entry_point = "main";
+
+    };
+
+    enum struct RayTracingShaderType
+    {
+        RAYGEN,
+        MISS,
+        CLOSEST_HIT,
+        ANY_HIT,
+        INTERSECTION,
+        CALLABLE,
+        MAX_ENUM = 0x7fffffff,
     };
 
-    // TODO: find a better way to link shader groups to shaders than by index
+    enum struct RayTracingShaderGroup 
+    {
+        RAYGEN_GROUP,
+        MISS_GROUP,
+        HIT_GROUP,
+        CALLABLE_GROUP,
+        MAX_ENUM = 0x7fffffff,
+    };
+
+    struct RayTracingShaderInfo {
+        RayTracingShaderType type;
+        ShaderInfo info;
+    };
+
     struct RayTracingShaderGroupInfo
     {
-        ShaderGroup type;
+        ExtendedShaderGroupType type;
         u32 general_shader_index = (~0U);
         u32 closest_hit_shader_index = (~0U);
         u32 any_hit_shader_index = (~0U);
@@ -44,6 +71,12 @@ namespace daxa
         u64 size {};
     };
 
+    // Group region info where group shader handles are stored.
+    struct GroupRegionInfo {
+        ShaderGroupType type;
+        StridedDeviceAddressRegion region;
+    };
+
     struct RayTracingShaderBindingTable
     {
         StridedDeviceAddressRegion raygen_region = {};
@@ -54,18 +87,24 @@ namespace daxa
 
     struct RayTracingPipelineInfo
     {
-        Span<ShaderInfo const> ray_gen_shaders = {};
-        Span<ShaderInfo const> intersection_shaders = {};
-        Span<ShaderInfo const> any_hit_shaders = {};
-        Span<ShaderInfo const> callable_shaders = {};
-        Span<ShaderInfo const> closest_hit_shaders = {};
-        Span<ShaderInfo const> miss_hit_shaders = {};
-        Span<RayTracingShaderGroupInfo const> shader_groups = {};
+        Span<RayTracingShaderInfo const> stages = {};
+        Span<RayTracingShaderGroupInfo const> groups = {};
         u32 max_ray_recursion_depth;
         u32 push_constant_size = {};
         SmallString name = {};
     };
 
+    struct BuildShaderBindingTableInfo
+    {
+        std::span<u32 const> group_indices = {};
+    };
+
+    struct RayTracingShaderBindingTableEntries
+    {
+        daxa::BufferId buffer = {};
+        Span<GroupRegionInfo const> group_regions = {};
+    };
+
     /**
      * @brief   Represents a pipeline state object, usable in recording commands.
      *
@@ -83,9 +122,14 @@ namespace daxa
         /// @return reference to info of object.
         [[nodiscard]] auto info() const -> RayTracingPipelineInfo const &;
 
-        struct SbtPair { daxa::BufferId buffer; RayTracingShaderBindingTable table; };
+        struct SbtPair {
+            daxa::BufferId buffer;
+            RayTracingShaderBindingTableEntries entries;
+        };
         [[nodiscard]] auto create_default_sbt() const -> SbtPair;
+        [[nodiscard]] auto create_sbt(BuildShaderBindingTableInfo const & info) const -> SbtPair;
         void get_shader_group_handles(void *out_blob) const;
+        auto get_shader_group_count() const -> u32;
 
       protected:
         template <typename T, typename H_T>

include/daxa/types.hpp

@@ -338,6 +338,21 @@ namespace daxa
         }
         SmallString(SmallString const & other) = default;
         SmallString & operator=(SmallString const & other) = default;
+        bool operator<(SmallString const & other) const
+        {
+            for (FixedListSizeT i = 0; i < std::min(this->m_size, other.m_size); ++i)
+            {
+                if (this->m_data[i] < other.m_data[i])
+                {
+                    return true;
+                }
+                if (this->m_data[i] > other.m_data[i])
+                {
+                    return false;
+                }
+            }
+            return this->m_size < other.m_size;
+        }
         [[nodiscard]] auto view() const -> std::string_view
         {
             return {this->m_data.data(), static_cast<usize>(this->m_size)};
@@ -1833,12 +1848,22 @@ namespace daxa
         none = 1000165000,
     };
 
-    // TODO: distinguish between GENERAL(raygen, miss & callable) cause shader handles must be set in order (raygen, miss, hit, callable)?
-    enum struct ShaderGroup
+    enum struct ExtendedShaderGroupType
+    {
+        RAYGEN = 0,
+        MISS = 1,
+        TRIANGLES_HIT_GROUP = 2,
+        PROCEDURAL_HIT_GROUP = 3,
+        CALLABLE = 4,
+        MAX_ENUM = 0x7fffffff,
+    };
+
+    enum struct ShaderGroupType
     {
-        GENERAL = 0,
-        TRIANGLES_HIT_GROUP = 1,
-        PROCEDURAL_HIT_GROUP = 2,
+        RAYGEN = 0,
+        MISS = 1,
+        HIT = 2,
+        CALLABLE = 3,
         MAX_ENUM = 0x7fffffff,
     };
 

include/daxa/utils/pipeline_manager.hpp

@@ -64,15 +64,16 @@ namespace daxa
         ShaderCompileOptions compile_options = {};
     };
 
+    struct RayTracingShaderCompileInfo
+    {
+        ShaderCompileInfo shader_info = {};
+        RayTracingShaderType type = {};
+    };
+
     struct RayTracingPipelineCompileInfo
     {
-        std::vector<ShaderCompileInfo> ray_gen_infos = {};
-        std::vector<ShaderCompileInfo> intersection_infos = {};
-        std::vector<ShaderCompileInfo> any_hit_infos = {};
-        std::vector<ShaderCompileInfo> callable_infos = {};
-        std::vector<ShaderCompileInfo> closest_hit_infos = {};
-        std::vector<ShaderCompileInfo> miss_hit_infos = {};
-        std::vector<RayTracingShaderGroupInfo> shader_groups_infos = {};
+        std::vector<RayTracingShaderCompileInfo> stages = {};
+        std::vector<RayTracingShaderGroupInfo> groups = {};
         u32 max_ray_recursion_depth = {};
         u32 push_constant_size = {};
         std::string name = {};

src/cpp_wrapper.cpp

@@ -861,19 +861,36 @@ namespace daxa
         auto result = SbtPair{};
         auto daxa_res = daxa_ray_tracing_pipeline_create_default_sbt(
             rc_cast<daxa_RayTracingPipeline>(this->object),
-            r_cast<daxa_RayTracingShaderBindingTable *>(&result.table),
+            r_cast<daxa_RayTracingShaderBindingTableEntries *>(&result.entries),
             r_cast<daxa_BufferId *>(&result.buffer));
         check_result(daxa_res, "failed in create_default_sbt");
         return result;
     }
 
+    auto RayTracingPipeline::create_sbt(BuildShaderBindingTableInfo const & info) const -> SbtPair
+    {
+        auto result = SbtPair{};
+        auto daxa_res = daxa_ray_tracing_pipeline_create_sbt(
+            rc_cast<daxa_RayTracingPipeline>(this->object),
+            r_cast<daxa_RayTracingShaderBindingTableEntries *>(&result.entries),
+            r_cast<daxa_BufferId *>(&result.buffer),
+            r_cast<daxa_BuildShaderBindingTableInfo const *>(&info));
+        check_result(daxa_res, "failed in create_default_sbt");
+        return result;
+    }
+
     void RayTracingPipeline::get_shader_group_handles(void * out_blob) const
     {
         auto daxa_res = daxa_ray_tracing_pipeline_get_shader_group_handles(
             rc_cast<daxa_RayTracingPipeline>(this->object), out_blob);
         check_result(daxa_res, "failed in get_shader_group_handles");
     }
 
+    auto RayTracingPipeline::get_shader_group_count() const -> u32
+    {
+        return daxa_ray_tracing_pipeline_get_shader_group_count(rc_cast<daxa_RayTracingPipeline>(this->object));
+    }
+
     auto RayTracingPipeline::inc_refcnt(ImplHandle const * object) -> u64
     {
         return daxa_ray_tracing_pipeline_inc_refcnt(rc_cast<daxa_RayTracingPipeline>(object));