Initial support for image stores, support texture sample on compute

5 files changed, 291 insertions, 144 deletions

diff --git a/Ryujinx.Graphics.Gpu/Image/Texture.cs b/Ryujinx.Graphics.Gpu/Image/Texture.cs
index 32db8688..7ebf01b8 100644
--- a/Ryujinx.Graphics.Gpu/Image/Texture.cs
+++ b/Ryujinx.Graphics.Gpu/Image/Texture.cs
@@ -217,8 +217,6 @@ namespace Ryujinx.Graphics.Gpu.Image
 
             ulong rangeSize = (EndAddress - Address + pageMask) & ~pageMask;
 
-            _context.Methods.InvalidateRange(rangeAddress, rangeSize);
-
             Span<byte> data = _context.PhysicalMemory.Read(Address, Size);
 
             if (_info.IsLinear)
@@ -683,11 +681,6 @@ namespace Ryujinx.Graphics.Gpu.Image
             return Address < address + size && address < EndAddress;
         }
 
-        public void Invalidate()
-        {
-            // _hasData = false;
-        }
-
         public void IncrementReferenceCount()
         {
             _referenceCount++;
diff --git a/Ryujinx.Graphics.Gpu/Image/TextureBindingsManager.cs b/Ryujinx.Graphics.Gpu/Image/TextureBindingsManager.cs
new file mode 100644
index 00000000..2d0e828e
--- /dev/null
+++ b/Ryujinx.Graphics.Gpu/Image/TextureBindingsManager.cs
@@ -0,0 +1,231 @@
+using Ryujinx.Graphics.GAL;
+using Ryujinx.Graphics.Shader;
+using System;
+
+namespace Ryujinx.Graphics.Gpu.Image
+{
+    class TextureBindingsManager
+    {
+        private GpuContext _context;
+
+        private bool _isCompute;
+
+        private SamplerPool _samplerPool;
+
+        private ulong _texturePoolAddress;
+        private int   _texturePoolMaximumId;
+
+        private TexturePoolCache _texturePoolCache;
+
+        private TextureBindingInfo[][] _textureBindings;
+        private TextureBindingInfo[][] _imageBindings;
+
+        private struct TextureStatePerStage
+        {
+            public ITexture Texture;
+            public ISampler Sampler;
+        }
+
+        private TextureStatePerStage[][] _textureState;
+        private TextureStatePerStage[][] _imageState;
+
+        private int _textureBufferIndex;
+
+        private bool _rebind;
+
+        public TextureBindingsManager(GpuContext context, bool isCompute)
+        {
+            _context   = context;
+            _isCompute = isCompute;
+
+            _texturePoolCache = new TexturePoolCache(context);
+
+            int stages = isCompute ? 1 : Constants.TotalShaderStages;
+
+            _textureBindings = new TextureBindingInfo[stages][];
+            _imageBindings   = new TextureBindingInfo[stages][];
+
+            _textureState = new TextureStatePerStage[stages][];
+            _imageState   = new TextureStatePerStage[stages][];
+        }
+
+        public void SetTextures(int stage, TextureBindingInfo[] bindings)
+        {
+            _textureBindings[stage] = bindings;
+
+            _textureState[stage] = new TextureStatePerStage[bindings.Length];
+        }
+
+        public void SetImages(int stage, TextureBindingInfo[] bindings)
+        {
+            _imageBindings[stage] = bindings;
+
+            _imageState[stage] = new TextureStatePerStage[bindings.Length];
+        }
+
+        public void SetTextureBufferIndex(int index)
+        {
+            _textureBufferIndex = index;
+        }
+
+        public void SetSamplerPool(ulong gpuVa, int maximumId)
+        {
+            ulong address = _context.MemoryManager.Translate(gpuVa);
+
+            if (_samplerPool != null)
+            {
+                if (_samplerPool.Address == address)
+                {
+                    return;
+                }
+
+                _samplerPool.Dispose();
+            }
+
+            _samplerPool = new SamplerPool(_context, address, maximumId);
+        }
+
+        public void SetTexturePool(ulong gpuVa, int maximumId)
+        {
+            ulong address = _context.MemoryManager.Translate(gpuVa);
+
+            _texturePoolAddress   = address;
+            _texturePoolMaximumId = maximumId;
+        }
+
+        public void CommitBindings()
+        {
+            TexturePool texturePool = _texturePoolCache.FindOrCreate(
+                _texturePoolAddress,
+                _texturePoolMaximumId);
+
+            if (_isCompute)
+            {
+                CommitTextureBindings(texturePool, ShaderStage.Compute, 0);
+                CommitImageBindings  (texturePool, ShaderStage.Compute, 0);
+            }
+            else
+            {
+                for (ShaderStage stage = ShaderStage.Vertex; stage <= ShaderStage.Fragment; stage++)
+                {
+                    int stageIndex = (int)stage - 1;
+
+                    CommitTextureBindings(texturePool, stage, stageIndex);
+                    CommitImageBindings  (texturePool, stage, stageIndex);
+                }
+            }
+
+            _rebind = false;
+        }
+
+        private void CommitTextureBindings(TexturePool pool, ShaderStage stage, int stageIndex)
+        {
+            if (_textureBindings[stageIndex] == null)
+            {
+                return;
+            }
+
+            for (int index = 0; index < _textureBindings[stageIndex].Length; index++)
+            {
+                TextureBindingInfo binding = _textureBindings[stageIndex][index];
+
+                int packedId = ReadPackedId(stageIndex, binding.Handle);
+
+                int textureId = UnpackTextureId(packedId);
+                int samplerId = UnpackSamplerId(packedId);
+
+                Texture texture = pool.Get(textureId);
+
+                ITexture hostTexture = texture?.GetTargetTexture(binding.Target);
+
+                if (_textureState[stageIndex][index].Texture != hostTexture || _rebind)
+                {
+                    _textureState[stageIndex][index].Texture = hostTexture;
+
+                    _context.Renderer.Pipeline.BindTexture(index, stage, hostTexture);
+                }
+
+                Sampler sampler = _samplerPool.Get(samplerId);
+
+                ISampler hostSampler = sampler?.HostSampler;
+
+                if (_textureState[stageIndex][index].Sampler != hostSampler || _rebind)
+                {
+                    _textureState[stageIndex][index].Sampler = hostSampler;
+
+                    _context.Renderer.Pipeline.BindSampler(index, stage, hostSampler);
+                }
+            }
+        }
+
+        private void CommitImageBindings(TexturePool pool, ShaderStage stage, int stageIndex)
+        {
+            if (_imageBindings[stageIndex] == null)
+            {
+                return;
+            }
+
+            for (int index = 0; index < _imageBindings[stageIndex].Length; index++)
+            {
+                TextureBindingInfo binding = _imageBindings[stageIndex][index];
+
+                int packedId = ReadPackedId(stageIndex, binding.Handle);
+
+                int textureId = UnpackTextureId(packedId);
+
+                Texture texture = pool.Get(textureId);
+
+                ITexture hostTexture = texture?.GetTargetTexture(binding.Target);
+
+                if (_imageState[stageIndex][index].Texture != hostTexture || _rebind)
+                {
+                    _imageState[stageIndex][index].Texture = hostTexture;
+
+                    _context.Renderer.Pipeline.BindImage(index, stage, hostTexture);
+                }
+            }
+        }
+
+        private int ReadPackedId(int stage, int wordOffset)
+        {
+            ulong address;
+
+            var bufferManager = _context.Methods.BufferManager;
+
+            if (_isCompute)
+            {
+                address = bufferManager.GetComputeUniformBufferAddress(_textureBufferIndex);
+            }
+            else
+            {
+                address = bufferManager.GetGraphicsUniformBufferAddress(stage, _textureBufferIndex);
+            }
+
+            address += (uint)wordOffset * 4;
+
+            return BitConverter.ToInt32(_context.PhysicalMemory.Read(address, 4));
+        }
+
+        private static int UnpackTextureId(int packedId)
+        {
+            return (packedId >> 0) & 0xfffff;
+        }
+
+        private static int UnpackSamplerId(int packedId)
+        {
+            return (packedId >> 20) & 0xfff;
+        }
+
+        public void InvalidatePoolRange(ulong address, ulong size)
+        {
+            _samplerPool?.InvalidateRange(address, size);
+
+            _texturePoolCache.InvalidateRange(address, size);
+        }
+
+        public void Rebind()
+        {
+            _rebind = true;
+        }
+    }
+}
\ No newline at end of file
diff --git a/Ryujinx.Graphics.Gpu/Image/TextureManager.cs b/Ryujinx.Graphics.Gpu/Image/TextureManager.cs
index 23416ddd..91a9cfd1 100644
--- a/Ryujinx.Graphics.Gpu/Image/TextureManager.cs
+++ b/Ryujinx.Graphics.Gpu/Image/TextureManager.cs
@@ -4,7 +4,6 @@ using Ryujinx.Graphics.GAL.Texture;
 using Ryujinx.Graphics.Gpu.Image;
 using Ryujinx.Graphics.Gpu.Memory;
 using Ryujinx.Graphics.Gpu.State;
-using Ryujinx.Graphics.Shader;
 using Ryujinx.Graphics.Texture;
 using System;
 
@@ -12,15 +11,10 @@ namespace Ryujinx.Graphics.Gpu.Image
 {
     class TextureManager
     {
-        private GpuContext    _context;
-        private BufferManager _bufferManager;
+        private GpuContext _context;
 
-        private SamplerPool _samplerPool;
-
-        private ulong _texturePoolAddress;
-        private int   _texturePoolMaximumId;
-
-        private TexturePoolCache _texturePoolCache;
+        private TextureBindingsManager _cpBindingsManager;
+        private TextureBindingsManager _gpBindingsManager;
 
         private Texture[] _rtColors;
         private Texture   _rtColor3D;
@@ -35,24 +29,12 @@ namespace Ryujinx.Graphics.Gpu.Image
 
         private AutoDeleteCache _cache;
 
-        private TextureBindingInfo[][] _bindings;
-
-        private struct TextureStatePerStage
-        {
-            public ITexture Texture;
-            public ISampler Sampler;
-        }
-
-        private TextureStatePerStage[][] _textureState;
-
-        private int _textureBufferIndex;
-
-        public TextureManager(GpuContext context, BufferManager bufferManager)
+        public TextureManager(GpuContext context)
         {
-            _context       = context;
-            _bufferManager = bufferManager;
+            _context = context;
 
-            _texturePoolCache = new TexturePoolCache(context, this);
+            _cpBindingsManager = new TextureBindingsManager(context, isCompute: true);
+            _gpBindingsManager = new TextureBindingsManager(context, isCompute: false);
 
             _rtColors = new Texture[Constants.TotalRenderTargets];
 
@@ -61,47 +43,56 @@ namespace Ryujinx.Graphics.Gpu.Image
             _textures = new RangeList<Texture>();
 
             _cache = new AutoDeleteCache();
+        }
 
-            _bindings = new TextureBindingInfo[Constants.TotalShaderStages][];
-
-            _textureState = new TextureStatePerStage[Constants.TotalShaderStages][];
+        public void SetComputeTextures(TextureBindingInfo[] bindings)
+        {
+            _cpBindingsManager.SetTextures(0, bindings);
         }
 
-        public void BindTextures(int stage, TextureBindingInfo[] bindings)
+        public void SetGraphicsTextures(int stage, TextureBindingInfo[] bindings)
         {
-            _bindings[stage] = bindings;
+            _gpBindingsManager.SetTextures(stage, bindings);
+        }
 
-            _textureState[stage] = new TextureStatePerStage[bindings.Length];
+        public void SetComputeImages(TextureBindingInfo[] bindings)
+        {
+            _cpBindingsManager.SetImages(0, bindings);
         }
 
-        public void SetTextureBufferIndex(int index)
+        public void SetGraphicsImages(int stage, TextureBindingInfo[] bindings)
         {
-            _textureBufferIndex = index;
+            _gpBindingsManager.SetImages(stage, bindings);
         }
 
-        public void SetSamplerPool(ulong gpuVa, int maximumId)
+        public void SetComputeTextureBufferIndex(int index)
         {
-            ulong address = _context.MemoryManager.Translate(gpuVa);
+            _cpBindingsManager.SetTextureBufferIndex(index);
+        }
 
-            if (_samplerPool != null)
-            {
-                if (_samplerPool.Address == address)
-                {
-                    return;
-                }
+        public void SetGraphicsTextureBufferIndex(int index)
+        {
+            _gpBindingsManager.SetTextureBufferIndex(index);
+        }
 
-                _samplerPool.Dispose();
-            }
+        public void SetComputeSamplerPool(ulong gpuVa, int maximumId)
+        {
+            _cpBindingsManager.SetSamplerPool(gpuVa, maximumId);
+        }
 
-            _samplerPool = new SamplerPool(_context, address, maximumId);
+        public void SetGraphicsSamplerPool(ulong gpuVa, int maximumId)
+        {
+            _gpBindingsManager.SetSamplerPool(gpuVa, maximumId);
         }
 
-        public void SetTexturePool(ulong gpuVa, int maximumId)
+        public void SetComputeTexturePool(ulong gpuVa, int maximumId)
         {
-            ulong address = _context.MemoryManager.Translate(gpuVa);
+            _cpBindingsManager.SetTexturePool(gpuVa, maximumId);
+        }
 
-            _texturePoolAddress   = address;
-            _texturePoolMaximumId = maximumId;
+        public void SetGraphicsTexturePool(ulong gpuVa, int maximumId)
+        {
+            _gpBindingsManager.SetTexturePool(gpuVa, maximumId);
         }
 
         public void SetRenderTargetColor(int index, Texture color)
@@ -121,59 +112,22 @@ namespace Ryujinx.Graphics.Gpu.Image
             _rtDepthStencil = depthStencil;
         }
 
-        public void CommitBindings()
+        public void CommitComputeBindings()
         {
-            UpdateTextures();
-            UpdateRenderTargets();
+            // Evert time we switch between graphics and compute work,
+            // we must rebind everything.
+            // Since compute work happens less often, we always do that
+            // before and after the compute dispatch.
+            _cpBindingsManager.Rebind();
+            _cpBindingsManager.CommitBindings();
+            _gpBindingsManager.Rebind();
         }
 
-        private void UpdateTextures()
+        public void CommitGraphicsBindings()
         {
-            TexturePool texturePool = _texturePoolCache.FindOrCreate(
-                _texturePoolAddress,
-                _texturePoolMaximumId);
-
-            for (ShaderStage stage = ShaderStage.Vertex; stage <= ShaderStage.Fragment; stage++)
-            {
-                int stageIndex = (int)stage - 1;
+            _gpBindingsManager.CommitBindings();
 
-                if (_bindings[stageIndex] == null)
-                {
-                    continue;
-                }
-
-                for (int index = 0; index < _bindings[stageIndex].Length; index++)
-                {
-                    TextureBindingInfo binding = _bindings[stageIndex][index];
-
-                    int packedId = ReadPackedId(stageIndex, binding.Handle);
-
-                    int textureId = (packedId >> 0)  & 0xfffff;
-                    int samplerId = (packedId >> 20) & 0xfff;
-
-                    Texture texture = texturePool.Get(textureId);
-
-                    ITexture hostTexture = texture?.GetTargetTexture(binding.Target);
-
-                    if (_textureState[stageIndex][index].Texture != hostTexture)
-                    {
-                        _textureState[stageIndex][index].Texture = hostTexture;
-
-                        _context.Renderer.GraphicsPipeline.BindTexture(index, stage, hostTexture);
-                    }
-
-                    Sampler sampler = _samplerPool.Get(samplerId);
-
-                    ISampler hostSampler = sampler?.HostSampler;
-
-                    if (_textureState[stageIndex][index].Sampler != hostSampler)
-                    {
-                        _textureState[stageIndex][index].Sampler = hostSampler;
-
-                        _context.Renderer.GraphicsPipeline.BindSampler(index, stage, hostSampler);
-                    }
-                }
-            }
+            UpdateRenderTargets();
         }
 
         private void UpdateRenderTargets()
@@ -203,7 +157,7 @@ namespace Ryujinx.Graphics.Gpu.Image
 
                 if (anyChanged)
                 {
-                    _context.Renderer.GraphicsPipeline.SetRenderTargets(_rtHostColors, _rtHostDs);
+                    _context.Renderer.Pipeline.SetRenderTargets(_rtHostColors, _rtHostDs);
                 }
             }
             else
@@ -217,20 +171,11 @@ namespace Ryujinx.Graphics.Gpu.Image
 
                 if (anyChanged)
                 {
-                    _context.Renderer.GraphicsPipeline.SetRenderTargets(_rtColor3D.HostTexture, _rtHostDs);
+                    _context.Renderer.Pipeline.SetRenderTargets(_rtColor3D.HostTexture, _rtHostDs);
                 }
             }
         }
 
