¿Qué es un Compute Shader?

Un shader es un programita que se ejecuta en la GPU. Los tenemos de distintos sabores: vertex shader, pixel shader, geometry shader, tessellation shader, etc.,

Algunos seguro ya conoces, sin ir más lejos, cada vez que creas un material con el editor de materiales estás escribiendo un vertex/pixel shader.

¿Qué es un compute shader? Como su propio nombre indica, un compute shader sirve para hacer cálculos aprovechando la arquitectura masivamente paralela de la GPU.

¿Qué clase de cálculos? ¡De todo tipo! Desde cálculos para render, IA, efectos post process, físicas, etc.,

¿Cómo funciona?

Primero, escribes el fichero USF (Unreal Shader Format) en lenguaje HLSL (el editor de materiales no nos sirve en este caso) que contiene el código de tu compute shader.

Segundo, tendrás que indicarle a UE4 qué tienes una ruta con ficheros USF para compilar. Esto tendrás que hacerlo usando un módulo.

¿Qué es un módulo? UE4 organiza todo su código en bloques lógicos llamados módulos. Tenemos el módulo de render, el de entrada al usuario, etc., ¡de hecho cuando programas tu juego realmente estás escribiendo un módulo!

Tercero, en tu compute shader habrás declarado una serie de parámetros de entrada y salida. Dichos parámetros deben ser escritos/leídos desde código, desde la CPU. Para ello tendrás que declarar una clase especial que hereda de FGlobalShader que servirá como enlace a tu compute shader.

Por último tendrás que invocar tu compute shader y leer los parámetros de salida.

Vamos a empezar con el primer punto, crear el módulo.

Crear un módulo

Antes de crear el módulo, voy a crear un proyecto vacío llamado ComputeShaderExample.

Genera el proyecto en Visual Studio. Si te resulta más cómodo crea una clase cualquier desde el editor, por ejemplo un GameMode, de este modo UE4 generará el proyecto C++ por ti.

Ahora vamos a añadir un nuevo módulo.

Quizás el modo más fácil de crear un nuevo módulo a fecha de hoy es copiar y pegar la carpeta fuente del proyecto, en este caso \Source\ComputeShaderExample, y renombrar tanto la carpeta como los ficheros con el nombre del nuevo módulo.

Haz clic derecho sobre el fichero del proyecto y Generate Visual Studio Project files.

Edita el contenido de cada fichero.

Para el caso de MyComputeShader.Build.cs no olvides renombrar la clase:

using UnrealBuildTool;

public class MyComputeShader : ModuleRules
{
	public MyComputeShader(ReadOnlyTargetRules Target) : base(Target)
	{
		PCHUsage = PCHUsageMode.UseExplicitOrSharedPCHs;

		PublicDependencyModuleNames.AddRange(new string[] { });
		PrivateDependencyModuleNames.AddRange(new string[]
		{
				"Core",
				"CoreUObject",
				"Engine",
				"Renderer",
				"RenderCore",
				"RHI",
				"Projects"
		});
	}
}

En cuanto a las clases responsables del módulo:

#pragma once

#include "CoreMinimal.h"

#include "Modules/ModuleInterface.h"
#include "Modules/ModuleManager.h"

class MYCOMPUTESHADER_API FMyComputeShaderModule : public IModuleInterface
{
public:
	virtual void StartupModule() override;
	virtual void ShutdownModule() override;
};

#include "MyComputeShaderModule.h"
#include "Modules/ModuleManager.h"
#include "Misc/Paths.h"
#include "GlobalShader.h"

IMPLEMENT_GAME_MODULE(FMyComputeShaderModule, MyComputeShader);

void FMyComputeShaderModule::StartupModule()
{
	// Aquí indicaremos a UE4 que cargue nuestro shader
}

void FMyComputeShaderModule::ShutdownModule()
{
}

No olvides dar de alta al módulo en el fichero que describe el proyecto ComputeShaderExample.uproject:

