2
\$\begingroup\$

EDIT

Per request I've provided a simple compute shader that reproduces the problem. The compute shader below is kind of a contrived example meant to compute the rolling average over an input structured buffer.

By defining "USE_FUNCTION" certain parts of the algorithm will be split up into a separate function. This results in additional temporary registers being used by the compiled code as described in the original question.

The Unity version used to compile the compute shader is 2020.3.0f1 Personal.

#pragma kernel ComputeRollingAverage

// Comment out "USE_FUNCTION" define to inline code manually.
#define USE_FUNCTION
#define SAMPLE_COUNT 10
#define CODE \
count = 0; \
 \
for (uint offset = 0; offset < SAMPLE_COUNT; offset++) \
{ \
    if (index + offset >= size) \
    { \
        break; \
    } \
    samples[count++] += input[index + offset]; \
}

uint size;                                         

StructuredBuffer<float> input;
RWStructuredBuffer<float> output;

void InitializeSamples(out float samples[SAMPLE_COUNT])
{
    for (int iterator = 0; iterator < SAMPLE_COUNT; iterator++)
    {
        samples[iterator] = 0.0;
    }
}

void SampleInput(uint index, inout float samples[SAMPLE_COUNT], out uint count)
{
    CODE
}

[numthreads(64, 1, 1)]
void ComputeRollingAverage(uint3 id : SV_DispatchThreadID)
{
    uint index = id.x;
    float samples[SAMPLE_COUNT];
    uint count;
    InitializeSamples(samples);
    
#ifdef USE_FUNCTION
    SampleInput(index, samples, count);
#else
    CODE
#endif
    
    float rollingAverage = 0.0f;
    
    for (uint iterator = 0; iterator < count; iterator++)
    {
        rollingAverage += samples[iterator];
    }
    output[id.x] = rollingAverage / count;
}

Original Question Below

I'm currently working on a compute shader in Unity (implementing Cubical Marching Squares) and noticed that splitting up certain parts of the algorithm into separate functions causes the number of needed temporary registers (in below's case dcl_indexableTemp registers) to increase.

The only workaround I found is manually inlining everything, basically going back to one single, giant kernel function. That works, but the code gets harder to read and maintain.

Below I've provided a small snippet of my code (with parts omitted) illustrating the problem. At the beginning I create a local array of Segment structs of size 12:

[numthreads(4, 4, 4)]
void GenerateMesh(uint3 id : SV_DispatchThreadID)
{
    if (IsOutOfVoxelBounds(id))
    {
        return;
    }

    Segment segments[12];
    uint segmentsCount = 0;

    for (uint faceIndex = 0; faceIndex < 6; faceIndex++)
    {
        // >>> Calculate face segments.
        uint4 faceSampleIndices = voxelFaceSampleIndices[faceIndex];
        uint segmentsIndex = 0;

        for (uint sampleIndex = 0; sampleIndex < 4; sampleIndex++)
        {
            uint3 position = voxelStride * (id + voxelCorners[faceSampleIndices[sampleIndex]]);
            HermiteData sample = hermiteVolume[CalculateHermiteIndex(position)];
            segmentsIndex |= (GetHermiteDataDensity(sample) < 0.0) << sampleIndex;
        }

        uint4 voxelFace = voxelFaces[faceIndex];
        uint4 faceSegments = marchingSquaresSegments[segmentsIndex];

        for (uint faceSegmentsIndex = 0; faceSegmentsIndex < 4; faceSegmentsIndex += 2)
        {
            if (faceSegments[faceSegmentsIndex] == -1)
            {
                break;
            }

            Segment segment;
            InitializeSegment(segment, 0, 0, 0, 0, 0);
            segments[segmentsCount++] = segment;
        }
        // <<<
    }

    if (segmentsCount == 0)
    {
        return;
    }
    
    // additional code omitted
}

A single Segment struct uses up two 4-component registers. So the whole local array uses up 2 * 12 = 24 4-component registers. When looking at the compiled code through Unity (snippet below) the local array shows up as dcl_indexableTemp x0[24], 4 as expected. Everything fine so far.

dcl_uav_structured u0, 16
dcl_uav_structured_opc u1, 24
dcl_uav_structured u2, 12
dcl_input vThreadID.xyz
dcl_temps 22
dcl_indexableTemp x0[24], 4
dcl_indexableTemp x1[8], 4
dcl_indexableTemp x2[2], 4
dcl_indexableTemp x3[8], 4
dcl_thread_group 4, 4, 4

I now put the parts between // >>> Calculate face segments. and // <<< into a separate function and call it from the kernel like so:

void GenerateSegments(uint3 id, uint faceIndex, inout Segment segments[12], inout uint segmentsCount)
{
    uint4 faceSampleIndices = voxelFaceSampleIndices[faceIndex];
    uint segmentsIndex = 0;

    for (uint sampleIndex = 0; sampleIndex < 4; sampleIndex++)
    {
        uint3 position = voxelStride * (id + voxelCorners[faceSampleIndices[sampleIndex]]);
        HermiteData sample = hermiteVolume[CalculateHermiteIndex(position)];
        segmentsIndex |= (GetHermiteDataDensity(sample) < 0.0) << sampleIndex;
    }

    uint4 voxelFace = voxelFaces[faceIndex];
    uint4 faceSegments = marchingSquaresSegments[segmentsIndex];

    for (uint faceSegmentsIndex = 0; faceSegmentsIndex < 4; faceSegmentsIndex += 2)
    {
        if (faceSegments[faceSegmentsIndex] == -1)
        {
            break;
        }

        Segment segment;
        InitializeSegment(segment, 0, 0, 0, 0, 0);
        segments[segmentsCount++] = segment;
    }
}

