Add support for shader constant buffer slot indexing (#1608)

* Add support for shader constant buffer slot indexing

* Fix typo

Ryujinx.Graphics.Gpu/Memory/BufferManager.cs

@@ -147,6 +147,13 @@ namespace Ryujinx.Graphics.Gpu.Memory
             }
         }
 
+        /// <summary>
+        /// Sets a transform feedback buffer on the graphics pipeline.
+        /// The output from the vertex transformation stages are written into the feedback buffer.
+        /// </summary>
+        /// <param name="index">Index of the transform feedback buffer</param>
+        /// <param name="gpuVa">Start GPU virtual address of the buffer</param>
+        /// <param name="size">Size in bytes of the transform feedback buffer</param>
         public void SetTransformFeedbackBuffer(int index, ulong gpuVa, ulong size)
         {
             ulong address = TranslateAndCreateBuffer(gpuVa, size);
@@ -264,6 +271,25 @@ namespace Ryujinx.Graphics.Gpu.Memory
             _cpUniformBuffers.EnableMask = mask;
         }
 
+        /// <summary>
+        /// Gets a bit mask indicating which compute uniform buffers are currently bound.
+        /// </summary>
+        /// <returns>Mask where each bit set indicates a bound constant buffer</returns>
+        public uint GetComputeUniformBufferUseMask()
+        {
+            uint mask = 0;
+
+            for (int i = 0; i < _cpUniformBuffers.Buffers.Length; i++)
+            {
+                if (_cpUniformBuffers.Buffers[i].Address != 0)
+                {
+                    mask |= 1u << i;
+                }
+            }
+
+            return mask;
+        }
+
         /// <summary>
         /// Sets the enabled uniform buffers mask on the graphics pipeline.
         /// Each bit set on the mask indicates that the respective buffer index is enabled.
@@ -277,6 +303,26 @@ namespace Ryujinx.Graphics.Gpu.Memory
             _gpUniformBuffersDirty = true;
         }
 
+        /// <summary>
+        /// Gets a bit mask indicating which graphics uniform buffers are currently bound.
+        /// </summary>
+        /// <param name="stage">Index of the shader stage</param>
+        /// <returns>Mask where each bit set indicates a bound constant buffer</returns>
+        public uint GetGraphicsUniformBufferUseMask(int stage)
+        {
+            uint mask = 0;
+
+            for (int i = 0; i < _gpUniformBuffers[stage].Buffers.Length; i++)
+            {
+                if (_gpUniformBuffers[stage].Buffers[i].Address != 0)
+                {
+                    mask |= 1u << i;
+                }
+            }
+
+            return mask;
+        }
+
         /// <summary>
         /// Performs address translation of the GPU virtual address, and creates a
         /// new buffer, if needed, for the specified range.

Ryujinx.Graphics.Gpu/Shader/GpuAccessor.cs

@@ -123,6 +123,17 @@ namespace Ryujinx.Graphics.Gpu.Shader
         /// <returns>Shared Memory size in bytes</returns>
         public int QueryComputeSharedMemorySize() => _sharedMemorySize;
 
+        /// <summary>
+        /// Queries Constant Buffer usage information.
+        /// </summary>
+        /// <returns>A mask where each bit set indicates a bound constant buffer</returns>
+        public uint QueryConstantBufferUse()
+        {
+            return _compute
+                ? _context.Methods.BufferManager.GetComputeUniformBufferUseMask()
+                : _context.Methods.BufferManager.GetGraphicsUniformBufferUseMask(_stageIndex);
+        }
+
         /// <summary>
         /// Queries texture target information.
         /// </summary>

Ryujinx.Graphics.Shader/CodeGen/Glsl/CodeGenContext.cs

@@ -12,6 +12,8 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         public ShaderConfig Config { get; }
 
+        public bool CbIndexable { get; }
+
         public List<BufferDescriptor>  CBufferDescriptors { get; }
         public List<BufferDescriptor>  SBufferDescriptors { get; }
         public List<TextureDescriptor> TextureDescriptors { get; }
@@ -25,9 +27,10 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         private string _indentation;
 
