4 files changed, 655 insertions, 0 deletions
diff --git a/Ryujinx.Graphics.Vulkan/Queries/BufferedQuery.cs b/Ryujinx.Graphics.Vulkan/Queries/BufferedQuery.cs
new file mode 100644
index 00000000..a53e02a1
--- /dev/null
+++ b/Ryujinx.Graphics.Vulkan/Queries/BufferedQuery.cs
@@ -0,0 +1,206 @@
+using Ryujinx.Common.Logging;
+using Ryujinx.Graphics.GAL;
+using Silk.NET.Vulkan;
+using System;
+using System.Runtime.InteropServices;
+using System.Threading;
+
+namespace Ryujinx.Graphics.Vulkan.Queries
+{
+    class BufferedQuery : IDisposable
+    {
+        private const int MaxQueryRetries = 5000;
+        private const long DefaultValue = -1;
+        private const long DefaultValueInt = 0xFFFFFFFF;
+
+        private readonly Vk _api;
+        private readonly Device _device;
+        private readonly PipelineFull _pipeline;
+
+        private QueryPool _queryPool;
+        private bool _isReset;
+
+        private readonly BufferHolder _buffer;
+        private readonly IntPtr _bufferMap;
+        private readonly CounterType _type;
+        private bool _result32Bit;
+        private bool _isSupported;
+
+        private long _defaultValue;
+
+        public unsafe BufferedQuery(VulkanRenderer gd, Device device, PipelineFull pipeline, CounterType type, bool result32Bit)
+        {
+            _api = gd.Api;
+            _device = device;
+            _pipeline = pipeline;
+            _type = type;
+            _result32Bit = result32Bit;
+
+            _isSupported = QueryTypeSupported(gd, type);
+
+            if (_isSupported)
+            {
+                QueryPipelineStatisticFlags flags = type == CounterType.PrimitivesGenerated ?
+                    QueryPipelineStatisticFlags.QueryPipelineStatisticGeometryShaderPrimitivesBit : 0;
+
+                var queryPoolCreateInfo = new QueryPoolCreateInfo()
+                {
+                    SType = StructureType.QueryPoolCreateInfo,
+                    QueryCount = 1,
+                    QueryType = GetQueryType(type),
+                    PipelineStatistics = flags
+                };
+
+                gd.Api.CreateQueryPool(device, queryPoolCreateInfo, null, out _queryPool).ThrowOnError();
+            }
+
+            var buffer = gd.BufferManager.Create(gd, sizeof(long), forConditionalRendering: true);
+
+            _bufferMap = buffer.Map(0, sizeof(long));
+            _defaultValue = result32Bit ? DefaultValueInt : DefaultValue;
+            Marshal.WriteInt64(_bufferMap, _defaultValue);
+            _buffer = buffer;
+        }
+
+        private bool QueryTypeSupported(VulkanRenderer gd, CounterType type)
+        {
+            return type switch
+            {
+                CounterType.SamplesPassed => true,
+                CounterType.TransformFeedbackPrimitivesWritten => gd.Capabilities.SupportsTransformFeedbackQueries,
+                CounterType.PrimitivesGenerated => gd.Capabilities.SupportsGeometryShader,
+                _ => false
+            };
+        }
+
+        private static QueryType GetQueryType(CounterType type)
+        {
+            return type switch
+            {
+                CounterType.SamplesPassed => QueryType.Occlusion,
+                CounterType.PrimitivesGenerated => QueryType.PipelineStatistics,
+                CounterType.TransformFeedbackPrimitivesWritten => QueryType.TransformFeedbackStreamExt,
+                _ => QueryType.Occlusion
+            };
+        }
+
+        public Auto<DisposableBuffer> GetBuffer()
+        {
+            return _buffer.GetBuffer();
+        }
+
+        public void Reset()
+        {
+            End(false);
+            Begin();
+        }
+
+        public void Begin()
+        {
+            if (_isSupported)
+            {
+                _pipeline.BeginQuery(this, _queryPool, !_isReset);
+            }
+            _isReset = false;
+        }
+
+        public unsafe void End(bool withResult)
+        {
+            if (_isSupported)
+            {
+                _pipeline.EndQuery(_queryPool);
+            }
+
+            if (withResult && _isSupported)
+            {
+                Marshal.WriteInt64(_bufferMap, _defaultValue);
+                _pipeline.CopyQueryResults(this);
+            }
+            else
+            {
+                // Dummy result, just return 0.
+                Marshal.WriteInt64(_bufferMap, 0);
+            }
+        }
+
+        public bool TryGetResult(out long result)
+        {
+            result = Marshal.ReadInt64(_bufferMap);
+
+            return result != _defaultValue;
+        }
+
+        public long AwaitResult(AutoResetEvent wakeSignal = null)
+        {
+            long data = _defaultValue;
+
+            if (wakeSignal == null)
+            {
+                while (data == _defaultValue)
+                {
+                    data = Marshal.ReadInt64(_bufferMap);
+                }
+            }
+            else
+            {
+                int iterations = 0;
+                while (data == _defaultValue && iterations++ < MaxQueryRetries)
+                {
+                    data = Marshal.ReadInt64(_bufferMap);
+                    if (data == _defaultValue)
+                    {
+                        wakeSignal.WaitOne(1);
+                    }
+                }
+
+                if (iterations >= MaxQueryRetries)
+                {
+                    Logger.Error?.Print(LogClass.Gpu, $"Error: Query result {_type} timed out. Took more than {MaxQueryRetries} tries.");
+                }
+            }
+
+            return data;
+        }
+
+        public void PoolReset(CommandBuffer cmd)
+        {
+            if (_isSupported)
+            {
+                _api.CmdResetQueryPool(cmd, _queryPool, 0, 1);
+            }
+            _isReset = true;
+        }
+
+        public void PoolCopy(CommandBufferScoped cbs)
+        {
+            var buffer = _buffer.GetBuffer(cbs.CommandBuffer, true).Get(cbs, 0, sizeof(long)).Value;
+
+            QueryResultFlags flags = QueryResultFlags.QueryResultWaitBit;
+
+            if (!_result32Bit)
+            {
+                flags |= QueryResultFlags.QueryResult64Bit;
+            }
+
+            _api.CmdCopyQueryPoolResults(
+                cbs.CommandBuffer,
+                _queryPool,
+                0,
+                1,
+                buffer,
+                0,
+                (ulong)(_result32Bit ? sizeof(int) : sizeof(long)),
+                flags);
+        }
+
+        public unsafe void Dispose()
+        {
+            _buffer.Dispose();
+            if (_isSupported)
+            {
+                _api.DestroyQueryPool(_device, _queryPool, null);
+            }
+            _queryPool = default;
+        }
+    }
+}
diff --git a/Ryujinx.Graphics.Vulkan/Queries/CounterQueue.cs b/Ryujinx.Graphics.Vulkan/Queries/CounterQueue.cs
new file mode 100644
index 00000000..7ee3c15a
--- /dev/null
+++ b/Ryujinx.Graphics.Vulkan/Queries/CounterQueue.cs
@@ -0,0 +1,224 @@
+using Ryujinx.Graphics.GAL;
+using Silk.NET.Vulkan;
+using System;
+using System.Collections.Generic;
+using System.Threading;
+
+namespace Ryujinx.Graphics.Vulkan.Queries
+{
+    class CounterQueue : IDisposable
+    {
+        private const int QueryPoolInitialSize = 100;
+
+        private readonly VulkanRenderer _gd;
+        private readonly Device _device;
+        private readonly PipelineFull _pipeline;
+
+        public CounterType Type { get; }
+        public bool Disposed { get; private set; }
+
+        private Queue<CounterQueueEvent> _events = new Queue<CounterQueueEvent>();
+        private CounterQueueEvent _current;
+
+        private ulong _accumulatedCounter;
+        private int _waiterCount;
+
+        private object _lock = new object();
+
+        private Queue<BufferedQuery> _queryPool;
+        private AutoResetEvent _queuedEvent = new AutoResetEvent(false);
+        private AutoResetEvent _wakeSignal = new AutoResetEvent(false);
+        private AutoResetEvent _eventConsumed = new AutoResetEvent(false);
+
+        private Thread _consumerThread;
+
+        internal CounterQueue(VulkanRenderer gd, Device device, PipelineFull pipeline, CounterType type)
+        {
+            _gd = gd;
+            _device = device;
+            _pipeline = pipeline;
+
+            Type = type;
+
+            _queryPool = new Queue<BufferedQuery>(QueryPoolInitialSize);
+            for (int i = 0; i < QueryPoolInitialSize; i++)
+            {
+                // AMD Polaris GPUs on Windows seem to have issues reporting 64-bit query results.
+                _queryPool.Enqueue(new BufferedQuery(_gd, _device, _pipeline, type, gd.IsAmdWindows));
+            }
+
+            _current = new CounterQueueEvent(this, type, 0);
+
+            _consumerThread = new Thread(EventConsumer);
+            _consumerThread.Start();
+        }
+
+        private void EventConsumer()
+        {
+            while (!Disposed)
+            {
+                CounterQueueEvent evt = null;
+                lock (_lock)
+                {
+                    if (_events.Count > 0)
+                    {
+                        evt = _events.Dequeue();
+                    }
+                }
+
+                if (evt == null)
+                {
+                    _queuedEvent.WaitOne(); // No more events to go through, wait for more.
+                }
+                else
+                {
+                    // Spin-wait rather than sleeping if there are any waiters, by passing null instead of the wake signal.
+                    evt.TryConsume(ref _accumulatedCounter, true, _waiterCount == 0 ? _wakeSignal : null);
+                }
+
+                if (_waiterCount > 0)
+                {
+                    _eventConsumed.Set();
+                }
+            }
+        }
+
+        internal BufferedQuery GetQueryObject()
+        {
+            // Creating/disposing query objects on a context we're sharing with will cause issues.
+            // So instead, make a lot of query objects on the main thread and reuse them.
+
+            lock (_lock)
+            {
+                if (_queryPool.Count > 0)
+                {
+                    BufferedQuery result = _queryPool.Dequeue();
+                    return result;
+                }
+                else
+                {
+                    return new BufferedQuery(_gd, _device, _pipeline, Type, _gd.IsAmdWindows);
+                }
+            }
+        }
+
+        internal void ReturnQueryObject(BufferedQuery query)
+        {
+            lock (_lock)
+            {
+                _pipeline.ResetQuery(query);
+                _queryPool.Enqueue(query);
+            }
+        }
+
+        public CounterQueueEvent QueueReport(EventHandler<ulong> resultHandler, ulong lastDrawIndex, bool hostReserved)
+        {
+            CounterQueueEvent result;
+            ulong draws = lastDrawIndex - _current.DrawIndex;
+
+            lock (_lock)
+            {
+                // A query's result only matters if more than one draw was performed during it.
+                // Otherwise, dummy it out and return 0 immediately.
+
+                if (hostReserved)
+                {
+                    // This counter event is guaranteed to be available for host conditional rendering.
+                    _current.ReserveForHostAccess();
+                }
+
+                _current.Complete(draws > 0 && Type != CounterType.TransformFeedbackPrimitivesWritten);
+                _events.Enqueue(_current);
+
+                _current.OnResult += resultHandler;
+
+                result = _current;
+
+                _current = new CounterQueueEvent(this, Type, lastDrawIndex);
+            }
+
+            _queuedEvent.Set();
+
+            return result;
+        }
+
+        public void QueueReset(ulong lastDrawIndex)
+        {
+            ulong draws = lastDrawIndex - _current.DrawIndex;
+
+            lock (_lock)
+            {
+                _current.Clear(draws != 0);
+            }
+        }
+
+        public void Flush(bool blocking)
+        {
+            if (!blocking)
+            {
+                // Just wake the consumer thread - it will update the queries.
+                _wakeSignal.Set();
+                return;
+            }
+
+            lock (_lock)
+            {
+                // Tell the queue to process all events.
+                while (_events.Count > 0)
+                {
+                    CounterQueueEvent flush = _events.Peek();
+                    if (!flush.TryConsume(ref _accumulatedCounter, true))
+                    {
+                        return; // If not blocking, then return when we encounter an event that is not ready yet.
+                    }
+                    _events.Dequeue();
+                }
+            }
+        }
+
+        public void FlushTo(CounterQueueEvent evt)
+        {
+            // Flush the counter queue on the main thread.
+            Interlocked.Increment(ref _waiterCount);
+
+            _wakeSignal.Set();
+
+            while (!evt.Disposed)
+            {
+                _eventConsumed.WaitOne(1);
+            }
+
+            Interlocked.Decrement(ref _waiterCount);
+        }
+
+        public void Dispose()
+        {
+            lock (_lock)
+            {
+                while (_events.Count > 0)
+                {
+                    CounterQueueEvent evt = _events.Dequeue();
+
+                    evt.Dispose();
+                }
+
+                Disposed = true;
+            }
+
+            _queuedEvent.Set();
+
+            _consumerThread.Join();
+
+            _current?.Dispose();
+
+            foreach (BufferedQuery query in _queryPool)
+            {
+                query.Dispose();
+            }
+
+            _queuedEvent.Dispose();
+            _wakeSignal.Dispose();
+            _eventConsumed.Dispose();
+        }
+    }
+}
diff --git a/Ryujinx.Graphics.Vulkan/Queries/CounterQueueEvent.cs b/Ryujinx.Graphics.Vulkan/Queries/CounterQueueEvent.cs
new file mode 100644
index 00000000..241fe1ee
--- /dev/null
+++ b/Ryujinx.Graphics.Vulkan/Queries/CounterQueueEvent.cs
@@ -0,0 +1,167 @@
+using Ryujinx.Graphics.GAL;
+using System;
+using System.Threading;
+
+namespace Ryujinx.Graphics.Vulkan.Queries
+{
+    class CounterQueueEvent : ICounterEvent
+    {
+        public event EventHandler<ulong> OnResult;
+
+        public CounterType Type { get; }
+        public bool ClearCounter { get; private set; }
+
+        public bool Disposed { get; private set; }
+        public bool Invalid { get; set; }
+
+        public ulong DrawIndex { get; }
+
+        private CounterQueue _queue;
+        private BufferedQuery _counter;
+
+        private bool _hostAccessReserved = false;
+        private int _refCount = 1; // Starts with a reference from the counter queue.
+
+        private object _lock = new object();
+        private ulong _result = ulong.MaxValue;
+
+        public CounterQueueEvent(CounterQueue queue, CounterType type, ulong drawIndex)
+        {
+            _queue = queue;
+
+            _counter = queue.GetQueryObject();
+            Type = type;
+
+            DrawIndex = drawIndex;
+
+            _counter.Begin();
+        }
+
+        public Auto<DisposableBuffer> GetBuffer()
+        {
+            return _counter.GetBuffer();
+        }
+
+        internal void Clear(bool counterReset)
+        {
+            if (counterReset)
+            {
+                _counter.Reset();
+            }
+
+            ClearCounter = true;
+        }
+
+        internal void Complete(bool withResult)
+        {
+            _counter.End(withResult);
+        }
+
+        internal bool TryConsume(ref ulong result, bool block, AutoResetEvent wakeSignal = null)
+        {
+            lock (_lock)
+            {
+                if (Disposed)
+                {
+                    return true;
+                }
+
+                if (ClearCounter)
+                {
+                    result = 0;
+                }
+
+                long queryResult;
+
+                if (block)
+                {
+                    queryResult = _counter.AwaitResult(wakeSignal);
+                }
+                else
+                {
+                    if (!_counter.TryGetResult(out queryResult))
+                    {
+                        return false;
+                    }
+                }
+
+                result += (ulong)queryResult;
+
+                _result = result;
+
+                OnResult?.Invoke(this, result);
+
+                Dispose(); // Return the our resources to the pool.
+
+                return true;
+            }
+        }
+
+        public void Flush()
+        {
+            if (Disposed)
+            {
+                return;
+            }
+
+            // Tell the queue to process all events up to this one.
+            _queue.FlushTo(this);
+        }
+
+        public void DecrementRefCount()
+        {
+            if (Interlocked.Decrement(ref _refCount) == 0)
+            {
+                DisposeInternal();
+            }
+        }
+
+        public bool ReserveForHostAccess()
+        {
+            if (_hostAccessReserved)
+            {
+                return true;
+            }
+
+            if (IsValueAvailable())
+            {
+                return false;
+            }
+
+            if (Interlocked.Increment(ref _refCount) == 1)
+            {
+                Interlocked.Decrement(ref _refCount);
+
+                return false;
+            }
+
+            _hostAccessReserved = true;
+
+            return true;
+        }
+
+        public void ReleaseHostAccess()
+        {
+            _hostAccessReserved = false;
+
+            DecrementRefCount();
+        }
+
+        private void DisposeInternal()
+        {
+            _queue.ReturnQueryObject(_counter);
+        }
+
+        private bool IsValueAvailable()
+        {
+            return _result != ulong.MaxValue || _counter.TryGetResult(out _);
+        }
+
+        public void Dispose()
+        {
+            Disposed = true;
+
+            DecrementRefCount();
+        }
+    }
+}
diff --git a/Ryujinx.Graphics.Vulkan/Queries/Counters.cs b/Ryujinx.Graphics.Vulkan/Queries/Counters.cs
new file mode 100644
index 00000000..63581e42
--- /dev/null
+++ b/Ryujinx.Graphics.Vulkan/Queries/Counters.cs
@@ -0,0 +1,58 @@
+using Ryujinx.Graphics.GAL;
+using Silk.NET.Vulkan;
+using System;
+
+namespace Ryujinx.Graphics.Vulkan.Queries
+{
+    class Counters : IDisposable
+    {
+        private readonly CounterQueue[] _counterQueues;
+        private readonly PipelineFull _pipeline;
+
+        public Counters(VulkanRenderer gd, Device device, PipelineFull pipeline)
+        {
+            _pipeline = pipeline;
+
+            int count = Enum.GetNames(typeof(CounterType)).Length;
+
+            _counterQueues = new CounterQueue[count];
+
+            for (int index = 0; index < count; index++)
+            {
+                CounterType type = (CounterType)index;
+                _counterQueues[index] = new CounterQueue(gd, device, pipeline, type);
+            }
+        }
+
+        public CounterQueueEvent QueueReport(CounterType type, EventHandler<ulong> resultHandler, bool hostReserved)
+        {
+            return _counterQueues[(int)type].QueueReport(resultHandler, _pipeline.DrawCount, hostReserved);
+        }
+
+        public void QueueReset(CounterType type)
+        {
+            _counterQueues[(int)type].QueueReset(_pipeline.DrawCount);
+        }
+
+        public void Update()
+        {
+            foreach (var queue in _counterQueues)
+            {
+                queue.Flush(false);
+            }
+        }
+
+        public void Flush(CounterType type)
+        {
+            _counterQueues[(int)type].Flush(true);
+        }
+
+        public void Dispose()
+        {
+            foreach (var queue in _counterQueues)
+            {
+                queue.Dispose();
+            }
+        }
+    }
+}