shader_ir: Implement BRX & BRA.CC

6 files changed, 76 insertions, 4 deletions

diff --git a/src/video_core/engines/shader_bytecode.h b/src/video_core/engines/shader_bytecode.h
index 404d4f5aa..c3055602b 100644
--- a/src/video_core/engines/shader_bytecode.h
+++ b/src/video_core/engines/shader_bytecode.h
@@ -1368,6 +1368,20 @@ union Instruction {
     } bra;
 
     union {
+        BitField<20, 24, u64> target;
+        BitField<5, 1, u64> constant_buffer;
+
+        s32 GetBranchExtend() const {
+            // Sign extend the branch target offset
+            u32 mask = 1U << (24 - 1);
+            u32 value = static_cast<u32>(target);
+            // The branch offset is relative to the next instruction and is stored in bytes, so
+            // divide it by the size of an instruction and add 1 to it.
+            return static_cast<s32>((value ^ mask) - mask) / sizeof(Instruction) + 1;
+        }
+    } brx;
+
+    union {
         BitField<39, 1, u64> emit; // EmitVertex
         BitField<40, 1, u64> cut;  // EndPrimitive
     } out;
@@ -1464,6 +1478,7 @@ public:
         BFE_IMM,
         BFI_IMM_R,
         BRA,
+        BRX,
         PBK,
         LD_A,
         LD_L,
@@ -1738,6 +1753,7 @@ private:
             INST("111000101001----", Id::SSY, Type::Flow, "SSY"),
             INST("111000101010----", Id::PBK, Type::Flow, "PBK"),
             INST("111000100100----", Id::BRA, Type::Flow, "BRA"),
+            INST("111000100101----", Id::BRX, Type::Flow, "BRX"),
             INST("1111000011111---", Id::SYNC, Type::Flow, "SYNC"),
             INST("111000110100---", Id::BRK, Type::Flow, "BRK"),
             INST("111000110000----", Id::EXIT, Type::Flow, "EXIT"),
diff --git a/src/video_core/renderer_opengl/gl_shader_decompiler.cpp b/src/video_core/renderer_opengl/gl_shader_decompiler.cpp
index 5f2f1510c..cedfe30b1 100644
--- a/src/video_core/renderer_opengl/gl_shader_decompiler.cpp
+++ b/src/video_core/renderer_opengl/gl_shader_decompiler.cpp
@@ -1555,6 +1555,14 @@ private:
         return {};
     }
 
+    std::string BranchIndirect(Operation operation) {
+        const std::string op_a = VisitOperand(operation, 0, Type::Uint);
+
+        code.AddLine("jmp_to = {};", op_a);
+        code.AddLine("break;");
+        return {};
+    }
+
     std::string PushFlowStack(Operation operation) {
         const auto stack = std::get<MetaStackClass>(operation.GetMeta());
         const auto target = std::get_if<ImmediateNode>(&*operation[0]);
@@ -1789,6 +1797,7 @@ private:
         &GLSLDecompiler::ImageStore,
 
         &GLSLDecompiler::Branch,
+        &GLSLDecompiler::BranchIndirect,
         &GLSLDecompiler::PushFlowStack,
         &GLSLDecompiler::PopFlowStack,
         &GLSLDecompiler::Exit,
diff --git a/src/video_core/renderer_vulkan/vk_shader_decompiler.cpp b/src/video_core/renderer_vulkan/vk_shader_decompiler.cpp
index 97ce214b1..1bb04607b 100644
--- a/src/video_core/renderer_vulkan/vk_shader_decompiler.cpp
+++ b/src/video_core/renderer_vulkan/vk_shader_decompiler.cpp
@@ -949,6 +949,14 @@ private:
         return {};
     }
 
