Optimizing compute shader with thread group shared memory

Question

I'm currently learning compute shaders and I'm trying to write an optimized Game Of Life. I have a first version working that uses a Shader Storage Buffer Object. I dispatch a thread per cell I want to update and that thread samples the SSBO 8 times to gather the cell's neighbors. This works fine.

I'm now trying to optimize this by using shared memory. Every thread in a work group will now load a single cell in shared memory wait for the memory and execution barrier to resolve and then sample the shared memory 8 times to compute its cell's state. Of course some threads need to load more than 1 cell to shared memory because if they're on the 'border' of a workgroup no thread will load the data for some of its neighbors. ( See picture )

I got this working as well. However, it is less performant than the brute force version. ( On an RTX 2070 and GTX 660M ) I'm very surprised. I'm using Vulkan and gpu queries to estimate how long the compute pass takes. Here is for example the numbers for a 4096x4096 grid. I dispatch 256x256 thread groups of size 16x16. On my old GTX 660M compute takes ~12ms in brute force mode but it takes ~17ms with TGSM. I can't test on the RTX right now but the TGSM version was less performant as well.

I was definitely expecting to see some improvement with TGSM. I need to look into it more with tools like Nvidia Nsight but I would assume this shader is limited by memory and definitely not ALU. I guess my questions are:

Am I correct in assuming the TGSM version should run faster somehow? Search for Neighborhood Processing here
If my assumption is correct, what went wrong?

Here is the brute force shader:

#version 450

layout (constant_id = 0) const uint CELLS_COUNT          = 4096;
layout (constant_id = 1) const uint GRID_SIZE            = 64;        
layout (constant_id = 2) const uint THREADS_PER_GROUP_X  = 8;
layout (constant_id = 3) const uint THREADS_PER_GROUP_Y  = 8;
layout (constant_id = 4) const uint THREADS_PER_GROUP    = 64;

layout(local_size_x_id = 2, local_size_y_id = 3) in;

layout(std430, set = 0, binding = 0) buffer SrcGrid {
    uint state[CELLS_COUNT];
} srcGrid;

layout(std430, set = 0, binding = 1) buffer DstGrid {
    uint state[CELLS_COUNT];
} dstGrid;

const ivec2 sampleXYOffsets[] = { 
    ivec2(-1, -1),   ivec2(0, -1),  ivec2(1, -1),
    ivec2(-1,  0),                  ivec2(1,  0),
    ivec2(-1,  1),   ivec2(0,  1),  ivec2(1,  1),
};

void main() {
    
    const uint maxIdX = gl_WorkGroupSize.x * gl_NumWorkGroups.x;
    const uint maxIdY = gl_WorkGroupSize.y * gl_NumWorkGroups.y;

    uint aliveNeighbors = 0;
    // Convert dispatch IDs into storage buffer id
    const uint currentCellIndex = gl_GlobalInvocationID.y * maxIdX + gl_GlobalInvocationID.x;
    // Convert storage buffer id into grid coordinates (x, y)
    const uvec2 currentCoords = uvec2( currentCellIndex % GRID_SIZE, currentCellIndex / GRID_SIZE);
    for( uint i = 0; i < 8; i++ ){

        uvec2 coords = currentCoords;
        // Bring everything above 0 to be able to use the modulo operator
        coords = (coords + sampleXYOffsets[i] + GRID_SIZE) % GRID_SIZE;
        // Convert grid coordinates (x, y) into storage buffer id
        uint neighborIndex = coords.x + coords.y * GRID_SIZE;
        aliveNeighbors += srcGrid.state[neighborIndex];
    }

    uint currentCellState = srcGrid.state[currentCellIndex];

    if( currentCellState < 1.0 && aliveNeighbors == 3 )
    {
        // Dead cell comes back to life
        dstGrid.state[currentCellIndex] = 1;
        return;
    }

    // Alive cell dies
    if( aliveNeighbors < 2.0 || aliveNeighbors > 3.0)
    {
        dstGrid.state[currentCellIndex] = 0;
        return;
    }

    dstGrid.state[currentCellIndex] = currentCellState;
}

And the one using TGSM:

#version 450