{
	"FileVersion": 3,
	"EngineAssociation": "4.25",
	"Category": "",
	"Description": "",
	"Modules": [
		{
			"Name": "ComputeShaderExample",
			"Type": "Runtime",
			"LoadingPhase": "Default",
			"AdditionalDependencies": [
				"Engine"
			]
		},
		{
			"Name": "MyComputeShader",
			"Type": "Runtime",
			"LoadingPhase": "PostConfigInit"
		}
	]
}

Nota como se usa como LoadingPhase el valor PostConfigInit porque UE4 requiere que todos los shaders sean cargados en dicha fase.

Vamos ahora a por el shader propiamente dicho.

Global Shader

Los shaders en UE4 vienen en dos sabores: Global y Material.

Tal y como su propio nombre sugiere, los shader de tipo global son únicos y no son instanciados, mientras que los tipo material son shaders que usualmente creas en el editor de materiales, son instanciados y existen varias copias de ellos.

La inmensa mayoría de shaders que usas son de tipo material.

Sin embargo, los compute shader son un ejemplo notorio de global shader. Otro ejemplo serían los shaders para postprocesamiento.

Aunque lo haremos enseguida, es de gran ayuda esta breve guía rápida sobre cómo añadir global shaders para tener un overview del proceso.

También tienes más información sobre desarrollo de shaders en UE4 aquí.

Declarar el shader

Para usar un shader en UE4 debes:

1. Colocar el shader en una carpeta especial. Normalmente en la carpeta Shaders.
2. Al cargar el módulo, índicarle a UE4 la ruta dónde encontrar tus shaders.
3. Crear una clase que herede de FGlobalShader. Usa esta clase para declarar los recursos que serán pasados al shader.

Empezamos por el punto 2 que es el más sencillo. La forma de indicar desde tu módulo a UE4 la ruta dónde encontrar tus shaders es la siguiente:

void FMyComputeShaderModule::StartupModule()
{
	FString ShaderDirectory = FPaths::Combine(FPaths::ProjectDir(), TEXT("Shaders/Private"));
	AddShaderSourceDirectoryMapping("/MyComputeShader", ShaderDirectory);	
}

En este caso la ruta dónde buscará los shaders será en /Shaders/Private.

Vamos a crear dicha ruta y escribir el shader.

Código del shader

Para nuestro proyecto de ejemplo, vamos a usar el shader que usamos aquí para generar un perlin noise.

Crea un fichero en la carpeta Shaders/Private llamado PerlinNoiseCS.usf (la extensión USF significa Unreal Shader File, básicamente es un fichero HLSL).

El shader es una adaptación de éste que ya usamos aquí.

Éste es el código fuente de PerlinNoiseCS.usf que procederemos a comentar inmediatamente:

#include "/Engine/Public/Platform.ush"

RWTexture2D<float> OutTexture;
float2 Size;

float2 hash22(float2 x)
{
    const float2 k = float2(0.3183099, 0.3678794);
    x = x * k + k.yx;
    return -1.0 + 2.0 * frac(16.0 * k * frac(x.x * x.y * (x.x + x.y)));
}

float pnoise(float2 p)
{
    float2 id = floor(p);
    float2 st = frac(p);
    
    float2 f = st;
    f = f * f * f * (f * (f * 6.0 - 15.0) + 10.0);
    
    return
        lerp(
        	lerp(dot(hash22(id + float2(0.0, 0.0)), float2(0.0, 0.0) - st),
                dot(hash22(id + float2(1.0, 0.0)), float2(1.0, 0.0) - st),
                f.x),
            lerp(dot(hash22(id + float2(0.0, 1.0)), float2(0.0, 1.0) - st),
                dot(hash22(id + float2(1.0, 1.0)), float2(1.0, 1.0) - st),
                f.x),
        	f.y
		);
}

[numthreads(32, 32, 1)]
void mainCS(uint3 GroupId : SV_GroupID, uint3 DispatchThreadId : SV_DispatchThreadID, uint3 GroupThreadId : SV_GroupThreadID, uint GroupIndex : SV_GroupIndex)
{
    float2 pos = (float2(DispatchThreadId.xy) - 0.5 * Size) / Size;
    float noise = pnoise(30.0 * pos);
    OutTexture[DispatchThreadId.xy] = noise;
}