+    Id BranchIndirect(Operation operation) {
+        const Id op_a = VisitOperand<Type::Uint>(operation, 0);
+
+        Emit(OpStore(jmp_to, op_a));
+        BranchingOp([&]() { Emit(OpBranch(continue_label)); });
+        return {};
+    }
+
     Id PushFlowStack(Operation operation) {
         const auto target = std::get_if<ImmediateNode>(&*operation[0]);
         ASSERT(target);
@@ -1334,6 +1342,7 @@ private:
         &SPIRVDecompiler::ImageStore,
 
         &SPIRVDecompiler::Branch,
+        &SPIRVDecompiler::BranchIndirect,
         &SPIRVDecompiler::PushFlowStack,
         &SPIRVDecompiler::PopFlowStack,
         &SPIRVDecompiler::Exit,
diff --git a/src/video_core/shader/control_flow.cpp b/src/video_core/shader/control_flow.cpp
index fcf22c7f2..a9de8f814 100644
--- a/src/video_core/shader/control_flow.cpp
+++ b/src/video_core/shader/control_flow.cpp
@@ -284,6 +284,9 @@ ParseResult ParseCode(CFGRebuildState& state, u32 address, ParseInfo& parse_info
             state.pbk_labels.emplace(offset, target);
             break;
         }
+        case OpCode::Id::BRX: {
+            return ParseResult::AbnormalFlow;
+        }
         default:
             break;
         }
diff --git a/src/video_core/shader/decode/other.cpp b/src/video_core/shader/decode/other.cpp
index d46a8ab82..ed3c63781 100644
--- a/src/video_core/shader/decode/other.cpp
+++ b/src/video_core/shader/decode/other.cpp
@@ -91,11 +91,45 @@ u32 ShaderIR::DecodeOther(NodeBlock& bb, u32 pc) {
         break;
     }
     case OpCode::Id::BRA: {
-        UNIMPLEMENTED_IF_MSG(instr.bra.constant_buffer != 0,
-                             "BRA with constant buffers are not implemented");
+        Node branch;
+        if (instr.bra.constant_buffer == 0) {
+            const u32 target = pc + instr.bra.GetBranchTarget();
+            branch = Operation(OperationCode::Branch, Immediate(target));
+        } else {
+            const u32 target = pc + 1;
+            const Node op_a = GetConstBuffer(instr.cbuf36.index, instr.cbuf36.GetOffset());
+            const Node convert = SignedOperation(OperationCode::IArithmeticShiftRight,
+                                               true, PRECISE, op_a, Immediate(3));
+            const Node operand = Operation(OperationCode::IAdd, PRECISE, convert, Immediate(target));
+            branch = Operation(OperationCode::BranchIndirect, convert);
+        }
 
-        const u32 target = pc + instr.bra.GetBranchTarget();
-        const Node branch = Operation(OperationCode::Branch, Immediate(target));
+        const Tegra::Shader::ConditionCode cc = instr.flow_condition_code;
+        if (cc != Tegra::Shader::ConditionCode::T) {
+            bb.push_back(Conditional(GetConditionCode(cc), {branch}));
+        } else {
+            bb.push_back(branch);
+        }
+        break;
+    }
+    case OpCode::Id::BRX: {
+        Node operand;
+        if (instr.brx.constant_buffer != 0) {
+            const s32 target = pc + 1;
+            const Node index = GetRegister(instr.gpr8);
+            const Node op_a =
+                GetConstBufferIndirect(instr.cbuf36.index, instr.cbuf36.GetOffset() + 0, index);
+            const Node convert = SignedOperation(OperationCode::IArithmeticShiftRight,
+                                               true, PRECISE, op_a, Immediate(3));
+            operand = Operation(OperationCode::IAdd, PRECISE, convert, Immediate(target));
+        } else {
+            const s32 target = pc + instr.brx.GetBranchExtend();
+            const Node op_a = GetRegister(instr.gpr8);
+            const Node convert = SignedOperation(OperationCode::IArithmeticShiftRight,
+                                               true, PRECISE, op_a, Immediate(3));
+            operand = Operation(OperationCode::IAdd, PRECISE, convert, Immediate(target));
+        }
+        const Node branch = Operation(OperationCode::BranchIndirect, operand);
 
         const Tegra::Shader::ConditionCode cc = instr.flow_condition_code;
         if (cc != Tegra::Shader::ConditionCode::T) {
diff --git a/src/video_core/shader/node.h b/src/video_core/shader/node.h
index 0ac83fcf0..e468758a6 100644
--- a/src/video_core/shader/node.h
+++ b/src/video_core/shader/node.h
@@ -149,6 +149,7 @@ enum class OperationCode {
     ImageStore, /// (MetaImage, float[N] coords) -> void
 
     Branch,        /// (uint branch_target) -> void
+    BranchIndirect,/// (uint branch_target) -> void
     PushFlowStack, /// (uint branch_target) -> void
     PopFlowStack,  /// () -> void
     Exit,          /// () -> void