layout (constant_id = 0) const uint CELLS_COUNT          = 4096;
layout (constant_id = 1) const uint GRID_SIZE            = 64;        
layout (constant_id = 2) const uint THREADS_PER_GROUP_X  = 8;
layout (constant_id = 3) const uint THREADS_PER_GROUP_Y  = 8;
layout (constant_id = 4) const uint THREADS_PER_GROUP    = 64;

layout(local_size_x_id = 2, local_size_y_id = 3) in;

layout(std430, set = 0, binding = 0) buffer SrcGrid {
    uint state[CELLS_COUNT];
} srcGrid;

layout(std430, set = 0, binding = 1) buffer DstGrid {
    uint state[CELLS_COUNT];
} dstGrid;

uint getSSBODataFromWorldGridCoords( uint worldGridX, uint worldGridY )
{
    uvec2 worldGridCoords = uvec2( worldGridX, worldGridY );
    worldGridCoords = (worldGridCoords + GRID_SIZE) % GRID_SIZE;
    return srcGrid.state[ worldGridCoords.x + worldGridCoords.y  * GRID_SIZE ];
}

shared uint sharedData[ THREADS_PER_GROUP_X + 2 ][ THREADS_PER_GROUP_Y + 2 ];

void main() 
{

    const uvec2 workGroupSize = uvec2( THREADS_PER_GROUP_X, THREADS_PER_GROUP_Y );

    // Grid coords within a thead group
    const uvec2 localGridCoords = gl_LocalInvocationID.xy;
    // Coordinates of each thread group tile
    const uvec2 worldGridOffset = gl_WorkGroupID.xy * workGroupSize.xy;
    // Grid coords within the whole game 
    const uvec2 worldGridCoords = worldGridOffset + localGridCoords; 

    // Early out if not in the grid
    if( worldGridCoords.x >= GRID_SIZE || worldGridCoords.y >= GRID_SIZE )
    {
        return;
    }

    // Load data into TGSM
    const uvec2 tgsmCoords = localGridCoords + uvec2( 1, 1 );

    // Top left corner 
    if( localGridCoords.x == 0 && localGridCoords.y == 0 )
    {
        // Top
        sharedData[ tgsmCoords.x     ][ tgsmCoords.y - 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x,     worldGridCoords.y - 1 );
        // Top Left
        sharedData[ tgsmCoords.x - 1 ][ tgsmCoords.y - 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x - 1, worldGridCoords.y - 1 );
        // Left
        sharedData[ tgsmCoords.x - 1 ][ tgsmCoords.y     ] = getSSBODataFromWorldGridCoords( worldGridCoords.x - 1, worldGridCoords.y     );
    }
    // Top right corner 
    else if( localGridCoords.x == workGroupSize.x - 1 && localGridCoords.y == 0 )
    {
        // Top
        sharedData[ tgsmCoords.x     ][ tgsmCoords.y - 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x,     worldGridCoords.y - 1 );
        // Top Right
        sharedData[ tgsmCoords.x + 1 ][ tgsmCoords.y - 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x + 1, worldGridCoords.y - 1 );
        // Right
        sharedData[ tgsmCoords.x + 1 ][ tgsmCoords.y     ] = getSSBODataFromWorldGridCoords( worldGridCoords.x + 1, worldGridCoords.y     );
    }
    // Bottom left corner 
    else if( localGridCoords.x == 0 && localGridCoords.y == workGroupSize.y - 1 )
    {
        // Bottom
        sharedData[ tgsmCoords.x     ][ tgsmCoords.y + 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x,     worldGridCoords.y + 1 );
        // Bottom Left
        sharedData[ tgsmCoords.x - 1 ][ tgsmCoords.y + 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x - 1, worldGridCoords.y + 1 );
        // Left
        sharedData[ tgsmCoords.x - 1 ][ tgsmCoords.y     ] = getSSBODataFromWorldGridCoords( worldGridCoords.x - 1, worldGridCoords.y     );
    }
    // Bottom right corner 
    else if( localGridCoords.x == workGroupSize.x - 1 && localGridCoords.y == workGroupSize.y - 1 )
    {
        // Bottom
        sharedData[ tgsmCoords.x     ][ tgsmCoords.y + 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x,     worldGridCoords.y + 1 );
        // Bottom Right
        sharedData[ tgsmCoords.x + 1 ][ tgsmCoords.y + 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x + 1, worldGridCoords.y + 1 );
        // Right
        sharedData[ tgsmCoords.x + 1 ][ tgsmCoords.y     ] = getSSBODataFromWorldGridCoords( worldGridCoords.x + 1, worldGridCoords.y     );
    }
    // Left Edge
    else if( localGridCoords.x == 0 )
    {
        sharedData[ tgsmCoords.x - 1 ][ tgsmCoords.y     ] = getSSBODataFromWorldGridCoords( worldGridCoords.x - 1, worldGridCoords.y     );
    }
    // Right Edge
    else if( localGridCoords.x == workGroupSize.x - 1 )
    {
        sharedData[ tgsmCoords.x + 1 ][ tgsmCoords.y     ] = getSSBODataFromWorldGridCoords( worldGridCoords.x + 1, worldGridCoords.y     );
    }
    // Top Edge
    else if( localGridCoords.y == 0 )
    {
        sharedData[ tgsmCoords.x     ][ tgsmCoords.y - 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x,     worldGridCoords.y - 1 );
    }
    // Bottom Edge
    else if( localGridCoords.y == workGroupSize.y - 1 )
    {
        sharedData[ tgsmCoords.x     ][ tgsmCoords.y + 1 ] = getSSBODataFromWorldGridCoords( worldGridCoords.x,     worldGridCoords.y + 1 );
    }

    // SSBO index for this thread
    const uint ssboIndex = worldGridCoords.x + worldGridCoords.y * GRID_SIZE;
    uint currentCellState = getSSBODataFromWorldGridCoords( worldGridCoords.x, worldGridCoords.y );

    // Load current cell into TGSM.
    sharedData[ tgsmCoords.x ][ tgsmCoords.y ] = currentCellState;

    // Barrier
    memoryBarrierShared();
    barrier();

    // Count alive neighbors from TGSM
    uint aliveNeighbors = 0;
    const uvec2 sharedMemoryGridCoordinates = uvec2( localGridCoords.x + 1, localGridCoords.y + 1 );
    aliveNeighbors += sharedData[ sharedMemoryGridCoordinates.x - 1 ][ sharedMemoryGridCoordinates.y - 1 ];
    aliveNeighbors += sharedData[ sharedMemoryGridCoordinates.x     ][ sharedMemoryGridCoordinates.y - 1 ];
    aliveNeighbors += sharedData[ sharedMemoryGridCoordinates.x + 1 ][ sharedMemoryGridCoordinates.y - 1 ];
    aliveNeighbors += sharedData[ sharedMemoryGridCoordinates.x - 1 ][ sharedMemoryGridCoordinates.y     ];
    aliveNeighbors += sharedData[ sharedMemoryGridCoordinates.x + 1 ][ sharedMemoryGridCoordinates.y     ];
    aliveNeighbors += sharedData[ sharedMemoryGridCoordinates.x - 1 ][ sharedMemoryGridCoordinates.y + 1 ];
    aliveNeighbors += sharedData[ sharedMemoryGridCoordinates.x     ][ sharedMemoryGridCoordinates.y + 1 ];
    aliveNeighbors += sharedData[ sharedMemoryGridCoordinates.x + 1 ][ sharedMemoryGridCoordinates.y + 1 ];

    if( currentCellState < 1.0 && aliveNeighbors == 3 )
    {
        // Dead cell comes back to life
        dstGrid.state[ssboIndex] = 1;
        return;
    }

    // Alive cell dies
    if( aliveNeighbors < 2.0 || aliveNeighbors > 3.0)
    {
        dstGrid.state[ssboIndex] = 0;
        return;
    }

    dstGrid.state[ssboIndex] = currentCellState;
}

