using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using Unity.Mathematics;
using UnityEngine.Rendering;
using UnityEngine.Rendering.Sampling;
using UnityEngine.Rendering.UnifiedRayTracing;
using static UnityEngine.PathTracing.Core.World;
namespace UnityEngine.PathTracing.Core
{
// Interface for many light sampling
internal interface IManyLightSampling : IDisposable
{
void Build(CommandBuffer cmd, World.LightState lightState, Bounds sceneBounds, SamplingResources samplingResources);
void Bind(CommandBuffer cmd, IRayTracingShader shader);
}
internal enum GridMemLayout { Sparse, Dense };
internal enum GridSizingStrategy { Uniform, FitToSceneBounds };
internal static class LightGridUtils
{
public static Vector3Int ComputeLightGridDims(Vector3 sceneBounds, int maxLightGridCellCount, GridSizingStrategy lightGridSizingStrategy)
{
if (lightGridSizingStrategy == GridSizingStrategy.Uniform)
{
int volumeSide = (int)Math.Pow((double)maxLightGridCellCount, 1.0 / 3.0);
if ((volumeSide + 1) * (volumeSide + 1) * (volumeSide + 1) <= maxLightGridCellCount)
volumeSide++;
return new Vector3Int(volumeSide, volumeSide, volumeSide);
}
// Fix scene bounds if the ratio between 2 dims is too important
float maxSceneDim = math.max(sceneBounds.x, math.max(sceneBounds.y, sceneBounds.z));
if (sceneBounds.x * Mathf.Sqrt(maxLightGridCellCount) < maxSceneDim)
sceneBounds.x = maxSceneDim / Mathf.Sqrt(maxLightGridCellCount);
if (sceneBounds.y * Mathf.Sqrt(maxLightGridCellCount) < maxSceneDim)
sceneBounds.y = maxSceneDim / Mathf.Sqrt(maxLightGridCellCount);
if (sceneBounds.z * Mathf.Sqrt(maxLightGridCellCount) < maxSceneDim)
sceneBounds.z = maxSceneDim / Mathf.Sqrt(maxLightGridCellCount);
// Compute ideal cell width (we aim for for cells having the same width along all 3 axes)
float idealCellWidth = Mathf.Pow(sceneBounds.x * sceneBounds.y * sceneBounds.z / ((float)maxLightGridCellCount), 1.0f / 3.0f);
// Use ideal cell width to compute grid dims
Vector3Int gridDims = new Vector3Int(
Math.Max((int)(sceneBounds.x / idealCellWidth), 1),
Math.Max((int)(sceneBounds.y / idealCellWidth), 1),
Math.Max((int)(sceneBounds.z / idealCellWidth), 1));
Debug.Assert(gridDims.x * gridDims.y * gridDims.z <= maxLightGridCellCount);
return gridDims;
}
}
internal class ConservativeLightGrid : IManyLightSampling
{
// Light grid parameters
public int LightGridCellCount = 64 * 64 * 64;
public int MaxLightsPerCell = 64; // Only used with GridMemLayout.Dense
public GridSizingStrategy LightGridSizingStrategy = GridSizingStrategy.FitToSceneBounds;
public GridMemLayout GridMemLayout = GridMemLayout.Sparse;
public ConservativeLightGrid(ComputeShader shader)
{
_shader = shader;
_buildLightGridlKernel = _shader.FindKernel("BuildConservativeLightGrid");
}
public void Init()
{
if (GridMemLayout == GridMemLayout.Dense && (_lightGridCellsDataBuffer == null || _lightGridCellsDataBuffer.count <= 1))
{
int count = LightGridCellCount * MaxLightsPerCell;
int stride = Marshal.SizeOf<World.ThinReservoir>();
_lightGridCellsDataBuffer?.Dispose();
_lightGridCellsDataBuffer = new ComputeBuffer(count, stride);
}
if (_lightGridBuffer == null || _lightGridBuffer.count <= 1)
{
int count = LightGridCellCount;
int stride = Marshal.SizeOf<int2>();
_lightGridBuffer?.Dispose();
_lightGridBuffer = new ComputeBuffer(count, stride);
}
if (_totalLightsInGridCountBuffer == null)
{
_totalLightsInGridCountBuffer = new ComputeBuffer(1, sizeof(int));
}
}
protected void BindComputeResources(CommandBuffer cmd, World.LightState lightState, Bounds sceneBounds, SamplingResources samplingResources)
{
SamplingResources.Bind(cmd, samplingResources);
// Set the input lighting state.Note that this is a subset, as we evaluate without light cookies
cmd.SetComputeIntParam(_shader, ShaderProperties.NumLights, lightState.LightCount);
cmd.SetComputeIntParam(_shader, ShaderProperties.MaxLightsPerCell, MaxLightsPerCell);
cmd.SetComputeIntParam(_shader, ShaderProperties.NumEmissiveMeshes, lightState.MeshLightCount);
cmd.SetComputeIntParam(_shader, ShaderProperties.GridDimX, _lightGridDims.x);
cmd.SetComputeIntParam(_shader, ShaderProperties.GridDimY, _lightGridDims.y);
cmd.SetComputeIntParam(_shader, ShaderProperties.GridDimZ, _lightGridDims.z);
cmd.SetComputeVectorParam(_shader, ShaderProperties.GridMin, sceneBounds.min);
cmd.SetComputeVectorParam(_shader, ShaderProperties.GridSize, sceneBounds.size);
cmd.SetComputeVectorParam(_shader, ShaderProperties.CellSize, _cellSize);
cmd.SetComputeVectorParam(_shader, ShaderProperties.InvCellSize, _invCellSize);
cmd.SetComputeBufferParam(_shader, _buildLightGridlKernel, ShaderProperties.LightList, lightState.LightListBuffer);
// Set the output buffer
cmd.SetComputeBufferParam(_shader, _buildLightGridlKernel, ShaderProperties.LightGrid, _lightGridBuffer);
cmd.SetComputeBufferParam(_shader, _buildLightGridlKernel, ShaderProperties.TotalReservoirCount, _totalLightsInGridCountBuffer);
}
public void Build(CommandBuffer cmd, World.LightState lightState, Bounds sceneBounds, SamplingResources samplingResources)
{
if (lightState.LightListBuffer == null)
return;
Init();
_sceneBounds = sceneBounds;
_lightGridDims = LightGridUtils.ComputeLightGridDims(sceneBounds.size, LightGridCellCount, LightGridSizingStrategy);
Vector3 div = new Vector3(1.0f / _lightGridDims.x, 1.0f / _lightGridDims.y, 1.0f / _lightGridDims.z);
Vector3 cellSize = Vector3.Scale(sceneBounds.size, div);
_cellSize = cellSize;
// The length of the diagonal
_cellSize.w = Mathf.Sqrt(cellSize.x * cellSize.x + cellSize.y * cellSize.y + cellSize.z * cellSize.z);
_invCellSize = new Vector4(1.0f / _cellSize.x, 1.0f / _cellSize.y, 1.0f / _cellSize.z, 1.0f / _cellSize.w);
BindComputeResources(cmd, lightState, sceneBounds, samplingResources);
// If the grid is sparse, do a first dispatch to determine the total light count for all cells
// And allocate the _lightGridBuffer based on that number
if (GridMemLayout == GridMemLayout.Sparse)
{
_shader.EnableKeyword("SPARSE_GRID");
cmd.SetComputeBufferParam(_shader, _buildLightGridlKernel, ShaderProperties.LightGridCellsData, _lightGridBuffer); // dummy bind
DispatchBuild(cmd, 0);
GraphicsHelpers.Flush(cmd);
var requiredLightCount = new int[1];
_totalLightsInGridCountBuffer.GetData(requiredLightCount);
if (_lightGridCellsDataBuffer == null || _lightGridCellsDataBuffer.count < requiredLightCount[0])
{
_lightGridCellsDataBuffer?.Dispose();
_lightGridCellsDataBuffer = new ComputeBuffer(math.max(requiredLightCount[0], 1), Marshal.SizeOf<World.ThinReservoir>());
}
// Need to re-bind everything after flush
BindComputeResources(cmd, lightState, sceneBounds, samplingResources);
}
else
{
_shader.DisableKeyword("SPARSE_GRID");
}
// Build the grid
cmd.SetComputeBufferParam(_shader, _buildLightGridlKernel, ShaderProperties.LightGridCellsData, _lightGridCellsDataBuffer);
DispatchBuild(cmd, 1);
}
public void Bind(CommandBuffer cmd, IRayTracingShader shader)
{
if (_lightGridCellsDataBuffer == null)
{
// dummy buffer, when the feature is disabled
int stride = Marshal.SizeOf<World.ThinReservoir>();
_lightGridCellsDataBuffer = new ComputeBuffer(1, stride);
}
if (_lightGridBuffer == null)
{
// dummy buffer, when the feature is disabled
_lightGridBuffer = new ComputeBuffer(1, Marshal.SizeOf<int2>());
}
shader.SetIntParam(cmd, ShaderProperties.GridDimX, _lightGridDims.x);
shader.SetIntParam(cmd, ShaderProperties.GridDimY, _lightGridDims.y);
shader.SetIntParam(cmd, ShaderProperties.GridDimZ, _lightGridDims.z);
shader.SetIntParam(cmd, ShaderProperties.NumReservoirs, MaxLightsPerCell);
shader.SetVectorParam(cmd, ShaderProperties.GridMin, _sceneBounds.min);
shader.SetVectorParam(cmd, ShaderProperties.GridSize, _sceneBounds.size);
shader.SetVectorParam(cmd, ShaderProperties.CellSize, _cellSize);
shader.SetVectorParam(cmd, ShaderProperties.InvCellSize, _invCellSize);
shader.SetBufferParam(cmd, ShaderProperties.LightGridCellsData, _lightGridCellsDataBuffer);
shader.SetBufferParam(cmd, ShaderProperties.LightGrid, _lightGridBuffer);
}
public void Dispose()
{
_lightGridCellsDataBuffer?.Dispose();
_lightGridBuffer?.Dispose();
_totalLightsInGridCountBuffer?.Dispose();
}
void DispatchBuild(CommandBuffer cmd, int buildPass)
{
const int groupDim = 4;
cmd.SetComputeIntParam(_shader, ShaderProperties.BuildPass, buildPass);
cmd.SetBufferData(_totalLightsInGridCountBuffer, new uint[] { 0 });
cmd.DispatchCompute(_shader, _buildLightGridlKernel,
GraphicsHelpers.DivUp(_lightGridDims.x, groupDim),
GraphicsHelpers.DivUp(_lightGridDims.y, groupDim),
GraphicsHelpers.DivUp(_lightGridDims.z, groupDim));
}
readonly ComputeShader _shader;
readonly int _buildLightGridlKernel;
ComputeBuffer _lightGridCellsDataBuffer;
ComputeBuffer _lightGridBuffer;
ComputeBuffer _totalLightsInGridCountBuffer;
Bounds _sceneBounds;
Vector4 _cellSize;
Vector4 _invCellSize;
Vector3Int _lightGridDims;
}
internal class RegirLightGrid : IManyLightSampling
{
// Light grid parameters
public int LightGridCellCount = 64 * 64 * 64;
public int MaxLightsPerCell = 64;
public int NumCandidates = -1; // -1 means we iterate over all the lights
public GridSizingStrategy LightGridSizingStrategy = GridSizingStrategy.Uniform;
public RegirLightGrid(ComputeShader shader)
{
_shader = shader;
_buildRegirLightGridlKernel = _shader.FindKernel("BuildRegirLightGrid");
}
public void Init()
{
if (_lightGridCellsDataBuffer == null || _lightGridCellsDataBuffer.count <= 1)
{
int count = LightGridCellCount * MaxLightsPerCell;
int stride = Marshal.SizeOf<World.ThinReservoir>();
_lightGridCellsDataBuffer?.Dispose();
_lightGridCellsDataBuffer = new ComputeBuffer(count, stride);
}
if (_lightGridBuffer == null || _lightGridBuffer.count <= 1)
{
int count = LightGridCellCount;
int stride = Marshal.SizeOf<int2>();
_lightGridBuffer?.Dispose();
_lightGridBuffer = new ComputeBuffer(count, stride);
}
}
public void Build(CommandBuffer cmd, World.LightState lightState, Bounds sceneBounds, SamplingResources samplingResources)
{
if (lightState.LightListBuffer == null)
return;
Init();
_sceneBounds = sceneBounds;
// The number of RIS candidates cannot exceed the number of light sources
int activeCandidates = NumCandidates == -1 ? lightState.LightCount : Mathf.Min(NumCandidates, lightState.LightCount);
_lightGridDims = LightGridUtils.ComputeLightGridDims(sceneBounds.size, LightGridCellCount, LightGridSizingStrategy);
Vector3 div = new Vector3(1.0f / _lightGridDims.x, 1.0f / _lightGridDims.y, 1.0f / _lightGridDims.z);
Vector3 cellSize = Vector3.Scale(sceneBounds.size, div);
_cellSize = cellSize;
// The length of the diagonal
_cellSize.w = Mathf.Sqrt(cellSize.x * cellSize.x + cellSize.y * cellSize.y + cellSize.z * cellSize.z);
_invCellSize = new Vector4(1.0f / _cellSize.x, 1.0f / _cellSize.y, 1.0f / _cellSize.z, 1.0f / _cellSize.w);
SamplingResources.Bind(cmd, samplingResources);
// Set the input lighting state.Note that this is a subset, as we evaluate without light cookies
cmd.SetComputeIntParam(_shader, ShaderProperties.NumLights, lightState.LightCount);
cmd.SetComputeIntParam(_shader, ShaderProperties.NumCandidates, activeCandidates);
cmd.SetComputeIntParam(_shader, ShaderProperties.NumReservoirs, MaxLightsPerCell);
cmd.SetComputeIntParam(_shader, ShaderProperties.NumEmissiveMeshes, lightState.MeshLightCount);
cmd.SetComputeIntParam(_shader, ShaderProperties.GridDimX, _lightGridDims.x);
cmd.SetComputeIntParam(_shader, ShaderProperties.GridDimY, _lightGridDims.y);
cmd.SetComputeIntParam(_shader, ShaderProperties.GridDimZ, _lightGridDims.z);
cmd.SetComputeVectorParam(_shader, ShaderProperties.GridMin, sceneBounds.min);
cmd.SetComputeVectorParam(_shader, ShaderProperties.GridSize, sceneBounds.size);
cmd.SetComputeVectorParam(_shader, ShaderProperties.CellSize, _cellSize);
cmd.SetComputeVectorParam(_shader, ShaderProperties.InvCellSize, _invCellSize);
cmd.SetComputeBufferParam(_shader, _buildRegirLightGridlKernel, ShaderProperties.LightList, lightState.LightListBuffer);
// Set the output buffer
cmd.SetComputeBufferParam(_shader, _buildRegirLightGridlKernel, ShaderProperties.LightGrid, _lightGridBuffer);
cmd.SetComputeBufferParam(_shader, _buildRegirLightGridlKernel, ShaderProperties.LightGridCellsData, _lightGridCellsDataBuffer);
// Build the grid
const int groupDim = 4;
cmd.DispatchCompute(_shader, _buildRegirLightGridlKernel,
GraphicsHelpers.DivUp(_lightGridDims.x, groupDim),
GraphicsHelpers.DivUp(_lightGridDims.y, groupDim),
GraphicsHelpers.DivUp(_lightGridDims.z, groupDim));
}
public void Bind(CommandBuffer cmd, IRayTracingShader shader)
{
if (_lightGridCellsDataBuffer == null)
{
// dummy buffer, when the feature is disabled
int stride = Marshal.SizeOf<World.ThinReservoir>();
_lightGridCellsDataBuffer = new ComputeBuffer(1, stride);
}
if (_lightGridBuffer == null)
{
// dummy buffer, when the feature is disabled
_lightGridBuffer = new ComputeBuffer(1, Marshal.SizeOf<int2>());
}
shader.SetIntParam(cmd, ShaderProperties.GridDimX, _lightGridDims.x);
shader.SetIntParam(cmd, ShaderProperties.GridDimY, _lightGridDims.y);
shader.SetIntParam(cmd, ShaderProperties.GridDimZ, _lightGridDims.z);
shader.SetIntParam(cmd, ShaderProperties.NumReservoirs, MaxLightsPerCell);
shader.SetVectorParam(cmd, ShaderProperties.GridMin, _sceneBounds.min);
shader.SetVectorParam(cmd, ShaderProperties.GridSize, _sceneBounds.size);
shader.SetVectorParam(cmd, ShaderProperties.CellSize, _cellSize);
shader.SetVectorParam(cmd, ShaderProperties.InvCellSize, _invCellSize);
shader.SetBufferParam(cmd, ShaderProperties.LightGridCellsData, _lightGridCellsDataBuffer);
shader.SetBufferParam(cmd, ShaderProperties.LightGrid, _lightGridBuffer);
}
public void Dispose()
{
_lightGridCellsDataBuffer?.Dispose();
_lightGridBuffer?.Dispose();
}
readonly ComputeShader _shader;
readonly int _buildRegirLightGridlKernel;
ComputeBuffer _lightGridCellsDataBuffer;
ComputeBuffer _lightGridBuffer;
Bounds _sceneBounds;
Vector4 _cellSize;
Vector4 _invCellSize;
Vector3Int _lightGridDims;
}
}