98 lines
3.1 KiB
HLSL
98 lines
3.1 KiB
HLSL
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
/*=============================================================================
|
|
NeuralPostProcessCommon.usf: Shaders for neural post process
|
|
=============================================================================*/
|
|
|
|
#ifndef MATERIAL_NEURAL_POST_PROCESS
|
|
#define MATERIAL_NEURAL_POST_PROCESS 0
|
|
#endif
|
|
|
|
#ifndef NEURAL_POSTPROCESS_PREPASS
|
|
#define NEURAL_POSTPROCESS_PREPASS 0
|
|
#endif
|
|
|
|
#if MATERIAL_NEURAL_POST_PROCESS
|
|
|
|
// Input buffer/textue set by the prepass
|
|
RWTexture2D<float4> RWNeuralTexture;
|
|
RWBuffer<float> InputNeuralBuffer;
|
|
RWBuffer<uint> NeuralSourceType;
|
|
float4 InputNeuralBufferDimension;
|
|
|
|
// Buffer/textures read by the main pass
|
|
Texture2D<float4> NeuralTexture;
|
|
Buffer<float> OutputNeuralBuffer;
|
|
float4 OutputNeuralBufferDimension;
|
|
|
|
#define SOURCE_TYPE_TEXTURE 0
|
|
#define SOURCE_TYPE_BUFFER 1
|
|
|
|
#endif
|
|
|
|
void UpdateNeuralProfileSourceType(int InputSourceType)
|
|
{
|
|
#if MATERIAL_NEURAL_POST_PROCESS && NEURAL_POSTPROCESS_PREPASS
|
|
NeuralSourceType[0] = InputSourceType;
|
|
#endif
|
|
}
|
|
|
|
|
|
// Save a 3 channel data to OutputNeuralBuffer from StartChannel to StartChannel+2
|
|
void SaveToNeuralBuffer(int Batch, int StartChannel, float2 BufferWH, float3 Value)
|
|
{
|
|
#if MATERIAL_NEURAL_POST_PROCESS && NEURAL_POSTPROCESS_PREPASS
|
|
float2 HWIndex = BufferWH.yx * InputNeuralBufferDimension.zw - 0.5f;
|
|
|
|
// point sampling
|
|
const int IX = (int)clamp(HWIndex.y, 0, InputNeuralBufferDimension.w - 1);
|
|
const int IY = (int)clamp(HWIndex.x, 0, InputNeuralBufferDimension.z - 1);
|
|
|
|
int BufferHeight = InputNeuralBufferDimension.z;
|
|
int BufferWidth = InputNeuralBufferDimension.w;
|
|
int BufferGridSize = BufferHeight * BufferWidth;
|
|
int NumChannels = InputNeuralBufferDimension.y;
|
|
int BatchSize = NumChannels * BufferGridSize;
|
|
|
|
int Idx = (Batch * BatchSize + StartChannel * BufferGridSize + IY * BufferWidth + IX);
|
|
int Offset = BufferGridSize;
|
|
|
|
// TODO: array of struct vs struct of array performance evaluation.
|
|
InputNeuralBuffer[Idx] = Value[0];
|
|
const int MaxOffset = min(3, InputNeuralBufferDimension.y - StartChannel) - 1;
|
|
for (int i = 1; i <= MaxOffset; ++i)
|
|
{
|
|
InputNeuralBuffer[Idx + i * Offset] = Value[i];
|
|
}
|
|
#endif
|
|
|
|
}
|
|
|
|
void ReadFromNeuralBuffer(int Batch, int StartChannel, float2 BufferWH, inout float3 Value)
|
|
{
|
|
#if MATERIAL_NEURAL_POST_PROCESS && MATERIAL_NEURAL_POST_PROCESS
|
|
float2 HWIndex = BufferWH.yx * OutputNeuralBufferDimension.zw - 0.5f;
|
|
|
|
// point sampling
|
|
const int IX = (int)clamp(HWIndex.y, 0, OutputNeuralBufferDimension.w - 1);
|
|
const int IY = (int)clamp(HWIndex.x, 0, OutputNeuralBufferDimension.z - 1);
|
|
|
|
int BufferHeight = OutputNeuralBufferDimension.z;
|
|
int BufferWidth = OutputNeuralBufferDimension.w;
|
|
int BufferGridSize = BufferHeight * BufferWidth;
|
|
int NumChannels = OutputNeuralBufferDimension.y;
|
|
int BatchSize = NumChannels * BufferGridSize;
|
|
|
|
int Idx = (Batch * BatchSize + StartChannel * BufferGridSize + IY * BufferWidth + IX);
|
|
int Offset = BufferGridSize;
|
|
|
|
// TODO: array of struct vs struct of array performance evaluation.
|
|
const int MaxOffset = min(3, OutputNeuralBufferDimension.y- StartChannel) - 1;
|
|
for (int i = 0; i <= MaxOffset; ++i)
|
|
{
|
|
Value[i] = OutputNeuralBuffer[Idx + i * Offset];
|
|
}
|
|
#else
|
|
Value = 0;
|
|
#endif
|
|
} |