Skip to content

bitonic_sort #940

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 14 commits into from
+568 −0
Closed

bitonic_sort #940

Changes from 1 commit
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
5d6322a
created a structured like fft
CrabExtra Oct 2, 2025
264650c
config question
CrabExtra Oct 2, 2025
268949e
subgroupsort
CrabExtra Oct 5, 2025
73aa820
added bitonic_sort name space
CrabExtra Oct 5, 2025
dcb5e6e
Merge branch 'Devsh-Graphics-Programming:master' into master
CrabExtra Oct 7, 2025
b84a4bd
subgroup changes
CrabExtra Oct 7, 2025
779815e
removed unused
CrabExtra Oct 7, 2025
78a307a
Merge branch 'master' of https://github.com/CrabExtra/Nabla
CrabExtra Oct 7, 2025
ad7a4c5
added last merge step as a function
CrabExtra Oct 19, 2025
b80283a
uncomplete workgroup fn
CrabExtra Oct 19, 2025
7c91744
complete the logic for some pr questions
CrabExtra Oct 19, 2025
4d253f3
Refactor bitonic sort for workgroup + Accessor support
CrabExtra Oct 22, 2025
f03b8b2
Update bitonic_sort.hlsl
CrabExtra Oct 22, 2025
555dcbe
VT implumentation
CrabExtra Oct 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
211 changes: 146 additions & 65 deletions include/nbl/builtin/hlsl/workgroup/bitonic_sort.hlsl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,74 +1,155 @@
#ifndef NBL_BUILTIN_HLSL_SUBGROUP_BITONIC_SORT_INCLUDED
#define NBL_BUILTIN_HLSL_SUBGROUP_BITONIC_SORT_INCLUDED
#ifndef NBL_BUILTIN_HLSL_WORKGROUP_BITONIC_SORT_INCLUDED
#define NBL_BUILTIN_HLSL_WORKGROUP_BITONIC_SORT_INCLUDED
#include "nbl/builtin/hlsl/bitonic_sort/common.hlsl"
#include "nbl/builtin/hlsl/memory_accessor.hlsl"
#include "nbl/builtin/hlsl/functional.hlsl"
#include "nbl/builtin/hlsl/subgroup/bitonic_sort.hlsl"
#include "nbl/builtin/hlsl/bit.hlsl"
#include "nbl/builtin/hlsl/workgroup/shuffle.hlsl"
#include "nbl/builtin/hlsl/workgroup/basic.hlsl"

