A game about forced loneliness, made by TACStudios
tacstudios.tngl.sh
1using System;
2
3namespace UnityEngine.Rendering
4{
5 /// <summary>
6 /// Utility class for computing inclusive or exclusive prefix sums, directly or indirectly dispatched on the GPU.
7 /// </summary>
8 public partial struct GPUPrefixSum
9 {
10 private SystemResources resources;
11
12 /// <summary>
13 /// Initializes a re-usable GPU prefix sum instance.
14 /// </summary>
15 /// <param name="resources">The required system resources.</param>
16 public GPUPrefixSum(SystemResources resources)
17 {
18 this.resources = resources;
19 this.resources.LoadKernels();
20 }
21
22 Vector4 PackPrefixSumArgs(int a, int b, int c, int d)
23 {
24 unsafe
25 {
26 return new Vector4(
27 *((float*)&a),
28 *((float*)&b),
29 *((float*)&c),
30 *((float*)&d));
31 }
32 }
33
34 internal void ExecuteCommonIndirect(CommandBuffer cmdBuffer, GraphicsBuffer inputBuffer, in SupportResources supportResources, bool isExclusive)
35 {
36 int sumOnGroupKernel = isExclusive ? resources.kernelPrefixSumOnGroupExclusive : resources.kernelPrefixSumOnGroup;
37 int sumResolveParentKernel = isExclusive ? resources.kernelPrefixSumResolveParentExclusive : resources.kernelPrefixSumResolveParent;
38
39 //hierarchy up
40 for (int levelId = 0; levelId < supportResources.maxLevelCount; ++levelId)
41 {
42 var packedArgs = PackPrefixSumArgs(0, 0, 0, levelId);
43 cmdBuffer.SetComputeVectorParam(resources.computeAsset, ShaderIDs._PrefixSumIntArgs, packedArgs);
44
45 if (levelId == 0)
46 cmdBuffer.SetComputeBufferParam(resources.computeAsset, sumOnGroupKernel, ShaderIDs._InputBuffer, inputBuffer);
47 else
48 cmdBuffer.SetComputeBufferParam(resources.computeAsset, sumOnGroupKernel, ShaderIDs._InputBuffer, supportResources.prefixBuffer1);
49
50 cmdBuffer.SetComputeBufferParam(resources.computeAsset, sumOnGroupKernel, ShaderIDs._TotalLevelsBuffer, supportResources.totalLevelCountBuffer);
51 cmdBuffer.SetComputeBufferParam(resources.computeAsset, sumOnGroupKernel, ShaderIDs._LevelsOffsetsBuffer, supportResources.levelOffsetBuffer);
52 cmdBuffer.SetComputeBufferParam(resources.computeAsset, sumOnGroupKernel, ShaderIDs._OutputBuffer, supportResources.prefixBuffer0);
53 cmdBuffer.DispatchCompute(resources.computeAsset, sumOnGroupKernel, supportResources.indirectDispatchArgsBuffer, (uint)(levelId * ShaderDefs.ArgsBufferStride * 4));
54
55 if (levelId == supportResources.maxLevelCount - 1)
56 continue;
57
58 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelPrefixSumNextInput, ShaderIDs._InputBuffer, supportResources.prefixBuffer0);
59 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelPrefixSumNextInput, ShaderIDs._LevelsOffsetsBuffer, supportResources.levelOffsetBuffer);
60 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelPrefixSumNextInput, ShaderIDs._OutputBuffer, supportResources.prefixBuffer1);
61 cmdBuffer.DispatchCompute(resources.computeAsset, resources.kernelPrefixSumNextInput, supportResources.indirectDispatchArgsBuffer, (uint)((levelId + 1) * ShaderDefs.ArgsBufferStride * 4));
62 }
63
64 //down the hierarchy
65 for (int levelId = supportResources.maxLevelCount - 1; levelId >= 1; --levelId)
66 {
67 var packedArgs = PackPrefixSumArgs(0, 0, 0, levelId);
68 cmdBuffer.SetComputeVectorParam(resources.computeAsset, ShaderIDs._PrefixSumIntArgs, packedArgs);
69 cmdBuffer.SetComputeBufferParam(resources.computeAsset, sumResolveParentKernel, ShaderIDs._InputBuffer, inputBuffer);
70 cmdBuffer.SetComputeBufferParam(resources.computeAsset, sumResolveParentKernel, ShaderIDs._OutputBuffer, supportResources.prefixBuffer0);
71 cmdBuffer.SetComputeBufferParam(resources.computeAsset, sumResolveParentKernel, ShaderIDs._LevelsOffsetsBuffer, supportResources.levelOffsetBuffer);
72 cmdBuffer.DispatchCompute(resources.computeAsset, sumResolveParentKernel, supportResources.indirectDispatchArgsBuffer, (uint)(((levelId - 1) * ShaderDefs.ArgsBufferStride + ShaderDefs.ArgsBufferLower) * 4));
73 }
74 }
75
76 /// <summary>
77 /// Prefix sum a list of data from a CPU-defined count.
78 /// </summary>
79 /// <param name="cmdBuffer">Command Buffer for recording the prefix sum commands.</param>
80 /// <param name="arguments">Runtime arguments for the prefix sum.</param>
81 /// <exception cref="Exception">When the input data is invalid.</exception>
82 public void DispatchDirect(CommandBuffer cmdBuffer, in DirectArgs arguments)
83 {
84 if (arguments.supportResources.prefixBuffer0 == null || arguments.supportResources.prefixBuffer1 == null)
85 throw new Exception("Support resources are not valid.");
86
87 if (arguments.input == null)
88 throw new Exception("Input source buffer cannot be null.");
89
90 if (arguments.inputCount > arguments.supportResources.alignedElementCount)
91 throw new Exception("Input count exceeds maximum count of support resources. Ensure to create support resources with enough space.");
92
93 //Generate level offsets first, from const value.
94 var packedArgs = PackPrefixSumArgs(arguments.inputCount, arguments.supportResources.maxLevelCount, 0, 0);
95 cmdBuffer.SetComputeVectorParam(resources.computeAsset, ShaderIDs._PrefixSumIntArgs, packedArgs);
96 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromConst, ShaderIDs._OutputLevelsOffsetsBuffer, arguments.supportResources.levelOffsetBuffer);
97 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromConst, ShaderIDs._OutputDispatchLevelArgsBuffer, arguments.supportResources.indirectDispatchArgsBuffer);
98 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromConst, ShaderIDs._OutputTotalLevelsBuffer, arguments.supportResources.totalLevelCountBuffer);
99 cmdBuffer.DispatchCompute(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromConst, 1, 1, 1);
100
101 ExecuteCommonIndirect(cmdBuffer, arguments.input, arguments.supportResources, arguments.exclusive);
102 }
103
104 /// <summary>
105 /// Prefix sum a list of data from a GPU-defined count.
106 /// </summary>
107 /// <param name="cmdBuffer">Command Buffer for recording the prefix sum commands.</param>
108 /// <param name="arguments">Runtime arguments for the prefix sum.</param>
109 /// <exception cref="Exception">When the input data is invalid.</exception>
110 public void DispatchIndirect(CommandBuffer cmdBuffer, in IndirectDirectArgs arguments)
111 {
112 if (arguments.supportResources.prefixBuffer0 == null || arguments.supportResources.prefixBuffer1 == null)
113 throw new Exception("Support resources are not valid.");
114
115 if (arguments.input == null || arguments.inputCountBuffer == null)
116 throw new Exception("Input source buffer and inputCountBuffer cannot be null.");
117
118 //Generate level offsets first, from const value.
119 var packedArgs = PackPrefixSumArgs(0, arguments.supportResources.maxLevelCount, arguments.inputCountBufferByteOffset, 0);
120 cmdBuffer.SetComputeVectorParam(resources.computeAsset, ShaderIDs._PrefixSumIntArgs, packedArgs);
121 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromBuffer, ShaderIDs._InputCountBuffer, arguments.inputCountBuffer);
122 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromBuffer, ShaderIDs._OutputLevelsOffsetsBuffer, arguments.supportResources.levelOffsetBuffer);
123 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromBuffer, ShaderIDs._OutputDispatchLevelArgsBuffer, arguments.supportResources.indirectDispatchArgsBuffer);
124 cmdBuffer.SetComputeBufferParam(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromBuffer, ShaderIDs._OutputTotalLevelsBuffer, arguments.supportResources.totalLevelCountBuffer);
125 cmdBuffer.DispatchCompute(resources.computeAsset, resources.kernelCalculateLevelDispatchArgsFromBuffer, 1, 1, 1);
126
127 ExecuteCommonIndirect(cmdBuffer, arguments.input, arguments.supportResources, arguments.exclusive);
128 }
129 }
130}