Monday, September 15, 2014

Finding Two Max Values in an Array in Shader

Go to Index

One of the problems I was faced was to, given four floats, find the two highest values and their indices in the array. This had to be done in a shader program so ideally no loops, arrays indexing or things like that. Fortunately my sort-of-a-boss, KriS, quickly found a solution that I'm presenting here.

Here comes the code:
void GetTwoMaxValues( in float4 values, out float2 weights, out float2 indices )
{
    float v1 = max( max( max( values.x, values.y ), values.z ), values.w );
    float i1 = 0;
    i1 = CndMask( values.y == v1, 1, i1 );
    i1 = CndMask( values.z == v1, 2, i1 );
    i1 = CndMask( values.w == v1, 3, i1 );

    //

    values.x = CndMask( i1 == 0, 0, values.x );
    values.y = CndMask( i1 == 1, 0, values.y );
    values.z = CndMask( i1 == 2, 0, values.z );
    values.w = CndMask( i1 == 3, 0, values.w );

    float v2 = max( max( max( values.x, values.y ), values.z ), values.w );
    float i2 = 0;
    i2 = CndMask( values.y == v2, 1, i2 );
    i2 = CndMask( values.z == v2, 2, i2 );
    i2 = CndMask( values.w == v2, 3, i2 );

    //

    weights.x = v1;
    weights.y = v2;
    indices.x = i1;
    indices.y = i2;
}
The code finds the first maximum value simply by "iterating" over elements and "max'ing" them. That was simple. But we also need to know the index of the value in the array. We assume it's the 0th index and compare the found value with other values from the array using CndMask statement. If the compared value is equal to the value of the element in the array, we update the index.

To find the second highest values we first zeroe the actual highest element found in the previous pass and do the same procedure as in the previous pass to find the second highest element.

The code can actually be understood much quicker by actually reading it rather than reading explanation trying to explain it :). It doesn't do anything magical in algorithmic terms but shows how to solve (efficiently) this simple problem of two max values in a shader program.