yum-mirror/slang

Making it easier to work with shaders

git clone https://git.yummers.dev/yum-mirror/slang

Ellie Hermaszewskaformatf65d756bf

master
branchesmasterall refs →
5.8 KiB155 linesraw
1// vk-shader-table.cpp
2#include "vk-shader-table.h"
3
4#include "vk-device.h"
5#include "vk-helper-functions.h"
6#include "vk-transient-heap.h"
7
8namespace gfx
9{
10
11using namespace Slang;
12
13namespace vk
14{
15
16RefPtr<BufferResource> ShaderTableImpl::createDeviceBuffer(
17    PipelineStateBase* pipeline,
18    TransientResourceHeapBase* transientHeap,
19    IResourceCommandEncoder* encoder)
20{
21    auto vkApi = m_device->m_api;
22    auto rtProps = vkApi.m_rtProperties;
23    uint32_t handleSize = rtProps.shaderGroupHandleSize;
24    m_raygenTableSize = m_rayGenShaderCount * rtProps.shaderGroupBaseAlignment;
25    m_missTableSize = (uint32_t)VulkanUtil::calcAligned(
26        m_missShaderCount * handleSize,
27        rtProps.shaderGroupBaseAlignment);
28    m_hitTableSize = (uint32_t)VulkanUtil::calcAligned(
29        m_hitGroupCount * handleSize,
30        rtProps.shaderGroupBaseAlignment);
31    m_callableTableSize = (uint32_t)VulkanUtil::calcAligned(
32        m_callableShaderCount * handleSize,
33        rtProps.shaderGroupBaseAlignment);
34    uint32_t tableSize = m_raygenTableSize + m_missTableSize + m_hitTableSize + m_callableTableSize;
35
36    auto pipelineImpl = static_cast<RayTracingPipelineStateImpl*>(pipeline);
37    ComPtr<IBufferResource> bufferResource;
38    IBufferResource::Desc bufferDesc = {};
39    bufferDesc.memoryType = MemoryType::DeviceLocal;
40    bufferDesc.defaultState = ResourceState::General;
41    bufferDesc.allowedStates =
42        ResourceStateSet(ResourceState::General, ResourceState::CopyDestination);
43    bufferDesc.type = IResource::Type::Buffer;
44    bufferDesc.sizeInBytes = tableSize;
45    static_cast<vk::DeviceImpl*>(m_device)->createBufferResourceImpl(
46        bufferDesc,
47        VK_BUFFER_USAGE_SHADER_BINDING_TABLE_BIT_KHR,
48        nullptr,
49        bufferResource.writeRef());
50
51    TransientResourceHeapImpl* transientHeapImpl =
52        static_cast<TransientResourceHeapImpl*>(transientHeap);
53
54    IBufferResource* stagingBuffer = nullptr;
55    Offset stagingBufferOffset = 0;
56    transientHeapImpl
57        ->allocateStagingBuffer(tableSize, stagingBuffer, stagingBufferOffset, MemoryType::Upload);
58
59    assert(stagingBuffer);
60    void* stagingPtr = nullptr;
61    stagingBuffer->map(nullptr, &stagingPtr);
62
63    List<uint8_t> handles;
64    auto handleCount = pipelineImpl->shaderGroupCount;
65    auto totalHandleSize = handleSize * handleCount;
66    handles.setCount(totalHandleSize);
67    auto result = vkApi.vkGetRayTracingShaderGroupHandlesKHR(
68        m_device->m_device,
69        pipelineImpl->m_pipeline,
70        0,
71        (uint32_t)handleCount,
72        totalHandleSize,
73        handles.getBuffer());
74
75    uint8_t* stagingBufferPtr = (uint8_t*)stagingPtr + stagingBufferOffset;
76    auto subTablePtr = stagingBufferPtr;
77    Int shaderTableEntryCounter = 0;
78
79    // Each loop calculates the copy source and destination locations by fetching the name
80    // of the shader group from the list of shader group names and getting its corresponding
81    // index in the buffer of handles.
82    for (uint32_t i = 0; i < m_rayGenShaderCount; i++)
83    {
84        auto dstHandlePtr = subTablePtr + i * rtProps.shaderGroupBaseAlignment;
85        auto shaderGroupName = m_shaderGroupNames[shaderTableEntryCounter++];
86        auto shaderGroupIndexPtr =
87            pipelineImpl->shaderGroupNameToIndex.tryGetValue(shaderGroupName);
88        if (!shaderGroupIndexPtr)
89            continue;
90
91        auto shaderGroupIndex = *shaderGroupIndexPtr;
92        auto srcHandlePtr = handles.getBuffer() + shaderGroupIndex * handleSize;
93        memcpy(dstHandlePtr, srcHandlePtr, handleSize);
94        memset(dstHandlePtr + handleSize, 0, rtProps.shaderGroupBaseAlignment - handleSize);
95    }
96    subTablePtr += m_raygenTableSize;
97
98    for (uint32_t i = 0; i < m_missShaderCount; i++)
99    {
100        auto dstHandlePtr = subTablePtr + i * handleSize;
101        auto shaderGroupName = m_shaderGroupNames[shaderTableEntryCounter++];
102        auto shaderGroupIndexPtr =
103            pipelineImpl->shaderGroupNameToIndex.tryGetValue(shaderGroupName);
104        if (!shaderGroupIndexPtr)
105            continue;
106
107        auto shaderGroupIndex = *shaderGroupIndexPtr;
108        auto srcHandlePtr = handles.getBuffer() + shaderGroupIndex * handleSize;
109        memcpy(dstHandlePtr, srcHandlePtr, handleSize);
110    }
111    subTablePtr += m_missTableSize;
112
113    for (uint32_t i = 0; i < m_hitGroupCount; i++)
114    {
115        auto dstHandlePtr = subTablePtr + i * handleSize;
116        auto shaderGroupName = m_shaderGroupNames[shaderTableEntryCounter++];
117        auto shaderGroupIndexPtr =
118            pipelineImpl->shaderGroupNameToIndex.tryGetValue(shaderGroupName);
119        if (!shaderGroupIndexPtr)
120            continue;
121
122        auto shaderGroupIndex = *shaderGroupIndexPtr;
123        auto srcHandlePtr = handles.getBuffer() + shaderGroupIndex * handleSize;
124        memcpy(dstHandlePtr, srcHandlePtr, handleSize);
125    }
126    subTablePtr += m_hitTableSize;
127
128    for (uint32_t i = 0; i < m_callableShaderCount; i++)
129    {
130        auto dstHandlePtr = subTablePtr + i * handleSize;
131        auto shaderGroupName = m_shaderGroupNames[shaderTableEntryCounter++];
132        auto shaderGroupIndexPtr =
133            pipelineImpl->shaderGroupNameToIndex.tryGetValue(shaderGroupName);
134        if (!shaderGroupIndexPtr)
135            continue;
136
137        auto shaderGroupIndex = *shaderGroupIndexPtr;
138        auto srcHandlePtr = handles.getBuffer() + shaderGroupIndex * handleSize;
139        memcpy(dstHandlePtr, srcHandlePtr, handleSize);
140    }
141    subTablePtr += m_callableTableSize;
142
143    stagingBuffer->unmap(nullptr);
144    encoder->copyBuffer(bufferResource, 0, stagingBuffer, stagingBufferOffset, tableSize);
145    encoder->bufferBarrier(
146        1,
147        bufferResource.readRef(),
148        gfx::ResourceState::CopyDestination,
149        gfx::ResourceState::ShaderResource);
150    RefPtr<BufferResource> resultPtr = static_cast<BufferResource*>(bufferResource.get());
151    return _Move(resultPtr);
152}
153
154} // namespace vk
155} // namespace gfx