descriptor heap rewrite

This commit is contained in:
Jorijn van der Graaf 2026-04-05 22:53:59 +02:00
commit f8e142fb06
31 changed files with 429 additions and 1017 deletions

View file

@ -30,26 +30,16 @@ import :ShaderBindingTableVulkan;
import :Types;
export namespace Crafter {
class PipelineRTVulkan {
public:
struct PipelineRTVulkan {
VkPipeline pipeline;
VkPipelineLayout pipelineLayout;
std::vector<std::uint8_t> shaderHandles;
VulkanBuffer<std::uint8_t, true, true, false> sbtBuffer;
VulkanBuffer<std::uint8_t, true, true> sbtBuffer;
VkStridedDeviceAddressRegionKHR raygenRegion;
VkStridedDeviceAddressRegionKHR missRegion;
VkStridedDeviceAddressRegionKHR hitRegion;
VkStridedDeviceAddressRegionKHR callableRegion;
void Init(VkCommandBuffer cmd, std::span<VkDescriptorSetLayout> setLayouts, std::span<VkRayTracingShaderGroupCreateInfoKHR> raygenGroups, std::span<VkRayTracingShaderGroupCreateInfoKHR> missGroups, std::span<VkRayTracingShaderGroupCreateInfoKHR> hitGroups, ShaderBindingTableVulkan& shaderTable) {
VkPipelineLayoutCreateInfo pipelineLayoutInfo {
.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
.setLayoutCount = static_cast<std::uint32_t>(setLayouts.size()),
.pSetLayouts = setLayouts.data()
};
Device::CheckVkResult(vkCreatePipelineLayout(Device::device, &pipelineLayoutInfo, nullptr, &pipelineLayout));
void Init(VkCommandBuffer cmd, std::span<VkRayTracingShaderGroupCreateInfoKHR> raygenGroups, std::span<VkRayTracingShaderGroupCreateInfoKHR> missGroups, std::span<VkRayTracingShaderGroupCreateInfoKHR> hitGroups, ShaderBindingTableVulkan& shaderTable) {
std::vector<VkRayTracingShaderGroupCreateInfoKHR> groups;
groups.reserve(raygenGroups.size() + missGroups.size() + hitGroups.size());
@ -57,14 +47,21 @@ export namespace Crafter {
groups.insert(groups.end(), missGroups.begin(), missGroups.end());
groups.insert(groups.end(), hitGroups.begin(), hitGroups.end());
VkRayTracingPipelineCreateInfoKHR rtPipelineInfo{
VkPipelineCreateFlags2CreateInfo flags2 = {
.sType = VK_STRUCTURE_TYPE_PIPELINE_CREATE_FLAGS_2_CREATE_INFO,
.flags = VK_PIPELINE_CREATE_2_DESCRIPTOR_HEAP_BIT_EXT
};
VkRayTracingPipelineCreateInfoKHR rtPipelineInfo {
.sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR,
.pNext = &flags2,
.flags = 0,
.stageCount = static_cast<std::uint32_t>(shaderTable.shaderStages.size()),
.pStages = shaderTable.shaderStages.data(),
.groupCount = static_cast<std::uint32_t>(groups.size()),
.pGroups = groups.data(),
.maxPipelineRayRecursionDepth = 1,
.layout = pipelineLayout
.layout = VK_NULL_HANDLE
};
Device::CheckVkResult(Device::vkCreateRayTracingPipelinesKHR(Device::device, {}, {}, 1, &rtPipelineInfo, nullptr, &pipeline));
@ -114,189 +111,6 @@ export namespace Crafter {
callableRegion.size = 0;
}
};
template <std::uint32_t GeneralShader, std::uint32_t ClosestHitShader, std::uint32_t AnyHitShader, std::uint32_t IntersectionShader>
struct ShaderGroup {
static constexpr std::uint32_t generalShader = GeneralShader;
static constexpr std::uint32_t closestHitShader = ClosestHitShader;
static constexpr std::uint32_t anyHitShader = AnyHitShader;
static constexpr std::uint32_t intersectionShader = IntersectionShader;
};
template <typename Shaders, typename ShaderGroups>
class PipelineRTVulkanConst {
public:
inline static VkPipeline pipeline;
inline static VkPipelineLayout pipelineLayout;
inline static std::vector<std::uint8_t> shaderHandles;
inline static VulkanBuffer<std::uint8_t, true, true, false> sbtBuffer;
inline static VkStridedDeviceAddressRegionKHR raygenRegion;
inline static VkStridedDeviceAddressRegionKHR missRegion;
inline static VkStridedDeviceAddressRegionKHR hitRegion;
inline static VkStridedDeviceAddressRegionKHR callableRegion;
static void Init(VkCommandBuffer cmd, const std::span<const VkDescriptorSetLayout> setLayouts) {
VkPipelineLayoutCreateInfo pipelineLayoutInfo {
.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
.setLayoutCount = static_cast<std::uint32_t>(setLayouts.size()),
.pSetLayouts = setLayouts.data()
};
Device::CheckVkResult(vkCreatePipelineLayout(Device::device, &pipelineLayoutInfo, nullptr, &pipelineLayout));
constexpr auto groupIndexSeq = std::make_index_sequence<std::tuple_size_v<ShaderGroups>>{};
constexpr std::array<VkRayTracingShaderGroupCreateInfoKHR, std::tuple_size_v<ShaderGroups>> groups = GetShaderGroups(groupIndexSeq);
VkRayTracingPipelineCreateInfoKHR rtPipelineInfo{
.sType = VK_STRUCTURE_TYPE_RAY_TRACING_PIPELINE_CREATE_INFO_KHR,
.stageCount = static_cast<std::uint32_t>(ShaderBindingTableVulkanConst<Shaders>::shaderStages.size()),
.pStages = ShaderBindingTableVulkanConst<Shaders>::shaderStages.data(),
.groupCount = static_cast<std::uint32_t>(groups.size()),
.pGroups = groups.data(),
.maxPipelineRayRecursionDepth = 1,
.layout = pipelineLayout
};
Device::CheckVkResult(Device::vkCreateRayTracingPipelinesKHR(Device::device, {}, {}, 1, &rtPipelineInfo, nullptr, &pipeline));
std::size_t dataSize = Device::rayTracingProperties.shaderGroupHandleSize * rtPipelineInfo.groupCount;
shaderHandles.resize(dataSize);
Device::CheckVkResult(Device::vkGetRayTracingShaderGroupHandlesKHR(Device::device, pipeline, 0, rtPipelineInfo.groupCount, dataSize, shaderHandles.data()));
std::uint32_t sbtStride = AlignUp(Device::rayTracingProperties.shaderGroupHandleSize, Device::rayTracingProperties.shaderGroupHandleAlignment);
raygenRegion.stride = sbtStride;
raygenRegion.deviceAddress = 0;
raygenRegion.size = GetGroupCount<VK_SHADER_STAGE_RAYGEN_BIT_KHR>(groupIndexSeq) * sbtStride;
missRegion.stride = sbtStride;
missRegion.deviceAddress = AlignUp(raygenRegion.size, Device::rayTracingProperties.shaderGroupBaseAlignment);
missRegion.size = GetGroupCount<VK_SHADER_STAGE_MISS_BIT_KHR>(groupIndexSeq) * sbtStride;
hitRegion.stride = sbtStride;
hitRegion.deviceAddress = AlignUp(missRegion.deviceAddress + missRegion.size, Device::rayTracingProperties.shaderGroupBaseAlignment);
hitRegion.size = (GetGroupCount<VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR>(groupIndexSeq) * sbtStride) + (GetGroupCount<VK_SHADER_STAGE_ANY_HIT_BIT_KHR>(groupIndexSeq) * sbtStride);
std::size_t bufferSize = hitRegion.deviceAddress + hitRegion.size;
sbtBuffer.Create(VK_BUFFER_USAGE_2_SHADER_BINDING_TABLE_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, bufferSize);
AddShaderGroupsToBuffer(sbtStride, groupIndexSeq);
sbtBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR);
raygenRegion.deviceAddress += sbtBuffer.address;
missRegion.deviceAddress += sbtBuffer.address;
hitRegion.deviceAddress += sbtBuffer.address;
callableRegion.deviceAddress = 0;
callableRegion.stride = 0;
callableRegion.size = 0;
}
private:
template<std::size_t index>
consteval static void AddShaderGroup(std::array<VkRayTracingShaderGroupCreateInfoKHR, std::tuple_size_v<ShaderGroups>>& groups) {
using groupTemplate = std::tuple_element_t<index, ShaderGroups>;
VkRayTracingShaderGroupTypeKHR groupType;
if constexpr(groupTemplate::generalShader != VK_SHADER_UNUSED_KHR) {
groupType = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
} else if constexpr(groupTemplate::closestHitShader != VK_SHADER_UNUSED_KHR || groupTemplate::anyHitShader != VK_SHADER_UNUSED_KHR) {
groupType = VK_RAY_TRACING_SHADER_GROUP_TYPE_TRIANGLES_HIT_GROUP_KHR;
} else {
static_assert(
groupTemplate::generalShader != VK_SHADER_UNUSED_KHR ||
groupTemplate::closestHitShader != VK_SHADER_UNUSED_KHR,
"Shader group must define either a general or closest-hit shader"
);
}
groups[index] = {
.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR,
.type = groupType,
.generalShader = groupTemplate::generalShader,
.closestHitShader = groupTemplate::closestHitShader,
.anyHitShader = groupTemplate::anyHitShader,
.intersectionShader = groupTemplate::intersectionShader
};
}
template<std::size_t... Is>
consteval static std::array<VkRayTracingShaderGroupCreateInfoKHR, std::tuple_size_v<ShaderGroups>> GetShaderGroups(std::index_sequence<Is...>) {
std::array<VkRayTracingShaderGroupCreateInfoKHR, std::tuple_size_v<ShaderGroups>> groups{};
(AddShaderGroup<Is>(groups), ...);
return groups;
}
template<std::size_t index, VkShaderStageFlagBits stage>
consteval static void GetGroupCountImpl(std::uint32_t& count) {
using groupTemplate = std::tuple_element_t<index, ShaderGroups>;
if constexpr(groupTemplate::generalShader != VK_SHADER_UNUSED_KHR) {
using shaderTemplate = std::tuple_element_t<groupTemplate::generalShader, Shaders>;
if constexpr(shaderTemplate::_stage == stage) {
count++;
}
} else if constexpr(groupTemplate::closestHitShader != VK_SHADER_UNUSED_KHR) {
using shaderTemplate = std::tuple_element_t<groupTemplate::closestHitShader, Shaders>;
if constexpr(shaderTemplate::_stage == stage) {
count++;
}
} else if constexpr(groupTemplate::anyHitShader != VK_SHADER_UNUSED_KHR) {
using shaderTemplate = std::tuple_element_t<groupTemplate::anyHitShader, Shaders>;
if constexpr(shaderTemplate::_stage == stage) {
count++;
}
} else {
static_assert(
groupTemplate::generalShader != VK_SHADER_UNUSED_KHR ||
groupTemplate::closestHitShader != VK_SHADER_UNUSED_KHR,
"Shader group must define either a general or closest-hit shader"
);
}
}
template<VkShaderStageFlagBits stage, std::size_t... Is>
consteval static std::uint32_t GetGroupCount(std::index_sequence<Is...>) {
std::uint32_t count = 0;
(GetGroupCountImpl<Is, stage>(count), ...);
return count;
}
template<std::size_t index, VkShaderStageFlagBits stage>
static void AddShaderGroupToBuffer(std::uint32_t sbtStride, std::uint32_t& offset) {
using groupTemplate = std::tuple_element_t<index, ShaderGroups>;
if constexpr(groupTemplate::generalShader != VK_SHADER_UNUSED_KHR) {
using shaderTemplate = std::tuple_element_t<groupTemplate::generalShader, Shaders>;
if constexpr(shaderTemplate::_stage == stage) {
std::memcpy(sbtBuffer.value + offset, shaderHandles.data() + index * Device::rayTracingProperties.shaderGroupHandleSize, Device::rayTracingProperties.shaderGroupHandleSize);
offset += sbtStride;
}
} else if constexpr(groupTemplate::closestHitShader != VK_SHADER_UNUSED_KHR) {
using shaderTemplate = std::tuple_element_t<groupTemplate::closestHitShader, Shaders>;
if constexpr(shaderTemplate::_stage == stage) {
std::memcpy(sbtBuffer.value + offset, shaderHandles.data() + index * Device::rayTracingProperties.shaderGroupHandleSize, Device::rayTracingProperties.shaderGroupHandleSize);
offset += sbtStride;
}
} else if constexpr(groupTemplate::anyHitShader != VK_SHADER_UNUSED_KHR) {
using shaderTemplate = std::tuple_element_t<groupTemplate::anyHitShader, Shaders>;
if constexpr(shaderTemplate::_stage == stage) {
std::memcpy(sbtBuffer.value + offset, shaderHandles.data() + index * Device::rayTracingProperties.shaderGroupHandleSize, Device::rayTracingProperties.shaderGroupHandleSize);
offset += sbtStride;
}
} else {
static_assert(
groupTemplate::generalShader != VK_SHADER_UNUSED_KHR ||
groupTemplate::closestHitShader != VK_SHADER_UNUSED_KHR,
"Shader group must define either a general or closest-hit shader"
);
}
}
template<std::size_t... Is>
static void AddShaderGroupsToBuffer(std::uint32_t sbtStride, std::index_sequence<Is...>) {
std::uint32_t offset = 0;
(AddShaderGroupToBuffer<Is, VK_SHADER_STAGE_RAYGEN_BIT_KHR>(sbtStride, offset), ...);
offset = AlignUp(offset, Device::rayTracingProperties.shaderGroupBaseAlignment);
(AddShaderGroupToBuffer<Is, VK_SHADER_STAGE_MISS_BIT_KHR>(sbtStride, offset), ...);
offset = AlignUp(offset, Device::rayTracingProperties.shaderGroupBaseAlignment);
(AddShaderGroupToBuffer<Is, VK_SHADER_STAGE_CLOSEST_HIT_BIT_KHR>(sbtStride, offset), ...);
(AddShaderGroupToBuffer<Is, VK_SHADER_STAGE_ANY_HIT_BIT_KHR>(sbtStride, offset), ...);
}
};
}
#endif