Files
2026-03-05 00:14:42 +08:00

429 lines
19 KiB
C#

// Crest Water System
// Copyright © 2024 Wave Harmonic. All rights reserved.
// Inspired by https://github.com/speps/GX-EncinoWaves
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
namespace WaveHarmonic.Crest
{
/// <summary>
/// Runs FFT to generate water surface displacements
/// </summary>
sealed class FFTCompute
{
// Must match 'SIZE' param of first kernel in FFTCompute.compute
const int k_Kernel0Resolution = 8;
// Must match CASCADE_COUNT in FFTCompute.compute
const int k_CascadeCount = 16;
bool _Initialized = false;
RenderTexture _SpectrumInitial;
/// <summary>
/// Generated 'raw', uncombined, wave data. Input for putting into AnimWaves data before combine pass.
/// </summary>
public RenderTexture WaveBuffers { get; private set; }
bool _SpectrumInitialized = false;
ComputeShader _ShaderSpectrum;
ComputeShader _ShaderFFT;
int _KernelSpectrumInitial;
int _KernelSpectrumUpdate;
Parameters _Parameters;
float _GenerationTime = -1f;
static readonly bool s_SupportsRandomWriteRGFloat =
SystemInfo.SupportsRandomWriteOnRenderTextureFormat(RenderTextureFormat.RGFloat);
public static class ShaderIDs
{
public static readonly int s_Size = Shader.PropertyToID("_Crest_Size");
public static readonly int s_WindSpeed = Shader.PropertyToID("_Crest_WindSpeed");
public static readonly int s_Turbulence = Shader.PropertyToID("_Crest_Turbulence");
public static readonly int s_Alignment = Shader.PropertyToID("_Crest_Alignment");
public static readonly int s_Gravity = Shader.PropertyToID("_Crest_Gravity");
public static readonly int s_Period = Shader.PropertyToID("_Crest_Period");
public static readonly int s_WindDir = Shader.PropertyToID("_Crest_WindDir");
public static readonly int s_SpectrumControls = Shader.PropertyToID("_Crest_SpectrumControls");
public static readonly int s_ResultInit = Shader.PropertyToID("_Crest_ResultInit");
public static readonly int s_Time = Shader.PropertyToID("_Crest_Time");
public static readonly int s_Chop = Shader.PropertyToID("_Crest_Chop");
public static readonly int s_Init0 = Shader.PropertyToID("_Crest_Init0");
public static readonly int s_ResultHeight = Shader.PropertyToID("_Crest_ResultHeight");
public static readonly int s_ResultDisplaceX = Shader.PropertyToID("_Crest_ResultDisplaceX");
public static readonly int s_ResultDisplaceZ = Shader.PropertyToID("_Crest_ResultDisplaceZ");
public static readonly int s_InputH = Shader.PropertyToID("_Crest_InputH");
public static readonly int s_InputX = Shader.PropertyToID("_Crest_InputX");
public static readonly int s_InputZ = Shader.PropertyToID("_Crest_InputZ");
public static readonly int s_InputButterfly = Shader.PropertyToID("_Crest_InputButterfly");
public static readonly int s_Output1 = Shader.PropertyToID("_Crest_Output1");
public static readonly int s_Output2 = Shader.PropertyToID("_Crest_Output2");
public static readonly int s_Output3 = Shader.PropertyToID("_Crest_Output3");
public static readonly int s_Output = Shader.PropertyToID("_Crest_Output");
public static readonly int s_TemporaryFFT1 = Shader.PropertyToID("_Crest_TemporaryFFT1");
public static readonly int s_TemporaryFFT2 = Shader.PropertyToID("_Crest_TemporaryFFT2");
public static readonly int s_TemporaryFFT3 = Shader.PropertyToID("_Crest_TemporaryFFT3");
}
internal readonly struct Parameters
{
public readonly WaveSpectrum _Spectrum;
public readonly int _Resolution;
public readonly float _LoopPeriod;
public readonly float _WindSpeed;
public readonly float _WindDirectionRadians;
public readonly float _WindTurbulence;
public readonly float _WindAlignment;
public readonly float _Gravity;
public Parameters(WaveSpectrum spectrum, int resolution, float period, float speed, float direction, float turbulence, float alignment, float gravity)
{
_Spectrum = spectrum;
_Resolution = resolution;
_LoopPeriod = period;
_WindSpeed = speed;
_WindDirectionRadians = direction;
_WindTurbulence = turbulence;
_WindAlignment = alignment;
_Gravity = gravity;
}
// Implement custom or incur allocations.
public override int GetHashCode()
{
return System.HashCode.Combine(_Spectrum, _LoopPeriod, _WindSpeed, _WindDirectionRadians, _WindTurbulence, _WindAlignment, _Gravity, _Resolution);
}
public int GetHashCode(int resolution)
{
return System.HashCode.Combine(_Spectrum, _LoopPeriod, _WindSpeed, _WindDirectionRadians, _WindTurbulence, _WindAlignment, _Gravity, resolution);
}
}
public FFTCompute(Parameters parameters)
{
Debug.Assert(Mathf.NextPowerOfTwo(parameters._Resolution) == parameters._Resolution, "Crest: FFTCompute resolution must be power of 2");
_Parameters = parameters;
}
public void Release()
{
if (_SpectrumInitial != null)
{
_SpectrumInitial.Release();
}
if (WaveBuffers != null)
{
WaveBuffers.Release();
}
Helpers.Destroy(_SpectrumInitial);
Helpers.Destroy(WaveBuffers);
_SpectrumInitialized = false;
_Initialized = false;
}
internal static void CleanUpAll()
{
foreach (var generator in s_Generators)
{
generator.Value.Release();
}
s_Generators?.Clear();
foreach (var texture in s_ButterflyTextures?.Values)
{
Helpers.Destroy(texture);
}
s_ButterflyTextures?.Clear();
}
static readonly Dictionary<int, FFTCompute> s_Generators = new();
static readonly Dictionary<int, Texture2D> s_ButterflyTextures = new();
[RuntimeInitializeOnLoadMethod(RuntimeInitializeLoadType.SubsystemRegistration)]
static void InitStatics()
{
CleanUpAll();
}
/// <summary>
/// Computes water surface displacement, with wave components split across slices of the output texture array
/// </summary>
public static RenderTexture GenerateDisplacements(CommandBuffer buf, float time, Parameters parameters, bool updateSpectrum)
{
var conditionsHash = parameters.GetHashCode();
// All static data arguments should be hashed here and passed to the generator constructor
if (!s_Generators.TryGetValue(conditionsHash, out var generator))
{
// No generator for these params - create one
generator = new(parameters);
s_Generators.Add(conditionsHash, generator);
}
// The remaining dynamic data arguments should be passed in to the generation here
return generator.GenerateDisplacementsInternal(buf, time, updateSpectrum);
}
RenderTexture GenerateDisplacementsInternal(CommandBuffer buffer, float time, bool updateSpectrum)
{
// Check if already generated, and we're not being asked to re-update the spectrum
if (_GenerationTime == time && !updateSpectrum)
{
return WaveBuffers;
}
var resolution = _Parameters._Resolution;
var period = _Parameters._LoopPeriod;
// Initialize.
if (!_Initialized || _SpectrumInitial == null)
{
Release();
_ShaderSpectrum = WaterResources.Instance.Compute._FFTSpectrum;
_KernelSpectrumInitial = _ShaderSpectrum.FindKernel("SpectrumInitalize");
_KernelSpectrumUpdate = _ShaderSpectrum.FindKernel("SpectrumUpdate");
_ShaderFFT = WaterResources.Instance.Compute._FFT;
var rtd = new RenderTextureDescriptor(0, 0);
rtd.width = rtd.height = resolution;
rtd.dimension = TextureDimension.Tex2DArray;
rtd.enableRandomWrite = true;
rtd.depthBufferBits = 0;
rtd.volumeDepth = k_CascadeCount;
rtd.colorFormat = RenderTextureFormat.ARGBFloat;
rtd.msaaSamples = 1;
Helpers.SafeCreateRenderTexture(ref _SpectrumInitial, rtd);
_SpectrumInitial.name = "_Crest_FFTSpectrumInit";
_SpectrumInitial.Create();
// Raw wave data buffer
WaveBuffers = new(resolution, resolution, 0, GraphicsFormat.R16G16B16A16_SFloat)
{
wrapMode = TextureWrapMode.Repeat,
antiAliasing = 1,
filterMode = FilterMode.Bilinear,
anisoLevel = 0,
useMipMap = false,
name = "_Crest_FFTCascades",
dimension = TextureDimension.Tex2DArray,
volumeDepth = k_CascadeCount,
enableRandomWrite = true,
};
WaveBuffers.Create();
// Initialize bufferfly. Cached per resolution.
if (!s_ButterflyTextures.ContainsKey(resolution))
{
// Computes the offsets used for the FFT calculation.
var size = Mathf.RoundToInt(Mathf.Log(resolution, 2));
var colors = new Color[resolution * size];
int offset = 1, iterations = resolution >> 1;
for (var index = 0; index < size; index++)
{
var rowOffset = index * resolution;
{
int start = 0, end = 2 * offset;
for (var iteration = 0; iteration < iterations; iteration++)
{
var bigK = 0f;
for (var k = start; k < end; k += 2)
{
var phase = 2.0f * Mathf.PI * bigK * iterations / resolution;
var cos = Mathf.Cos(phase);
var sin = Mathf.Sin(phase);
colors[rowOffset + k / 2] = new(cos, -sin, 0, 1);
colors[rowOffset + k / 2 + offset] = new(-cos, sin, 0, 1);
bigK += 1f;
}
start += 4 * offset;
end = start + 2 * offset;
}
}
iterations >>= 1;
offset <<= 1;
}
var texture = new Texture2D(resolution, Mathf.RoundToInt(Mathf.Log(resolution, 2)), TextureFormat.RGBAFloat, false, true);
texture.SetPixels(colors);
texture.Apply();
s_ButterflyTextures.Add(resolution, texture);
}
_Initialized = true;
}
// Initialize spectrum.
// Computes base spectrum values based on wind speed and turbulence and spectrum controls.
if (!_SpectrumInitialized || updateSpectrum)
{
var wrapper = new PropertyWrapperCompute(buffer, _ShaderSpectrum, _KernelSpectrumInitial);
wrapper.SetInteger(ShaderIDs.s_Size, resolution);
wrapper.SetFloat(ShaderIDs.s_WindSpeed, _Parameters._WindSpeed);
wrapper.SetFloat(ShaderIDs.s_Turbulence, _Parameters._WindTurbulence);
wrapper.SetFloat(ShaderIDs.s_Alignment, _Parameters._WindAlignment);
wrapper.SetFloat(ShaderIDs.s_Gravity, _Parameters._Gravity);
wrapper.SetFloat(ShaderIDs.s_Period, period < Mathf.Infinity ? period : -1);
wrapper.SetVector(ShaderIDs.s_WindDir, new(Mathf.Cos(_Parameters._WindDirectionRadians), Mathf.Sin(_Parameters._WindDirectionRadians)));
wrapper.SetTexture(ShaderIDs.s_SpectrumControls, _Parameters._Spectrum.ControlsTexture);
wrapper.SetTexture(ShaderIDs.s_ResultInit, _SpectrumInitial);
wrapper.Dispatch(resolution / 8, resolution / 8, k_CascadeCount);
_SpectrumInitialized = true;
}
// Update Spectrum.
// Computes a spectrum for the current time which can be FFT'd into the final surface.
{
var wrapper = new PropertyWrapperCompute(buffer, _ShaderSpectrum, _KernelSpectrumUpdate);
var descriptor = _SpectrumInitial.descriptor;
if (s_SupportsRandomWriteRGFloat)
{
descriptor.colorFormat = RenderTextureFormat.RGFloat;
}
// No need to clear as overwritten.
buffer.GetTemporaryRT(ShaderIDs.s_TemporaryFFT1, descriptor);
buffer.GetTemporaryRT(ShaderIDs.s_TemporaryFFT2, descriptor);
buffer.GetTemporaryRT(ShaderIDs.s_TemporaryFFT3, descriptor);
wrapper.SetInteger(ShaderIDs.s_Size, resolution);
wrapper.SetFloat(ShaderIDs.s_Time, time * _Parameters._Spectrum._GravityScale);
wrapper.SetFloat(ShaderIDs.s_Chop, _Parameters._Spectrum._Chop);
wrapper.SetFloat(ShaderIDs.s_Period, period < Mathf.Infinity ? period : -1);
wrapper.SetTexture(ShaderIDs.s_Init0, _SpectrumInitial);
wrapper.SetTexture(ShaderIDs.s_ResultHeight, ShaderIDs.s_TemporaryFFT1);
wrapper.SetTexture(ShaderIDs.s_ResultDisplaceX, ShaderIDs.s_TemporaryFFT2);
wrapper.SetTexture(ShaderIDs.s_ResultDisplaceZ, ShaderIDs.s_TemporaryFFT3);
wrapper.Dispatch(resolution / 8, resolution / 8, k_CascadeCount);
}
// Dispatch FFT.
// FFT the spectrum into surface displacements.
{
var kernel = 2 * Mathf.RoundToInt(Mathf.Log(resolution / k_Kernel0Resolution, 2f));
var wrapper = new PropertyWrapperCompute(buffer, _ShaderFFT, kernel);
var butterfly = s_ButterflyTextures[resolution];
wrapper.SetTexture(ShaderIDs.s_InputButterfly, butterfly);
wrapper.SetTexture(ShaderIDs.s_Output1, ShaderIDs.s_TemporaryFFT1);
wrapper.SetTexture(ShaderIDs.s_Output2, ShaderIDs.s_TemporaryFFT2);
wrapper.SetTexture(ShaderIDs.s_Output3, ShaderIDs.s_TemporaryFFT3);
wrapper.Dispatch(1, resolution, k_CascadeCount);
wrapper = new PropertyWrapperCompute(buffer, _ShaderFFT, kernel + 1);
wrapper.SetTexture(ShaderIDs.s_InputH, ShaderIDs.s_TemporaryFFT1);
wrapper.SetTexture(ShaderIDs.s_InputX, ShaderIDs.s_TemporaryFFT2);
wrapper.SetTexture(ShaderIDs.s_InputZ, ShaderIDs.s_TemporaryFFT3);
wrapper.SetTexture(ShaderIDs.s_InputButterfly, butterfly);
wrapper.SetTexture(ShaderIDs.s_Output, WaveBuffers);
wrapper.Dispatch(resolution, 1, k_CascadeCount);
buffer.ReleaseTemporaryRT(ShaderIDs.s_TemporaryFFT1);
buffer.ReleaseTemporaryRT(ShaderIDs.s_TemporaryFFT2);
buffer.ReleaseTemporaryRT(ShaderIDs.s_TemporaryFFT3);
}
_GenerationTime = time;
return WaveBuffers;
}
/// <summary>
/// Changing wave gen data can result in creating lots of new generators. This gives a way to notify
/// that a parameter has changed. If there is no existing generator for the new param values, but there
/// is one for the old param values, this old generator is repurposed.
/// </summary>
public static void OnGenerationDataUpdated(Parameters oldParameters, Parameters newParameters)
{
// If multiple wave components share one FFT, then one of them changes its settings, it will
// actually steal the generator from the rest. Then the first from the rest which request the
// old settings will trigger creation of a new generator, and the remaining ones will use this
// new generator. In the end one new generator is created, but it's created for the old settings.
// Generators are requested single threaded so there should not be a race condition. Odd pattern
// but I don't think any other way works without ugly checks to see if old generators are still
// used, or other complicated things.
// Check if no generator exists for new values
var newHash = newParameters.GetHashCode();
if (!s_Generators.TryGetValue(newHash, out var oldGenerator))
{
// Try to adapt an existing generator rather than default to creating a new one
// Adapting requires the resolution to the same.
var oldHash = oldParameters.GetHashCode(newParameters._Resolution);
if (s_Generators.TryGetValue(oldHash, out var generator))
{
// Hash will change for this generator, so remove the current one
s_Generators.Remove(oldHash);
// Update params
generator._Parameters = newParameters;
// Trigger generator to re-init the spectrum
generator._SpectrumInitialized = false;
// Re-add with new hash
s_Generators.Add(newHash, generator);
}
}
else
{
// There is already a new generator which will be used. Remove the previous one - if it really is needed
// then it will be created later.
oldGenerator.Release();
s_Generators.Remove(oldParameters.GetHashCode());
}
}
/// <summary>
/// Number of FFT generators
/// </summary>
public static int GeneratorCount => s_Generators != null ? s_Generators.Count : 0;
public static FFTCompute GetInstance(Parameters parameters)
{
return s_Generators.GetValueOrDefault(parameters.GetHashCode(), null);
}
public bool HasData()
{
return WaveBuffers != null && WaveBuffers.IsCreated();
}
internal void OnGUI()
{
if (WaveBuffers != null && WaveBuffers.IsCreated())
{
DebugGUI.DrawTextureArray(WaveBuffers, 8, 0.5f, 20f);
}
if (_Parameters._Spectrum != null)
{
_Parameters._Spectrum.OnGUI();
}
}
}
}