bitonic_sort

[CrabExtra]

Description

Testing

TODO list:

[devshgraphicsprogrammingjenkins]

[CI]: Can one of the admins verify this patch?

[Fletterio]

The compiler is probably dumb and might not realize the right term is the negation of the left term. Ternaries in SPIR-V usually get compiled to an OpSelect which treats both terms after the ? not as branches to conditionally execute, but as operands whose result must be evaluated before the select operation runs. That is to say, if the compiler is stupid you're going to run two comparisons. If you make the right term the negation of the left one, CSE is likely to kick in and evaluate the comparison only once.

[Fletterio]

don't add indentation after compiler directives

[Fletterio]

I get that Ascending is used because when moving onto workgroup you're going to need to call alternating subgroup sorts. However, as a front-facing API if I wanted a single subgroup shuffle I'd usually want it in the order specified by the Comparator. Maybe push it after the Config and give it a default value of true. Or better yet, since Ascending can be confusing, consider calling it ReverseOrder or something simpler that conveys the intent better

[Fletterio]

Ascending and later names like takeLarger implicitly assume the comparator is less (lo and hi don't, those are related to the "lane" order in the bitonic sort diagram). That's fine on its own, it makes the code more readable vs naming them with a more generic option. However, there should be comments mentioning that names assume this implicitly so there's no confusion.

[Fletterio]

Are you sure this isn't reversed? Assume a less comparator, bitonicAscending = true for the current stage and upperHalf = true for the current thread. Then takeLarger semantically conveys that this thread wants to keep the larger values. And yet this code assigns the smaller values.

[Fletterio]

Again, if the compiler is dumb this code is very costly: half your threads in a subgroup will have upperHalf = true and the other half will have it set to false. Parallel code execution needs to be uniform across threads in the same SM, so this section of code will run twice: first some half of your threads (say, those in the upper half) will run, then the other half. This kills your throughput.

Inside each branch, the inner ifs will likely get compiled down to two OpSelects each. You can make this whole code branchless by doing

loKey = loCondition ? loKey : pLoKey;
loVal = loCondition ? loVal : pLoVal;
hiKey = hiCondition ? hiKey : pHiKey;
hiVal = hiCondition ? hiVal : pHiVal;

where loCondition and hiCondition are predicates that depend on both takeLarger and the result of the key comparison

[Fletterio]

In the end this is just mergeStage with bitonicAscending = true, right? I think you can just have mergeStage and avoid having this function duplicated

[Fletterio]

unlike the other method, both threads keep the min elements? Like upperHalf is not being considred here, so I'm inclined to believe this function is going to fail. I'd delete this method and just use mergeStage, since this is just that method but with a forced bitonicAscending = true.

[Fletterio]

Test this code to make sure it's right, but it feels wrong. Either way, just use mergeStage and avoid having this duped.

[Fletterio]

stage == subgroupSizeLog2is never true in this loop, so just assign the term for the false clause.

[Fletterio]

We already have a workgroup shuffle:

Nabla/include/nbl/builtin/hlsl/workgroup/shuffle.hlsl

Line 94 in aa5d6cb

void shuffleXor(NBL_REF_ARG(T) value, uint32_t mask, NBL_REF_ARG(SharedMemoryAdaptor) sharedmemAdaptor)

You need to template this bitonic sort workgroup struct on a shared memory accessor. You can either do the shuffle in two rounds (shuffle keys -> barrier -> shuffle values -> barrier) or in a single round, by shuffling a pair of both key and value together. Ideally this would be a setting you can choose as well. For now settle on one and leave a comment on how we can also consider the other case down the line (two rounds is probably easier at this stage, and useful for the bigger array sizes)

[Fletterio]

Ascending should be a parameter you pass to __call and not a template parameter. That way, you can control whether this starting subgroup sort is ascending or descending (notice that whenever you do a bigger-than-subgroup sort, some subgroup sorts are descending and some ascending, depending on the parity of the subgroupID. This is a condition you can't control from compiletime since the subgroupID is only known at runtime

[Fletterio]

This is also true of the workgroup struct btw. When you try to do virtual threading down the line, you will have to do the same thing based on "virtual workgroup ID", which will be known at runtime. So nbl::hlsl::workgroup::bitonic_sort::bitonic_sort should not be templated on Ascensing, it should instead be a parameter you pass to __call

[Fletterio]

On that note, nbl::hlsl::workgroup::bitonic_sort::bitonic_sort is an ugly namespace. The struct you use to run a bitonic sort should be nbl::hlsl::workgroup::BitonicSort (and similarly for subgroup). nbl::hlsl::workgroup::bitonic_sort should have structs related to the bitonic sort, but not the functional struct itself.

In the FFT, for example, we have the struct nbl::hlsl::workgroup::FFT to run the FFT, and nbl::hlsl::workgroup::fft is a namespace hat has structs useful for running an FFT, such as the config struct it's templated on

[Fletterio]

This is already in the workgroup namespace, you can just call it mergeStage

[Fletterio]

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?

[Fletterio]

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

[Fletterio]

this would be a shuffle using the key_adaptor

[Fletterio]

and this one would be a shuffle using the value_adaptor

[Fletterio]

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

[Fletterio]

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)

[Fletterio]

pull this one out of the loop

[Fletterio]

Make this branchless like you did the swaps in the subgroup branch

[Fletterio]

It feels like both the lo and hi update can be expressed using the compareAndSwap method above, using !takeLarger (or maybe takeLarger but I feel it's negated) instead of ascending

[Fletterio]

You are using the same accessor here (the adaptors might be different, but notice you create both from the same SharedMemoryAccessor and therefore they both address the same memory), these shuffles are therefore not independent. Imagine this scenario: thread 0 is running this code.

On the first shuffleXor, it will write loKey to position 0 into the sharedmem array, wait on a workgroup barrier. Meanwhile, thread threadStride will write its own loKey (which is pLoKey for thread 0, since 0 XOR threadStride = threadStride and similarly block on the same barrier.

The thing that thread 0 wants to do next is therefore to read pLoKey from position threadStride. But what if after the barrier, it's thread threadStride that gets to run first? Well, this thread wants to read its own pLoKey from position 0 into the array to complete the first shuffleXor, and immediately afterwards it's going to call the next shuffleXor, which will cause it to write its loVal into position threadStride of the array. If all of this happens before thread 0 had read its own pLoKey from that position, you've overwritten what that thread wanted to get

[Fletterio]

So what do we do? There's two options:

1. You ask the user to provide two sharedmem accessors that access different memory, one will be used to shuffle keys and the other will be used to shuffle values.

2. The user provides a single shared memory accessor with enough space for all keys and all values (max(sizeof(key_t), sizeof(value_t)) * WorkgroupSize) but you MUST barrier between each consecutive shuffle. You only call the barrier on the adaptor that was used last. Also, if pass > 0 you also barrier before the first shuffleXor, because you might have aliased memory from the last shuffleXor of the previous loop iteration

[Fletterio]

Ideally you'd want this to also be something you can choose in the Config struct (use more sharedmem with less barriers or use a single array with more barriers) but that's an optimization for the future. Right now just do the second one.

[Fletterio]

Don't barrier at the end of an iteration, but rather at the start of every iteration past the first. This saves one barrier at the end of the call, which might be unnecessary.

[Fletterio]

This is again compareAndSwap, identical to the one in the subgroup code. Make that method branchless like you did the inter-thread swaps in the subgroup sort, pull it out into bitonic_sort/common.hlsl then reuse it

[Fletterio]

Since loIdx is even you can get away with doing this addition as hiIdx = loIdx | 1

[Fletterio]

Doesn't the last pass of subgroup::bitonic_sort<SortConfig>::mergeStage ensure this already?

[Fletterio]

make these const

[Fletterio]

Make these const as well