namespace nbl
{
namespace hlsl
namespace hlsl
{
namespace workgroup
{
namespace bitonic_sort
{
// Reorder: non-type parameters FIRST, then typename parameters with defaults
// This matches FFT's pattern and avoids DXC bugs
template<uint16_t _ElementsPerInvocationLog2, uint16_t _WorkgroupSizeLog2, typename KeyType, typename ValueType, typename Comparator = less<KeyType> >
struct bitonic_sort_config
{
using key_t = KeyType;
using value_t = ValueType;
using comparator_t = Comparator;

NBL_CONSTEXPR_STATIC_INLINE uint16_t ElementsPerInvocationLog2 = _ElementsPerInvocationLog2;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NBL_CONSTEXPR_STATIC_INLINE resolves to const static when preprocessed. DXC is stupid and when it sees a static it WILL initialize a variable. But if it sees const on its own it does compile it down to a constant, which is the behaviour you would expect (and what would happen in C++). So for now just replace these usages with just const.

@devshgraphicsprogramming do we just change this macro to resolve to const in HLSL in master?

NBL_CONSTEXPR_STATIC_INLINE uint16_t WorkgroupSizeLog2 = _WorkgroupSizeLog2;

NBL_CONSTEXPR_STATIC_INLINE uint32_t ElementsPerInvocation = 1u << ElementsPerInvocationLog2;
NBL_CONSTEXPR_STATIC_INLINE uint32_t WorkgroupSize = 1u << WorkgroupSizeLog2;
};
}

template<typename Config, class device_capabilities = void>
struct BitonicSort;


template<uint16_t ElementsPerInvocationLog2, uint16_t WorkgroupSizeLog2, typename KeyType, typename ValueType, typename Comparator, class device_capabilities>
struct BitonicSort<bitonic_sort::bitonic_sort_config<ElementsPerInvocationLog2, WorkgroupSizeLog2, KeyType, ValueType, Comparator>, device_capabilities>
{
using config_t = bitonic_sort::bitonic_sort_config<ElementsPerInvocationLog2, WorkgroupSizeLog2, KeyType, ValueType, Comparator>;
using key_t = KeyType;
using value_t = ValueType;
using comparator_t = Comparator;

using SortConfig = subgroup::bitonic_sort_config<key_t, value_t, comparator_t>;

template<typename SharedMemoryAccessor>
static void mergeStage(NBL_REF_ARG(SharedMemoryAccessor) sharedmemAccessor, uint32_t stage, bool bitonicAscending, uint32_t invocationID, NBL_REF_ARG(key_t) loKey, NBL_REF_ARG(key_t) hiKey,
NBL_REF_ARG(value_t) loVal, NBL_REF_ARG(value_t) hiVal)
{
namespace workgroup
NBL_CONSTEXPR_STATIC_INLINE uint32_t WorkgroupSize = config_t::WorkgroupSize;
using adaptor_t = accessor_adaptors::StructureOfArrays<SharedMemoryAccessor, key_t, value_t, 1, WorkgroupSize>;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks wrong. You're making the index type for the array be value_t. If your values were floats, for example, this would blow up. You are getting lucky here because (I guess) the types you're testing with are all integers.

You'll want the shared memory accessor to satisfy this concept:

Nabla/include/nbl/builtin/hlsl/concepts/accessors/generic_shared_data.hlsl

Line 15 in aa5d6cb

#define NBL_CONCEPT_NAME GenericSharedMemoryAccessor

That concept basically states that the accessor can read and write uint32_ts. This is because shared memory (in most architectures) works with certain restrictions due to memory banking and size per transaction for each bank. It is the adaptor that is later in charge of reading/writing from/to shared memory with your actual type.

What you want here is to have a SharedMemoryAccessor accessing shared memory that is at least max(sizeof(key_t), sizeof(value_t)) * ArraySize bytes (this is unenforceable via concepts but you can make a utility for the config that returns this value so the user can allocate such an array).

Then you want TWO different adaptors: one is going to be
key_adaptor = accessor_adaptors::StructureOfArrays<SharedMemoryAccessor, key_t, uint32_t, 1, WorkgroupSize>
and the other is going to be
value_adaptor = accessor_adaptors::StructureOfArrays<SharedMemoryAccessor, value_t, uint32_t, 1, WorkgroupSize>

You would then shuffle the keys using the key adaptor, barrier, then shuffle the values using the value adaptor

adaptor_t sharedmemAdaptor;
sharedmemAdaptor.accessor = sharedmemAccessor;

const uint32_t subgroupSizeLog2 = glsl::gl_SubgroupSizeLog2();

[unroll]
for (uint32_t pass = 0; pass <= stage; pass++)
{
namespace bitonic_sort
{
template<typename KeyType, typename ValueType, typename Comparator = less<KeyType> >
struct bitonic_sort_config
{
using key_t = KeyType;
using value_t = ValueType;
using comparator_t = Comparator;
};

template<bool Ascending, typename Config, class device_capabilities = void>
struct bitonic_sort;
template<bool Ascending, typename KeyType, typename ValueType, typename Comparator, class device_capabilities>
struct bitonic_sort<Ascending, bitonic_sort_config<KeyType, ValueType, Comparator>, device_capabilities>
{
using config_t = bitonic_sort_config<KeyType, ValueType, Comparator>;
using key_t = typename config_t::key_t;
using value_t = typename config_t::value_t;
using comparator_t = typename config_t::comparator_t;

using SortConfig = subgroup::bitonic_sort_config<uint32_t, uint32_t, less<uint32_t> >;


static void mergeWGStage(uint32_t stage, bool bitonicAscending, uint32_t invocationID, NBL_REF_ARG(key_t) loKey, NBL_REF_ARG(key_t) hiKey,
NBL_REF_ARG(value_t) loVal, NBL_REF_ARG(value_t) hiVal)
{
[unroll]
for (uint32_t pass = 0; pass <= stage; pass++)
{
const uint32_t stride = 1u << ((stage - pass) + subgroupSizeLog2); // Element stride shifts to inter-subgroup scale
// Shuffle from partner using WG XOR need to implument

}
}


static void __call(NBL_REF_ARG(key_t) loKey, NBL_REF_ARG(key_t) hiKey,
NBL_REF_ARG(value_t) loVal, NBL_REF_ARG(value_t) hiVal)
{
const uint32_t invocationID = glsl::gl_SubgroupInvocationID();
const uint32_t subgroupSizeLog2 = glsl::gl_SubgroupSizeLog2();

//first sort all subgroup inside wg
subgroup::bitonic_sort<true, SortConfig>::__call(loKey, hiKey, loVal, hiVal);
//then we go over first work group shuffle
//we have n = log2(x), where n is how many wgshuffle we have to do on x(subgroup num)

[unroll]
for (uint32_t stage = 1; stage <= n; ++stage)
{
mergeWGStage(stage, Ascending, invocationID, hiKey, loKey, loVal, hiVal);
subgroup::bitonic_sort<true, SortConfig>::lastMergeStage(subgroupSizeLog2, invocationIDloKey, hiKey, loKey,loVal, hiVal);
workgroupExecutionAndMemoryBarrier();
}


}
};

}
// Stride calculation: stage S merges 2^(S+1) subgroups
const uint32_t stridePower = (stage - pass + 1) + subgroupSizeLog2;
const uint32_t stride = 1u << stridePower;
const uint32_t threadStride = stride >> 1;

// Separate shuffles for lo/hi streams (two-round shuffle as per PR review)
// TODO: Consider single-round shuffle of key-value pairs for better performance
key_t pLoKey = loKey;
shuffleXor(pLoKey, threadStride, sharedmemAdaptor);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this would be a shuffle using the key_adaptor

value_t pLoVal = loVal;
shuffleXor(pLoVal, threadStride, sharedmemAdaptor);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

and this one would be a shuffle using the value_adaptor


key_t pHiKey = hiKey;
shuffleXor(pHiKey, threadStride, sharedmemAdaptor);
value_t pHiVal = hiVal;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inbetween shuffles, your array is aliased. The first shuffle has the following behaviour: all threads write values, then there's a barrier (so all threads are done writing before they start reading) then they start reading. On the next shuffle, they immediately start writing again. If you don't barrier inbetween these shuffles, you risk writing before some other thread was done reading, overwriting what needed to be read. Between shuffles, therefore, you need to barrier to unalias the memory

shuffleXor(pHiVal, threadStride, sharedmemAdaptor);

const bool isUpper = (invocationID & threadStride) != 0;
const bool takeLarger = isUpper == bitonicAscending;

comparator_t comp;

// lo update
const bool loSelfSmaller = comp(loKey, pLoKey);
const bool takePartnerLo = takeLarger ? loSelfSmaller : !loSelfSmaller;
loKey = takePartnerLo ? pLoKey : loKey;
loVal = takePartnerLo ? pLoVal : loVal;

// hi update
const bool hiSelfSmaller = comp(hiKey, pHiKey);
const bool takePartnerHi = takeLarger ? hiSelfSmaller : !hiSelfSmaller;
hiKey = takePartnerHi ? pHiKey : hiKey;
hiVal = takePartnerHi ? pHiVal : hiVal;

sharedmemAdaptor.workgroupExecutionAndMemoryBarrier();
}
}

template<typename Accessor, typename SharedMemoryAccessor>
static void __call(
NBL_REF_ARG(Accessor) accessor, NBL_REF_ARG(SharedMemoryAccessor) sharedmemAccessor,
NBL_REF_ARG(key_t) loKey, NBL_REF_ARG(key_t) hiKey,
NBL_REF_ARG(value_t) loVal, NBL_REF_ARG(value_t) hiVal)
{
NBL_CONSTEXPR_STATIC_INLINE uint32_t WorkgroupSize = config_t::WorkgroupSize;

const uint32_t invocationID = glsl::gl_LocalInvocationID().x;
const uint32_t subgroupSizeLog2 = glsl::gl_SubgroupSizeLog2();
const uint32_t subgroupSize = 1u << subgroupSizeLog2;
const uint32_t subgroupID = glsl::gl_SubgroupID();
const uint32_t numSubgroups = WorkgroupSize / subgroupSize;
const uint32_t numSubgroupsLog2 = findMSB(numSubgroups);


const bool subgroupAscending = (subgroupID & 1) == 0;
subgroup::bitonic_sort<SortConfig>::__call(subgroupAscending, loKey, hiKey, loVal, hiVal);


[unroll]
for (uint32_t stage = 0; stage < numSubgroupsLog2; ++stage)
{
const bool isLastStage = (stage == numSubgroupsLog2 - 1);
const bool bitonicAscending = isLastStage ? true : !bool(invocationID & (subgroupSize << (stage + 1)));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

might be wrong here but it feels like the formula on the right yields true even for the last stage (at that point the single 1 in the bitmask is too far to the left, so the result of the & is a 0)


mergeStage(sharedmemAccessor, stage, bitonicAscending, invocationID, loKey, hiKey, loVal, hiVal);

const uint32_t subgroupInvocationID = glsl::gl_SubgroupInvocationID();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pull this one out of the loop

subgroup::bitonic_sort<SortConfig>::mergeStage(subgroupSizeLog2, bitonicAscending, subgroupInvocationID, loKey, hiKey, loVal, hiVal);
}


// Final: ensure lo <= hi within each thread (for ascending sort)
comparator_t comp;
if (comp(hiKey, loKey))
{
// Swap keys
key_t tempKey = loKey;
loKey = hiKey;
hiKey = tempKey;
// Swap values
value_t tempVal = loVal;
loVal = hiVal;
hiVal = tempVal;
}
}
};

}
}
#endif
}

#endif