SHA256
345 lines
12 KiB
C#
345 lines
12 KiB
C#
using System;
|
|
using System.Buffers;
|
|
using System.Collections.Generic;
|
|
using System.Linq;
|
|
using ADOFAIRenderer.Core;
|
|
using ADOFAIRenderer.Util.FFmpeg;
|
|
using FFmpeg.AutoGen;
|
|
using MelonLoader;
|
|
using Unity.Collections;
|
|
using Unity.Collections.LowLevel.Unsafe;
|
|
using UnityEngine;
|
|
using UnityEngine.Rendering;
|
|
|
|
namespace ADOFAIRenderer.Util
|
|
{
|
|
public class FrameReader : IDisposable
|
|
{
|
|
private static readonly int Height = Shader.PropertyToID("Height");
|
|
private static readonly int Width = Shader.PropertyToID("Width");
|
|
private static readonly int DestinationBuffer = Shader.PropertyToID("DestinationBuffer");
|
|
private static readonly int SourceTexture = Shader.PropertyToID("SourceTexture");
|
|
private const int InitialBufferSize = 16;
|
|
|
|
private readonly Stack<ReadbackRequestState> _statePool = new();
|
|
private readonly List<ReadbackRequestState> _allocatedStates = new();
|
|
private readonly object _poolLock = new();
|
|
private bool _isInitialized;
|
|
private volatile bool _isDisposed;
|
|
|
|
public FrameReader()
|
|
{
|
|
InitializePool();
|
|
}
|
|
|
|
private void InitializePool()
|
|
{
|
|
lock (_poolLock)
|
|
{
|
|
if (_isInitialized) return;
|
|
for (var i = 0; i < InitialBufferSize; i++)
|
|
{
|
|
var state = new ReadbackRequestState(this);
|
|
_statePool.Push(state);
|
|
_allocatedStates.Add(state);
|
|
}
|
|
|
|
_isInitialized = true;
|
|
}
|
|
}
|
|
|
|
private static unsafe AVPixelFormat PredictPixelFormat(Codec codec)
|
|
{
|
|
var candidates = GPUUtil.GetPrioritizedEncoderCandidates(codec);
|
|
// ReSharper disable once ForeachCanBePartlyConvertedToQueryUsingAnotherGetEnumerator
|
|
foreach (var encoderName in candidates)
|
|
{
|
|
var pCodec = ffmpeg.avcodec_find_encoder_by_name(encoderName);
|
|
if (pCodec == null) continue;
|
|
|
|
if (encoderName.Contains("qsv") || encoderName.Contains("nvenc") || encoderName.Contains("amf"))
|
|
{
|
|
return AVPixelFormat.AV_PIX_FMT_NV12;
|
|
}
|
|
return encoderName == "libvvenc" ? AVPixelFormat.AV_PIX_FMT_YUV420P10LE : AVPixelFormat.AV_PIX_FMT_YUV420P;
|
|
}
|
|
return AVPixelFormat.AV_PIX_FMT_YUV420P;
|
|
}
|
|
|
|
public void RequestFrameCapture(Texture texture, long currentFrame)
|
|
{
|
|
if (!RenderManager.IsCapturingActive) return;
|
|
ReadbackRequestState state;
|
|
|
|
lock (_poolLock)
|
|
{
|
|
if (_statePool.Count == 0)
|
|
{
|
|
state = new ReadbackRequestState(this);
|
|
_allocatedStates.Add(state);
|
|
}
|
|
else
|
|
{
|
|
state = _statePool.Pop();
|
|
}
|
|
|
|
state.IsInFlight = true;
|
|
}
|
|
|
|
// GPU 컬러 변환 활성화 여부 검사
|
|
bool useGPUConversion = RenderMod.ColorConversionShader;
|
|
AVPixelFormat expectedFmt = AVPixelFormat.AV_PIX_FMT_RGBA;
|
|
|
|
if (useGPUConversion && RenderManager.Instance)
|
|
{
|
|
expectedFmt = PredictPixelFormat(RenderManager.Instance.Configuration.Codec);
|
|
// Compute Shader는 NV12와 YUV420P만 지원하므로 그 외 포맷은 CPU 폴백 처리
|
|
if (expectedFmt != AVPixelFormat.AV_PIX_FMT_NV12 && expectedFmt != AVPixelFormat.AV_PIX_FMT_YUV420P)
|
|
{
|
|
useGPUConversion = false;
|
|
}
|
|
}
|
|
|
|
var bufferSize = texture.width * texture.height * 4; // RGBA = 4 channel
|
|
state.Setup(currentFrame, bufferSize, useGPUConversion, expectedFmt, texture.width, texture.height);
|
|
|
|
if (useGPUConversion && state.GPUBuffer != null && RenderMod.ColorConversionShader)
|
|
{
|
|
int kernelIndex = -1;
|
|
if (expectedFmt == AVPixelFormat.AV_PIX_FMT_NV12)
|
|
{
|
|
kernelIndex = RenderMod.ColorConversionShader.FindKernel("CSMain_NV12");
|
|
}
|
|
else if (expectedFmt == AVPixelFormat.AV_PIX_FMT_YUV420P)
|
|
{
|
|
kernelIndex = RenderMod.ColorConversionShader.FindKernel("CSMain_YUV420P");
|
|
}
|
|
|
|
if (kernelIndex >= 0)
|
|
{
|
|
RenderMod.ColorConversionShader.SetTexture(kernelIndex, SourceTexture, texture);
|
|
RenderMod.ColorConversionShader.SetBuffer(kernelIndex, DestinationBuffer, state.GPUBuffer);
|
|
RenderMod.ColorConversionShader.SetInt(Width, texture.width);
|
|
RenderMod.ColorConversionShader.SetInt(Height, texture.height);
|
|
|
|
var threadGroupsX = expectedFmt == AVPixelFormat.AV_PIX_FMT_NV12
|
|
? (texture.width + 31) / 32
|
|
: (texture.width + 63) / 64;
|
|
var threadGroupsY = (texture.height + 15) / 16;
|
|
|
|
RenderMod.ColorConversionShader.Dispatch(kernelIndex, threadGroupsX, threadGroupsY, 1);
|
|
|
|
AsyncGPUReadback.Request(
|
|
state.GPUBuffer,
|
|
state.CachedCallback
|
|
);
|
|
}
|
|
else
|
|
{
|
|
// 커널을 못 찾은 경우 기존 RGBA 텍스처 백업 기법으로 폴백
|
|
AsyncGPUReadback.RequestIntoNativeArray(
|
|
ref state.NativeBuffer,
|
|
texture,
|
|
0,
|
|
state.CachedCallback
|
|
);
|
|
}
|
|
}
|
|
else
|
|
{
|
|
AsyncGPUReadback.RequestIntoNativeArray(
|
|
ref state.NativeBuffer,
|
|
texture,
|
|
0,
|
|
state.CachedCallback
|
|
);
|
|
}
|
|
}
|
|
|
|
internal void OnReadbackComplete(ReadbackRequestState state, AsyncGPUReadbackRequest request)
|
|
{
|
|
var frame = state.FrameIndex;
|
|
|
|
lock (_poolLock)
|
|
{
|
|
state.IsInFlight = false;
|
|
if (_isDisposed)
|
|
{
|
|
state.Dispose();
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (RenderManager.Instance && RenderManager.Instance.IsRenderCancelled)
|
|
{
|
|
RecycleState(state);
|
|
return;
|
|
}
|
|
|
|
if (request.hasError)
|
|
{
|
|
if (RenderManager.Instance)
|
|
{
|
|
RenderManager.Instance.readbackFailureCount++;
|
|
}
|
|
Melon<RenderMod>.Logger.Warning($"[OOM Prevent] GPU Readback failed for frame {frame}. Queue deadlock prevented by sending a black frame.");
|
|
|
|
var byteSize = state.IsGPUConversion
|
|
? state.Width * state.Height * 3 / 2
|
|
: state.Width * state.Height * 4;
|
|
var blackFrame = ArrayPool<byte>.Shared.Rent(byteSize);
|
|
Array.Clear(blackFrame, 0, byteSize);
|
|
|
|
var capFrame = new RenderManager.CapturedFrame
|
|
{
|
|
ManagedBuffer = blackFrame,
|
|
Length = byteSize
|
|
};
|
|
PushToQueue(frame, capFrame);
|
|
RecycleState(state);
|
|
}
|
|
else
|
|
{
|
|
unsafe
|
|
{
|
|
if (state.IsGPUConversion)
|
|
{
|
|
var data = request.GetData<byte>();
|
|
// 비동기 요청 버퍼 데이터를 지속성(NativeBuffer) 메모리로 복사
|
|
NativeArray<byte>.Copy(data, state.NativeBuffer, data.Length);
|
|
|
|
var ptr = (IntPtr)state.NativeBuffer.GetUnsafeReadOnlyPtr();
|
|
var capFrame = new RenderManager.CapturedFrame
|
|
{
|
|
NativePointer = ptr,
|
|
Length = state.NativeBuffer.Length,
|
|
ReadbackState = state
|
|
};
|
|
PushToQueue(frame, capFrame);
|
|
}
|
|
else
|
|
{
|
|
var ptr = (IntPtr)state.NativeBuffer.GetUnsafeReadOnlyPtr();
|
|
var capFrame = new RenderManager.CapturedFrame
|
|
{
|
|
NativePointer = ptr,
|
|
Length = state.NativeBuffer.Length,
|
|
ReadbackState = state
|
|
};
|
|
PushToQueue(frame, capFrame);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
internal void RecycleState(ReadbackRequestState state)
|
|
{
|
|
lock (_poolLock)
|
|
{
|
|
if (_isDisposed || !_isInitialized)
|
|
{
|
|
state.Dispose();
|
|
return;
|
|
}
|
|
|
|
state.Recycle();
|
|
_statePool.Push(state);
|
|
}
|
|
}
|
|
|
|
private static void PushToQueue(long frame, RenderManager.CapturedFrame capturedFrame)
|
|
{
|
|
if (!RenderManager.Instance) return;
|
|
lock (RenderManager.Instance.ImageQueue)
|
|
{
|
|
RenderManager.Instance.ImageQueue[frame] = capturedFrame;
|
|
}
|
|
RenderManager.Instance.FrameAvailableEvent.Set();
|
|
}
|
|
|
|
public void Dispose()
|
|
{
|
|
lock (_poolLock)
|
|
{
|
|
if (_isDisposed) return;
|
|
_isDisposed = true;
|
|
foreach (var state in _allocatedStates.Where(state => !state.IsInFlight))
|
|
{
|
|
state.Dispose();
|
|
}
|
|
_statePool.Clear();
|
|
_allocatedStates.Clear();
|
|
_isInitialized = false;
|
|
}
|
|
}
|
|
}
|
|
|
|
internal class ReadbackRequestState : IDisposable
|
|
{
|
|
private readonly FrameReader _owner;
|
|
public long FrameIndex { get; private set; }
|
|
public NativeArray<byte> NativeBuffer;
|
|
public bool IsInFlight { get; internal set; }
|
|
|
|
public ComputeBuffer? GPUBuffer;
|
|
public bool IsGPUConversion { get; private set; }
|
|
public int Width { get; private set; }
|
|
public int Height { get; private set; }
|
|
|
|
public readonly Action<AsyncGPUReadbackRequest> CachedCallback;
|
|
|
|
public ReadbackRequestState(FrameReader owner)
|
|
{
|
|
_owner = owner;
|
|
CachedCallback = OnComplete;
|
|
FrameIndex = -1;
|
|
}
|
|
|
|
public void Setup(long frameIndex, int bufferSize, bool isGPUConversion, AVPixelFormat pixelFormat, int width, int height)
|
|
{
|
|
FrameIndex = frameIndex;
|
|
IsGPUConversion = isGPUConversion;
|
|
Width = width; Height = height;
|
|
|
|
var readbackSize = isGPUConversion ? (width * height * 3 / 2) : bufferSize;
|
|
if (!NativeBuffer.IsCreated || NativeBuffer.Length != readbackSize)
|
|
{
|
|
if (NativeBuffer.IsCreated) NativeBuffer.Dispose();
|
|
|
|
NativeBuffer = new NativeArray<byte>(readbackSize, Allocator.Persistent,
|
|
NativeArrayOptions.UninitializedMemory);
|
|
}
|
|
|
|
if (isGPUConversion)
|
|
{
|
|
var computeBufferSize = readbackSize / 4;
|
|
if (GPUBuffer != null && GPUBuffer.count == computeBufferSize) return;
|
|
|
|
GPUBuffer?.Dispose();
|
|
GPUBuffer = new ComputeBuffer(computeBufferSize, 4, ComputeBufferType.Raw);
|
|
}
|
|
else
|
|
{
|
|
GPUBuffer?.Dispose();
|
|
GPUBuffer = null;
|
|
}
|
|
}
|
|
|
|
private void OnComplete(AsyncGPUReadbackRequest request)
|
|
{
|
|
_owner.OnReadbackComplete(this, request);
|
|
}
|
|
|
|
public void Recycle()
|
|
{
|
|
FrameIndex = -1;
|
|
}
|
|
|
|
public void Dispose()
|
|
{
|
|
if (NativeBuffer.IsCreated) NativeBuffer.Dispose();
|
|
GPUBuffer?.Dispose();
|
|
GPUBuffer = null;
|
|
}
|
|
}
|
|
} |