Implement Load/Store Local/Shared and Atomic shared using new instructions (#5241)

* Implement Load/Store Local/Shared and Atomic shared using new instructions

* Remove now unused code

* Fix base offset register overwrite

* Fix missing storage buffer set index when generating GLSL for Vulkan

* Shader cache version bump

* Remove more unused code

* Some PR feedback

12 files changed, 105 insertions, 412 deletions

diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
index 958f1cef..08e8eb19 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs
@@ -71,40 +71,10 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             context.AppendLine($"const int {DefaultNames.UndefinedName} = 0;");
             context.AppendLine();
 
-            if (context.Config.Stage == ShaderStage.Compute)
-            {
-                int localMemorySize = BitUtils.DivRoundUp(context.Config.GpuAccessor.QueryComputeLocalMemorySize(), 4);
-
-                if (localMemorySize != 0)
-                {
-                    string localMemorySizeStr = NumberFormatter.FormatInt(localMemorySize);
-
-                    context.AppendLine($"uint {DefaultNames.LocalMemoryName}[{localMemorySizeStr}];");
-                    context.AppendLine();
-                }
-
-                int sharedMemorySize = BitUtils.DivRoundUp(context.Config.GpuAccessor.QueryComputeSharedMemorySize(), 4);
-
-                if (sharedMemorySize != 0)
-                {
-                    string sharedMemorySizeStr = NumberFormatter.FormatInt(sharedMemorySize);
-
-                    context.AppendLine($"shared uint {DefaultNames.SharedMemoryName}[{sharedMemorySizeStr}];");
-                    context.AppendLine();
-                }
-            }
-            else if (context.Config.LocalMemorySize != 0)
-            {
-                int localMemorySize = BitUtils.DivRoundUp(context.Config.LocalMemorySize, 4);
-
-                string localMemorySizeStr = NumberFormatter.FormatInt(localMemorySize);
-
-                context.AppendLine($"uint {DefaultNames.LocalMemoryName}[{localMemorySizeStr}];");
-                context.AppendLine();
-            }
-
             DeclareConstantBuffers(context, context.Config.Properties.ConstantBuffers.Values);
             DeclareStorageBuffers(context, context.Config.Properties.StorageBuffers.Values);
+            DeclareMemories(context, context.Config.Properties.LocalMemories.Values, isShared: false);
+            DeclareMemories(context, context.Config.Properties.SharedMemories.Values, isShared: true);
 
             var textureDescriptors = context.Config.GetTextureDescriptors();
             if (textureDescriptors.Length != 0)
@@ -238,11 +208,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 context.AppendLine();
             }
 
-            if ((info.HelperFunctionsMask & HelperFunctionsMask.AtomicMinMaxS32Shared) != 0)
-            {
-                AppendHelperFunction(context, "Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/AtomicMinMaxS32Shared.glsl");
-            }
-
             if ((info.HelperFunctionsMask & HelperFunctionsMask.MultiplyHighS32) != 0)
             {
                 AppendHelperFunction(context, "Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/MultiplyHighS32.glsl");
@@ -273,11 +238,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                 AppendHelperFunction(context, "Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/ShuffleXor.glsl");
             }
 
-            if ((info.HelperFunctionsMask & HelperFunctionsMask.StoreSharedSmallInt) != 0)
-            {
-                AppendHelperFunction(context, "Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/StoreSharedSmallInt.glsl");
-            }
-
             if ((info.HelperFunctionsMask & HelperFunctionsMask.SwizzleAdd) != 0)
             {
                 AppendHelperFunction(context, "Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/SwizzleAdd.glsl");
@@ -358,7 +318,14 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                     _ => "std430"
                 };
 
-                context.AppendLine($"layout (binding = {buffer.Binding}, {layout}) {declType} _{buffer.Name}");
+                string set = string.Empty;
+
+                if (context.Config.Options.TargetApi == TargetApi.Vulkan)
+                {
+                    set = $"set = {buffer.Set}, ";
+                }
+
+                context.AppendLine($"layout ({set}binding = {buffer.Binding}, {layout}) {declType} _{buffer.Name}");
                 context.EnterScope();
 
                 foreach (StructureField field in buffer.Type.Fields)
@@ -391,6 +358,27 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             }
         }
 
