SDFTexture.compute

#include "UnityCG.cginc"

#pragma kernel _SeedKernel
#pragma kernel _FloodKernel
#pragma kernel _DistKernel
#pragma kernel _ShowSeedsKernel

#pragma multi_compile_local __ _ADD_BORDERS

#define THREAD_GROUP_WIDTH 8

Texture2D _SeedTexRead;
RWTexture2D<uint4> _FloodTex;
Texture2D<uint4> _FloodTexRead;
RWTexture2D<float> _SdfTex;

CBUFFER_START( Often )
	float _SeedThreshold;
	float2 _TexelSize;
	uint2 _Resolution;
CBUFFER_END

CBUFFER_START( Always )
	int _StepSize;
CBUFFER_END

SamplerState _LinearClamp;

#define NEIGHBOUR_COUNT 9
static const int2 offsets[ NEIGHBOUR_COUNT ] = {
	int2( -1, -1 ), int2( 0, -1 ), int2( 1, -1 ),
	int2( -1, 0 ), int2( 0, 0 ), int2( 1, 0 ),
	int2( -1, 1 ), int2( 0, 1 ), int2( 1, 1 ),
};

[numthreads( THREAD_GROUP_WIDTH, THREAD_GROUP_WIDTH, 1 )]
void _SeedKernel( uint2 coord : SV_DispatchThreadID )
{
	if( coord.x >= _Resolution.x || coord.y >= _Resolution.y ) return;

	// Sample and test for seed.
	float2 uv = ( coord + 0.5 ) * _TexelSize;
	bool isInsideSeed = _SeedTexRead.SampleLevel( _LinearClamp, uv, 0 ).r > _SeedThreshold;

	// Since reading outside texture bounds returns zero, we shift all coordinates (1,1) so that (0,0) can be used to indicate "no seed coordinate".
	int2 seed = coord + int2( 1, 1 );

	// TODO add border
	#ifdef _ADD_BORDERS
		if( coord.x == 0 || coord.y == 0 || coord.x >= _Resolution.x-1 || coord.y >= _Resolution.y-1 ) isInsideSeed = true;
	#endif

	// Store inside seeds at xy and outside seeds at zw.
	_FloodTex[ coord ] = isInsideSeed ? uint4( seed, 0, 0 ) : uint4( 0, 0, seed );
}

groupshared uint2x3 sharedCoordDist[ NEIGHBOUR_COUNT ];

[numthreads( 1, 1, NEIGHBOUR_COUNT )]
void _FloodKernel( uint2 coord : SV_DispatchThreadID, uint3 gtId : SV_GroupThreadID )
{
	// Get the jump offset.
	int2 sampleCoord = coord + offsets[ gtId.z ] * _StepSize;

	// Sample.
	int4 seed = _FloodTex[ sampleCoord ]; // Reads outside bounds will return zeros.
	int2 shiftedCoord = coord + int2( 1, 1 );
	int4 diff = seed - shiftedCoord.xyxy;
	sharedCoordDist[ gtId.z ] = int2x3(
		seed.xy, any( seed.xy ) ? dot( diff.xy, diff.xy ) : 9999999, // enough to be ignored
		seed.zw, any( seed.zw ) ? dot( diff.zw, diff.zw ) : 9999999
	);

	// Wait until all threads in this group has executed issued reads/writes to groupshared memory.
	GroupMemoryBarrier();

	// If first thread in group, read from shared memory and find coordinate with smallest distance.
	if( gtId.z == 0 )
	{
		uint bestOuterDist = 9999999;
		uint bestInnerDist = 9999999;
		uint2 bestOuterCoord = int2( 0, 0 );
		uint2 bestInnerCoord = int2( 0, 0 );
		uint2x3 data;
		[unroll]
		for( int i = 0; i < NEIGHBOUR_COUNT; i++ )
		{
			data = sharedCoordDist[ i ];
			if( data._m02 < bestInnerDist ){
				bestInnerDist = data._m02;
				bestInnerCoord = data._m00_m01;
			}
			if( data._m12 < bestOuterDist ){
				bestOuterDist = data._m12;
				bestOuterCoord = data._m10_m11;
			}
		}

		// Write.
		_FloodTex[ coord ] = uint4( bestInnerCoord, bestOuterCoord );
	}
}

[numthreads( THREAD_GROUP_WIDTH, THREAD_GROUP_WIDTH, 1 )]
void _DistKernel( uint2 coord : SV_DispatchThreadID )
{
	if( coord.x >= _Resolution.x || coord.y >= _Resolution.y ) return;

	int4 closestCoords = _FloodTexRead[ coord ];
	int2 shiftedCoord = coord + int2( 1, 1 );
	float4 diff = closestCoords - shiftedCoord.xyxy;

	float outerDist = any( closestCoords.xy ) ? length( diff.xy ) : 0;
	float innerDist = any( closestCoords.zw ) ? length( diff.zw ) : 0;

	float sd = outerDist - innerDist; // If coord is inside then outerDist will be zero, else innerDist will be zero.
	
	// Use same convension as Unity's SDF Bake Tool
	// "the underlying surface scales such that the largest side of the Texture is of length 1".
	// https://docs.unity3d.com/Packages/[email protected]/manual/sdf-in-vfx-graph.html
	sd *= min( _TexelSize.x, _TexelSize.y );
	
	_SdfTex[ coord ] = sd;
}

[numthreads( THREAD_GROUP_WIDTH, THREAD_GROUP_WIDTH, 1 )]
void _ShowSeedsKernel( uint2 id : SV_DispatchThreadID )
{
	if( id.x >= _Resolution.x || id.y >= _Resolution.y ) return;

	_SdfTex[ id ] = any( _FloodTexRead[ id ].xy ) ? 1 : 0;
}