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 →
16.1 KiB535 linesraw
1#pragma once
2
3
4#include "core/slang-basic.h"
5#include "core/slang-short-list.h"
6#include "core/slang-virtual-object-pool.h"
7#include "slang-com-ptr.h"
8
9#include <d3d12.h>
10#include <dxgi.h>
11
12namespace gfx
13{
14
15/*! \brief A simple class to manage an underlying Dx12 Descriptor Heap. Allocations are made
16linearly in order. It is not possible to free individual allocations, but all allocations can be
17deallocated with 'deallocateAll'. */
18class D3D12DescriptorHeap
19{
20public:
21    typedef D3D12DescriptorHeap ThisType;
22
23    /// Initialize
24    Slang::Result init(
25        ID3D12Device* device,
26        int size,
27        D3D12_DESCRIPTOR_HEAP_TYPE type,
28        D3D12_DESCRIPTOR_HEAP_FLAGS flags);
29    /// Initialize with an array of handles copying over the representation
30    Slang::Result init(
31        ID3D12Device* device,
32        const D3D12_CPU_DESCRIPTOR_HANDLE* handles,
33        int numHandles,
34        D3D12_DESCRIPTOR_HEAP_TYPE type,
35        D3D12_DESCRIPTOR_HEAP_FLAGS flags);
36
37    /// Returns the number of slots that have been used
38    SLANG_FORCE_INLINE int getUsedSize() const { return m_currentIndex; }
39
40    /// Get the total amount of descriptors possible on the heap
41    SLANG_FORCE_INLINE int getTotalSize() const { return m_totalSize; }
42    /// Allocate a descriptor. Returns the index, or -1 if none left.
43    SLANG_FORCE_INLINE int allocate();
44    /// Allocate a number of descriptors. Returns the start index (or -1 if not possible)
45    SLANG_FORCE_INLINE int allocate(int numDescriptors);
46
47    ///
48    SLANG_FORCE_INLINE int placeAt(int index);
49
50    /// Deallocates all allocations, and starts allocation from the start of the underlying heap
51    /// again
52    SLANG_FORCE_INLINE void deallocateAll() { m_currentIndex = 0; }
53
54    /// Get the size of each
55    SLANG_FORCE_INLINE int getDescriptorSize() const { return m_descriptorSize; }
56
57    /// Get the GPU heap start
58    SLANG_FORCE_INLINE D3D12_GPU_DESCRIPTOR_HANDLE getGpuStart() const
59    {
60        return m_heap->GetGPUDescriptorHandleForHeapStart();
61    }
62    /// Get the CPU heap start
63    SLANG_FORCE_INLINE D3D12_CPU_DESCRIPTOR_HANDLE getCpuStart() const
64    {
65        return m_heap->GetCPUDescriptorHandleForHeapStart();
66    }
67
68    /// Get the GPU handle at the specified index
69    SLANG_FORCE_INLINE D3D12_GPU_DESCRIPTOR_HANDLE getGpuHandle(int index) const;
70    /// Get the CPU handle at the specified index
71    SLANG_FORCE_INLINE D3D12_CPU_DESCRIPTOR_HANDLE getCpuHandle(int index) const;
72
73    /// Get the underlying heap
74    SLANG_FORCE_INLINE ID3D12DescriptorHeap* getHeap() const { return m_heap; }
75
76    /// Ctor
77    D3D12DescriptorHeap();
78
79protected:
80    Slang::ComPtr<ID3D12Device> m_device;
81    Slang::ComPtr<ID3D12DescriptorHeap> m_heap; ///< The underlying heap being allocated from
82    int m_totalSize;                         ///< Total amount of allocations available on the heap
83    int m_currentIndex;                      ///< The current descriptor
84    int m_descriptorSize;                    ///< The size of each descriptor
85    D3D12_DESCRIPTOR_HEAP_FLAGS m_heapFlags; ///< The flags of the heap
86};
87
88/// A d3d12 descriptor, used as "backing storage" for a view.
89///
90/// This type is intended to be used to represent descriptors that
91/// are allocated and freed through a `D3D12GeneralDescriptorHeap`.
92struct D3D12Descriptor
93{
94    D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle;
95};
96
97/// An allocator for host-visible descriptors.
98///
99/// Unlike the `D3D12DescriptorHeap` type, this class allows for both
100/// allocation and freeing of descriptors, by maintaining a free list.
101///
102class D3D12GeneralDescriptorHeap : public Slang::RefObject
103{
104    ID3D12Device* m_device;
105    int m_chunkSize;
106    D3D12_DESCRIPTOR_HEAP_TYPE m_type;
107
108    D3D12DescriptorHeap m_heap;
109    Slang::VirtualObjectPool m_allocator;
110
111public:
112    int getSize() { return m_chunkSize; }
113
114    Slang::Result init(
115        ID3D12Device* device,
116        int chunkSize,
117        D3D12_DESCRIPTOR_HEAP_TYPE type,
118        D3D12_DESCRIPTOR_HEAP_FLAGS flag)
119    {
120        m_device = device;
121        m_chunkSize = chunkSize;
122        m_type = type;
123
124        SLANG_RETURN_ON_FAIL(m_heap.init(m_device, m_chunkSize, m_type, flag));
125        m_allocator.initPool(m_chunkSize);
126        return SLANG_OK;
127    }
128
129    SLANG_FORCE_INLINE D3D12_CPU_DESCRIPTOR_HANDLE getCpuHandle(int index) const
130    {
131        return m_heap.getCpuHandle(index);
132    }
133
134    SLANG_FORCE_INLINE D3D12_GPU_DESCRIPTOR_HANDLE getGpuHandle(int index) const
135    {
136        return m_heap.getGpuHandle(index);
137    }
138
139    int allocate(int count) { return m_allocator.alloc(count); }
140
141    Slang::Result allocate(D3D12Descriptor* outDescriptor)
142    {
143        // TODO: this allocator would take some work to make thread-safe
144
145        int index = m_allocator.alloc(1);
146        if (index < 0)
147        {
148            assert(!"descriptor allocation failed");
149            return SLANG_FAIL;
150        }
151
152        D3D12Descriptor descriptor;
153        descriptor.cpuHandle = m_heap.getCpuHandle(index);
154
155        *outDescriptor = descriptor;
156        return SLANG_OK;
157    }
158
159    void free(int index, int count) { m_allocator.free(index, count); }
160
161    void free(D3D12Descriptor descriptor)
162    {
163        auto index =
164            (int)(descriptor.cpuHandle.ptr - m_heap.getCpuStart().ptr) / m_heap.getDescriptorSize();
165        free(index, 1);
166    }
167};
168
169class D3D12GeneralExpandingDescriptorHeap : public Slang::RefObject
170{
171    ID3D12Device* m_device;
172    D3D12_DESCRIPTOR_HEAP_TYPE m_type;
173    D3D12_DESCRIPTOR_HEAP_FLAGS m_flag;
174    int m_chunkSize;
175    Slang::List<Slang::RefPtr<D3D12GeneralDescriptorHeap>> m_subHeaps;
176    Slang::List<int> m_subHeapStartingIndex;
177
178public:
179    Slang::Result newSubHeap()
180    {
181        Slang::RefPtr<D3D12GeneralDescriptorHeap> subHeap = new D3D12GeneralDescriptorHeap();
182        SLANG_RETURN_ON_FAIL(subHeap->init(m_device, m_chunkSize, m_type, m_flag));
183        m_subHeaps.add(subHeap);
184        if (m_subHeapStartingIndex.getCount())
185        {
186            m_subHeapStartingIndex.add(
187                m_subHeapStartingIndex.getLast() + m_subHeaps.getLast()->getSize());
188        }
189        else
190        {
191            m_subHeapStartingIndex.add(0);
192        }
193        return SLANG_OK;
194    }
195
196    int getSubHeapIndex(int descriptorIndex) const
197    {
198        Slang::Index l = 0;
199        Slang::Index r = m_subHeapStartingIndex.getCount();
200        while (l < r - 1)
201        {
202            Slang::Index m = l + (r - l) / 2;
203            if (m_subHeapStartingIndex[m] < descriptorIndex)
204                l = m;
205            else if (m_subHeapStartingIndex[m] > descriptorIndex)
206                r = m;
207            else
208                return (int)m;
209        }
210        assert(
211            m_subHeapStartingIndex[l] <= descriptorIndex &&
212            m_subHeapStartingIndex[l] + m_subHeaps[l]->getSize() > descriptorIndex);
213        return (int)l;
214    }
215
216    Slang::Result init(
217        ID3D12Device* device,
218        int chunkSize,
219        D3D12_DESCRIPTOR_HEAP_TYPE type,
220        D3D12_DESCRIPTOR_HEAP_FLAGS flag)
221    {
222        m_device = device;
223        m_chunkSize = chunkSize;
224        m_type = type;
225        m_flag = flag;
226
227        return newSubHeap();
228    }
229
230    SLANG_FORCE_INLINE D3D12_CPU_DESCRIPTOR_HANDLE getCpuHandle(int index) const
231    {
232        auto subHeapIndex = getSubHeapIndex(index);
233        return m_subHeaps[subHeapIndex]->getCpuHandle(index - m_subHeapStartingIndex[subHeapIndex]);
234    }
235
236    SLANG_FORCE_INLINE D3D12_GPU_DESCRIPTOR_HANDLE getGpuHandle(int index) const
237    {
238        auto subHeapIndex = getSubHeapIndex(index);
239        return m_subHeaps[subHeapIndex]->getGpuHandle(index - m_subHeapStartingIndex[subHeapIndex]);
240    }
241
242    int allocate(int count)
243    {
244        auto result = m_subHeaps.getLast()->allocate(count);
245        if (result == -1)
246        {
247            newSubHeap();
248            return allocate(count);
249        }
250        return result + m_subHeapStartingIndex.getLast();
251    }
252
253    Slang::Result allocate(D3D12Descriptor* outDescriptor)
254    {
255        int index = allocate(1);
256        if (index < 0)
257        {
258            assert(!"descriptor allocation failed");
259            return SLANG_FAIL;
260        }
261
262        D3D12Descriptor descriptor;
263        descriptor.cpuHandle = getCpuHandle(index);
264
265        *outDescriptor = descriptor;
266        return SLANG_OK;
267    }
268
269    void free(int index, int count)
270    {
271        auto subHeapIndex = getSubHeapIndex(index);
272        m_subHeaps[subHeapIndex]->free(index - m_subHeapStartingIndex[subHeapIndex], count);
273    }
274
275    void free(D3D12Descriptor descriptor)
276    {
277        for (auto& subHeap : m_subHeaps)
278        {
279            if (descriptor.cpuHandle.ptr >= subHeap->getCpuHandle(0).ptr)
280            {
281                auto subIndex = descriptor.cpuHandle.ptr - subHeap->getCpuHandle(0).ptr;
282                if (subIndex < (SIZE_T)subHeap->getSize())
283                {
284                    subHeap->free(descriptor);
285                    break;
286                }
287            }
288        }
289    }
290};
291
292class D3D12LinearExpandingDescriptorHeap : public Slang::RefObject
293{
294    ID3D12Device* m_device;
295    D3D12_DESCRIPTOR_HEAP_TYPE m_type;
296    D3D12_DESCRIPTOR_HEAP_FLAGS m_flag;
297    int m_chunkSize;
298    Slang::ShortList<D3D12DescriptorHeap, 4> m_subHeaps;
299    int32_t m_subHeapIndex;
300
301public:
302    Slang::Result newSubHeap()
303    {
304        m_subHeapIndex++;
305        if (m_subHeapIndex <= m_subHeaps.getCount())
306        {
307            D3D12DescriptorHeap subHeap;
308            SLANG_RETURN_ON_FAIL(subHeap.init(m_device, m_chunkSize, m_type, m_flag));
309            m_subHeaps.add(Slang::_Move(subHeap));
310        }
311        return SLANG_OK;
312    }
313
314    Slang::Result init(
315        ID3D12Device* device,
316        int chunkSize,
317        D3D12_DESCRIPTOR_HEAP_TYPE type,
318        D3D12_DESCRIPTOR_HEAP_FLAGS flag)
319    {
320        m_device = device;
321        m_chunkSize = chunkSize;
322        m_type = type;
323        m_flag = flag;
324        m_subHeapIndex = -1;
325        return newSubHeap();
326    }
327
328    int allocate(int count)
329    {
330        auto result = m_subHeaps[m_subHeapIndex].allocate(count);
331        if (result == -1)
332        {
333            newSubHeap();
334            return allocate(count);
335        }
336        assert(result <= 0xFFFFFF);
337        assert(m_subHeapIndex <= 255);
338        return (m_subHeapIndex << 24) + result;
339    }
340
341    SLANG_FORCE_INLINE D3D12_CPU_DESCRIPTOR_HANDLE getCpuHandle(int index) const
342    {
343        auto subHeapIndex = ((uint32_t)(index >> 24) & 0xFF);
344        return m_subHeaps[subHeapIndex].getCpuHandle(index & 0xFFFFFF);
345    }
346
347    void free(int index, int count) { assert(0 && "not supported"); }
348
349    void free(D3D12Descriptor descriptor) { assert(0 && "not supported"); }
350
351    void freeAll()
352    {
353        for (auto& subHeap : m_subHeaps)
354            subHeap.deallocateAll();
355        m_subHeapIndex = 0;
356    }
357};
358
359struct DescriptorHeapReference
360{
361    enum class Type
362    {
363        Linear,
364        General,
365        ExpandingGeneral,
366        ExpandingLinear
367    };
368    union Ptr
369    {
370        D3D12DescriptorHeap* linearHeap;
371        D3D12GeneralDescriptorHeap* generalHeap;
372        D3D12GeneralExpandingDescriptorHeap* generalExpandingHeap;
373        D3D12LinearExpandingDescriptorHeap* linearExpandingHeap;
374    };
375    Type type;
376    Ptr ptr;
377    DescriptorHeapReference() = default;
378    DescriptorHeapReference(D3D12DescriptorHeap* heap)
379    {
380        type = Type::Linear;
381        ptr.linearHeap = heap;
382    }
383    DescriptorHeapReference(D3D12GeneralDescriptorHeap* heap)
384    {
385        type = Type::General;
386        ptr.generalHeap = heap;
387    }
388    DescriptorHeapReference(D3D12GeneralExpandingDescriptorHeap* heap)
389    {
390        type = Type::ExpandingGeneral;
391        ptr.generalExpandingHeap = heap;
392    }
393    DescriptorHeapReference(D3D12LinearExpandingDescriptorHeap* heap)
394    {
395        type = Type::ExpandingLinear;
396        ptr.linearExpandingHeap = heap;
397    }
398    D3D12_CPU_DESCRIPTOR_HANDLE getCpuHandle(int index) const
399    {
400        switch (type)
401        {
402        case Type::Linear:
403            return ptr.linearHeap->getCpuHandle(index);
404        case Type::General:
405            return ptr.generalHeap->getCpuHandle(index);
406        case Type::ExpandingGeneral:
407            return ptr.generalExpandingHeap->getCpuHandle(index);
408        case Type::ExpandingLinear:
409            return ptr.linearExpandingHeap->getCpuHandle(index);
410        default:
411            return D3D12_CPU_DESCRIPTOR_HANDLE();
412        }
413    }
414    D3D12_GPU_DESCRIPTOR_HANDLE getGpuHandle(int index) const
415    {
416        switch (type)
417        {
418        case Type::Linear:
419            return ptr.linearHeap->getGpuHandle(index);
420        case Type::General:
421            return ptr.generalHeap->getGpuHandle(index);
422        case Type::ExpandingGeneral:
423            return ptr.generalExpandingHeap->getGpuHandle(index);
424        default:
425            return D3D12_GPU_DESCRIPTOR_HANDLE();
426        }
427    }
428    int allocate(int numDescriptors)
429    {
430        switch (type)
431        {
432        case Type::Linear:
433            return ptr.linearHeap->allocate(numDescriptors);
434        case Type::General:
435            return ptr.generalHeap->allocate(numDescriptors);
436        case Type::ExpandingGeneral:
437            return ptr.generalExpandingHeap->allocate(numDescriptors);
438        default:
439            return ptr.linearExpandingHeap->allocate(numDescriptors);
440        }
441    }
442    void free(int index, int count)
443    {
444        switch (type)
445        {
446        default:
447        case Type::Linear:
448            SLANG_ASSERT(!"Linear heap does not support free().");
449            break;
450        case Type::General:
451            return ptr.generalHeap->free(index, count);
452        case Type::ExpandingGeneral:
453            return ptr.generalExpandingHeap->free(index, count);
454        }
455    }
456    void freeIfSupported(int index, int count)
457    {
458        switch (type)
459        {
460        case Type::Linear:
461            return;
462        case Type::General:
463            return ptr.generalHeap->free(index, count);
464        case Type::ExpandingGeneral:
465            return ptr.generalExpandingHeap->free(index, count);
466        default:
467            break;
468        }
469    }
470};
471
472// ---------------------------------------------------------------------------
473int D3D12DescriptorHeap::allocate()
474{
475    return allocate(1);
476}
477// ---------------------------------------------------------------------------
478int D3D12DescriptorHeap::allocate(int numDescriptors)
479{
480    if (m_currentIndex + numDescriptors <= m_totalSize)
481    {
482        const int index = m_currentIndex;
483        m_currentIndex += numDescriptors;
484        return index;
485    }
486    if (m_heapFlags & D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE)
487    {
488        // No automatic resizing for GPU visible heaps.
489        return -1;
490    }
491    // We don't have enough heap size, resize the heap.
492    auto oldHeap = m_heap;
493    auto oldSize = m_totalSize;
494    auto currentIndex = m_currentIndex;
495    auto desc = m_heap->GetDesc();
496    this->init(m_device, (int)desc.NumDescriptors * 2, desc.Type, desc.Flags);
497    m_device->CopyDescriptorsSimple(
498        (UINT)currentIndex,
499        m_heap->GetCPUDescriptorHandleForHeapStart(),
500        oldHeap->GetCPUDescriptorHandleForHeapStart(),
501        desc.Type);
502    m_currentIndex = currentIndex;
503    // Now allocate again.
504    const int index = m_currentIndex;
505    m_currentIndex += numDescriptors;
506    return index;
507}
508// ---------------------------------------------------------------------------
509SLANG_FORCE_INLINE int D3D12DescriptorHeap::placeAt(int index)
510{
511    assert(index >= 0 && index < m_totalSize);
512    m_currentIndex = index + 1;
513    return index;
514}
515
516// ---------------------------------------------------------------------------
517SLANG_FORCE_INLINE D3D12_CPU_DESCRIPTOR_HANDLE D3D12DescriptorHeap::getCpuHandle(int index) const
518{
519    assert(index >= 0 && index < m_totalSize);
520    D3D12_CPU_DESCRIPTOR_HANDLE start = m_heap->GetCPUDescriptorHandleForHeapStart();
521    D3D12_CPU_DESCRIPTOR_HANDLE dst;
522    dst.ptr = start.ptr + m_descriptorSize * index;
523    return dst;
524}
525// ---------------------------------------------------------------------------
526SLANG_FORCE_INLINE D3D12_GPU_DESCRIPTOR_HANDLE D3D12DescriptorHeap::getGpuHandle(int index) const
527{
528    assert(index >= 0 && index < m_totalSize);
529    D3D12_GPU_DESCRIPTOR_HANDLE start = m_heap->GetGPUDescriptorHandleForHeapStart();
530    D3D12_GPU_DESCRIPTOR_HANDLE dst;
531    dst.ptr = start.ptr + m_descriptorSize * index;
532    return dst;
533}
534
535} // namespace gfx