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?