Al igual que un pixel/vertex shader, un compute shader tiene parámetros de entrada.

En este caso los parámetros son Size y OutNoise, de tipo flaot2 y UAV respectivamente.

Algunos de los tipos de recursos más habituales:

• Unordered access view (UAV), para escribir recursos (por ejemplo una textura)
• Shader resource view (SRV) para leer recursos (por ejemplo una textura)
• Constant buffer view (CBV), para leer datos constantes con respecto al pipeline (por ejemplo matrices de transformación)

Aquí una lista de recursos que puedes pasar a un shader. Si tienes interés en profundizar más, y de manera opcional, aquí una guía sobre DirectX.

A diferencia de un pixel/vertex shader, un compute shader tiene como método principal:

[numthreads(nX, nY, nZ)]
void mainCS(uint3 GroupId : SV_GroupID, uint3 DispatchId : SV_DispatchThreadID, uint3 GroupThreadId : SV_GroupThreadID, uint GroupIndex : SV_GroupIndex)
{   
    .......
}

¿Qué son esos argumentos de entrada?

Cuando invocas (técnicamente se dice Dispatch) un compute shader le pasas como argumento cuántos grupos de hilos quieres lanzar por cada dimensión.

FIntVector GroupThreadCount = /*cuántos grupos de hilos */;
FComputeShaderUtils::Dispatch(
	RHICmdList, 
	TheComputeShader, 
	Parameters, 
	GroupThreadCount
);

Para cada grupo de hilo se lanzan tantos hilos como indiques con el decorador numthreads(X, Y, Z).

El método mainCS se ejecuta por cada hilo.

Por tanto, los argumentos de entrada de mainCS sirven para identificar a qué grupo de hilos pertenece dicho hilo, y dentro de ese grupo sus coordenadas.

En concreto el argumento SV_DispatchThreadID es muy útil ya que aglutina la información del siguiente modo:

SV_DispatchThreadID = SV_GroupID * numthreads + SV_GroupThreadID

Este argumento lo usamos en nuestro shader para indexar la textura:

[numthreads(32, 32, 1)]
void mainCS(uint3 GroupId : SV_GroupID, uint3 DispatchThreadId : SV_DispatchThreadID, uint3 GroupThreadId : SV_GroupThreadID, uint GroupIndex : SV_GroupIndex)
{
    // ...
    OutTexture[DispatchThreadId.xy] = noise;
}

Por tanto, si vamos a usar 32x32 hilos por grupo de trabajo. Y la textura es de 1024x1024, ¿cuantos grupos de trabajo necesitamos lanzar?

Tendremos que hacer Dispatch con $1024/32 = 32$ grupos de trabajo.

Parámetros

Ahora desde el lado de la CPU vamos a declarar nuestro FGlobalShader.

Una clase FGlobalShader hace uso extensivo de macros para definir los parámetros que usa el shader.

En nuestro caso, la clase en cuestión es ésta:

class FPerlinNoiseCS : public FGlobalShader
{
public:
	DECLARE_GLOBAL_SHADER(FPerlinNoiseCS);
	
	SHADER_USE_PARAMETER_STRUCT(FPerlinNoiseCS, FGlobalShader);
	
	BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )		
		SHADER_PARAMETER_UAV(RWTexture2D<float>, OutTexture)
		SHADER_PARAMETER(FVector2D, Size)
	END_SHADER_PARAMETER_STRUCT()

	static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters)
	{
		return IsFeatureLevelSupported(Parameters.Platform, ERHIFeatureLevel::SM5);
	}

	static inline void ModifyCompilationEnvironment(
    	const FGlobalShaderPermutationParameters& Parameters,
        FShaderCompilerEnvironment& OutEnvironment
	)
	{
		FGlobalShader::ModifyCompilationEnvironment(Parameters, OutEnvironment);
		// aquí puedes setear los defines del shader
	}

};

IMPLEMENT_GLOBAL_SHADER(FPerlinNoiseCS, "/MyComputeShader/PerlinNoiseCS.usf", "mainCS", SF_Compute);