[numthreads(4, 4, 4)]
void GenerateMesh(uint3 id : SV_DispatchThreadID)
{
    if (IsOutOfVoxelBounds(id))
    {
        return;
    }

    Segment segments[12];
    uint segmentsCount = 0;

    for (uint faceIndex = 0; faceIndex < 6; faceIndex++)
    {
        GenerateSegments(id, faceIndex, segments, segmentsCount);
    }

    if (segmentsCount == 0)
    {
        return;
    }
    
    // additional code omitted
}

Again, looking at the compiled code through Unity (snippet below), an additional dcl_indexableTemp register is used just because of the function call that should be inlined regardless.

dcl_uav_structured u0, 16
dcl_uav_structured_opc u1, 24
dcl_uav_structured u2, 12
dcl_input vThreadID.xyz
dcl_temps 22
dcl_indexableTemp x0[24], 4
dcl_indexableTemp x1[24], 4
dcl_indexableTemp x2[8], 4
dcl_indexableTemp x3[2], 4
dcl_indexableTemp x4[8], 4
dcl_thread_group 4, 4, 4

My understanding is that using the input modifier inout copies the function argument in before and out after the function call. So in that sense the additional dcl_indexableTemp register makes sense. But shouldn't the compiler notice that this copying can be omitted in aboves case and save up on those temporary registers?

Am I forced to manually inline everything or is there another way (in Unity) which would actually allow for the usage of functions to split up the larger algorithm?

\$\endgroup\$
3
  • 1
    \$\begingroup\$ It might be easier for someone to help you with this if you could post an example shader that reproduces the problem and compiles. In addition to that, have you verified that the additional temporary registers actually affect performance? \$\endgroup\$
    – Adam
    Mar 25, 2021 at 21:21
  • \$\begingroup\$ I had to inline everything as well, and it seemed that using classes and inheritance did not work as expected either, got name collisions, so had to refactor even more. And I got compiler errors that I used more registers than allowed, so had to introduce magic numbers to get down variable count. So would be nice with solution... but other than that remember to check unity window often so you don't spend weeks on something that must be refactored. \$\endgroup\$
    – Emil
    Mar 26, 2021 at 6:50
  • \$\begingroup\$ @Adam, I've provided an example compute shader. No, I haven't checked whether performance is affected but Unity warns about too many temporary registers being used for my compute shader if I split my algorithm up in a way I would prefer. \$\endgroup\$ Mar 28, 2021 at 14:34

1 Answer 1

0
\$\begingroup\$

I found a way to get the compiler to do the right thing while using functions. It's not ideal, but better than inlining everything manually.

What I did was use static global variables instead of inout parameters. You may want to do some more experimentation to find out exactly which things need to be global variables instead of parameters.

Here's a modified version of your example shader that compiles to exactly the same assembly regardless of the USE_FUNCTION define being set or not:

#pragma kernel ComputeRollingAverage

// Comment out "USE_FUNCTION" define to inline code manually.
#define USE_FUNCTION
#define SAMPLE_COUNT 10
#define CODE \
count = 0; \
 \
for (uint offset = 0; offset < SAMPLE_COUNT; offset++) \
{ \
    if (index + offset >= size) \
    { \
        break; \
    } \
    samples[count++] += input[index + offset]; \
}

uint size;                                         

StructuredBuffer<float> input;
RWStructuredBuffer<float> output;

static float samples[SAMPLE_COUNT];

void InitializeSamples()
{
    for (int iterator = 0; iterator < SAMPLE_COUNT; iterator++)
    {
        samples[iterator] = 0.0;
    }
}

void SampleInput(uint index, out uint count)
{
    CODE
}

[numthreads(64, 1, 1)]
void ComputeRollingAverage(uint3 id : SV_DispatchThreadID)
{
    uint index = id.x;
    uint count;
    InitializeSamples();
    
#ifdef USE_FUNCTION
    SampleInput(index, count);
#else
    CODE
#endif
    
    float rollingAverage = 0.0f;
    
    for (uint iterator = 0; iterator < count; iterator++)
    {
        rollingAverage += samples[iterator];
    }
    output[id.x] = rollingAverage / count;
}

Note that I tested this by compiling with fxc.exe directly, and not using Unity, but I'd expect the results to be the same.

Using fxc.exe directly allowed me to compile it twice, and pass /D USE_FUNCTION on the command line for one compilation, output the results to a file with /Fc, and then diff the two results with fc. For reference, this is what the batch file I used looked like:

fxc.exe /E ComputeRollingAverage /T cs_5_0 ComputeShader.txt /Fc inline.txt
fxc.exe /E ComputeRollingAverage /T cs_5_0 ComputeShader.txt /Fc function.txt /D USE_FUNCTION
fc function.txt inline.txt > diff.txt
pause

Note that you may need to specify the full path to fxc.exe.

\$\endgroup\$
1
  • \$\begingroup\$ I tested your code with Unity and I can confirm your results. Furthermore, I was able to use your approach to split up my initial algorithm (Cubical Marching Squares) into separate functions without increasing the number of temporary registers. Using static variables is a great compromise. Thank you! :) \$\endgroup\$ Mar 29, 2021 at 10:02

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .