115 lines
3.0 KiB
HLSL
115 lines
3.0 KiB
HLSL
// Copyright Epic Games, Inc. All Rights Reserved.
|
|
|
|
#pragma once
|
|
|
|
// the define is set by FGPUWorkGroupLoadBalancer::SetShaderDefines which should be used by any CS that uses this.
|
|
#ifdef WGLB_ENABLE
|
|
|
|
#include "Common.ush"
|
|
#include "WaveOpUtil.ush"
|
|
#include "ComputeShaderUtils.ush"
|
|
|
|
struct FWorkGroupInfo
|
|
{
|
|
uint2 WorkGroupWorkBoundary;
|
|
uint FirstItem;
|
|
uint NumItems;
|
|
uint CarryOverStartOffset;
|
|
uint Payload; // aribitrary 32-bit payload for each workgroup
|
|
};
|
|
|
|
struct FPackedItem
|
|
{
|
|
uint Packed;
|
|
};
|
|
|
|
StructuredBuffer<FWorkGroupInfo> WorkGroupInfoBuffer;
|
|
StructuredBuffer<FPackedItem> ItemBuffer;
|
|
uint NumWorkGroupInfos;
|
|
uint NumItems;
|
|
|
|
struct FWorkGroupItem
|
|
{
|
|
bool bHasCarryOver;
|
|
uint NumChildren;
|
|
uint Payload;
|
|
uint BatchPrefixSum;
|
|
};
|
|
|
|
FWorkGroupItem UnpackItem(FPackedItem PackedItem)
|
|
{
|
|
FWorkGroupItem Result;
|
|
|
|
Result.bHasCarryOver = (PackedItem.Packed & 1u) != 0u;
|
|
Result.NumChildren = 1u + ((PackedItem.Packed >> 1u) & WGLB_NUM_ITEM_MASK);
|
|
Result.BatchPrefixSum = (PackedItem.Packed >> (1u + WGLB_NUM_ITEM_BITS)) & WGLB_PREFIX_BIT_MASK;
|
|
Result.Payload = (PackedItem.Packed >> (1u + WGLB_NUM_ITEM_BITS + WGLB_PREFIX_BITS));
|
|
|
|
return Result;
|
|
}
|
|
|
|
groupshared FPackedItem Items[WGLB_NUM_THREADS_PER_GROUP];
|
|
|
|
struct FWorkGroupSetup
|
|
{
|
|
bool bValid;
|
|
FWorkGroupItem Item;
|
|
FWorkGroupInfo WorkGroupInfo;
|
|
// Work out who I am in expansion of the item in range [0,WorkItem.Count], may be negative when no work at end of batch
|
|
int LocalItemIndex;
|
|
uint WorkItemIndex;
|
|
uint ChildIndex;
|
|
uint DispatchThreadId;
|
|
};
|
|
|
|
FWorkGroupSetup WorkGroupLoadBalancer_Setup(uint3 GroupId, uint GroupThreadIndex)
|
|
{
|
|
FWorkGroupSetup Setup = (FWorkGroupSetup)0;
|
|
uint WorkGroupIndex = GetUnWrappedDispatchGroupId(GroupId);
|
|
Setup.DispatchThreadId = WorkGroupIndex * WGLB_NUM_THREADS_PER_GROUP + GroupThreadIndex;
|
|
|
|
if (WorkGroupIndex >= NumWorkGroupInfos)
|
|
{
|
|
return Setup;
|
|
}
|
|
|
|
// Load work description for whole work group
|
|
Setup.WorkGroupInfo = WorkGroupInfoBuffer[WorkGroupIndex];
|
|
|
|
FPackedItem PackedItem = (FPackedItem)0u;
|
|
if (GroupThreadIndex < Setup.WorkGroupInfo.NumItems)
|
|
{
|
|
PackedItem = ItemBuffer[Setup.WorkGroupInfo.FirstItem + GroupThreadIndex];
|
|
}
|
|
Items[GroupThreadIndex] = PackedItem;
|
|
|
|
GroupMemoryBarrierWithGroupSync();
|
|
Setup.WorkItemIndex = MaskedBitCount(Setup.WorkGroupInfo.WorkGroupWorkBoundary, GroupThreadIndex);
|
|
Setup.Item = UnpackItem(Items[Setup.WorkItemIndex]);
|
|
// Work out who I am in expansion of the item in range [0,WorkItem.Count], may be negative when no work at end of batch
|
|
Setup.LocalItemIndex = int(GroupThreadIndex - Setup.Item.BatchPrefixSum);
|
|
Setup.ChildIndex = uint(Setup.LocalItemIndex);
|
|
|
|
if (Setup.Item.bHasCarryOver)
|
|
{
|
|
Setup.ChildIndex += Setup.WorkGroupInfo.CarryOverStartOffset;
|
|
}
|
|
Setup.bValid = Setup.LocalItemIndex >= 0 && Setup.LocalItemIndex < int(Setup.Item.NumChildren);
|
|
|
|
return Setup;
|
|
}
|
|
|
|
|
|
uint WorkGroupLoadBalancer_GetNumWorkGroups()
|
|
{
|
|
return NumWorkGroupInfos;
|
|
}
|
|
|
|
|
|
uint WorkGroupLoadBalancer_GetNumItems()
|
|
{
|
|
return NumItems;
|
|
}
|
|
|
|
#endif // WGLB_ENABLE
|