Las MACROS son bastante descriptivas.

La primera macro declara el shader como FGlobalShader. Piensa en esta macro como la típica GENERATED_BODY() pero para shaders.

DECLARE_GLOBAL_SHADER(FPerlinNoiseCS);

El siguiente paso declarar una estructura, llamada FParameters y nos apoyaremos en macros para declararla.

SHADER_USE_PARAMETER_STRUCT(FPerlinNoiseCS, FGlobalShader);

BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )		    
    SHADER_PARAMETER(UINT, Seed)		
    SHADER_PARAMETER_UAV(RWTexture2D<float>, OutNoise)
END_SHADER_PARAMETER_STRUCT()

Esta estructura debe corresponder uno-a-uno con los parámetros del shader.

Si olvidas alguno de los parámetros obtendrás el error:

Shader FPerlinNoiseCS has unbound parameters not represented in the parameter struct: [nombre del parámetro]

Todas las macros para declarar parámetros las puedes encontrar en: /Engine/Source/Runtime/RenderCore/Public/ShaderParameterStruct.h

Por último, y paso crucial, asociar dicha declaración con el fichero USF:

IMPLEMENT_GLOBAL_SHADER(
	FPerlinNoiseCS, 
    "/MyComputeShader/PerlinNoiseCS.usf",
    "mainCS",
    SF_Compute
);

Ya tenemos el shader listo para ser invocado.

Invocación del shader

Para invocar un compute shader:

1. Rellenamos la estructura FParameters que definimos anteriormente,
2. Obtenemos una referencia al shader usando GetGlobalShaderMap,
3. Definimos cuantos grupos de hilos queremos lanzar,
4. Usamos FComputeShaderUtils para lanzar el compute shader.

Aquí en código:

FPerlinNoiseCS::FParameters Parameters;	
Parameters.OutTexture = /* ... */;
Parameters.Size = /* ... */;

TShaderMapRef<FPerlinNoiseCS> PerlinNoiseCS(
	GetGlobalShaderMap(GMaxRHIFeatureLevel)
);
	
FIntVector WorkgroupCount = /* ... ie FIntVector(32, 32, 1) .... */;

FComputeShaderUtils::Dispatch(
	RHICmdList,
    PerlinNoiseCS,
    Parameters,
    WorkgroupCount
);	

Dónde la textura de salida puede ser cualquier render target.

El problema es que no puedes invocar un compute shader desde cualquier lugar.

En concreto, solo puedes invocar un compute shader desde el hilo de render.

Si quieres saber más información sobre cómo funciona el tema de los hilos aquí hay un post.

La buena noticia es que hay una macro que nos facilita la tarea de ejecutar un trozo de código en el hilo de render:

ENQUEUE_RENDER_COMMAND(DispatchComputeShader)(
    [this](FRHICommandList& RHICmdList)
    {
    	...
        FComputeShaderUtils::Dispatch(...);
        ...
    }
);

Si necesitas invocar tu compute shader en cada frame, quizás sea mejor subscribirse a un evento del sistema de render que tener que encolar el comando en cada frame.

Es decir, en vez de usar en cada frame ENQUEUE_RENDER_COMMAND hacer lo siguiente:

IRendererModule* RendererModule = FModuleManager::GetModulePtr<IRendererModule>(TEXT("Renderer"));
FDelegateHandle ResolvedSceneColorHandle = RendererModule->GetResolvedSceneColorCallbacks().AddRaw(this, &FMyComputeShaderHelper::ExecuteInRenderThread);

Cuando quieras dejar de ejecutar el computer shader en cada frame:

RendererModule->GetResolvedSceneColorCallbacks().Remove(ResolvedSceneColorHandle);
OnPostResolvedSceneColorHandle.Reset();

Y recuerda que ahora el método con el que te has suscrito es llamado en cada frame y desde el render thread:

// Este metodo es llamado en render thread
void FMyComputeShaderHelper::ExecuteInRenderThread(
	FRHICommandListImmediate& RHICmdList, 
    class FSceneRenderTargets& SceneContext
)
{
    ...
    FComputeShaderUtils::Dispatch(...);
    ...
}