Any pointers towards articles or tools that could help me figure this out would be greatly appreciated :).

misdake · Accepted Answer · 2022-02-08 09:39:21Z

1

Your second shader reads 10x10 pixels of data and writes 8x8 (which is correct), however, it handles corner and edge cases with lots of "if"s. This leads to divergence and kills performance.

One way to avoid divergence is that, a workgroup reads NxN (1 thread to 1 pixel), and only writes (N-2)x(N-2) pixels in the middle, discard corner and edge cases. Dispatch with N-2 stride.

answered Feb 8, 2022 at 9:39

misdake

111 bronze badge

$\begingroup$ Wouldn't the if statements be the same, just for writing instead of reading? I'm also not sure I understand what you mean with the N-2 stride. $\endgroup$
– François Guthmann
Feb 8, 2022 at 14:08
$\begingroup$ There are 8 different if blocks for each special case, all threads in a subgroup will have to wait for every cases to finish. Like, all threads (32 threads for example) in subgroup 0 will wait for thread0 (corner) to fetch 3 pixels, then wait for thread1-6 to fetch 1 pixel, ... . Texture fetches are quite serialized. $\endgroup$
– misdake
Feb 9, 2022 at 3:36
$\begingroup$ An optmization can be done to your second shader. Instead of letting corner threads to fetch 3 more pixels, see those extra corner/edge pixels as jobs and assign them to threads (store jobs in a constant array), so that each thread will fetch zero or one more pixels with minimal divergence (one uniform "if" instead of eight). pseudo code "fetch(x, y); if (hasJob[localIndex]) { fetch(x + job[localIndex].x, y + job[localIndex].y); }" $\endgroup$
– misdake
Feb 9, 2022 at 8:00
$\begingroup$ The key is to use one uniform "if" instead of 8 different ones. My method is to read uniformly and write with a mask. pseudo code "if (localcoord.xy > 0 and < N-1) { write data; }". Let N=16, Map=256x256, workgroups read 16x16 pixels, write 14x14 pixels => dispatch (w=ceil(256/14), h=ceil(256/14)) workgroups, and change coord/cellindex calculation accordingly. $\endgroup$
– misdake
Feb 9, 2022 at 8:08
$\begingroup$ Ah right, that makes a lot of sense! ( both solutions ) I was confused by the N-2 stride thinking you meant a workgroup of size N-2xN-2. I haven't touched this code in a while but I might find some time to look into it again. Thanks for answering this old post! $\endgroup$
– François Guthmann
Feb 10, 2022 at 15:47

