/**
 * This file contains the compute shaders used to sort the index map buffer such that our individual particles 
 * will be rendered in a back-to-front fashion relative to the rendering camera.
 * These shader programs implement a bitonic sorting network with individual GPU threads acting as nodes and is
 * based on the "ComputeShaderSort11" example included with the DirectX June 2010 SDK release.
 * It is implemented in such a way that (assuming the default block sizes are preserved) it can run on DX10 / SM4
 * hardware with the only tweak being to replace the Buffer / RWBuffer declarations below with 
 * StructuredBuffer / RWStructuredBuffer and making the corresponding change in the associated DBPro code.
 * Beware that the DBPro example program requires DX11 because it uses a set of simultaneously accessed append/consume
 * buffers (DX10 / SM4 only permits a single writable resource at any one time from a shader program, be it a RWBuffer 
 * or a render target image) so further tweaks would be required to get the entire example to run in DX10 mode.
 **/

// These defines should correspond to the constants having the same name in the DBPro source.
// If you wish to change them you must change both these and the DBP #constants.
// The BITONIC_BLOCK_SIZE corresponds to how many threads will perform sorting per thread group; 
// TRANSPOSE_BLOCK_SIZE ^ 2 threads will be issued for the MatrixTranspose shader program which is used
// to shift the data being sorted by envisioning it as matrices that are transposed between sorting 
// its rows and columns. See the DBPro code for how the transposes and partial sorting passes come together
// to finally yield a completely sorted buffer, as well as the MatrixTranspose program and the node sorting
// program below.
#define BITONIC_BLOCK_SIZE		512
#define TRANSPOSE_BLOCK_SIZE	16


// Include common declarations
#include "Declarations.hlsl"


// This constant buffer is used to provide per-pass data to the BitonicSort and MatrixTranspose shader programs alike
cbuffer PassData : register(b0) {
    unsigned int Level;
    unsigned int LevelMask;
    unsigned int MatrixWidth;
    unsigned int MatrixHeight;
};

// This constant buffer provides the (used) particle buffer size on a per-frame basis
cbuffer ParticleData : register(b1) {
	uint ParticleCount;
};


// The current sorted index map will be transferred between two buffers as a means of simplifying the MatrixTranspose
// shader program. The Input and Data buffers are then swapped by the DBPro source between each such transpose such that 
// the Input buffer contains the previous state that is being transposed into the Data buffer.
// The BitonicSort program only uses the (writable) Data buffer.
Buffer<unsigned int>	Input : register(t0);
RWBuffer<unsigned int>	Data  : register(u0);

// The particle buffer is needed to read the particle depth data which is used as the sorting criterion
StructuredBuffer<Particle> ParticleBuffer : register(t1);


// We store the current segment of data being sorted in groupshared memory.
// This is a small, very fast cache memory (usually L1 or L2 depending on graphics card) that is limited to 32Kb per thread group
// on DX11-capable hardware, or 16Kb on DX10 hardware. As the name implies, this memory block can be accessed by all threads in the same group.
// This increases the efficiency of the sorting algorithm, which is essentially a comparison-based sort, significantly compared to
// constantly reading the data from the buffers (which reside in VRAM).
// We need to store two values per element being sorted; its view space depth, which is what we sort based on, and the index into 
// the particle buffer to the particle having this particular depth value, which is what we are storing in the index lookup table buffer
// we are sorting.
struct sSortData {
	uint	id;
	float	depth;
};
groupshared sSortData shared_data[BITONIC_BLOCK_SIZE];

/**
 * The BitonicSort program will perform partial sorting of sub-lists of the entire buffer to be sorted of BITONIC_BLOCK_SIZE elements 
 * per thread group. To yield a sorted list which exceeds the BITONIC_BLOCK_SIZE number of elements, the data is viewed as a matrix
 * for which sub-matrices are identified and transposed by the MatrixTranspose program defined below. The DBPro source will perform 
 * the control flow that allows the sorting of partial lists and the subsequent matrix transpositions and further sortings that 
 * eventually results in the entire list being in sorted order.
 * Because this sorting algorithm is comparative, memory access speed will be the biggest performance bottleneck; as such we use 
 * groupshared memory, which is usually implemented as on-chip L1 or L2 cache by GPU manufacturers, for making efficient comparisons.
 **/
