shader: add support for const buffer indirect addressing

2 files changed, 68 insertions, 18 deletions

diff --git a/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp b/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
index 8ea730c80..1cfe1d49f 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv_context_get_set.cpp
@@ -124,25 +124,56 @@ std::optional<OutAttr> OutputAttrPointer(EmitContext& ctx, IR::Attribute attr) {
 
 Id GetCbuf(EmitContext& ctx, Id result_type, Id UniformDefinitions::*member_ptr, u32 element_size,
            const IR::Value& binding, const IR::Value& offset) {
-    if (!binding.IsImmediate()) {
-        throw NotImplementedException("Constant buffer indexing");
-    }
-    const Id cbuf{ctx.cbufs[binding.U32()].*member_ptr};
+    std::array<Id, 2> indexes;
+
     const Id uniform_type{ctx.uniform_types.*member_ptr};
-    if (!offset.IsImmediate()) {
+    if (offset.IsImmediate()) {
+        // Hardware been proved to read the aligned offset (e.g. LDC.U32 at 6 will read offset 4)
+        const Id imm_offset{ctx.Const(offset.U32() / element_size)};
+        indexes = {ctx.u32_zero_value, imm_offset};
+    } else {
         Id index{ctx.Def(offset)};
         if (element_size > 1) {
             const u32 log2_element_size{static_cast<u32>(std::countr_zero(element_size))};
             const Id shift{ctx.Const(log2_element_size)};
             index = ctx.OpShiftRightArithmetic(ctx.U32[1], ctx.Def(offset), shift);
         }
-        const Id access_chain{ctx.OpAccessChain(uniform_type, cbuf, ctx.u32_zero_value, index)};
+        indexes = {ctx.u32_zero_value, index};
+    }
+
+    if (binding.IsImmediate()) {
+        const Id cbuf{ctx.cbufs[binding.U32()].*member_ptr};
+        const Id access_chain{ctx.OpAccessChain(uniform_type, cbuf, indexes)};
         return ctx.OpLoad(result_type, access_chain);
+    } else {
+        const Id index{ctx.Def(binding)};
+        const Id ptr{ctx.TypePointer(spv::StorageClass::Function, result_type)};
+        const Id value{ctx.AddLocalVariable(ptr, spv::StorageClass::Function)};
+        const Id merge_label = ctx.OpLabel();
+
+        std::array<Id, Info::MAX_CBUFS> buf_labels;
+        std::array<Sirit::Literal, Info::MAX_CBUFS> buf_literals;
+        for (u32 i = 0; i < Info::MAX_CBUFS; i++) {
+            buf_labels[i] = ctx.OpLabel();
+            buf_literals[i] = Sirit::Literal{i};
+        }
+
+        ctx.OpSelectionMerge(merge_label, spv::SelectionControlMask::MaskNone);
+        ctx.OpSwitch(index, buf_labels[0], buf_literals, buf_labels);
+
+        for (u32 i = 0; i < Info::MAX_CBUFS; i++) {
+            ctx.AddLabel(buf_labels[i]);
+            const Id cbuf{ctx.cbufs[i].*member_ptr};
+            const Id access_chain{ctx.OpAccessChain(uniform_type, cbuf, indexes)};
+            const Id result = ctx.OpLoad(result_type, access_chain);
+            ctx.OpStore(value, result);
+            ctx.OpBranch(merge_label);
+        }
+
+        ctx.AddLabel(merge_label);
+
+        return ctx.OpLoad(result_type, value);
     }
-    // Hardware been proved to read the aligned offset (e.g. LDC.U32 at 6 will read offset 4)
-    const Id imm_offset{ctx.Const(offset.U32() / element_size)};
-    const Id access_chain{ctx.OpAccessChain(uniform_type, cbuf, ctx.u32_zero_value, imm_offset)};
-    return ctx.OpLoad(result_type, access_chain);
 }
 
 Id GetCbufU32(EmitContext& ctx, const IR::Value& binding, const IR::Value& offset) {
diff --git a/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp b/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
index bfd2ae650..1a50dd382 100644
--- a/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
+++ b/src/shader_recompiler/ir_opt/collect_shader_info_pass.cpp
@@ -29,6 +29,20 @@ void AddConstantBufferDescriptor(Info& info, u32 index, u32 count) {
                  });
 }
 
+void AddRegisterIndexedLdc(Info& info) {
+    // The shader can use any possible constant buffer
+    info.constant_buffer_mask = (1 << Info::MAX_CBUFS) - 1;
+
+    auto& cbufs{info.constant_buffer_descriptors};
+    cbufs.clear();
+    for (u32 i = 0; i < Info::MAX_CBUFS; i++) {
+        cbufs.push_back(ConstantBufferDescriptor{.index = i, .count = 1});
+
+        // The shader can use any possible access size
+        info.constant_buffer_used_sizes[i] = 0x10'000;
+    }
+}
+
 void GetPatch(Info& info, IR::Patch patch) {
     if (!IR::IsGeneric(patch)) {
         throw NotImplementedException("Reading non-generic patch {}", patch);
@@ -463,10 +477,12 @@ void VisitUsages(Info& info, IR::Inst& inst) {
     case IR::Opcode::GetCbufU32x2: {
         const IR::Value index{inst.Arg(0)};
         const IR::Value offset{inst.Arg(1)};
-        if (!index.IsImmediate()) {
-            throw NotImplementedException("Constant buffer with non-immediate index");
+        if (index.IsImmediate()) {
+            AddConstantBufferDescriptor(info, index.U32(), 1);
+        } else {
+            AddRegisterIndexedLdc(info);
         }
-        AddConstantBufferDescriptor(info, index.U32(), 1);
+
         u32 element_size{};
         switch (inst.GetOpcode()) {
         case IR::Opcode::GetCbufU8:
@@ -494,11 +510,14 @@ void VisitUsages(Info& info, IR::Inst& inst) {
         default:
             break;
         }
-        u32& size{info.constant_buffer_used_sizes[index.U32()]};
-        if (offset.IsImmediate()) {
-            size = Common::AlignUp(std::max(size, offset.U32() + element_size), 16u);
-        } else {
-            size = 0x10'000;
+
+        if (index.IsImmediate()) {
+            u32& size{info.constant_buffer_used_sizes[index.U32()]};
+            if (offset.IsImmediate()) {
+                size = Common::AlignUp(std::max(size, offset.U32() + element_size), 16u);
+            } else {
+                size = 0x10'000;
+            }
         }
         break;
     }