Add a comment |

Антон Васин · Accepted Answer · 2021-08-20 18:22:11Z

Just a guess here.

The minimal physical workgroup on Nvidia GPUs is so-called 'warp', which is 32 threads wide. AMD's so-called 'wavefront' is 64 threads wide. This means, if one of these 32/64 virtual threads goes to memory (boundary condition), all of the rest ones have to wait, because warp performs a single instruction for all of the threads. It is a vector instruction. It's called 'thread divergency"

Your workgroup is 8x8 which is 64 virtual scalar threads or 1 or 2 real vector threads. This means the warp/wavefront goes to memory twice. The second fetch is less effective, than the first one.

You can try either expand the workgroup, for example 32x32 threads, or redesign your code to make if fetch the data only once. The out-of-boundary threads may be disabled after the prefetch stage.

I'd recommend having the amount of threads to be between 128 and 512 for better occupancy. Also, in brute force case instead of shared memory the cache is being used, possibly L1 cache, which on nVidia cards uses same memory as shared. Your software "cache" would only work if there is some known behaviour that you could exploit. You could try expanding the shared memory and do multiple fetches before calculaing results with multiple passes. — user369070, Aug 26, 2021 at 9:31

Stack Exchange Network

Optimizing compute shader with thread group shared memory

2 Answers 2

Your Answer

Not the answer you're looking for? Browse other questions tagged
gpu
performance
compute-shader
simulation
or ask your own question.

Hot Network Questions

Optimizing compute shader with thread group shared memory

2 Answers 2

Your Answer

Sign up or log in

Post as a guest

Not the answer you're looking for? Browse other questions tagged gpuperformancecompute-shadersimulation or ask your own question.

Related

Hot Network Questions

Not the answer you're looking for? Browse other questions tagged
gpu
performance
compute-shader
simulation
or ask your own question.