[numthreads(BITONIC_BLOCK_SIZE, 1, 1)]
void BitonicSort(uint3 threadId : SV_DispatchThreadID, uint groupId : SV_GroupIndex) {
    // Load shared data into groupshared (cache) memory
	uint localId				= Data[threadId.x];
    shared_data[groupId].id		= localId;
	// Since we are always sorting the full buffer we need to handle the case where we have unused particle indices being sorted.
	// This is the case when localId >= ParticleCount. To solve this, we simply set the depth used to sort the element to a very 
	// low value; this will result in the unused particle indices being sorted at the end of the index lookup buffer, from where
	// only the first ParticleCount number of elements will be read when rendering.
	shared_data[groupId].depth	= localId < ParticleCount ? ParticleBuffer[localId].depth : -9999999.0;	
    
	// This intrinsic function will cause the thread to be suspended until all other threads in the group reach this point of the 
	// shader program. This is required to ensure that we won't read from the shared data block until all threads have finished 
	// writing to it.
	GroupMemoryBarrierWithGroupSync();
    
    // Sort the shared data
    for (uint j = Level >> 1; j > 0; j >>= 1) {
		// Determine index into the Particle buffer based on minimum/maximum particle depth. Each set of elements will be considered by two threads 
		// where one choses the largest and the other the smallest.
		// On DX11 hardware we could technically swap both values directly in the shared memory, however for compatibility with DX10 hardware 
		// (where groupshared memory does not support unordered write operations) we will instead use this approach where two threads will be
		// used and each will chose one element to write. Linear memory writing like this is also prone to better optimization possibilities which 
		// makes it quite likely that it will be faster to write the data like this rather than having half the amount of threads performing two 
		// non-linear writes each.
        uint id		= ((shared_data[groupId & ~j].depth >= shared_data[groupId | j].depth) == (bool)(LevelMask & threadId.x))? groupId ^ j : groupId;
		// Store index and depth on the stack
		uint index	= shared_data[id].id;
		float depth = shared_data[id].depth;

		// Again we need to ensure that all threads in the group have finished reading from the shared data block before we continue to write
		// the swap-sorted elements back to it.
        GroupMemoryBarrierWithGroupSync();
        shared_data[groupId].id		= index;
		shared_data[groupId].depth	= depth;

		// And again, wait for all threads to finish writing before continuing with reading in the next iteration again (if any)
        GroupMemoryBarrierWithGroupSync();
    }
    
    // Store the partial sorted data stored in shared memory to the data buffer
    Data[threadId.x] = shared_data[groupId].id;
}


/**
 * This shader program will transpose imaginary matrices by viewing the data buffer as a 2D data set (a matrix) and flip its rows and columns.
 * By coupling this with sorting rows (and thus by extension, transposing to sort columns) using the BitonicSort program defined above we can
 * gradually piece together separatedly sorted sub-lists into one, cohesive sorted list (the data buffer).
 * We use groupshared memory for efficient memory operations here as well; as with all other declarations in a HLSL source, it cannot share
 * the name of another variable.
 * Here we actually use two dimensions for specifying threads and thread groups since it maps well onto the problem at hand (one thread per 
 * element in the imaginary 2D matrix to transpose) and allows for simpler code.
 **/
groupshared uint transpose_data[TRANSPOSE_BLOCK_SIZE * TRANSPOSE_BLOCK_SIZE];

[numthreads(TRANSPOSE_BLOCK_SIZE, TRANSPOSE_BLOCK_SIZE, 1)]
void MatrixTranspose(uint3 threadId : SV_DispatchThreadID, uint3 gtId : SV_GroupThreadID, uint groupId : SV_GroupIndex) {
    // This works quite similar to the BitonicSort program defined above. Here each thread will read a single element into
	// shared memory for temporary storage and then write it out, transposing the imaginary sub-matrix identified by the 
	// group index such that rows data is swapped with column data.
	transpose_data[groupId] = Input[threadId.y * MatrixWidth + threadId.x];
	// Wait for all threads to finish reading data
    GroupMemoryBarrierWithGroupSync();
	// And then write the transpose of the data to the Data buffer such that X is swapped with Y
    uint2 XY = threadId.yx - gtId.yx + gtId.xy;
    Data[XY.y * MatrixHeight + XY.x] = transpose_data[gtId.x * TRANSPOSE_BLOCK_SIZE + gtId.y];
}

/** 
 * This simple compute shader is called at the beginning of each sorting pass to reset the offset buffer.
 * It simply sets the initial state of the index lookup (ĞDatağ) buffer to be a list of incrementing indices starting at 0.
 **/
[numthreads(BITONIC_BLOCK_SIZE, 1, 1)]
void ResetOffsetBuffer(uint3 threadId : SV_DispatchThreadId) {
	Data[threadId.x] = threadId.x;
}