Nabla-Examples-and-Tests/14_BitonicSort/app_resources/common.hlsl at 5a4e46400a04ceb67063a5289059dd8137f7ba34 · CrabExtra/Nabla-Examples-and-Tests · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
#ifndef _BITONIC_SORT_COMMON_INCLUDED_
#define _BITONIC_SORT_COMMON_INCLUDED_
#include "nbl/builtin/hlsl/cpp_compat.hlsl"
#include "nbl/builtin/hlsl/workgroup/basic.hlsl"
#include "nbl/builtin/hlsl/functional.hlsl"

struct PushConstantData
{
	uint64_t deviceBufferAddress;
};

NBL_CONSTEXPR uint32_t WorkgroupSizeLog2 = 10;
NBL_CONSTEXPR uint32_t ElementsPerThreadLog2 = 2;
NBL_CONSTEXPR uint32_t elementCount = uint32_t(1) << (WorkgroupSizeLog2 + ElementsPerThreadLog2);

template<typename KeyType>
struct WorkgroupType
{
    KeyType key;
    uint32_t workgroupRelativeIndex;
};

template<typename KeyType, uint32_t KeyBits, typename StorageType = uint32_t>
struct SubgroupType
{
    static const StorageType KeyMask = (StorageType(1) << KeyBits) - 1;
    StorageType packed;

    static SubgroupType create(KeyType key, uint32_t subgroupRelativeIndex)
    {
        SubgroupType st;
        st.packed = (StorageType(key) & KeyMask) | (StorageType(subgroupRelativeIndex) << KeyBits);
        return st;
    }

    KeyType getKey() { return KeyType(packed & KeyMask); }
    uint32_t getSubgroupRelativeIndex() { return packed >> KeyBits; }

    WorkgroupType<KeyType> toWorkgroupType()
    {
        WorkgroupType<KeyType> wt;
        wt.key = getKey();
        wt.workgroupRelativeIndex = nbl::hlsl::workgroup::SubgroupContiguousIndex();
        return wt;
    }
};

template<typename KeyType>
struct KeyComparator
{
    bool operator()(KeyType a, KeyType b)
    {
        return a < b;
    }

    bool operator()(nbl::hlsl::pair<KeyType, KeyType> a, nbl::hlsl::pair<KeyType, KeyType> b)
    {
        return a.first < b.first;
    }

    bool operator()(WorkgroupType<KeyType> a, WorkgroupType<KeyType> b)
    {
        return a.key < b.key;
    }

    template<uint32_t KeyBits, typename StorageType>
    bool operator()(SubgroupType<KeyType, KeyBits, StorageType> a,
                    SubgroupType<KeyType, KeyBits, StorageType> b)
    {
        return a.getKey() < b.getKey();
    }
};

#endif