Files
UnrealEngine/Engine/Plugins/Experimental/NNERuntimeRDG/Shaders/Private/NNEHlslShaders/NNEHlslShadersInstanceNormalization.usf
2025-05-18 13:04:45 +08:00

36 lines
1.0 KiB
HLSL

// Copyright Epic Games, Inc. All Rights Reserved.
#include "/Engine/Public/Platform.ush"
#define WORK_TYPE float
#define BUFFER_TYPE float
#define READ(x) x
#define WRITE(x) x
Buffer<WORK_TYPE> Input;
Buffer<WORK_TYPE> InputScale;
Buffer<WORK_TYPE> InputBias;
Buffer<WORK_TYPE> InputMean;
Buffer<WORK_TYPE> InputInvStdDev;
RWBuffer<WORK_TYPE> Output;
uint Num;
uint ThreadCountX;
uint InstanceSize;
uint ChannelSize;
[numthreads(THREADGROUP_SIZE_X, 1, 1)]
void InstanceNormalization(in const uint3 DispatchThreadID : SV_DispatchThreadID)
{
const uint Index = DispatchThreadID.y * ThreadCountX + DispatchThreadID.x;
if (Index < Num)
{
const uint InstanceIdx = Index/InstanceSize;
const uint ChannelIdx = InstanceIdx%ChannelSize;
WORK_TYPE Mean = READ(InputMean[InstanceIdx]);
WORK_TYPE InvStdDev = READ(InputInvStdDev[InstanceIdx]);
WORK_TYPE Scale = READ(InputScale[ChannelIdx]);
WORK_TYPE Bias = READ(InputBias[ChannelIdx]);
Output[Index] = WRITE((READ(Input[Index]) - Mean) * Scale * InvStdDev + Bias);
}
}