-        private int ReadPackedId(int stage, int wordOffset)
-        {
-            ulong address = _bufferManager.GetGraphicsUniformBufferAddress(stage, _textureBufferIndex);
-
-            address += (uint)wordOffset * 4;
-
-            return BitConverter.ToInt32(_context.PhysicalMemory.Read(address, 4));
-        }
-
         public Texture FindOrCreateTexture(CopyTexture copyTexture)
         {
             ulong address = _context.MemoryManager.Translate(copyTexture.Address.Pack());
@@ -645,20 +590,6 @@ namespace Ryujinx.Graphics.Gpu.Image
             return ts[0];
         }
 
-        public void InvalidateRange(ulong address, ulong size)
-        {
-            Texture[] overlaps = _textures.FindOverlaps(address, size);
-
-            foreach (Texture overlap in overlaps)
-            {
-                overlap.Invalidate();
-            }
-
-            _samplerPool?.InvalidateRange(address, size);
-
-            _texturePoolCache.InvalidateRange(address, size);
-        }
-
         public void Flush()
         {
             foreach (Texture texture in _cache)
diff --git a/Ryujinx.Graphics.Gpu/Image/TexturePool.cs b/Ryujinx.Graphics.Gpu/Image/TexturePool.cs
index 558f4def..8512e370 100644
--- a/Ryujinx.Graphics.Gpu/Image/TexturePool.cs
+++ b/Ryujinx.Graphics.Gpu/Image/TexturePool.cs
@@ -9,8 +9,6 @@ namespace Ryujinx.Graphics.Gpu.Image
 {
     class TexturePool : Pool<Texture>
     {
-        private TextureManager _textureManager;
-
         public LinkedListNode<TexturePool> CacheNode { get; set; }
 
         private struct TextureContainer
@@ -20,13 +18,9 @@ namespace Ryujinx.Graphics.Gpu.Image
         }
 
         public TexturePool(
-            GpuContext     context,
-            TextureManager textureManager,
-            ulong          address,
-            int            maximumId) : base(context, address, maximumId)
-        {
-            _textureManager = textureManager;
-        }
+            GpuContext context,
+            ulong      address,
+            int        maximumId) : base(context, address, maximumId) { }
 
         public override Texture Get(int id)
         {
@@ -56,7 +50,7 @@ namespace Ryujinx.Graphics.Gpu.Image
                     return null;
                 }
 
-                texture = _textureManager.FindOrCreateTexture(info, TextureSearchFlags.Sampler);
+                texture = Context.Methods.TextureManager.FindOrCreateTexture(info, TextureSearchFlags.Sampler);
 
                 texture.IncrementReferenceCount();
 
diff --git a/Ryujinx.Graphics.Gpu/Image/TexturePoolCache.cs b/Ryujinx.Graphics.Gpu/Image/TexturePoolCache.cs
index 8e8313ae..9ab7e292 100644
--- a/Ryujinx.Graphics.Gpu/Image/TexturePoolCache.cs
+++ b/Ryujinx.Graphics.Gpu/Image/TexturePoolCache.cs
@@ -6,15 +6,13 @@ namespace Ryujinx.Graphics.Gpu.Image
     {
         private const int MaxCapacity = 4;
 
-        private GpuContext     _context;
-        private TextureManager _textureManager;
+        private GpuContext _context;
 
         private LinkedList<TexturePool> _pools;
 
-        public TexturePoolCache(GpuContext context, TextureManager textureManager)
+        public TexturePoolCache(GpuContext context)
         {
-            _context        = context;
-            _textureManager = textureManager;
+            _context = context;
 
             _pools = new LinkedList<TexturePool>();
         }
@@ -42,7 +40,7 @@ namespace Ryujinx.Graphics.Gpu.Image
             }
 
             // If not found, create a new one.
-            pool = new TexturePool(_context, _textureManager, address, maximumId);
+            pool = new TexturePool(_context, address, maximumId);
 
             pool.CacheNode = _pools.AddLast(pool);