-        public CodeGenContext(ShaderConfig config)
+        public CodeGenContext(ShaderConfig config, bool cbIndexable)
         {
             Config = config;
+            CbIndexable = cbIndexable;
 
             CBufferDescriptors = new List<BufferDescriptor>();
             SBufferDescriptors = new List<BufferDescriptor>();
@@ -85,9 +88,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             int cBufSlot = bindless ? operand.CbufSlot : 0;
             int cBufOffset = bindless ? operand.CbufOffset : 0;
 
-            return TextureDescriptors.FindIndex(descriptor => 
-                descriptor.Type == texOp.Type && 
-                descriptor.HandleIndex == texOp.Handle && 
+            return TextureDescriptors.FindIndex(descriptor =>
+                descriptor.Type == texOp.Type &&
+                descriptor.HandleIndex == texOp.Handle &&
                 descriptor.CbufSlot == cBufSlot &&
                 descriptor.CbufOffset == cBufOffset);
         }

Ryujinx.Graphics.Shader/CodeGen/Glsl/Declarations.cs

@@ -213,51 +213,57 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
 
         private static void DeclareUniforms(CodeGenContext context, StructuredProgramInfo info)
         {
-            foreach (int cbufSlot in info.CBuffers.OrderBy(x => x))
+            string ubSize = "[" + NumberFormatter.FormatInt(Constants.ConstantBufferSize / 16) + "]";
+
+            if (info.UsesCbIndexing)
             {
                 string ubName = OperandManager.GetShaderStagePrefix(context.Config.Stage);
 
-                ubName += "_" + DefaultNames.UniformNamePrefix + cbufSlot;
+                ubName += "_" + DefaultNames.UniformNamePrefix;
+
+                string blockName = $"{ubName}_{DefaultNames.BlockSuffix}";
+
+                int maxSlot = 0;
+
+                foreach (int cbufSlot in info.CBuffers.OrderBy(x => x))
+                {
+                    context.CBufferDescriptors.Add(new BufferDescriptor($"{blockName}[{cbufSlot}]", cbufSlot));
 
-                context.CBufferDescriptors.Add(new BufferDescriptor(ubName, cbufSlot));
+                    if (maxSlot < cbufSlot)
+                    {
+                        maxSlot = cbufSlot;
+                    }
+                }
 
-                context.AppendLine("layout (std140) uniform " + ubName);
+                context.AppendLine("layout (std140) uniform " + blockName);
 
                 context.EnterScope();
 
-                string ubSize = "[" + NumberFormatter.FormatInt(Constants.ConstantBufferSize / 16) + "]";
+                context.AppendLine("vec4 " + DefaultNames.DataName + ubSize + ";");
 
-                context.AppendLine("vec4 " + OperandManager.GetUbName(context.Config.Stage, cbufSlot) + ubSize + ";");
+                string arraySize = NumberFormatter.FormatInt(maxSlot + 1);
 
-                context.LeaveScope(";");
+                context.LeaveScope($" {ubName}[{arraySize}];");
             }
-        }
-
-        private static bool DeclareRenderScale(CodeGenContext context)
-        {
-            if ((context.Config.UsedFeatures & (FeatureFlags.FragCoordXY | FeatureFlags.IntegerSampling)) != 0)
+            else
             {
-                string stage = OperandManager.GetShaderStagePrefix(context.Config.Stage);
+                foreach (int cbufSlot in info.CBuffers.OrderBy(x => x))
+                {
+                    string ubName = OperandManager.GetShaderStagePrefix(context.Config.Stage);
 
-                int scaleElements = context.TextureDescriptors.Count;
+                    ubName += "_" + DefaultNames.UniformNamePrefix + cbufSlot;
 
-                if (context.Config.Stage == ShaderStage.Fragment)
-                {
-                    scaleElements++; // Also includes render target scale, for gl_FragCoord.
-                }
+                    context.CBufferDescriptors.Add(new BufferDescriptor(ubName, cbufSlot));
 
-                context.AppendLine($"uniform float {stage}_renderScale[{scaleElements}];");
+                    context.AppendLine("layout (std140) uniform " + ubName);
 
-                if (context.Config.UsedFeatures.HasFlag(FeatureFlags.IntegerSampling))
-                {
-                    context.AppendLine();
-                    AppendHelperFunction(context, $"Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/TexelFetchScale_{stage}.glsl");
-                }
+                    context.EnterScope();
 
-                return true;
-            }
+                    context.AppendLine("vec4 " + OperandManager.GetUbName(context.Config.Stage, cbufSlot, false) + ubSize + ";");
 
-            return false;
+                    context.LeaveScope(";");
+                }
+            }
         }
 
         private static void DeclareStorages(CodeGenContext context, StructuredProgramInfo info)
