A game about forced loneliness, made by TACStudios
tacstudios.tngl.sh
1using System;
2using UnityEngine.Rendering.RenderGraphModule;
3
4namespace UnityEngine.Rendering
5{
6 public partial struct GPUPrefixSum
7 {
8 [GenerateHLSL]
9 internal static class ShaderDefs
10 {
11 public const int GroupSize = 128;
12
13 // Stride of the indirect arguement buffer in uints, the buffer is split into two sections dispatch options ( a lower or upper arguement set )
14 public const int ArgsBufferStride = 16;
15 public const int ArgsBufferUpper = 0;
16 public const int ArgsBufferLower = 8;
17
18 public static int DivUpGroup(int value)
19 {
20 return (value + GroupSize - 1) / GroupSize;
21 }
22
23 public static int AlignUpGroup(int value)
24 {
25 return DivUpGroup(value) * GroupSize;
26 }
27
28 public static void CalculateTotalBufferSize(int maxElementCount, out int totalSize, out int levelCounts)
29 {
30 int alignedSupportMaxCount = AlignUpGroup(maxElementCount);
31 totalSize = alignedSupportMaxCount;
32 levelCounts = 1;
33 while (alignedSupportMaxCount > GroupSize)
34 {
35 alignedSupportMaxCount = AlignUpGroup(DivUpGroup(alignedSupportMaxCount));
36 totalSize += alignedSupportMaxCount;
37 ++levelCounts;
38 }
39 }
40 }
41
42 /// <summary>
43 /// Structure defining level offsets.
44 /// </summary>
45 [GenerateHLSL]
46 public struct LevelOffsets
47 {
48 /// <summary> Number of levels. </summary>
49 public uint count;
50
51 /// <summary> Level offset. </summary>
52 public uint offset;
53
54 /// <summary> Parent level offset. </summary>
55 public uint parentOffset;
56 }
57
58 /// <summary>
59 /// Utility for adapting to render graph usage.
60 /// </summary>
61 public struct RenderGraphResources
62 {
63 internal int alignedElementCount;
64 internal int maxBufferCount;
65 internal int maxLevelCount;
66
67 internal BufferHandle prefixBuffer0;
68 internal BufferHandle prefixBuffer1;
69 internal BufferHandle totalLevelCountBuffer;
70 internal BufferHandle levelOffsetBuffer;
71 internal BufferHandle indirectDispatchArgsBuffer;
72
73 /// <summary>The prefix sum result.</summary>
74 public BufferHandle output => prefixBuffer0;
75
76 /// <summary>
77 /// Creates the render graph buffer resources from an input count.
78 /// </summary>
79 /// <param name="newMaxElementCount">The maximum number of elements that the buffer will support.</param>
80 /// <param name="renderGraph">Render Graph</param>
81 /// <param name="builder">Render Graph Builder</param>
82 /// <param name="outputIsTemp">Whether or not to allocate a transient resource.</param>
83 /// <returns>The created Render Graph Resources.</returns>
84 public static RenderGraphResources Create(int newMaxElementCount, RenderGraph renderGraph, RenderGraphBuilder builder, bool outputIsTemp = false)
85 {
86 var resources = new RenderGraphResources();
87 resources.Initialize(newMaxElementCount, renderGraph, builder, outputIsTemp);
88 return resources;
89 }
90
91 void Initialize(int newMaxElementCount, RenderGraph renderGraph, RenderGraphBuilder builder, bool outputIsTemp = false)
92 {
93 newMaxElementCount = Math.Max(newMaxElementCount, 1);
94 ShaderDefs.CalculateTotalBufferSize(newMaxElementCount, out int totalSize, out int levelCounts);
95
96 var prefixBuffer0Desc = new BufferDesc(totalSize, 4, GraphicsBuffer.Target.Raw) { name = "prefixBuffer0" };
97 prefixBuffer0 = outputIsTemp ? builder.CreateTransientBuffer(prefixBuffer0Desc) : builder.WriteBuffer(renderGraph.CreateBuffer(prefixBuffer0Desc));
98 prefixBuffer1 = builder.CreateTransientBuffer(new BufferDesc(newMaxElementCount, 4, GraphicsBuffer.Target.Raw) { name = "prefixBuffer1" });
99 totalLevelCountBuffer = builder.CreateTransientBuffer(new BufferDesc(1, 4, GraphicsBuffer.Target.Raw) { name = "totalLevelCountBuffer" });
100 levelOffsetBuffer = builder.CreateTransientBuffer(new BufferDesc(levelCounts, System.Runtime.InteropServices.Marshal.SizeOf<LevelOffsets>(), GraphicsBuffer.Target.Structured) { name = "levelOffsetBuffer" });
101 indirectDispatchArgsBuffer = builder.CreateTransientBuffer(new BufferDesc(ShaderDefs.ArgsBufferStride * levelCounts, sizeof(uint), GraphicsBuffer.Target.Structured | GraphicsBuffer.Target.IndirectArguments) { name = "indirectDispatchArgsBuffer" });//3 arguments for upp dispatch, 3 arguments for lower dispatch
102 alignedElementCount = ShaderDefs.AlignUpGroup(newMaxElementCount);
103 maxBufferCount = totalSize;
104 maxLevelCount = levelCounts;
105 }
106 }
107
108 /// <summary>
109 /// Data structure containing the runtime resources that are bound by the command buffer.
110 /// </summary>
111 public struct SupportResources
112 {
113 internal bool ownsResources;
114 internal int alignedElementCount;
115 internal int maxBufferCount;
116 internal int maxLevelCount;
117
118 internal GraphicsBuffer prefixBuffer0;
119 internal GraphicsBuffer prefixBuffer1;
120 internal GraphicsBuffer totalLevelCountBuffer;
121 internal GraphicsBuffer levelOffsetBuffer;
122 internal GraphicsBuffer indirectDispatchArgsBuffer;
123
124 /// <summary>The prefix sum result.</summary>
125 public GraphicsBuffer output => prefixBuffer0;
126
127 /// <summary>
128 /// Allocate support resources to accomodate a max count.
129 /// </summary>
130 /// <param name="maxElementCount">The max element count.</param>
131 /// <returns>The created support resources.</returns>
132 public static SupportResources Create(int maxElementCount)
133 {
134 var resources = new SupportResources() { alignedElementCount = 0, ownsResources = true };
135 resources.Resize(maxElementCount);
136 return resources;
137 }
138
139 /// <summary>
140 /// Load supporting resources from Render Graph Resources.
141 /// </summary>
142 /// <param name="shaderGraphResources">Render Graph Resources</param>
143 /// <returns>The created support resources.</returns>
144 public static SupportResources Load(RenderGraphResources shaderGraphResources)
145 {
146 var resources = new SupportResources() { alignedElementCount = 0, ownsResources = false };
147 resources.LoadFromShaderGraph(shaderGraphResources);
148 return resources;
149 }
150
151 internal void Resize(int newMaxElementCount)
152 {
153 if (!ownsResources)
154 throw new Exception("Cannot resize resources unless they are owned. Use GpuPrefixSumSupportResources.Create() for this.");
155
156 newMaxElementCount = Math.Max(newMaxElementCount, 1); //at bare minimum support a single group.
157 if (alignedElementCount >= newMaxElementCount)
158 return;
159
160 Dispose();
161 ShaderDefs.CalculateTotalBufferSize(newMaxElementCount, out int totalSize, out int levelCounts);
162
163 alignedElementCount = ShaderDefs.AlignUpGroup(newMaxElementCount);
164 maxBufferCount = totalSize;
165 maxLevelCount = levelCounts;
166
167 prefixBuffer0 = new GraphicsBuffer(GraphicsBuffer.Target.Raw, totalSize, 4);
168 prefixBuffer1 = new GraphicsBuffer(GraphicsBuffer.Target.Raw, newMaxElementCount, 4);
169 totalLevelCountBuffer = new GraphicsBuffer(GraphicsBuffer.Target.Raw, 1, 4);
170 levelOffsetBuffer = new GraphicsBuffer(GraphicsBuffer.Target.Structured, levelCounts, System.Runtime.InteropServices.Marshal.SizeOf<LevelOffsets>());
171 indirectDispatchArgsBuffer = new GraphicsBuffer(GraphicsBuffer.Target.IndirectArguments, ShaderDefs.ArgsBufferStride * levelCounts, sizeof(uint));//3 arguments for upp dispatch, 3 arguments for lower dispatch
172 }
173
174 void LoadFromShaderGraph(RenderGraphResources shaderGraphResources)
175 {
176 alignedElementCount = shaderGraphResources.alignedElementCount;
177 maxBufferCount = shaderGraphResources.maxBufferCount;
178 maxLevelCount = shaderGraphResources.maxLevelCount;
179
180 prefixBuffer0 = (GraphicsBuffer)shaderGraphResources.prefixBuffer0;
181 prefixBuffer1 = (GraphicsBuffer)shaderGraphResources.prefixBuffer1;
182 totalLevelCountBuffer = (GraphicsBuffer)shaderGraphResources.totalLevelCountBuffer;
183 levelOffsetBuffer = (GraphicsBuffer)shaderGraphResources.levelOffsetBuffer;
184 indirectDispatchArgsBuffer = (GraphicsBuffer)shaderGraphResources.indirectDispatchArgsBuffer;
185 }
186
187 /// <summary>
188 /// Dispose the supporting resources.
189 /// </summary>
190 public void Dispose()
191 {
192 if (alignedElementCount == 0 || !ownsResources)
193 return;
194
195 alignedElementCount = 0;
196
197 void TryFreeBuffer(GraphicsBuffer resource)
198 {
199 if (resource != null)
200 {
201 resource.Dispose();
202 resource = null;
203 }
204 }
205
206 TryFreeBuffer(prefixBuffer0);
207 TryFreeBuffer(prefixBuffer1);
208 TryFreeBuffer(levelOffsetBuffer);
209 TryFreeBuffer(indirectDispatchArgsBuffer);
210 TryFreeBuffer(totalLevelCountBuffer);
211 }
212 }
213
214 /// <summary>
215 /// Arguments for a direct prefix sum.
216 /// </summary>
217 public struct DirectArgs
218 {
219 /// <summary>An inclusive or exclusive prefix sum.</summary>
220 public bool exclusive;
221 /// <summary>The size of the input list.</summary>
222 public int inputCount;
223 /// <summary>The input list.</summary>
224 public GraphicsBuffer input;
225 /// <summary>Required runtime resources.</summary>
226 public SupportResources supportResources;
227 }
228
229 /// <summary>
230 /// Arguments for an indirect prefix sum.
231 /// </summary>
232 public struct IndirectDirectArgs
233 {
234 /// <summary>An inclusive or exclusive prefix sum.</summary>
235 public bool exclusive;
236 /// <summary>Byte offset of the count inside the input count buffer.</summary>
237 public int inputCountBufferByteOffset;
238 /// <summary>GPU buffer defining the size of the input list.</summary>
239 public ComputeBuffer inputCountBuffer;
240 /// <summary>The input list.</summary>
241 public GraphicsBuffer input;
242 /// <summary>Required runtime resources.</summary>
243 public SupportResources supportResources;
244 }
245
246 /// <summary>
247 /// Structure defining any required assets used by the GPU sort.
248 /// </summary>
249 public struct SystemResources
250 {
251 /// <summary>
252 /// The compute asset that defines all of the kernels for the GPU prefix sum.
253 /// </summary>
254 public ComputeShader computeAsset;
255
256 internal int kernelCalculateLevelDispatchArgsFromConst;
257 internal int kernelCalculateLevelDispatchArgsFromBuffer;
258 internal int kernelPrefixSumOnGroup;
259 internal int kernelPrefixSumOnGroupExclusive;
260 internal int kernelPrefixSumNextInput;
261 internal int kernelPrefixSumResolveParent;
262 internal int kernelPrefixSumResolveParentExclusive;
263
264 internal void LoadKernels()
265 {
266 if (computeAsset == null)
267 return;
268
269 kernelCalculateLevelDispatchArgsFromConst = computeAsset.FindKernel("MainCalculateLevelDispatchArgsFromConst");
270 kernelCalculateLevelDispatchArgsFromBuffer = computeAsset.FindKernel("MainCalculateLevelDispatchArgsFromBuffer");
271 kernelPrefixSumOnGroup = computeAsset.FindKernel("MainPrefixSumOnGroup");
272 kernelPrefixSumOnGroupExclusive = computeAsset.FindKernel("MainPrefixSumOnGroupExclusive");
273 kernelPrefixSumNextInput = computeAsset.FindKernel("MainPrefixSumNextInput");
274 kernelPrefixSumResolveParent = computeAsset.FindKernel("MainPrefixSumResolveParent");
275 kernelPrefixSumResolveParentExclusive = computeAsset.FindKernel("MainPrefixSumResolveParentExclusive");
276 }
277 }
278 }
279}