Rasagar/Library/PackageCache/com.unity.rendering.light-transport/Runtime/UnifiedRayTracing/Compute/RadeonRays/RadeonRaysAPI.cs
2024-08-26 23:07:20 +03:00

465 lines
18 KiB
C#

using System;
using System.IO;
using System.Runtime.InteropServices;
using Unity.Mathematics;
namespace UnityEngine.Rendering.RadeonRays
{
enum IndexFormat { Int32 = 0, Int16 };
internal struct MeshBuildInfo
{
public GraphicsBuffer vertices;
public int verticesStartOffset; // in DWORD
public uint vertexCount;
public uint vertexStride; // in DWORD
public int baseVertex;
public GraphicsBuffer triangleIndices;
public int indicesStartOffset; // in DWORD
public int baseIndex;
public IndexFormat indexFormat;
public uint triangleCount;
}
internal struct MeshBuildMemoryRequirements
{
public ulong buildScratchSizeInDwords;
public ulong resultSizeInDwords;
}
internal struct SceneBuildMemoryRequirements
{
public ulong buildScratchSizeInDwords;
}
internal class SceneMemoryRequirements
{
public ulong buildScratchSizeInDwords;
public ulong[] bottomLevelBvhSizeInDwords;
public uint[] bottomLevelBvhOffsetInNodes;
public ulong totalBottomLevelBvhSizeInDwords;
}
[System.Flags]
internal enum BuildFlags
{
None = 0,
PreferFastBuild = 1 << 0,
MinimizeMemory = 1 << 1
}
internal enum RayQueryType
{
ClosestHit,
AnyHit
}
internal enum RayQueryOutputType
{
FullHitData,
InstanceID
}
[StructLayout(LayoutKind.Sequential)]
internal struct Transform
{
public float4 row0;
public float4 row1;
public float4 row2;
public Transform(float4 row0, float4 row1, float4 row2)
{
this.row0 = row0;
this.row1 = row1;
this.row2 = row2;
}
public static Transform Identity()
{
return new Transform(
new float4(1.0f, 0.0f, 0.0f, 0.0f),
new float4(0.0f, 1.0f, 0.0f, 0.0f),
new float4(0.0f, 0.0f, 1.0f, 0.0f));
}
public static Transform Translation(float3 translation)
{
return new Transform(
new float4(1.0f, 0.0f, 0.0f, translation.x),
new float4(0.0f, 1.0f, 0.0f, translation.y),
new float4(0.0f, 0.0f, 1.0f, translation.z));
}
public static Transform Scale(float3 scale)
{
return new Transform(
new float4(scale.x, 0.0f, 0.0f, 0.0f),
new float4(0.0f, scale.y, 0.0f, 0.0f),
new float4(0.0f, 0.0f, scale.z, 0.0f));
}
public static Transform TRS(float3 translation, float3 rotation, float3 scale)
{
var rot = float3x3.Euler(rotation);
rot.c0 *= scale.x;
rot.c1 *= scale.y;
rot.c2 *= scale.z;
return new Transform(
new float4(rot.c0.x, rot.c1.x, rot.c2.x, translation.x),
new float4(rot.c0.y, rot.c1.y, rot.c2.y, translation.y),
new float4(rot.c0.z, rot.c1.z, rot.c2.z, translation.z));
}
public Transform Inverse()
{
float4x4 m = new float4x4();
m[0] = new float4(row0.x, row1.x, row2.x, 0);
m[1] = new float4(row0.y, row1.y, row2.y, 0);
m[2] = new float4(row0.z, row1.z, row2.z, 0);
m[3] = new float4(row0.w, row1.w, row2.w, 1);
m = math.fastinverse(m);
Transform res;
res.row0 = new float4(m[0].x, m[1].x, m[2].x, m[3].x);
res.row1 = new float4(m[0].y, m[1].y, m[2].y, m[3].y);
res.row2 = new float4(m[0].z, m[1].z, m[2].z, m[3].z);
return res;
}
}
[StructLayout(LayoutKind.Sequential)]
internal struct BvhNode
{
public uint child0;
public uint child1;
public uint parent;
public uint update;
public uint3 aabb0_min;
public uint3 aabb0_max;
public uint3 aabb1_min;
public uint3 aabb1_max;
}
[StructLayout(LayoutKind.Sequential)]
internal struct BvhHeader
{
public uint totalNodeCount;
public uint leafNodeCount;
public uint root;
public uint unused;
public uint3 unused1;
public uint3 unused2;
public uint3 unused3;
public uint3 unused4;
}
internal struct Instance
{
public uint meshAccelStructOffset;
public uint instanceMask;
public uint vertexOffset;
public uint indexOffset;
public bool triangleCullingEnabled;
public bool invertTriangleCulling;
public uint userInstanceID;
public Transform localToWorldTransform;
}
[StructLayout(LayoutKind.Sequential)]
internal struct InstanceInfo
{
public int blasOffset;
public int instanceMask;
public int vertexOffset;
public int indexOffset;
public int triangleCullingEnabled;
public int invertTriangleCulling;
public uint userInstanceID;
public int padding2;
public Transform worldToLocalTransform;
public Transform localToWorldTransform;
}
internal class RadeonRaysShaders
{
public ComputeShader bitHistogram;
public ComputeShader blockReducePart;
public ComputeShader blockScan;
public ComputeShader buildHlbvh;
public ComputeShader reorderTriangleIndices;
public ComputeShader restructureBvh;
public ComputeShader scatter;
public ComputeShader topLevelIntersector;
public ComputeShader intersector;
#if UNITY_EDITOR
public static RadeonRaysShaders LoadFromPath(string kernelFolderPath)
{
var res = new RadeonRaysShaders();
res.bitHistogram = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "bit_histogram.compute"));
res.blockReducePart = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "block_reduce_part.compute"));
res.blockScan = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "block_scan.compute"));
res.buildHlbvh = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "build_hlbvh.compute"));
res.reorderTriangleIndices = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "reorder_triangle_indices.compute"));
res.restructureBvh = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "restructure_bvh.compute"));
res.scatter = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "scatter.compute"));
res.topLevelIntersector = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "top_level_intersector.compute"));
res.intersector = UnityEditor.AssetDatabase.LoadAssetAtPath<ComputeShader>(Path.Combine(kernelFolderPath, "intersector.compute"));
return res;
}
#endif
}
internal class RadeonRaysAPI : IDisposable
{
private HlbvhBuilder buildBvh;
private HlbvhTopLevelBuilder buildTopLevelBvh;
private RestructureBvh restructureBvh;
private TraceGeometry traceGeometry;
private TraceScene traceScene;
public const GraphicsBuffer.Target BufferTarget = GraphicsBuffer.Target.Structured;
public RadeonRaysAPI(RadeonRaysShaders shaders)
{
buildBvh = new HlbvhBuilder(shaders);
buildTopLevelBvh = new HlbvhTopLevelBuilder(shaders);
restructureBvh = new RestructureBvh(shaders);
traceGeometry = new TraceGeometry(shaders);
traceScene = new TraceScene(shaders);
}
public void Dispose()
{
restructureBvh.Dispose();
}
static public int BvhNodeSizeInDwords() { return Marshal.SizeOf<BvhNode>() / 4; }
static public int BvhNodeSizeInBytes() { return Marshal.SizeOf<BvhNode>(); }
public void BuildMeshAccelStruct(
CommandBuffer cmd,
MeshBuildInfo buildInfo, BuildFlags buildFlags,
GraphicsBuffer scratchBuffer,
GraphicsBuffer accelStructBuffer, uint accelStructOffsetInNodes, uint accelStructSizeInNodes,
Action<AsyncGPUReadbackRequest> callback = null)
{
if (SystemInfo.graphicsDeviceType == GraphicsDeviceType.Metal)
buildFlags |= BuildFlags.PreferFastBuild;
buildBvh.Execute(cmd,
buildInfo.vertices, buildInfo.verticesStartOffset, buildInfo.vertexStride,
buildInfo.triangleIndices, buildInfo.indicesStartOffset, buildInfo.triangleCount,
scratchBuffer,
accelStructBuffer, accelStructOffsetInNodes, accelStructSizeInNodes,
(buildFlags & BuildFlags.MinimizeMemory) != 0 ? 2u : 0);
if (callback != null && (buildFlags & BuildFlags.MinimizeMemory) != 0)
cmd.RequestAsyncReadback(accelStructBuffer, 4*sizeof(uint), (int)accelStructOffsetInNodes*Marshal.SizeOf<BvhNode>(), callback);
if ((buildFlags & BuildFlags.PreferFastBuild) == 0)
{
restructureBvh.Execute(cmd,
buildInfo.vertices, buildInfo.verticesStartOffset, buildInfo.vertexStride,
buildInfo.triangleIndices, buildInfo.indicesStartOffset, buildInfo.triangleCount,
scratchBuffer, accelStructBuffer, accelStructOffsetInNodes);
}
}
public MeshBuildMemoryRequirements GetMeshBuildMemoryRequirements(MeshBuildInfo buildInfo, BuildFlags buildFlags)
{
if (SystemInfo.graphicsDeviceType == GraphicsDeviceType.Metal)
buildFlags |= BuildFlags.PreferFastBuild;
var result = new MeshBuildMemoryRequirements();
if ((buildFlags & BuildFlags.MinimizeMemory) != 0)
result.resultSizeInDwords = buildBvh.GetResultDataSizeInDwordsPrediction(buildInfo.triangleCount);
else
result.resultSizeInDwords = buildBvh.GetResultDataSizeInDwords(buildInfo.triangleCount);
result.buildScratchSizeInDwords = buildBvh.GetScratchDataSizeInDwords(buildInfo.triangleCount);
ulong restructureScratchSize = ((buildFlags & BuildFlags.PreferFastBuild) == 0) ? restructureBvh.GetScratchDataSizeInDwords(buildInfo.triangleCount) : 0;
result.buildScratchSizeInDwords = math.max(result.buildScratchSizeInDwords, restructureScratchSize);
return result;
}
public TopLevelAccelStruct BuildSceneAccelStruct(
CommandBuffer cmd,
GraphicsBuffer meshAccelStructsBuffer,
Instance[] instances, BuildFlags buildFlags,
GraphicsBuffer scratchBuffer)
{
if (SystemInfo.graphicsDeviceType == GraphicsDeviceType.Metal)
buildFlags |= BuildFlags.PreferFastBuild;
var accelStruct = new TopLevelAccelStruct();
if (instances.Length == 0)
{
buildTopLevelBvh.CreateEmpty(ref accelStruct);
return accelStruct;
}
buildTopLevelBvh.AllocateResultBuffers((uint)instances.Length, ref accelStruct);
var instancesInfos = new InstanceInfo[instances.Length];
for (uint i = 0; i < instances.Length; ++i)
{
instancesInfos[i] = new InstanceInfo
{
blasOffset = (int)instances[i].meshAccelStructOffset,
instanceMask = (int)instances[i].instanceMask,
vertexOffset = (int)instances[i].vertexOffset,
indexOffset = (int)instances[i].indexOffset,
localToWorldTransform = instances[i].localToWorldTransform,
triangleCullingEnabled = instances[i].triangleCullingEnabled ? 1 : 0,
invertTriangleCulling = instances[i].invertTriangleCulling ? 1 : 0,
userInstanceID = instances[i].userInstanceID
// worldToLocal computed in the shader
};
}
accelStruct.instanceInfos.SetData(instancesInfos);
accelStruct.bottomLevelBvhs = meshAccelStructsBuffer;
accelStruct.instanceCount = (uint)instances.Length;
buildTopLevelBvh.Execute(cmd, scratchBuffer, ref accelStruct);
return accelStruct;
}
public TopLevelAccelStruct CreateSceneAccelStructBuffers(
GraphicsBuffer meshAccelStructsBuffer,
uint tlasSizeInDwords,
Instance[] instances)
{
var accelStruct = new TopLevelAccelStruct();
if (instances.Length == 0)
{
buildTopLevelBvh.CreateEmpty(ref accelStruct);
return accelStruct;
}
var instancesInfos = new InstanceInfo[instances.Length];
for (uint i = 0; i < instances.Length; ++i)
{
instancesInfos[i] = new InstanceInfo
{
blasOffset = (int)instances[i].meshAccelStructOffset,
instanceMask = (int)instances[i].instanceMask,
vertexOffset = (int)instances[i].vertexOffset,
indexOffset = (int)instances[i].indexOffset,
localToWorldTransform = instances[i].localToWorldTransform,
triangleCullingEnabled = instances[i].triangleCullingEnabled ? 1 : 0,
invertTriangleCulling = instances[i].invertTriangleCulling ? 1 : 0,
userInstanceID = instances[i].userInstanceID,
worldToLocalTransform = instances[i].localToWorldTransform.Inverse()
};
}
accelStruct.instanceInfos = new GraphicsBuffer(TopLevelAccelStruct.instanceInfoTarget, instances.Length, Marshal.SizeOf<InstanceInfo>());
accelStruct.instanceInfos.SetData(instancesInfos);
accelStruct.bottomLevelBvhs = meshAccelStructsBuffer;
accelStruct.topLevelBvh = new GraphicsBuffer(TopLevelAccelStruct.topLevelBvhTarget, (int)tlasSizeInDwords / BvhNodeSizeInDwords(), Marshal.SizeOf<BvhNode>());
accelStruct.instanceCount = (uint)instances.Length;
return accelStruct;
}
public SceneBuildMemoryRequirements GetSceneBuildMemoryRequirements(uint instanceCount)
{
var result = new SceneBuildMemoryRequirements();
result.buildScratchSizeInDwords = buildTopLevelBvh.GetScratchDataSizeInDwords(instanceCount);
return result;
}
public SceneMemoryRequirements GetSceneMemoryRequirements(MeshBuildInfo[] buildInfos, BuildFlags buildFlags)
{
if (SystemInfo.graphicsDeviceType == GraphicsDeviceType.Metal)
buildFlags |= BuildFlags.PreferFastBuild;
var requirements = new SceneMemoryRequirements();
requirements.bottomLevelBvhSizeInDwords = new ulong[buildInfos.Length];
requirements.bottomLevelBvhOffsetInNodes = new uint[buildInfos.Length];
requirements.buildScratchSizeInDwords = 0;
requirements.totalBottomLevelBvhSizeInDwords = 0;
int i = 0;
uint offset = 0;
foreach (var buildInfo in buildInfos)
{
var meshRequirements = GetMeshBuildMemoryRequirements(buildInfo, buildFlags);
requirements.buildScratchSizeInDwords = math.max(requirements.buildScratchSizeInDwords, meshRequirements.buildScratchSizeInDwords);
requirements.totalBottomLevelBvhSizeInDwords += meshRequirements.resultSizeInDwords;
requirements.bottomLevelBvhSizeInDwords[i] = meshRequirements.resultSizeInDwords;
requirements.bottomLevelBvhOffsetInNodes[i] = offset;
offset += (uint)(meshRequirements.resultSizeInDwords / (ulong)RadeonRaysAPI.BvhNodeSizeInDwords());
i++;
}
ulong topLevelScratchSize = buildTopLevelBvh.GetScratchDataSizeInDwords((uint)buildInfos.Length);
requirements.buildScratchSizeInDwords = math.max(requirements.buildScratchSizeInDwords, topLevelScratchSize);
return requirements;
}
static public ulong GetTraceMemoryRequirements(uint rayCount)
{
return math.max(TraceGeometry.GetScratchDataSizeInDwords(rayCount), TraceScene.GetScratchDataSizeInDwords(rayCount));
}
public void Intersect(
CommandBuffer cmd,
GraphicsBuffer raysBuffer, uint rayCount,
RayQueryType queryType, RayQueryOutputType queryOutputType,
GraphicsBuffer vertexBuffer, GraphicsBuffer indexBuffer,
GraphicsBuffer accelStructBuffer, uint accelStructOffset, GraphicsBuffer scratchBuffer,
GraphicsBuffer hitsBuffer)
{
traceGeometry.Execute(cmd, raysBuffer, rayCount, null, queryType, queryOutputType,
vertexBuffer, indexBuffer, accelStructBuffer, accelStructOffset, scratchBuffer, hitsBuffer);
}
public void IntersectIndirect(
CommandBuffer cmd,
GraphicsBuffer raysBuffer, GraphicsBuffer indirectRayCount,
RayQueryType queryType, RayQueryOutputType queryOutputType,
GraphicsBuffer vertexBuffer, GraphicsBuffer indexBuffer,
GraphicsBuffer accelStructBuffer, uint accelStructOffset, GraphicsBuffer scratchBuffer,
GraphicsBuffer hitsBuffer)
{
traceGeometry.Execute(cmd, raysBuffer, 0, indirectRayCount, queryType, queryOutputType,
vertexBuffer, indexBuffer, accelStructBuffer, accelStructOffset, scratchBuffer, hitsBuffer);
}
public void Intersect(
CommandBuffer cmd,
GraphicsBuffer raysBuffer, uint rayCount, GraphicsBuffer indirectRayCount,
RayQueryType queryType, RayQueryOutputType queryOutputType,
GraphicsBuffer vertexBuffer, int vertexStride, GraphicsBuffer indexBuffer,
TopLevelAccelStruct accelStruct, GraphicsBuffer scratchBuffer,
GraphicsBuffer hitsBuffer)
{
traceScene.Execute(
cmd,
raysBuffer, rayCount, indirectRayCount, queryType, queryOutputType,
vertexBuffer, vertexStride, indexBuffer,
accelStruct, scratchBuffer, hitsBuffer);
}
}
}