@@ -500,6 +506,33 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             }
         }
 
+        private static bool DeclareRenderScale(CodeGenContext context)
+        {
+            if ((context.Config.UsedFeatures & (FeatureFlags.FragCoordXY | FeatureFlags.IntegerSampling)) != 0)
+            {
+                string stage = OperandManager.GetShaderStagePrefix(context.Config.Stage);
+
+                int scaleElements = context.TextureDescriptors.Count;
+
+                if (context.Config.Stage == ShaderStage.Fragment)
+                {
+                    scaleElements++; // Also includes render target scale, for gl_FragCoord.
+                }
+
+                context.AppendLine($"uniform float {stage}_renderScale[{scaleElements}];");
+
+                if (context.Config.UsedFeatures.HasFlag(FeatureFlags.IntegerSampling))
+                {
+                    context.AppendLine();
+                    AppendHelperFunction(context, $"Ryujinx.Graphics.Shader/CodeGen/Glsl/HelperFunctions/TexelFetchScale_{stage}.glsl");
+                }
+
+                return true;
+            }
+
+            return false;
+        }
+
         private static void AppendHelperFunction(CodeGenContext context, string filename)
         {
             string code = EmbeddedResources.ReadAllText(filename);

Ryujinx.Graphics.Shader/CodeGen/Glsl/GlslGenerator.cs

@@ -12,7 +12,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
     {
         public static GlslProgram Generate(StructuredProgramInfo info, ShaderConfig config)
         {
-            CodeGenContext context = new CodeGenContext(config);
+            CodeGenContext context = new CodeGenContext(config, info.UsesCbIndexing);
 
             Declarations.Declare(context, info);
 

Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGen.cs

@@ -19,7 +19,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
             }
             else if (node is AstOperand operand)
             {
-                return context.OperandManager.GetExpression(operand, context.Config);
+                return context.OperandManager.GetExpression(operand, context.Config, context.CbIndexable);
             }
 
             throw new ArgumentException($"Invalid node type \"{node?.GetType().Name ?? "null"}\".");

Ryujinx.Graphics.Shader/CodeGen/Glsl/Instructions/InstGenMemory.cs

@@ -125,7 +125,16 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl.Instructions
 
             offsetExpr = Enclose(offsetExpr, src2, Instruction.ShiftRightS32, isLhs: true);
 
-            return OperandManager.GetConstantBufferName(src1, offsetExpr, context.Config.Stage);
+            if (src1 is AstOperand oper && oper.Type == OperandType.Constant)
+            {
+                return OperandManager.GetConstantBufferName(oper.Value, offsetExpr, context.Config.Stage, context.CbIndexable);
+            }
+            else
+            {
+                string slotExpr = GetSoureExpr(context, src1, GetSrcVarType(operation.Inst, 0));
+
+                return OperandManager.GetConstantBufferName(slotExpr, offsetExpr, context.Config.Stage);
+            }
         }
 
         public static string LoadLocal(CodeGenContext context, AstOperation operation)

Ryujinx.Graphics.Shader/CodeGen/Glsl/OperandManager.cs

@@ -92,7 +92,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return name;
         }
 
