85 lines
3.2 KiB
C++
85 lines
3.2 KiB
C++
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#pragma once
|
|
|
|
#include "GlobalShader.h"
|
|
#include "RenderGraphUtils.h"
|
|
#include "ShaderParameterUtils.h"
|
|
|
|
#define UE_API NNEDENOISERSHADERS_API
|
|
|
|
namespace UE::NNEDenoiserShaders::Internal
|
|
{
|
|
|
|
// Maps to ENNETensorDataType
|
|
enum class EDataType : uint8
|
|
{
|
|
None = 0,
|
|
// Char,
|
|
// Boolean,
|
|
Half = 3,
|
|
Float,
|
|
MAX
|
|
};
|
|
|
|
class FMappedCopyConstants
|
|
{
|
|
public:
|
|
static constexpr int32 THREAD_GROUP_SIZE{16};
|
|
static constexpr int32 MAX_NUM_MAPPED_CHANNELS{4};
|
|
};
|
|
|
|
// Note: The shader also supports Texture/Texture and Buffer/Buffer mapped copy, just add the Shader class.
|
|
|
|
class FTextureBufferMappedCopyCS : public FGlobalShader
|
|
{
|
|
DECLARE_EXPORTED_GLOBAL_SHADER(FTextureBufferMappedCopyCS, UE_API);
|
|
SHADER_USE_PARAMETER_STRUCT(FTextureBufferMappedCopyCS, FGlobalShader)
|
|
|
|
class FInputDataType : SHADER_PERMUTATION_ENUM_CLASS("INPUT_DATA_TYPE_INDEX", EDataType);
|
|
class FOutputDataType : SHADER_PERMUTATION_ENUM_CLASS("OUTPUT_DATA_TYPE_INDEX", EDataType);
|
|
class FNumMappedChannels : SHADER_PERMUTATION_RANGE_INT("NUM_MAPPED_CHANNELS", 0, FMappedCopyConstants::MAX_NUM_MAPPED_CHANNELS + 1);
|
|
using FPermutationDomain = TShaderPermutationDomain<FInputDataType, FOutputDataType, FNumMappedChannels>;
|
|
|
|
public:
|
|
|
|
BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
|
|
SHADER_PARAMETER(int32, Width)
|
|
SHADER_PARAMETER(int32, Height)
|
|
SHADER_PARAMETER_RDG_TEXTURE(Texture2D, InputTexture)
|
|
SHADER_PARAMETER_RDG_BUFFER_UAV(RWBuffer, OutputBuffer)
|
|
SHADER_PARAMETER_ARRAY(FIntVector4, OutputChannel_InputChannel_Unused_Unused, [FMappedCopyConstants::MAX_NUM_MAPPED_CHANNELS])
|
|
END_SHADER_PARAMETER_STRUCT()
|
|
|
|
static UE_API void ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& InParameters, FShaderCompilerEnvironment& OutEnvironment);
|
|
static UE_API bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters);
|
|
};
|
|
|
|
class FBufferTextureMappedCopyCS : public FGlobalShader
|
|
{
|
|
DECLARE_EXPORTED_GLOBAL_SHADER(FBufferTextureMappedCopyCS, UE_API);
|
|
SHADER_USE_PARAMETER_STRUCT(FBufferTextureMappedCopyCS, FGlobalShader)
|
|
|
|
class FInputDataType : SHADER_PERMUTATION_ENUM_CLASS("INPUT_DATA_TYPE_INDEX", EDataType);
|
|
class FOutputDataType : SHADER_PERMUTATION_ENUM_CLASS("OUTPUT_DATA_TYPE_INDEX", EDataType);
|
|
class FNumMappedChannels : SHADER_PERMUTATION_RANGE_INT("NUM_MAPPED_CHANNELS", 0, FMappedCopyConstants::MAX_NUM_MAPPED_CHANNELS + 1);
|
|
using FPermutationDomain = TShaderPermutationDomain<FInputDataType, FOutputDataType, FNumMappedChannels>;
|
|
|
|
public:
|
|
|
|
BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
|
|
SHADER_PARAMETER(int32, Width)
|
|
SHADER_PARAMETER(int32, Height)
|
|
SHADER_PARAMETER_RDG_BUFFER_SRV(Buffer, InputBuffer)
|
|
SHADER_PARAMETER_RDG_TEXTURE_UAV(RWTexture2D, OutputTexture)
|
|
SHADER_PARAMETER_ARRAY(FIntVector4, OutputChannel_InputChannel_Unused_Unused, [FMappedCopyConstants::MAX_NUM_MAPPED_CHANNELS])
|
|
END_SHADER_PARAMETER_STRUCT()
|
|
|
|
static UE_API void ModifyCompilationEnvironment(const FGlobalShaderPermutationParameters& InParameters, FShaderCompilerEnvironment& OutEnvironment);
|
|
static UE_API bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters);
|
|
};
|
|
|
|
} // namespace UE::NNEDenoiser::Private
|
|
|
|
#undef UE_API
|