Textura y UAV

El último ingrediente nos queda por añadir es ¿cómo conseguir un recurso UAV?

Anteriormente dijimos que se trata de un render target, pero requiere un poco más de trabajo.

Dado que estamos trabajando a bajo nivel debemos preocuparnos de ciertos aspectos como la transición de estados en un recurso.

En concreto, un recurso, como nuestro UAV, que es usado en un compute shader para ser escrito debemos tenerlo en un estado válido para escritura. Del mismo modo cuando va a ser leído (o copiado a otro render target) debe transicionar a un estado valido para lectura.

Estas transiciones de estados son propias de API de bajo nivel como Vulkan o DirectX. En este post tienes un tutorial de DirectX dónde hacemos uso extenso de transicionar recursos.

Dado que estamos trabajando a ese nivel debemos tener cuenta de ello.

Primero vamos a declarar una referencia al render target que vamos a usar como UAV. Es decir, el render target que vamos a pasarle a la GPU y transicionar su estado.

TRefCountPtr<IPooledRenderTarget> TextureOutput;  // nuestro UAV

Por otra parte, tenemos nuestro RenderTarget (de tipo UTextureRenderTarget2D) que será dónde copiemos nuestro TextureOutput.

Así que por un lado está nuestro asset, creado desde el propio editor en el Content Browser, de tipo UTextureRenderTarget2D qué será nuestro render target objetivo.

La forma de hacerlo es preguntarle directamente al engine que nos devuelva nuestro render target si ya existe o que cree uno si aún no existe:

if (TextureOutput.IsValid())
{
	return;
}

UTextureRenderTarget2D* RenderTarget = /* ... asset en el content browser ... */;
		
FIntPoint TextureSize{ RenderTarget->SizeX, RenderTarget->SizeY };

EPixelFormat Format = RenderTarget->GetRenderTargetResource()->TextureRHI->GetFormat();

FPooledRenderTargetDesc TextureOutputDesc =
    FPooledRenderTargetDesc::Create2DDesc(
        TextureSize, 
        Format, 
        FClearValueBinding::None, 
        TexCreate_None, // Flags
        TexCreate_ShaderResource | TexCreate_UAV, // Target flags
        false // separate target and shader resource
    );
TextureOutputDesc.DebugName = TEXT("PerlinNoiseCS_OutputTexture");

GRenderTargetPool.FindFreeElement(
    RHICmdList, 
    TextureOutputDesc, 
    TextureOutput, 
    TEXT("PerlinNoiseCS_OutputTexture")
);

Para transicionar el recurso:

RHICmdList.TransitionResource(
    
    /* Mostly for UAVs.  
    Transition to read/write state and 
    always insert a resource barrier */
	EResourceTransitionAccess::ERWBarrier, 
    
    /* from graphics to compute */
    EResourceTransitionPipeline::EGfxToCompute, 
    
    TextureOutput->GetRenderTargetItem().UAV
    
);

Para copiar el recurso TextureOutput a RenderTarget.

RHICmdList.CopyTexture(
	TextureOutput->GetRenderTargetItem().ShaderResourceTexture,
    RenderTarget->GetRenderTargetResource()->TextureRHI,
    FRHICopyTextureInfo()
);

Código final

El código final quedaría así:

#pragma once

#include "CoreMinimal.h"
#include "UObject/NoExportTypes.h"
#include "PerlinNoiseCSHelper.generated.h"

UCLASS()
class MYCOMPUTESHADER_API UPerlinNoiseCSHelper : public UObject
{
public:
	GENERATED_BODY()
	
	UPROPERTY(VisibleAnywhere, BlueprintReadOnly, Category = "Compute Helper")
	UTextureRenderTarget2D* RenderTarget;

	TRefCountPtr<IPooledRenderTarget> TextureOutput;

	void ExecuteInRenderThread(FRHICommandList& RHICmdList);
};

#include "PerlinNoiseCSHelper.h"
#include "GlobalShader.h"
#include "RenderGraphUtils.h"
#include "RenderTargetPool.h"
#include "Shader.h"
#include "RHICommandList.h"
#include "Engine/TextureRenderTarget2D.h"