-        public string GetExpression(AstOperand operand, ShaderConfig config)
+        public string GetExpression(AstOperand operand, ShaderConfig config, bool cbIndexable)
         {
             switch (operand.Type)
             {
@@ -103,7 +103,7 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
                     return NumberFormatter.FormatInt(operand.Value);
 
                 case OperandType.ConstantBuffer:
-                    return GetConstantBufferName(operand.CbufSlot, operand.CbufOffset, config.Stage);
+                    return GetConstantBufferName(operand.CbufSlot, operand.CbufOffset, config.Stage, cbIndexable);
 
                 case OperandType.LocalVariable:
                     return _locals[operand];
@@ -115,13 +115,9 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             throw new ArgumentException($"Invalid operand type \"{operand.Type}\".");
         }
 
-        public static string GetConstantBufferName(int slot, int offset, ShaderStage stage)
+        public static string GetConstantBufferName(int slot, int offset, ShaderStage stage, bool cbIndexable)
         {
-            string ubName = GetUbName(stage, slot);
-
-            ubName += "[" + (offset >> 2) + "]";
-
-            return ubName + "." + GetSwizzleMask(offset & 3);
+            return $"{GetUbName(stage, slot, cbIndexable)}[{offset >> 2}].{GetSwizzleMask(offset & 3)}";
         }
 
         private static string GetVec4Indexed(string vectorName, string indexExpr)
@@ -134,18 +130,14 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return $"({result})";
         }
 
-        public static string GetConstantBufferName(IAstNode slot, string offsetExpr, ShaderStage stage)
+        public static string GetConstantBufferName(int slot, string offsetExpr, ShaderStage stage, bool cbIndexable)
         {
-            // Non-constant slots are not supported.
-            // It is expected that upstream stages are never going to generate non-constant
-            // slot access.
-            AstOperand operand = (AstOperand)slot;
-
-            string ubName = GetUbName(stage, operand.Value);
-
-            string index0 = "[" + offsetExpr + " >> 2]";
+            return GetVec4Indexed(GetUbName(stage, slot, cbIndexable) + $"[{offsetExpr} >> 2]", offsetExpr + " & 3");
+        }
 
-            return GetVec4Indexed(ubName + index0, offsetExpr + " & 3");
+        public static string GetConstantBufferName(string slotExpr, string offsetExpr, ShaderStage stage)
+        {
+            return GetVec4Indexed(GetUbName(stage, slotExpr) + $"[{offsetExpr} >> 2]", offsetExpr + " & 3");
         }
 
         public static string GetOutAttributeName(AstOperand attr, ShaderConfig config)
@@ -228,13 +220,19 @@ namespace Ryujinx.Graphics.Shader.CodeGen.Glsl
             return isOutAttr ? "// bad_attr0x" + value.ToString("X") : "0.0";
         }
 
-        public static string GetUbName(ShaderStage stage, int slot)
+        public static string GetUbName(ShaderStage stage, int slot, bool cbIndexable)
         {
-            string ubName = GetShaderStagePrefix(stage);
+            if (cbIndexable)
+            {
+                return GetUbName(stage, NumberFormatter.FormatInt(slot, VariableType.S32));
+            }
 
-            ubName += "_" + DefaultNames.UniformNamePrefix + slot;
+            return $"{GetShaderStagePrefix(stage)}_{DefaultNames.UniformNamePrefix}{slot}_{DefaultNames.UniformNameSuffix}";
+        }
 
-            return ubName + "_" + DefaultNames.UniformNameSuffix;
+        private static string GetUbName(ShaderStage stage, string slotExpr)
+        {
+            return $"{GetShaderStagePrefix(stage)}_{DefaultNames.UniformNamePrefix}[{slotExpr}].{DefaultNames.DataName}";
         }
 
         public static string GetSamplerName(ShaderStage stage, AstTextureOperation texOp, string indexExpr)

Ryujinx.Graphics.Shader/Decoders/CbIndexMode.cs

@@ -0,0 +1,10 @@
+namespace Ryujinx.Graphics.Shader.Decoders
+{
+    enum CbIndexMode
+    {
+        Default = 0,
+        Il      = 1,
+        Is      = 2,
+        Isl     = 3
+    }
+}

Ryujinx.Graphics.Shader/Decoders/OpCodeLdc.cs

@@ -10,6 +10,7 @@ namespace Ryujinx.Graphics.Shader.Decoders
         public int Offset { get; }
         public int Slot   { get; }
 
+        public CbIndexMode IndexMode { get; }
         public IntegerSize Size { get; }
 
         public OpCodeLdc(InstEmitter emitter, ulong address, long opCode) : base(emitter, address, opCode)
