36 lines
1.0 KiB
HLSL
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);
|
|
}
|
|
} |