class FPerlinNoiseCS : public FGlobalShader
{
public:
	DECLARE_GLOBAL_SHADER(FPerlinNoiseCS);

	SHADER_USE_PARAMETER_STRUCT(FPerlinNoiseCS, FGlobalShader);

	BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
		SHADER_PARAMETER_UAV(RWTexture2D<float>, OutTexture)
		SHADER_PARAMETER(FVector2D, Size)
		END_SHADER_PARAMETER_STRUCT()

		static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters)
	{
		return IsFeatureLevelSupported(Parameters.Platform, ERHIFeatureLevel::SM5);
	}

	static inline void ModifyCompilationEnvironment(
		const FGlobalShaderPermutationParameters& Parameters,
		FShaderCompilerEnvironment& OutEnvironment
	)
	{
		FGlobalShader::ModifyCompilationEnvironment(Parameters, OutEnvironment);
		// aquí puedes setear los defines del shader
	}

};

IMPLEMENT_GLOBAL_SHADER(FPerlinNoiseCS, "/MyComputeShader/PerlinNoiseCS.usf", "mainCS", SF_Compute);

void UPerlinNoiseCSHelper::ExecuteInRenderThread(FRHICommandList& RHICmdList)
{
	if (!TextureOutput.IsValid()) {
		/* UTextureRenderTarget2D* RenderTarget =  ... */;
		
		FIntPoint TextureSize{ RenderTarget->SizeX, RenderTarget->SizeY };

		EPixelFormat Format = RenderTarget->GetRenderTargetResource()->TextureRHI->GetFormat();

		FPooledRenderTargetDesc TextureOutputDesc =
			FPooledRenderTargetDesc::Create2DDesc(
				TextureSize,
				Format,
				FClearValueBinding::None,
				TexCreate_None, // Flags
				TexCreate_ShaderResource | TexCreate_UAV, // Target flags
				false // separate target and shader resource
			);
		TextureOutputDesc.DebugName = TEXT("PerlinNoiseCS_OutputTexture");

		GRenderTargetPool.FindFreeElement(
			RHICmdList,
			TextureOutputDesc,
			TextureOutput,
			TEXT("PerlinNoiseCS_OutputTexture")
		);
	}

	FPerlinNoiseCS::FParameters Parameters;
	Parameters.OutTexture = TextureOutput->GetRenderTargetItem().UAV;
	Parameters.Size = FVector2D(RenderTarget->SizeX, RenderTarget->SizeY);

	TShaderMapRef<FPerlinNoiseCS> PerlinNoiseCS(
		GetGlobalShaderMap(GMaxRHIFeatureLevel)
	);

	FIntVector WorkgroupCount = FIntVector(32, 32, 1);

	FComputeShaderUtils::Dispatch(
		RHICmdList,
		PerlinNoiseCS,
		Parameters,
		WorkgroupCount
	);

	RHICmdList.TransitionResource(

		/* Mostly for UAVs.
		Transition to read/write state and
		always insert a resource barrier */
		EResourceTransitionAccess::ERWBarrier,

		/* from graphics to compute */
		EResourceTransitionPipeline::EGfxToCompute,

		TextureOutput->GetRenderTargetItem().UAV

	);
	
	RHICmdList.CopyTexture(
		TextureOutput->GetRenderTargetItem().ShaderResourceTexture,
		RenderTarget->GetRenderTargetResource()->TextureRHI,
		FRHICopyTextureInfo()
	);
}

Para hacer uso, basta con instanciar la clase:

PerlinNoiseCS = NewObject<UPerlinNoiseCSHelper>();

Y encolar el dispatch:

PerlinNoiseCS->RenderTarget = RenderTarget;
ENQUEUE_RENDER_COMMAND(DispatchComputeShader)(
    [this](FRHICommandList& RHICmdList) {
        PerlinNoiseCS->ExecuteInRenderThread(RHICmdList);
    }
);

Descargar proyecto completo

Puedes descargar el proyecto completo desde aquí.