@@ -20,7 +21,8 @@ namespace Ryujinx.Graphics.Shader.Decoders
             Offset = (opCode.Extract(20, 16) << 16) >> 16;
             Slot   = opCode.Extract(36, 5);
 
-            Size = (IntegerSize)opCode.Extract(48, 3);
+            IndexMode = (CbIndexMode)opCode.Extract(44, 2);
+            Size      = (IntegerSize)opCode.Extract(48, 3);
         }
     }
 }

Ryujinx.Graphics.Shader/IGpuAccessor.cs

@@ -39,6 +39,11 @@
             return 0xc000;
         }
 
+        public uint QueryConstantBufferUse()
+        {
+            return 0xffff;
+        }
+
         public bool QueryIsTextureBuffer(int handle)
         {
             return false;

Ryujinx.Graphics.Shader/Instructions/InstEmitMemory.cs

@@ -112,7 +112,7 @@ namespace Ryujinx.Graphics.Shader.Instructions
                     res = context.FPMultiply(res, Attribute(AttributeConsts.PositionW));
                 }
             }
- 
+
             if (op.Mode == InterpolationMode.Default)
             {
                 Operand srcB = GetSrcB(context);
@@ -152,7 +152,17 @@ namespace Ryujinx.Graphics.Shader.Instructions
 
             int count = op.Size == IntegerSize.B64 ? 2 : 1;
 
-            Operand addr = context.IAdd(GetSrcA(context), Const(op.Offset));
+            Operand slot = Const(op.Slot);
+            Operand srcA = GetSrcA(context);
+
+            if (op.IndexMode == CbIndexMode.Is ||
+                op.IndexMode == CbIndexMode.Isl)
+            {
+                slot = context.IAdd(slot, context.BitfieldExtractU32(srcA, Const(16), Const(16)));
+                srcA = context.BitwiseAnd(srcA, Const(0xffff));
+            }
+
+            Operand addr = context.IAdd(srcA, Const(op.Offset));
 
             Operand wordOffset = context.ShiftRightU32(addr, Const(2));
 
@@ -169,7 +179,7 @@ namespace Ryujinx.Graphics.Shader.Instructions
 
                 Operand offset = context.IAdd(wordOffset, Const(index));
 
-                Operand value = context.LoadConstant(Const(op.Slot), offset);
+                Operand value = context.LoadConstant(slot, offset);
 
                 if (isSmallInt)
                 {

Ryujinx.Graphics.Shader/StructuredIr/StructuredProgram.cs

@@ -2,6 +2,7 @@ using Ryujinx.Graphics.Shader.IntermediateRepresentation;
 using Ryujinx.Graphics.Shader.Translation;
 using System;
 using System.Collections.Generic;
+using System.Numerics;
 
 namespace Ryujinx.Graphics.Shader.StructuredIr
 {
@@ -73,12 +74,24 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
                 {
                     Operand slot = operation.GetSource(0);
 
-                    if (slot.Type != OperandType.Constant)
+                    if (slot.Type == OperandType.Constant)
                     {
-                        throw new InvalidOperationException("Found load with non-constant constant buffer slot.");
+                        context.Info.CBuffers.Add(slot.Value);
                     }
+                    else
+                    {
+                        // If the value is not constant, then we don't know
+                        // how many constant buffers are used, so we assume
+                        // all of them are used.
+                        int cbCount = 32 - BitOperations.LeadingZeroCount(context.Config.GpuAccessor.QueryConstantBufferUse());
+
+                        for (int index = 0; index < cbCount; index++)
+                        {
+                            context.Info.CBuffers.Add(index);
+                        }
 
-                    context.Info.CBuffers.Add(slot.Value);
+                        context.Info.UsesCbIndexing = true;
+                    }
                 }
                 else if (UsesStorage(inst))
                 {

Ryujinx.Graphics.Shader/StructuredIr/StructuredProgramInfo.cs

@@ -15,6 +15,7 @@ namespace Ryujinx.Graphics.Shader.StructuredIr
         public HashSet<int> OAttributes { get; }
 
         public bool UsesInstanceId { get; set; }
+        public bool UsesCbIndexing { get; set; }
 
         public HelperFunctionsMask HelperFunctionsMask { get; set; }