+        private static void DeclareMemories(CodeGenContext context, IEnumerable<MemoryDefinition> memories, bool isShared)
+        {
+            string prefix = isShared ? "shared " : string.Empty;
+
+            foreach (MemoryDefinition memory in memories)
+            {
+                string typeName = GetVarTypeName(context, memory.Type & ~AggregateType.Array);
+
+                if (memory.ArrayLength > 0)
+                {
+                    string arraySize = memory.ArrayLength.ToString(CultureInfo.InvariantCulture);
+
+                    context.AppendLine($"{prefix}{typeName} {memory.Name}[{arraySize}];");
+                }
+                else
+                {
+                    context.AppendLine($"{prefix}{typeName} {memory.Name}[];");
+                }
+            }
+        }
+
         private static void DeclareSamplers(CodeGenContext context, TextureDescriptor[] descriptors)
         {
             int arraySize = 0;
@@ -717,7 +705,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             string code = EmbeddedResources.ReadAllText(filename);
 
             code = code.Replace("\t", CodeGenContext.Tab);
-            code = code.Replace("$SHARED_MEM$", DefaultNames.SharedMemoryName);
 
             if (context.Config.GpuAccessor.QueryHostSupportsShaderBallot())
             {
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
index 5ee8259c..e909dcf0 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/DefaultNames.cs
@@ -11,9 +11,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
         public const string IAttributePrefix = "in_attr";
         public const string OAttributePrefix = "out_attr";
 
-        public const string LocalMemoryName  = "local_mem";
-        public const string SharedMemoryName = "shared_mem";
-
         public const string ArgumentNamePrefix = "a";
 
         public const string UndefinedName = "undef";
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/AtomicMinMaxS32Shared.glsl b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/AtomicMinMaxS32Shared.glsl
deleted file mode 100644
index 82b76bcc..00000000
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/AtomicMinMaxS32Shared.glsl
+++ /dev/null
@@ -1,21 +0,0 @@
-int Helper_AtomicMaxS32(int offset, int value)
-{
-    uint oldValue, newValue;
-    do
-    {
-        oldValue = $SHARED_MEM$[offset];
-        newValue = uint(max(int(oldValue), value));
-    } while (atomicCompSwap($SHARED_MEM$[offset], oldValue, newValue) != oldValue);
-    return int(oldValue);
-}
-
-int Helper_AtomicMinS32(int offset, int value)
-{
-    uint oldValue, newValue;
-    do
-    {
-        oldValue = $SHARED_MEM$[offset];
-        newValue = uint(min(int(oldValue), value));
-    } while (atomicCompSwap($SHARED_MEM$[offset], oldValue, newValue) != oldValue);
-    return int(oldValue);
-}
\ No newline at end of file
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/HelperFunctionNames.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/HelperFunctionNames.cs
index 54f35b15..21c43547 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/HelperFunctionNames.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/HelperFunctionNames.cs
@@ -2,9 +2,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 {
     static class HelperFunctionNames
     {
-        public static string AtomicMaxS32 = "Helper_AtomicMaxS32";
-        public static string AtomicMinS32 = "Helper_AtomicMinS32";
-
         public static string MultiplyHighS32 = "Helper_MultiplyHighS32";
         public static string MultiplyHighU32 = "Helper_MultiplyHighU32";
 
@@ -13,10 +10,5 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
         public static string ShuffleUp   = "Helper_ShuffleUp";
         public static string ShuffleXor  = "Helper_ShuffleXor";
         public static string SwizzleAdd  = "Helper_SwizzleAdd";
-
-        public static string StoreShared16  = "Helper_StoreShared16";
-        public static string StoreShared8   = "Helper_StoreShared8";
-        public static string StoreStorage16 = "Helper_StoreStorage16";
-        public static string StoreStorage8  = "Helper_StoreStorage8";
     }
 }
\ No newline at end of file
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/StoreSharedSmallInt.glsl b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/StoreSharedSmallInt.glsl
deleted file mode 100644
index 2f57b5ff..00000000
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/StoreSharedSmallInt.glsl
+++ /dev/null
@@ -1,23 +0,0 @@
-void Helper_StoreShared16(int offset, uint value)
-{
-    int wordOffset = offset >> 2;
-    int bitOffset = (offset & 3) * 8;
-    uint oldValue, newValue;
-    do
-    {
-        oldValue = $SHARED_MEM$[wordOffset];
-        newValue = bitfieldInsert(oldValue, value, bitOffset, 16);
-    } while (atomicCompSwap($SHARED_MEM$[wordOffset], oldValue, newValue) != oldValue);
-}
-
-void Helper_StoreShared8(int offset, uint value)
-{
-    int wordOffset = offset >> 2;
-    int bitOffset = (offset & 3) * 8;
-    uint oldValue, newValue;
-    do
-    {
-        oldValue = $SHARED_MEM$[wordOffset];
-        newValue = bitfieldInsert(oldValue, value, bitOffset, 8);
-    } while (atomicCompSwap($SHARED_MEM$[wordOffset], oldValue, newValue) != oldValue);
-}
\ No newline at end of file
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
index 01d8a6e7..b2577a99 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs
@@ -68,7 +68,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 
                 string args = string.Empty;
 
-                if (atomic && operation.StorageKind == StorageKind.StorageBuffer)
+                if (atomic && (operation.StorageKind == StorageKind.StorageBuffer || operation.StorageKind == StorageKind.SharedMemory))
                 {
                     args = GenerateLoadOrStore(context, operation, isStore: false);
 
@@ -81,23 +81,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                         args += ", " + GetSoureExpr(context, operation.GetSource(argIndex), dstType);
                     }
                 }
-                else if (atomic && operation.StorageKind == StorageKind.SharedMemory)
-                {
-                    args = LoadShared(context, operation);
-
-                    // For shared memory access, the second argument is unused and should be ignored.
-                    // It is there to make both storage and shared access have the same number of arguments.
-                    // For storage, both inputs are consumed when the argument index is 0, so we should skip it here.
-
-                    for (int argIndex = 2; argIndex < arity; argIndex++)
-                    {
-                        args += ", ";
-
-                        AggregateType dstType = GetSrcVarType(inst, argIndex);
-
-                        args += GetSoureExpr(context, operation.GetSource(argIndex), dstType);
-                    }
-                }
                 else
                 {
                     for (int argIndex = 0; argIndex < arity; argIndex++)
@@ -179,12 +162,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                     case Instruction.Load:
                         return Load(context, operation);
 
-                    case Instruction.LoadLocal:
-                        return LoadLocal(context, operation);
-
-                    case Instruction.LoadShared:
-                        return LoadShared(context, operation);
-
                     case Instruction.Lod:
                         return Lod(context, operation);
 
@@ -200,18 +177,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                     case Instruction.Store:
                         return Store(context, operation);
 
-                    case Instruction.StoreLocal:
-                        return StoreLocal(context, operation);
-
-                    case Instruction.StoreShared:
-                        return StoreShared(context, operation);
-
-                    case Instruction.StoreShared16:
-                        return StoreShared16(context, operation);
-
-                    case Instruction.StoreShared8:
-                        return StoreShared8(context, operation);
-
                     case Instruction.TextureSample:
                         return TextureSample(context, operation);
 
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
index f42d9898..8b0b744a 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenHelper.cs
@@ -17,9 +17,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             Add(Instruction.AtomicAdd,                InstType.AtomicBinary,   "atomicAdd");
             Add(Instruction.AtomicAnd,                InstType.AtomicBinary,   "atomicAnd");
             Add(Instruction.AtomicCompareAndSwap,     InstType.AtomicTernary,  "atomicCompSwap");
-            Add(Instruction.AtomicMaxS32,             InstType.CallTernary,    HelperFunctionNames.AtomicMaxS32);
             Add(Instruction.AtomicMaxU32,             InstType.AtomicBinary,   "atomicMax");
-            Add(Instruction.AtomicMinS32,             InstType.CallTernary,    HelperFunctionNames.AtomicMinS32);
             Add(Instruction.AtomicMinU32,             InstType.AtomicBinary,   "atomicMin");
             Add(Instruction.AtomicOr,                 InstType.AtomicBinary,   "atomicOr");
             Add(Instruction.AtomicSwap,               InstType.AtomicBinary,   "atomicExchange");
@@ -83,8 +81,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             Add(Instruction.ImageAtomic,              InstType.Special);
             Add(Instruction.IsNan,                    InstType.CallUnary,      "isnan");
             Add(Instruction.Load,                     InstType.Special);
-            Add(Instruction.LoadLocal,                InstType.Special);
-            Add(Instruction.LoadShared,               InstType.Special);
             Add(Instruction.Lod,                      InstType.Special);
             Add(Instruction.LogarithmB2,              InstType.CallUnary,      "log2");
             Add(Instruction.LogicalAnd,               InstType.OpBinaryCom,    "&&",              9);
@@ -118,10 +114,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             Add(Instruction.Sine,                     InstType.CallUnary,      "sin");
             Add(Instruction.SquareRoot,               InstType.CallUnary,      "sqrt");
             Add(Instruction.Store,                    InstType.Special);
-            Add(Instruction.StoreLocal,               InstType.Special);
-            Add(Instruction.StoreShared,              InstType.Special);
-            Add(Instruction.StoreShared16,            InstType.Special);
-            Add(Instruction.StoreShared8,             InstType.Special);
             Add(Instruction.Subtract,                 InstType.OpBinary,       "-",               2);
             Add(Instruction.SwizzleAdd,               InstType.CallTernary,    HelperFunctionNames.SwizzleAdd);
             Add(Instruction.TextureSample,            InstType.Special);
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
index c8084d9d..99376ffb 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs
@@ -191,25 +191,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             return GenerateLoadOrStore(context, operation, isStore: false);
         }
 
-        public static string LoadLocal(CodeGenContext context, AstOperation operation)
-        {
-            return LoadLocalOrShared(context, operation, DefaultNames.LocalMemoryName);
-        }
-
-        public static string LoadShared(CodeGenContext context, AstOperation operation)
-        {
-            return LoadLocalOrShared(context, operation, DefaultNames.SharedMemoryName);
-        }
-
-        private static string LoadLocalOrShared(CodeGenContext context, AstOperation operation, string arrayName)
-        {
-            IAstNode src1 = operation.GetSource(0);
-
-            string offsetExpr = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
-
-            return $"{arrayName}[{offsetExpr}]";
-        }
-
         public static string Lod(CodeGenContext context, AstOperation operation)
         {
             AstTextureOperation texOp = (AstTextureOperation)operation;
@@ -263,58 +244,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             return GenerateLoadOrStore(context, operation, isStore: true);
         }
 
-        public static string StoreLocal(CodeGenContext context, AstOperation operation)
-        {
-            return StoreLocalOrShared(context, operation, DefaultNames.LocalMemoryName);
-        }
-
-        public static string StoreShared(CodeGenContext context, AstOperation operation)
-        {
-            return StoreLocalOrShared(context, operation, DefaultNames.SharedMemoryName);
-        }
-
-        private static string StoreLocalOrShared(CodeGenContext context, AstOperation operation, string arrayName)
-        {
-            IAstNode src1 = operation.GetSource(0);
-            IAstNode src2 = operation.GetSource(1);
-
-            string offsetExpr = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
-
-            AggregateType srcType = OperandManager.GetNodeDestType(context, src2);
-
-            string src = TypeConversion.ReinterpretCast(context, src2, srcType, AggregateType.U32);
-
-            return $"{arrayName}[{offsetExpr}] = {src}";
-        }
-
-        public static string StoreShared16(CodeGenContext context, AstOperation operation)
-        {
-            IAstNode src1 = operation.GetSource(0);
-            IAstNode src2 = operation.GetSource(1);
-
-            string offsetExpr = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
-
-            AggregateType srcType = OperandManager.GetNodeDestType(context, src2);
-
-            string src = TypeConversion.ReinterpretCast(context, src2, srcType, AggregateType.U32);
-
-            return $"{HelperFunctionNames.StoreShared16}({offsetExpr}, {src})";
-        }
-
-        public static string StoreShared8(CodeGenContext context, AstOperation operation)
-        {
-            IAstNode src1 = operation.GetSource(0);
-            IAstNode src2 = operation.GetSource(1);
-
-            string offsetExpr = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
-
-            AggregateType srcType = OperandManager.GetNodeDestType(context, src2);
-
-            string src = TypeConversion.ReinterpretCast(context, src2, srcType, AggregateType.U32);
-
-            return $"{HelperFunctionNames.StoreShared8}({offsetExpr}, {src})";
-        }
-
         public static string TextureSample(CodeGenContext context, AstOperation operation)
         {
             AstTextureOperation texOp = (AstTextureOperation)operation;
@@ -675,6 +604,21 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
                     varType = field.Type;
                     break;
 
+                case StorageKind.LocalMemory:
+                case StorageKind.SharedMemory:
+                    if (!(operation.GetSource(srcIndex++) is AstOperand bindingId) || bindingId.Type != OperandType.Constant)
+                    {
+                        throw new InvalidOperationException($"First input of {operation.Inst} with {storageKind} storage must be a constant operand.");
+                    }
+
+                    MemoryDefinition memory = storageKind == StorageKind.LocalMemory
+                        ? context.Config.Properties.LocalMemories[bindingId.Value]
+                        : context.Config.Properties.SharedMemories[bindingId.Value];
+
+                    varName = memory.Name;
+                    varType = memory.Type;
+                    break;
+
                 case StorageKind.Input:
                 case StorageKind.InputPerPatch:
                 case StorageKind.Output:
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
index 4fd1d17c..4f6ca642 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs
@@ -113,7 +113,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
             if (node is AstOperation operation)
             {
-                if (operation.Inst == Instruction.Load)
+                if (operation.Inst == Instruction.Load || operation.Inst.IsAtomic())
                 {
                     switch (operation.StorageKind)
                     {
@@ -136,6 +136,19 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
                             return field.Type & AggregateType.ElementTypeMask;
 
+                        case StorageKind.LocalMemory:
+                        case StorageKind.SharedMemory:
+                            if (!(operation.GetSource(0) is AstOperand bindingId) || bindingId.Type != OperandType.Constant)
+                            {
+                                throw new InvalidOperationException($"First input of {operation.Inst} with {operation.StorageKind} storage must be a constant operand.");
+                            }
+
+                            MemoryDefinition memory = operation.StorageKind == StorageKind.LocalMemory
+                                ? context.Config.Properties.LocalMemories[bindingId.Value]
+                                : context.Config.Properties.SharedMemories[bindingId.Value];
+
+                            return memory.Type & AggregateType.ElementTypeMask;
+
                         case StorageKind.Input:
                         case StorageKind.InputPerPatch:
                         case StorageKind.Output:
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/CodeGenContext.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/CodeGenContext.cs
index 1f5167e6..a4daaa67 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/CodeGenContext.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/CodeGenContext.cs
@@ -25,8 +25,8 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
 
         public Dictionary<int, Instruction> ConstantBuffers { get; } = new Dictionary<int, Instruction>();
         public Dictionary<int, Instruction> StorageBuffers { get; } = new Dictionary<int, Instruction>();
-        public Instruction LocalMemory { get; set; }
-        public Instruction SharedMemory { get; set; }
+        public Dictionary<int, Instruction> LocalMemories { get; } = new Dictionary<int, Instruction>();
+        public Dictionary<int, Instruction> SharedMemories { get; } = new Dictionary<int, Instruction>();
         public Dictionary<TextureMeta, SamplerType> SamplersTypes { get; } = new Dictionary<TextureMeta, SamplerType>();
         public Dictionary<TextureMeta, (Instruction, Instruction, Instruction)> Samplers { get; } = new Dictionary<TextureMeta, (Instruction, Instruction, Instruction)>();
         public Dictionary<TextureMeta, (Instruction, Instruction)> Images { get; } = new Dictionary<TextureMeta, (Instruction, Instruction)>();
@@ -35,7 +35,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
         public Dictionary<IoDefinition, Instruction> InputsPerPatch { get; } = new Dictionary<IoDefinition, Instruction>();
         public Dictionary<IoDefinition, Instruction> OutputsPerPatch { get; } = new Dictionary<IoDefinition, Instruction>();
 
-        public Instruction CoordTemp { get; set; }
         public StructuredFunction CurrentFunction { get; set; }
         private readonly Dictionary<AstOperand, Instruction> _locals = new Dictionary<AstOperand, Instruction>();
         private readonly Dictionary<int, Instruction[]> _localForArgs = new Dictionary<int, Instruction[]>();
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs
index eb2db514..59acea4f 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Declarations.cs
@@ -6,7 +6,6 @@ using Spv.Generator;
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
-using System.Linq;
 using System.Numerics;
 using static Spv.Specification;
 using SpvInstruction = Spv.Generator.Instruction;
@@ -44,13 +43,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
                 context.AddLocalVariable(spvLocal);
                 context.DeclareLocal(local, spvLocal);
             }
-
-            var ivector2Type = context.TypeVector(context.TypeS32(), 2);
-            var coordTempPointerType = context.TypePointer(StorageClass.Function, ivector2Type);
-            var coordTemp = context.Variable(coordTempPointerType, StorageClass.Function);
-
-            context.AddLocalVariable(coordTemp);
-            context.CoordTemp = coordTemp;
         }
 
         public static void DeclareLocalForArgs(CodeGenContext context, List<StructuredFunction> functions)
@@ -77,54 +69,30 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
 
         public static void DeclareAll(CodeGenContext context, StructuredProgramInfo info)
         {
-            if (context.Config.Stage == ShaderStage.Compute)
-            {
-                int localMemorySize = BitUtils.DivRoundUp(context.Config.GpuAccessor.QueryComputeLocalMemorySize(), 4);
-
-                if (localMemorySize != 0)
-                {
-                    DeclareLocalMemory(context, localMemorySize);
-                }
-
-                int sharedMemorySize = BitUtils.DivRoundUp(context.Config.GpuAccessor.QueryComputeSharedMemorySize(), 4);
-
-                if (sharedMemorySize != 0)
-                {
-                    DeclareSharedMemory(context, sharedMemorySize);
-                }
-            }
-            else if (context.Config.LocalMemorySize != 0)
-            {
-                int localMemorySize = BitUtils.DivRoundUp(context.Config.LocalMemorySize, 4);
-                DeclareLocalMemory(context, localMemorySize);
-            }
-
             DeclareConstantBuffers(context, context.Config.Properties.ConstantBuffers.Values);
             DeclareStorageBuffers(context, context.Config.Properties.StorageBuffers.Values);
+            DeclareMemories(context, context.Config.Properties.LocalMemories, context.LocalMemories, StorageClass.Private);
+            DeclareMemories(context, context.Config.Properties.SharedMemories, context.SharedMemories, StorageClass.Workgroup);
             DeclareSamplers(context, context.Config.GetTextureDescriptors());
             DeclareImages(context, context.Config.GetImageDescriptors());
             DeclareInputsAndOutputs(context, info);
         }
 
-        private static void DeclareLocalMemory(CodeGenContext context, int size)
-        {
-            context.LocalMemory = DeclareMemory(context, StorageClass.Private, size);
-        }
-
-        private static void DeclareSharedMemory(CodeGenContext context, int size)
+        private static void DeclareMemories(
+            CodeGenContext context,
+            IReadOnlyDictionary<int, MemoryDefinition> memories,
+            Dictionary<int, SpvInstruction> dict,
+            StorageClass storage)
         {
-            context.SharedMemory = DeclareMemory(context, StorageClass.Workgroup, size);
-        }
-
-        private static SpvInstruction DeclareMemory(CodeGenContext context, StorageClass storage, int size)
-        {
-            var arrayType = context.TypeArray(context.TypeU32(), context.Constant(context.TypeU32(), size));
-            var pointerType = context.TypePointer(storage, arrayType);
-            var variable = context.Variable(pointerType, storage);
+            foreach ((int id, MemoryDefinition memory) in memories)
+            {
+                var pointerType = context.TypePointer(storage, context.GetType(memory.Type, memory.ArrayLength));
+                var variable = context.Variable(pointerType, storage);
 
-            context.AddGlobalVariable(variable);
+                context.AddGlobalVariable(variable);
 
-            return variable;
+                dict.Add(id, variable);
+            }
         }
 
         private static void DeclareConstantBuffers(CodeGenContext context, IEnumerable<BufferDefinition> buffers)
diff --git a/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs b/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs
index 6c115752..b451f7a4 100644
--- a/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs
+++ b/src/Ryujinx.Graphics.Shader/CodeGen/Spirv/Instructions.cs
@@ -97,8 +97,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
             Add(Instruction.ImageStore,               GenerateImageStore);
             Add(Instruction.IsNan,                    GenerateIsNan);
             Add(Instruction.Load,                     GenerateLoad);
-            Add(Instruction.LoadLocal,                GenerateLoadLocal);
-            Add(Instruction.LoadShared,               GenerateLoadShared);
             Add(Instruction.Lod,                      GenerateLod);
             Add(Instruction.LogarithmB2,              GenerateLogarithmB2);
             Add(Instruction.LogicalAnd,               GenerateLogicalAnd);
@@ -132,10 +130,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
             Add(Instruction.Sine,                     GenerateSine);
             Add(Instruction.SquareRoot,               GenerateSquareRoot);
             Add(Instruction.Store,                    GenerateStore);
-            Add(Instruction.StoreLocal,               GenerateStoreLocal);
-            Add(Instruction.StoreShared,              GenerateStoreShared);
-            Add(Instruction.StoreShared16,            GenerateStoreShared16);
-            Add(Instruction.StoreShared8,             GenerateStoreShared8);
             Add(Instruction.Subtract,                 GenerateSubtract);
             Add(Instruction.SwizzleAdd,               GenerateSwizzleAdd);
             Add(Instruction.TextureSample,            GenerateTextureSample);
@@ -871,30 +865,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
             return GenerateLoadOrStore(context, operation, isStore: false);
         }
 
-        private static OperationResult GenerateLoadLocal(CodeGenContext context, AstOperation operation)
-        {
-            return GenerateLoadLocalOrShared(context, operation, StorageClass.Private, context.LocalMemory);
-        }
-
-        private static OperationResult GenerateLoadShared(CodeGenContext context, AstOperation operation)
-        {
-            return GenerateLoadLocalOrShared(context, operation, StorageClass.Workgroup, context.SharedMemory);
-        }
-
-        private static OperationResult GenerateLoadLocalOrShared(
-            CodeGenContext context,
-            AstOperation operation,
-            StorageClass storageClass,
-            SpvInstruction memory)
-        {
-            var offset = context.Get(AggregateType.S32, operation.GetSource(0));
-
-            var elemPointer = context.AccessChain(context.TypePointer(storageClass, context.TypeU32()), memory, offset);
-            var value = context.Load(context.TypeU32(), elemPointer);
-
-            return new OperationResult(AggregateType.U32, value);
-        }
-
         private static OperationResult GenerateLod(CodeGenContext context, AstOperation operation)
         {
             AstTextureOperation texOp = (AstTextureOperation)operation;
@@ -1268,45 +1238,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
             return GenerateLoadOrStore(context, operation, isStore: true);
         }
 
-        private static OperationResult GenerateStoreLocal(CodeGenContext context, AstOperation operation)
-        {
-            return GenerateStoreLocalOrShared(context, operation, StorageClass.Private, context.LocalMemory);
-        }
-
-        private static OperationResult GenerateStoreShared(CodeGenContext context, AstOperation operation)
-        {
-            return GenerateStoreLocalOrShared(context, operation, StorageClass.Workgroup, context.SharedMemory);
-        }
-
-        private static OperationResult GenerateStoreLocalOrShared(
-            CodeGenContext context,
-            AstOperation operation,
-            StorageClass storageClass,
-            SpvInstruction memory)
-        {
-            var offset = context.Get(AggregateType.S32, operation.GetSource(0));
-            var value = context.Get(AggregateType.U32, operation.GetSource(1));
-
-            var elemPointer = context.AccessChain(context.TypePointer(storageClass, context.TypeU32()), memory, offset);
-            context.Store(elemPointer, value);
-
-            return OperationResult.Invalid;
-        }
-
-        private static OperationResult GenerateStoreShared16(CodeGenContext context, AstOperation operation)
-        {
-            GenerateStoreSharedSmallInt(context, operation, 16);
-
-            return OperationResult.Invalid;
-        }
-
-        private static OperationResult GenerateStoreShared8(CodeGenContext context, AstOperation operation)
-        {
-            GenerateStoreSharedSmallInt(context, operation, 8);
-
-            return OperationResult.Invalid;
-        }
-
         private static OperationResult GenerateSubtract(CodeGenContext context, AstOperation operation)
         {
             return GenerateBinary(context, operation, context.Delegates.FSub, context.Delegates.ISub);
@@ -1827,55 +1758,27 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
             AstOperation operation,
             Func<SpvInstruction, SpvInstruction, SpvInstruction, SpvInstruction, SpvInstruction, SpvInstruction> emitU)
         {
-            var value = context.GetU32(operation.GetSource(operation.SourcesCount - 1));
+            SpvInstruction elemPointer = GetStoragePointer(context, operation, out AggregateType varType);
 
-            SpvInstruction elemPointer;
-
-            if (operation.StorageKind == StorageKind.StorageBuffer)
-            {
-                elemPointer = GetStoragePointer(context, operation, out _);
-            }
-            else if (operation.StorageKind == StorageKind.SharedMemory)
-            {
-                var offset = context.GetU32(operation.GetSource(0));
-                elemPointer = context.AccessChain(context.TypePointer(StorageClass.Workgroup, context.TypeU32()), context.SharedMemory, offset);
-            }
-            else
-            {
-                throw new InvalidOperationException($"Invalid storage kind \"{operation.StorageKind}\".");
-            }
+            var value = context.Get(varType, operation.GetSource(operation.SourcesCount - 1));
 
             var one = context.Constant(context.TypeU32(), 1);
             var zero = context.Constant(context.TypeU32(), 0);
 
-            return new OperationResult(AggregateType.U32, emitU(context.TypeU32(), elemPointer, one, zero, value));
+            return new OperationResult(varType, emitU(context.GetType(varType), elemPointer, one, zero, value));
         }
 
         private static OperationResult GenerateAtomicMemoryCas(CodeGenContext context, AstOperation operation)
         {
-            var value0 = context.GetU32(operation.GetSource(operation.SourcesCount - 2));
-            var value1 = context.GetU32(operation.GetSource(operation.SourcesCount - 1));
+            SpvInstruction elemPointer = GetStoragePointer(context, operation, out AggregateType varType);
 
-            SpvInstruction elemPointer;
-
-            if (operation.StorageKind == StorageKind.StorageBuffer)
-            {
-                elemPointer = GetStoragePointer(context, operation, out _);
-            }
-            else if (operation.StorageKind == StorageKind.SharedMemory)
-            {
-                var offset = context.GetU32(operation.GetSource(0));
-                elemPointer = context.AccessChain(context.TypePointer(StorageClass.Workgroup, context.TypeU32()), context.SharedMemory, offset);
-            }
-            else
-            {
-                throw new InvalidOperationException($"Invalid storage kind \"{operation.StorageKind}\".");
-            }
+            var value0 = context.Get(varType, operation.GetSource(operation.SourcesCount - 2));
+            var value1 = context.Get(varType, operation.GetSource(operation.SourcesCount - 1));
 
             var one = context.Constant(context.TypeU32(), 1);
             var zero = context.Constant(context.TypeU32(), 0);
 
-            return new OperationResult(AggregateType.U32, context.AtomicCompareExchange(context.TypeU32(), elemPointer, one, zero, zero, value1, value0));
+            return new OperationResult(varType, context.AtomicCompareExchange(context.GetType(varType), elemPointer, one, zero, zero, value1, value0));
         }
 
         private static OperationResult GenerateLoadOrStore(CodeGenContext context, AstOperation operation, bool isStore)
@@ -1928,6 +1831,27 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
                         : context.StorageBuffers[bindingIndex.Value];
                     break;
 
+                case StorageKind.LocalMemory:
+                case StorageKind.SharedMemory:
+                    if (!(operation.GetSource(srcIndex++) is AstOperand bindingId) || bindingId.Type != OperandType.Constant)
+                    {
+                        throw new InvalidOperationException($"First input of {operation.Inst} with {storageKind} storage must be a constant operand.");
+                    }
+
+                    if (storageKind == StorageKind.LocalMemory)
+                    {
+                        storageClass = StorageClass.Private;
+                        varType = context.Config.Properties.LocalMemories[bindingId.Value].Type & AggregateType.ElementTypeMask;
+                        baseObj = context.LocalMemories[bindingId.Value];
+                    }
+                    else
+                    {
+                        storageClass = StorageClass.Workgroup;
+                        varType = context.Config.Properties.SharedMemories[bindingId.Value].Type & AggregateType.ElementTypeMask;
+                        baseObj = context.SharedMemories[bindingId.Value];
+                    }
+                    break;
+
                 case StorageKind.Input:
                 case StorageKind.InputPerPatch:
                 case StorageKind.Output:
@@ -2048,50 +1972,6 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Spirv
             return context.Load(context.GetType(varType), context.Inputs[ioDefinition]);
         }
 
-        private static void GenerateStoreSharedSmallInt(CodeGenContext context, AstOperation operation, int bitSize)
-        {
-            var offset = context.Get(AggregateType.U32, operation.GetSource(0));
-            var value = context.Get(AggregateType.U32, operation.GetSource(1));
-
-            var wordOffset = context.ShiftRightLogical(context.TypeU32(), offset, context.Constant(context.TypeU32(), 2));
-            var bitOffset = context.BitwiseAnd(context.TypeU32(), offset, context.Constant(context.TypeU32(), 3));
-            bitOffset = context.ShiftLeftLogical(context.TypeU32(), bitOffset, context.Constant(context.TypeU32(), 3));
-
-            var memory = context.SharedMemory;
-
-            var elemPointer = context.AccessChain(context.TypePointer(StorageClass.Workgroup, context.TypeU32()), memory, wordOffset);
-
-            GenerateStoreSmallInt(context, elemPointer, bitOffset, value, bitSize);
-        }
-
-        private static void GenerateStoreSmallInt(
-            CodeGenContext context,
-            SpvInstruction elemPointer,
-            SpvInstruction bitOffset,
-            SpvInstruction value,
-            int bitSize)
-        {
-            var loopStart = context.Label();
-            var loopEnd = context.Label();
-
-            context.Branch(loopStart);
-            context.AddLabel(loopStart);
-
-            var oldValue = context.Load(context.TypeU32(), elemPointer);
-            var newValue = context.BitFieldInsert(context.TypeU32(), oldValue, value, bitOffset, context.Constant(context.TypeU32(), bitSize));
-
-            var one = context.Constant(context.TypeU32(), 1);
-            var zero = context.Constant(context.TypeU32(), 0);
-
-            var result = context.AtomicCompareExchange(context.TypeU32(), elemPointer, one, zero, zero, newValue, oldValue);
-            var failed = context.INotEqual(context.TypeBool(), result, oldValue);
-
-            context.LoopMerge(loopEnd, loopStart, LoopControlMask.MaskNone);
-            context.BranchConditional(failed, loopStart, loopEnd);
-
-            context.AddLabel(loopEnd);
-        }
-
         private static OperationResult GetZeroOperationResult(
             CodeGenContext context,
             AstTextureOperation texOp,