yum-mirror/slang

Making it easier to work with shaders

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

Ellie HermaszewskaMove switch statement bodies to their own lines (#5493)b118451e3

master
branchesmasterall refs →
10.6 KiB352 linesraw
1// cuda-shader-object.cpp
2#include "cuda-shader-object.h"
3
4#include "cuda-helper-functions.h"
5#include "cuda-resource-views.h"
6#include "cuda-shader-object-layout.h"
7
8namespace gfx
9{
10#ifdef GFX_ENABLE_CUDA
11using namespace Slang;
12
13namespace cuda
14{
15Result ShaderObjectData::setCount(Index count)
16{
17    if (isHostOnly)
18    {
19        m_cpuBuffer.setCount(count);
20        if (!m_bufferView)
21        {
22            IResourceView::Desc viewDesc = {};
23            viewDesc.type = IResourceView::Type::UnorderedAccess;
24            m_bufferView = new ResourceViewImpl();
25            m_bufferView->proxyBuffer = m_cpuBuffer.getBuffer();
26            m_bufferView->m_desc = viewDesc;
27        }
28        return SLANG_OK;
29    }
30
31    if (!m_bufferResource)
32    {
33        IBufferResource::Desc desc;
34        desc.type = IResource::Type::Buffer;
35        desc.sizeInBytes = count;
36        m_bufferResource = new BufferResourceImpl(desc);
37        if (count)
38        {
39            SLANG_CUDA_RETURN_ON_FAIL(
40                cuMemAlloc((CUdeviceptr*)&m_bufferResource->m_cudaMemory, (size_t)count));
41        }
42        IResourceView::Desc viewDesc = {};
43        viewDesc.type = IResourceView::Type::UnorderedAccess;
44        m_bufferView = new ResourceViewImpl();
45        m_bufferView->memoryResource = m_bufferResource;
46        m_bufferView->m_desc = viewDesc;
47    }
48    auto oldSize = m_bufferResource->getDesc()->sizeInBytes;
49    if ((size_t)count != oldSize)
50    {
51        void* newMemory = nullptr;
52        if (count)
53        {
54            SLANG_CUDA_RETURN_ON_FAIL(cuMemAlloc((CUdeviceptr*)&newMemory, (size_t)count));
55        }
56        if (oldSize)
57        {
58            SLANG_CUDA_RETURN_ON_FAIL(cuMemcpy(
59                (CUdeviceptr)newMemory,
60                (CUdeviceptr)m_bufferResource->m_cudaMemory,
61                Math::Min((size_t)count, oldSize)));
62        }
63        cuMemFree((CUdeviceptr)m_bufferResource->m_cudaMemory);
64        m_bufferResource->m_cudaMemory = newMemory;
65        m_bufferResource->getDesc()->sizeInBytes = count;
66    }
67    return SLANG_OK;
68}
69
70Slang::Index ShaderObjectData::getCount()
71{
72    if (isHostOnly)
73        return m_cpuBuffer.getCount();
74    if (m_bufferResource)
75        return (Slang::Index)(m_bufferResource->getDesc()->sizeInBytes);
76    else
77        return 0;
78}
79
80void* ShaderObjectData::getBuffer()
81{
82    if (isHostOnly)
83        return m_cpuBuffer.getBuffer();
84
85    if (m_bufferResource)
86        return m_bufferResource->m_cudaMemory;
87    return nullptr;
88}
89
90/// Returns a resource view for GPU access into the buffer content.
91ResourceViewBase* ShaderObjectData::getResourceView(
92    RendererBase* device,
93    slang::TypeLayoutReflection* elementLayout,
94    slang::BindingType bindingType)
95{
96    SLANG_UNUSED(device);
97    m_bufferResource->getDesc()->elementSize = (int)elementLayout->getSize();
98    return m_bufferView.Ptr();
99}
100
101SlangResult ShaderObjectImpl::init(IDevice* device, ShaderObjectLayoutImpl* typeLayout)
102{
103    m_layout = typeLayout;
104
105    // If the layout tells us that there is any uniform data,
106    // then we need to allocate a constant buffer to hold that data.
107    //
108    // TODO: Do we need to allocate a shadow copy for use from
109    // the CPU?
110    //
111    // TODO: When/where do we bind this constant buffer into
112    // a descriptor set for later use?
113    //
114    auto slangLayout = getLayout()->getElementTypeLayout();
115    size_t uniformSize = slangLayout->getSize();
116    if (uniformSize)
117    {
118        m_data.setCount((Index)uniformSize);
119    }
120
121    // If the layout specifies that we have any resources or sub-objects,
122    // then we need to size the appropriate arrays to account for them.
123    //
124    // Note: the counts here are the *total* number of resources/sub-objects
125    // and not just the number of resource/sub-object ranges.
126    //
127    resources.setCount(typeLayout->getResourceCount());
128    m_objects.setCount(typeLayout->getSubObjectCount());
129
130    for (auto subObjectRange : getLayout()->subObjectRanges)
131    {
132        RefPtr<ShaderObjectLayoutImpl> subObjectLayout = subObjectRange.layout;
133
134        // In the case where the sub-object range represents an
135        // existential-type leaf field (e.g., an `IBar`), we
136        // cannot pre-allocate the object(s) to go into that
137        // range, since we can't possibly know what to allocate
138        // at this point.
139        //
140        if (!subObjectLayout)
141            continue;
142        //
143        // Otherwise, we will allocate a sub-object to fill
144        // in each entry in this range, based on the layout
145        // information we already have.
146
147        auto& bindingRangeInfo = getLayout()->m_bindingRanges[subObjectRange.bindingRangeIndex];
148        for (Index i = 0; i < bindingRangeInfo.count; ++i)
149        {
150            RefPtr<ShaderObjectImpl> subObject = new ShaderObjectImpl();
151            SLANG_RETURN_ON_FAIL(subObject->init(device, subObjectLayout));
152
153            ShaderOffset offset;
154            offset.uniformOffset = bindingRangeInfo.uniformOffset + sizeof(void*) * i;
155            offset.bindingRangeIndex = (GfxIndex)subObjectRange.bindingRangeIndex;
156            offset.bindingArrayIndex = (GfxIndex)i;
157
158            SLANG_RETURN_ON_FAIL(setObject(offset, subObject));
159        }
160    }
161    return SLANG_OK;
162}
163
164SLANG_NO_THROW GfxCount SLANG_MCALL ShaderObjectImpl::getEntryPointCount()
165{
166    return 0;
167}
168
169SLANG_NO_THROW Result SLANG_MCALL
170ShaderObjectImpl::getEntryPoint(GfxIndex index, IShaderObject** outEntryPoint)
171{
172    *outEntryPoint = nullptr;
173    return SLANG_OK;
174}
175
176SLANG_NO_THROW const void* SLANG_MCALL ShaderObjectImpl::getRawData()
177{
178    return m_data.getBuffer();
179}
180
181SLANG_NO_THROW Size SLANG_MCALL ShaderObjectImpl::getSize()
182{
183    return (Size)m_data.getCount();
184}
185
186SLANG_NO_THROW Result SLANG_MCALL
187ShaderObjectImpl::setData(ShaderOffset const& offset, void const* data, Size size)
188{
189    Size temp = m_data.getCount() - (Size)offset.uniformOffset;
190    size = Math::Min(size, temp);
191    SLANG_CUDA_RETURN_ON_FAIL(cuMemcpy(
192        (CUdeviceptr)((uint8_t*)m_data.getBuffer() + offset.uniformOffset),
193        (CUdeviceptr)data,
194        size));
195    return SLANG_OK;
196}
197
198SLANG_NO_THROW Result SLANG_MCALL
199ShaderObjectImpl::setResource(ShaderOffset const& offset, IResourceView* resourceView)
200{
201    if (!resourceView)
202        return SLANG_OK;
203
204    auto layout = getLayout();
205
206    auto bindingRangeIndex = offset.bindingRangeIndex;
207    SLANG_ASSERT(bindingRangeIndex >= 0);
208    SLANG_ASSERT(bindingRangeIndex < layout->m_bindingRanges.getCount());
209
210    auto& bindingRange = layout->m_bindingRanges[bindingRangeIndex];
211
212    auto viewIndex = bindingRange.baseIndex + offset.bindingArrayIndex;
213    auto cudaView = static_cast<ResourceViewImpl*>(resourceView);
214
215    resources[viewIndex] = cudaView;
216
217    if (cudaView->textureResource)
218    {
219        if (cudaView->m_desc.type == IResourceView::Type::UnorderedAccess)
220        {
221            auto handle = cudaView->textureResource->m_cudaSurfObj;
222            setData(offset, &handle, sizeof(uint64_t));
223        }
224        else
225        {
226            auto handle = cudaView->textureResource->getBindlessHandle();
227            setData(offset, &handle, sizeof(uint64_t));
228        }
229    }
230    else if (cudaView->memoryResource)
231    {
232        auto handle = cudaView->memoryResource->getBindlessHandle();
233        setData(offset, &handle, sizeof(handle));
234        auto sizeOffset = offset;
235        sizeOffset.uniformOffset += sizeof(handle);
236        auto& desc = *cudaView->memoryResource->getDesc();
237        size_t size = desc.sizeInBytes;
238        if (desc.elementSize > 1)
239            size /= desc.elementSize;
240        setData(sizeOffset, &size, sizeof(size));
241    }
242    else if (cudaView->proxyBuffer)
243    {
244        auto handle = cudaView->proxyBuffer;
245        setData(offset, &handle, sizeof(handle));
246        auto sizeOffset = offset;
247        sizeOffset.uniformOffset += sizeof(handle);
248        auto& desc = *cudaView->memoryResource->getDesc();
249        size_t size = desc.sizeInBytes;
250        if (desc.elementSize > 1)
251            size /= desc.elementSize;
252        setData(sizeOffset, &size, sizeof(size));
253    }
254    return SLANG_OK;
255}
256
257SLANG_NO_THROW Result SLANG_MCALL
258ShaderObjectImpl::setObject(ShaderOffset const& offset, IShaderObject* object)
259{
260    SLANG_RETURN_ON_FAIL(Super::setObject(offset, object));
261
262    auto bindingRangeIndex = offset.bindingRangeIndex;
263    auto& bindingRange = getLayout()->m_bindingRanges[bindingRangeIndex];
264
265    ShaderObjectImpl* subObject = static_cast<ShaderObjectImpl*>(object);
266    switch (bindingRange.bindingType)
267    {
268    default:
269        {
270            void* subObjectDataBuffer = subObject->getBuffer();
271            SLANG_RETURN_ON_FAIL(setData(offset, &subObjectDataBuffer, sizeof(void*)));
272        }
273        break;
274    case slang::BindingType::ExistentialValue:
275    case slang::BindingType::RawBuffer:
276    case slang::BindingType::MutableRawBuffer:
277        break;
278    }
279    return SLANG_OK;
280}
281
282SLANG_NO_THROW Result SLANG_MCALL
283ShaderObjectImpl::setSampler(ShaderOffset const& offset, ISamplerState* sampler)
284{
285    SLANG_UNUSED(sampler);
286    SLANG_UNUSED(offset);
287    return SLANG_OK;
288}
289
290SLANG_NO_THROW Result SLANG_MCALL ShaderObjectImpl::setCombinedTextureSampler(
291    ShaderOffset const& offset,
292    IResourceView* textureView,
293    ISamplerState* sampler)
294{
295    SLANG_UNUSED(sampler);
296    setResource(offset, textureView);
297    return SLANG_OK;
298}
299
300EntryPointShaderObjectImpl::EntryPointShaderObjectImpl()
301{
302    m_data.isHostOnly = true;
303}
304
305SLANG_NO_THROW uint32_t SLANG_MCALL RootShaderObjectImpl::addRef()
306{
307    return 1;
308}
309
310SLANG_NO_THROW uint32_t SLANG_MCALL RootShaderObjectImpl::release()
311{
312    return 1;
313}
314
315SlangResult RootShaderObjectImpl::init(IDevice* device, ShaderObjectLayoutImpl* typeLayout)
316{
317    SLANG_RETURN_ON_FAIL(ShaderObjectImpl::init(device, typeLayout));
318    auto programLayout = dynamic_cast<RootShaderObjectLayoutImpl*>(typeLayout);
319    for (auto& entryPoint : programLayout->entryPointLayouts)
320    {
321        RefPtr<EntryPointShaderObjectImpl> object = new EntryPointShaderObjectImpl();
322        SLANG_RETURN_ON_FAIL(object->init(device, entryPoint));
323        entryPointObjects.add(object);
324    }
325    return SLANG_OK;
326}
327
328SLANG_NO_THROW GfxCount SLANG_MCALL RootShaderObjectImpl::getEntryPointCount()
329{
330    return (GfxCount)entryPointObjects.getCount();
331}
332
333SLANG_NO_THROW Result SLANG_MCALL
334RootShaderObjectImpl::getEntryPoint(GfxIndex index, IShaderObject** outEntryPoint)
335{
336    returnComPtr(outEntryPoint, entryPointObjects[index]);
337    return SLANG_OK;
338}
339
340Result RootShaderObjectImpl::collectSpecializationArgs(ExtendedShaderObjectTypeList& args)
341{
342    SLANG_RETURN_ON_FAIL(ShaderObjectImpl::collectSpecializationArgs(args));
343    for (auto& entryPoint : entryPointObjects)
344    {
345        SLANG_RETURN_ON_FAIL(entryPoint->collectSpecializationArgs(args));
346    }
347    return SLANG_OK;
348}
349
350} // namespace cuda
351#endif
352} // namespace gfx