using System;
using System.IO;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEditor;
using UnityEngine.LightTransport;
using UnityEngine.PathTracing.Core;
using UnityEngine.Rendering;
using UnityEngine.Rendering.UnifiedRayTracing;
namespace UnityEngine.PathTracing.Integration
{
internal class UnityComputeProbeIntegrator : IProbeIntegrator
{
private readonly ProbeIntegrator _probeIntegrator;
private UnityComputeWorld _world;
private uint _bounceCount;
private uint _directLightingEvaluationCount;
private uint _numIndirectLightingEvaluations;
private uint _basePositionsOffset;
private static class ShaderProperties
{
public static readonly int MappingTable = Shader.PropertyToID("g_MappingTable");
public static readonly int PerProbeLightIndicesInput = Shader.PropertyToID("g_PerProbeLightIndicesInput");
public static readonly int PerProbeLightIndicesOutput = Shader.PropertyToID("g_PerProbeLightIndicesOutput");
public static readonly int PerProbeLightIndicesInputOffset = Shader.PropertyToID("g_PerProbeLightIndicesInputOffset");
public static readonly int MaxLightsPerProbe = Shader.PropertyToID("g_MaxLightsPerProbe");
public static readonly int ProbeCount = Shader.PropertyToID("g_ProbeCount");
}
private ComputeShader _probeOcclusionLightIndexMappingShader;
private int _probeOcclusionLightIndexMappingKernel;
private Rendering.Sampling.SamplingResources _samplingResources;
private ProbeIntegratorResources _integrationResources;
public UnityComputeProbeIntegrator(bool countNEERayAsPathSegment, Rendering.Sampling.SamplingResources samplingResources, ProbeIntegratorResources integrationResources, ComputeShader probeOcclusionLightIndexMappingShader)
{
_probeIntegrator = new ProbeIntegrator(countNEERayAsPathSegment);
_samplingResources = samplingResources;
_probeOcclusionLightIndexMappingShader = probeOcclusionLightIndexMappingShader;
_probeOcclusionLightIndexMappingKernel = _probeOcclusionLightIndexMappingShader.FindKernel("MapIndices");
_integrationResources = integrationResources;
}
public void Dispose()
{
_probeIntegrator.Dispose();
}
public IProbeIntegrator.Result IntegrateDirectRadiance(IDeviceContext context, int positionOffset, int positionCount, int sampleCount,
bool ignoreEnvironment, BufferSlice<SphericalHarmonicsL2> radianceEstimateOut)
{
UnityComputeDeviceContext unifiedContext = context as UnityComputeDeviceContext;
Debug.Assert(unifiedContext != null);
Debug.Assert(sampleCount > 0);
ProbeIntegrator.GetRadianceScratchBufferSizesInDwords((uint)positionCount, (uint)sampleCount, out uint expansionBufferSize, out uint reductionBufferSize);
var expansionBuffer = unifiedContext.GetTemporaryBuffer(expansionBufferSize, sizeof(float));
var reductionBuffer = unifiedContext.GetTemporaryBuffer(reductionBufferSize, sizeof(float));
uint sampleOffset = 0;
_probeIntegrator.EstimateDirectRadianceShl2(
unifiedContext.GetCommandBuffer(),
_world.PathTracingWorld,
_basePositionsOffset + (uint)positionOffset,
(uint)positionCount,
sampleOffset,
(uint)sampleCount,
_directLightingEvaluationCount,
ignoreEnvironment,
unifiedContext.GetComputeBuffer(radianceEstimateOut.Id),
(uint)radianceEstimateOut.Offset,
unifiedContext.GetComputeBuffer(expansionBuffer),
unifiedContext.GetComputeBuffer(reductionBuffer));
return new IProbeIntegrator.Result(IProbeIntegrator.ResultType.Success, string.Empty);
}
public IProbeIntegrator.Result IntegrateIndirectRadiance(IDeviceContext context, int positionOffset, int positionCount, int sampleCount,
bool ignoreEnvironment, BufferSlice<SphericalHarmonicsL2> radianceEstimateOut)
{
var unifiedContext = context as UnityComputeDeviceContext;
Debug.Assert(unifiedContext != null);
Debug.Assert(sampleCount > 0);
ProbeIntegrator.GetRadianceScratchBufferSizesInDwords((uint)positionCount, (uint)sampleCount, out uint expansionBufferSize, out uint reductionBufferSize);
var expansionBuffer = unifiedContext.GetTemporaryBuffer(expansionBufferSize, sizeof(float));
var reductionBuffer = unifiedContext.GetTemporaryBuffer(reductionBufferSize, sizeof(float));
uint sampleOffset = 0;
_probeIntegrator.EstimateIndirectRadianceShl2(
unifiedContext.GetCommandBuffer(),
_world.PathTracingWorld,
_basePositionsOffset + (uint)positionOffset,
(uint)positionCount,
_bounceCount,
sampleOffset,
(uint)sampleCount,
_numIndirectLightingEvaluations,
ignoreEnvironment,
unifiedContext.GetComputeBuffer(radianceEstimateOut.Id),
(uint)radianceEstimateOut.Offset,
unifiedContext.GetComputeBuffer(expansionBuffer),
unifiedContext.GetComputeBuffer(reductionBuffer));
return new IProbeIntegrator.Result(IProbeIntegrator.ResultType.Success, string.Empty);
}
public IProbeIntegrator.Result IntegrateValidity(IDeviceContext context, int positionOffset, int positionCount, int sampleCount, BufferSlice<float> validityEstimateOut)
{
UnityComputeDeviceContext unifiedContext = context as UnityComputeDeviceContext;
Debug.Assert(unifiedContext != null);
Debug.Assert(sampleCount > 0);
ProbeIntegrator.GetValidityScratchBufferSizesInDwords((uint)positionCount, (uint)sampleCount, out uint expansionBufferSize, out uint reductionBufferSize);
var expansionBuffer = unifiedContext.GetTemporaryBuffer(expansionBufferSize, sizeof(float));
var reductionBuffer = unifiedContext.GetTemporaryBuffer(reductionBufferSize, sizeof(float));
uint sampleOffset = 0;
_probeIntegrator.EstimateValidity(
unifiedContext.GetCommandBuffer(),
_world.PathTracingWorld,
_basePositionsOffset + (uint)positionOffset,
(uint)positionCount,
sampleOffset,
(uint)sampleCount,
unifiedContext.GetComputeBuffer(validityEstimateOut.Id),
(uint)validityEstimateOut.Offset,
unifiedContext.GetComputeBuffer(expansionBuffer),
unifiedContext.GetComputeBuffer(reductionBuffer));
return new IProbeIntegrator.Result(IProbeIntegrator.ResultType.Success, string.Empty);
}
public IProbeIntegrator.Result IntegrateOcclusion(IDeviceContext context, int positionOffset, int positionCount, int sampleCount,
int maxLightsPerProbe, BufferSlice<int> perProbeLightIndices, BufferSlice<float> probeOcclusionEstimateOut)
{
UnityComputeDeviceContext unifiedContext = context as UnityComputeDeviceContext;
Debug.Assert(unifiedContext != null);
Debug.Assert(maxLightsPerProbe > 0);
Debug.Assert(sampleCount > 0);
var cmd = unifiedContext.GetCommandBuffer();
// The input per-probe light indices refer to elements of the light list used by LightBaker (ie. BakeInput.lightData).
// This is by necessity, since that is the contract of the IProbeIntegrator interface, and the ordering of lights may be considered 'global'.
// However, ProbeIntegrator needs per-probe light indices that refer to elements of the light list used by the path tracer
// (ie. World.LightList) in order to access the light data in shader. We therefore need to convert the input indices.
int[] lightIndexMapping = new int[_world.LightHandles.Length];
for (int lightIndex = 0; lightIndex < lightIndexMapping.Length; lightIndex++)
{
var lightHandle = _world.LightHandles[lightIndex];
int worldLightIndex = _world.PathTracingWorld.LightHandleToLightListIndex[lightHandle];
lightIndexMapping[lightIndex] = worldLightIndex;
}
if (lightIndexMapping.Length == 0) // Avoid 0-sized buffer in case of no lights in scene
lightIndexMapping = new int[] { -1 };
using NativeArray<int> lightIndexMappingArray = new(lightIndexMapping, Allocator.Temp);
var lightIndexMappingBuffer = unifiedContext.GetTemporaryBuffer((ulong)lightIndexMapping.Length, sizeof(int));
unifiedContext.WriteBuffer(lightIndexMappingBuffer.Slice<int>(), lightIndexMappingArray);
var perProbeLightIndicesWorld = unifiedContext.GetTemporaryBuffer((ulong)(positionCount * maxLightsPerProbe), sizeof(int));
cmd.SetComputeBufferParam(_probeOcclusionLightIndexMappingShader, _probeOcclusionLightIndexMappingKernel, ShaderProperties.MappingTable, unifiedContext.GetComputeBuffer(lightIndexMappingBuffer));
cmd.SetComputeBufferParam(_probeOcclusionLightIndexMappingShader, _probeOcclusionLightIndexMappingKernel, ShaderProperties.PerProbeLightIndicesInput, unifiedContext.GetComputeBuffer(perProbeLightIndices.Id));
cmd.SetComputeBufferParam(_probeOcclusionLightIndexMappingShader, _probeOcclusionLightIndexMappingKernel, ShaderProperties.PerProbeLightIndicesOutput, unifiedContext.GetComputeBuffer(perProbeLightIndicesWorld));
cmd.SetComputeIntParam(_probeOcclusionLightIndexMappingShader, ShaderProperties.PerProbeLightIndicesInputOffset, (int)perProbeLightIndices.Offset);
cmd.SetComputeIntParam(_probeOcclusionLightIndexMappingShader, ShaderProperties.MaxLightsPerProbe, maxLightsPerProbe);
cmd.SetComputeIntParam(_probeOcclusionLightIndexMappingShader, ShaderProperties.ProbeCount, positionCount);
_probeOcclusionLightIndexMappingShader.GetKernelThreadGroupSizes(_probeOcclusionLightIndexMappingKernel, out uint threadGroupSizeX, out _, out _);
cmd.DispatchCompute(_probeOcclusionLightIndexMappingShader, _probeOcclusionLightIndexMappingKernel, GraphicsHelpers.DivUp(positionCount, threadGroupSizeX), 1, 1);
ProbeIntegrator.GetOcclusionScratchBufferSizesInDwords((uint)maxLightsPerProbe, (uint)positionCount, (uint)sampleCount, out uint expansionBufferSize, out uint reductionBufferSize);
var expansionBuffer = unifiedContext.GetTemporaryBuffer(expansionBufferSize, sizeof(float));
var reductionBuffer = unifiedContext.GetTemporaryBuffer(reductionBufferSize, sizeof(float));
uint sampleOffset = 0;
_probeIntegrator.EstimateLightOcclusion(
cmd,
_world.PathTracingWorld,
_basePositionsOffset + (uint)positionOffset,
(uint)positionCount,
sampleOffset,
(uint)sampleCount,
(uint)maxLightsPerProbe,
unifiedContext.GetComputeBuffer(perProbeLightIndicesWorld),
0u,
unifiedContext.GetComputeBuffer(probeOcclusionEstimateOut.Id),
(uint)probeOcclusionEstimateOut.Offset,
unifiedContext.GetComputeBuffer(expansionBuffer),
unifiedContext.GetComputeBuffer(reductionBuffer));
return new IProbeIntegrator.Result(IProbeIntegrator.ResultType.Success, string.Empty);
}
public void Prepare(IDeviceContext context, IWorld world, BufferSlice<Vector3> positions, float pushoff, int bounceCount)
{
_bounceCount = (uint)bounceCount;
_directLightingEvaluationCount = 4;
_numIndirectLightingEvaluations = 1;
_world = world as UnityComputeWorld;
Debug.Assert(world != null);
UnityComputeDeviceContext unifiedContext = context as UnityComputeDeviceContext;
Debug.Assert(unifiedContext != null);
_basePositionsOffset = (uint)positions.Offset;
Debug.Assert(_world != null, nameof(_world) + " != null");
_probeIntegrator.Prepare(unifiedContext.GetComputeBuffer(positions.Id), _integrationResources, _samplingResources);
}
public void SetProgressReporter(BakeProgressState progressState) => _probeIntegrator.SetProgressReporter(progressState);
}
}