1 files changed, 16 insertions, 3 deletions
diff --git a/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs b/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
index ec8fca1d..c280a6d8 100644
--- a/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
+++ b/Ryujinx.Graphics.Shader/Translation/Optimizations/GlobalToStorage.cs
@@ -8,11 +8,14 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
 {
     static class GlobalToStorage
     {
-        public static void RunPass(BasicBlock block, ShaderConfig config, ref int sbUseMask)
+        public static void RunPass(BasicBlock block, ShaderConfig config, ref int sbUseMask, ref int ubeUseMask)
         {
             int sbStart = GetStorageBaseCbOffset(config.Stage);
             int sbEnd = sbStart + StorageDescsSize;
 
+            int ubeStart = UbeBaseOffset;
+            int ubeEnd = UbeBaseOffset + UbeDescsSize;
+
             for (LinkedListNode<INode> node = block.Operations.First; node != null; node = node.Next)
             {
                 for (int index = 0; index < node.Value.SourcesCount; index++)
@@ -25,6 +28,16 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
                     {
                         sbUseMask |= 1 << storageIndex;
                     }
+
+                    if (config.Stage == ShaderStage.Compute)
+                    {
+                        int constantIndex = GetStorageIndex(src, ubeStart, ubeEnd);
+
+                        if (constantIndex >= 0)
+                        {
+                            ubeUseMask |= 1 << constantIndex;
+                        }
+                    }
                 }
 
                 if (!(node.Value is Operation operation))
@@ -54,7 +67,7 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
                         // so NVN "emulates" more constant buffers using global memory access.
                         // Here we try to replace the global access back to a constant buffer
                         // load.
-                        storageIndex = SearchForStorageBase(block, source, UbeBaseOffset, UbeBaseOffset + UbeDescsSize);
+                        storageIndex = SearchForStorageBase(block, source, ubeStart, ubeStart + ubeEnd);
 
                         if (storageIndex >= 0)
                         {
@@ -64,7 +77,7 @@ namespace Ryujinx.Graphics.Shader.Translation.Optimizations
                 }
             }
 
-            config.SetAccessibleStorageBuffersMask(sbUseMask);
+            config.SetAccessibleBufferMasks(sbUseMask, ubeUseMask);
         }
 
         private static LinkedListNode<INode> ReplaceGlobalWithStorage(BasicBlock block, LinkedListNode<INode> node, ShaderConfig config, int storageIndex)