shader: SSA and dominance

8 files changed, 223 insertions, 55 deletions

diff --git a/src/shader_recompiler/frontend/maxwell/control_flow.cpp b/src/shader_recompiler/frontend/maxwell/control_flow.cpp
index fc4dba826..21ee98137 100644
--- a/src/shader_recompiler/frontend/maxwell/control_flow.cpp
+++ b/src/shader_recompiler/frontend/maxwell/control_flow.cpp
@@ -36,6 +36,7 @@ static std::array<Block, 2> Split(Block&& block, Location pc, BlockId new_id) {
             .cond{true},
             .branch_true{new_id},
             .branch_false{UNREACHABLE_BLOCK_ID},
+            .imm_predecessors{},
         },
         Block{
             .begin{pc},
@@ -46,6 +47,7 @@ static std::array<Block, 2> Split(Block&& block, Location pc, BlockId new_id) {
             .cond{block.cond},
             .branch_true{block.branch_true},
             .branch_false{block.branch_false},
+            .imm_predecessors{},
         },
     };
 }
@@ -108,7 +110,7 @@ static bool HasFlowTest(Opcode opcode) {
     }
 }
 
-static std::string Name(const Block& block) {
+static std::string NameOf(const Block& block) {
     if (block.begin.IsVirtual()) {
         return fmt::format("\"Virtual {}\"", block.id);
     } else {
@@ -154,13 +156,127 @@ bool Block::Contains(Location pc) const noexcept {
 }
 
 Function::Function(Location start_address)
-    : entrypoint{start_address}, labels{Label{
+    : entrypoint{start_address}, labels{{
                                      .address{start_address},
                                      .block_id{0},
                                      .stack{},
                                  }} {}
 
+void Function::BuildBlocksMap() {
+    const size_t num_blocks{NumBlocks()};
+    blocks_map.resize(num_blocks);
+    for (size_t block_index = 0; block_index < num_blocks; ++block_index) {
+        Block& block{blocks_data[block_index]};
+        blocks_map[block.id] = &block;
+    }
+}
+
+void Function::BuildImmediatePredecessors() {
+    for (const Block& block : blocks_data) {
+        if (block.branch_true != UNREACHABLE_BLOCK_ID) {
+            blocks_map[block.branch_true]->imm_predecessors.push_back(block.id);
+        }
+        if (block.branch_false != UNREACHABLE_BLOCK_ID) {
+            blocks_map[block.branch_false]->imm_predecessors.push_back(block.id);
+        }
+    }
+}
+
+void Function::BuildPostOrder() {
+    boost::container::small_vector<BlockId, 0x110> block_stack;
+    post_order_map.resize(NumBlocks());
+
+    Block& first_block{blocks_data[blocks.front()]};
+    first_block.post_order_visited = true;
+    block_stack.push_back(first_block.id);
+
+    const auto visit_branch = [&](BlockId block_id, BlockId branch_id) {
+        if (branch_id == UNREACHABLE_BLOCK_ID) {
+            return false;
+        }
+        if (blocks_map[branch_id]->post_order_visited) {
+            return false;
+        }
+        blocks_map[branch_id]->post_order_visited = true;
+
+        // Calling push_back twice is faster than insert on msvc
+        block_stack.push_back(block_id);
+        block_stack.push_back(branch_id);
+        return true;
+    };
+    while (!block_stack.empty()) {
+        const Block* const block{blocks_map[block_stack.back()]};
+        block_stack.pop_back();
+
+        if (!visit_branch(block->id, block->branch_true) &&
+            !visit_branch(block->id, block->branch_false)) {
+            post_order_map[block->id] = static_cast<u32>(post_order_blocks.size());
+            post_order_blocks.push_back(block->id);
+        }
+    }
+}
+
+void Function::BuildImmediateDominators() {
+    auto transform_block_id{std::views::transform([this](BlockId id) { return blocks_map[id]; })};
+    auto reverse_order_but_first{std::views::reverse | std::views::drop(1) | transform_block_id};
+    auto has_idom{std::views::filter([](Block* block) { return block->imm_dominator; })};
+    auto intersect{[this](Block* finger1, Block* finger2) {
+        while (finger1 != finger2) {
+            while (post_order_map[finger1->id] < post_order_map[finger2->id]) {
+                finger1 = finger1->imm_dominator;
+            }
+            while (post_order_map[finger2->id] < post_order_map[finger1->id]) {
+                finger2 = finger2->imm_dominator;
+            }
+        }
+        return finger1;
+    }};
+    for (Block& block : blocks_data) {
+        block.imm_dominator = nullptr;
+    }
+    Block* const start_block{&blocks_data[blocks.front()]};
+    start_block->imm_dominator = start_block;
+
+    bool changed{true};
+    while (changed) {
+        changed = false;
+        for (Block* const block : post_order_blocks | reverse_order_but_first) {
+            Block* new_idom{};
+            for (Block* predecessor : block->imm_predecessors | transform_block_id | has_idom) {
+                new_idom = new_idom ? intersect(predecessor, new_idom) : predecessor;
+            }
+            changed |= block->imm_dominator != new_idom;
+            block->imm_dominator = new_idom;
+        }
+    }
+}
+
+void Function::BuildDominanceFrontier() {
+    auto transform_block_id{std::views::transform([this](BlockId id) { return blocks_map[id]; })};
+    auto has_enough_predecessors{[](Block& block) { return block.imm_predecessors.size() >= 2; }};
+    for (Block& block : blocks_data | std::views::filter(has_enough_predecessors)) {
+        for (Block* current : block.imm_predecessors | transform_block_id) {
+            while (current != block.imm_dominator) {
+                current->dominance_frontiers.push_back(current->id);
+                current = current->imm_dominator;
+            }
+        }
+    }
+}
+
 CFG::CFG(Environment& env_, Location start_address) : env{env_} {
+    VisitFunctions(start_address);
+
+    for (Function& function : functions) {
+        function.BuildBlocksMap();
+        function.BuildImmediatePredecessors();
+        function.BuildPostOrder();
+        function.BuildImmediateDominators();
+        function.BuildDominanceFrontier();
+    }
+}
+
+void CFG::VisitFunctions(Location start_address) {
     functions.emplace_back(start_address);
     for (FunctionId function_id = 0; function_id < functions.size(); ++function_id) {
         while (!functions[function_id].labels.empty()) {
@@ -202,6 +318,7 @@ void CFG::AnalyzeLabel(FunctionId function_id, Label& label) {
         .cond{true},
         .branch_true{UNREACHABLE_BLOCK_ID},
         .branch_false{UNREACHABLE_BLOCK_ID},
+        .imm_predecessors{},
     };
     // Analyze instructions until it reaches an already visited block or there's a branch
     bool is_branch{false};
@@ -310,7 +427,7 @@ CFG::AnalysisState CFG::AnalyzeInst(Block& block, FunctionId function_id, Locati
         // Technically CAL pushes into PRET, but that's implicit in the function call for us
         // Insert the function into the list if it doesn't exist
         if (std::ranges::find(functions, cal_pc, &Function::entrypoint) == functions.end()) {
-            functions.push_back(cal_pc);
+            functions.emplace_back(cal_pc);
         }
         // Handle CAL like a regular instruction
         break;
@@ -352,6 +469,7 @@ void CFG::AnalyzeCondInst(Block& block, FunctionId function_id, Location pc,
         .cond{cond},
         .branch_true{conditional_block_id},
         .branch_false{UNREACHABLE_BLOCK_ID},
+        .imm_predecessors{},
     })};
     // Set the end properties of the conditional instruction and give it a new identity
     Block& conditional_block{block};
@@ -465,14 +583,14 @@ std::string CFG::Dot() const {
         dot += fmt::format("\t\tnode [style=filled];\n");
         for (const u32 block_index : function.blocks) {
             const Block& block{function.blocks_data[block_index]};
-            const std::string name{Name(block)};
+            const std::string name{NameOf(block)};
             const auto add_branch = [&](BlockId branch_id, bool add_label) {
                 const auto it{std::ranges::find(function.blocks_data, branch_id, &Block::id)};
                 dot += fmt::format("\t\t{}->", name);
                 if (it == function.blocks_data.end()) {
                     dot += fmt::format("\"Unknown label {}\"", branch_id);
                 } else {
-                    dot += Name(*it);
+                    dot += NameOf(*it);
                 };
                 if (add_label && block.cond != true && block.cond != false) {
                     dot += fmt::format(" [label=\"{}\"]", block.cond);
@@ -520,7 +638,7 @@ std::string CFG::Dot() const {
         if (functions.front().blocks.empty()) {
             dot += "Start;\n";
         } else {
-            dot += fmt::format("\tStart -> {};\n", Name(functions.front().blocks_data.front()));
+            dot += fmt::format("\tStart -> {};\n", NameOf(functions.front().blocks_data.front()));
         }
         dot += fmt::format("\tStart [shape=diamond];\n");
     }
diff --git a/src/shader_recompiler/frontend/maxwell/control_flow.h b/src/shader_recompiler/frontend/maxwell/control_flow.h
index b2ab0cdc3..20ada8afd 100644
--- a/src/shader_recompiler/frontend/maxwell/control_flow.h
+++ b/src/shader_recompiler/frontend/maxwell/control_flow.h
@@ -70,6 +70,12 @@ struct Block {
     IR::Condition cond;
     BlockId branch_true;
     BlockId branch_false;
+    boost::container::small_vector<BlockId, 4> imm_predecessors;
+    boost::container::small_vector<BlockId, 8> dominance_frontiers;
+    union {
+        bool post_order_visited{false};
+        Block* imm_dominator;
+    };
 };
 
 struct Label {
@@ -81,11 +87,30 @@ struct Label {
 struct Function {
     Function(Location start_address);
 
+    void BuildBlocksMap();
+
+    void BuildImmediatePredecessors();
+
+    void BuildPostOrder();
+
+    void BuildImmediateDominators();
+
+    void BuildDominanceFrontier();
+
+    [[nodiscard]] size_t NumBlocks() const noexcept {
+        return static_cast<size_t>(current_block_id) + 1;
+    }
+
     Location entrypoint;
     BlockId current_block_id{0};
     boost::container::small_vector<Label, 16> labels;
     boost::container::small_vector<u32, 0x130> blocks;
     boost::container::small_vector<Block, 0x130> blocks_data;
+    // Translates from BlockId to block index
+    boost::container::small_vector<Block*, 0x130> blocks_map;
+
+    boost::container::small_vector<u32, 0x130> post_order_blocks;
+    boost::container::small_vector<BlockId, 0x130> post_order_map;
 };
 
 class CFG {
@@ -97,6 +122,12 @@ class CFG {
 public:
     explicit CFG(Environment& env, Location start_address);
 
+    CFG& operator=(const CFG&) = delete;
+    CFG(const CFG&) = delete;
+
+    CFG& operator=(CFG&&) = delete;
+    CFG(CFG&&) = delete;
+
     [[nodiscard]] std::string Dot() const;
 
     [[nodiscard]] std::span<const Function> Functions() const noexcept {
@@ -104,20 +135,22 @@ public:
     }
 
 private:
+    void VisitFunctions(Location start_address);
+
     void AnalyzeLabel(FunctionId function_id, Label& label);
 
     /// Inspect already visited blocks.
     /// Return true when the block has already been visited
-    [[nodiscard]] bool InspectVisitedBlocks(FunctionId function_id, const Label& label);
+    bool InspectVisitedBlocks(FunctionId function_id, const Label& label);
 
-    [[nodiscard]] AnalysisState AnalyzeInst(Block& block, FunctionId function_id, Location pc);
+    AnalysisState AnalyzeInst(Block& block, FunctionId function_id, Location pc);
 
     void AnalyzeCondInst(Block& block, FunctionId function_id, Location pc, EndClass insn_end_class,
                          IR::Condition cond);
 
     /// Return true when the branch instruction is confirmed to be a branch
-    [[nodiscard]] bool AnalyzeBranch(Block& block, FunctionId function_id, Location pc,
-                                     Instruction inst, Opcode opcode);
+    bool AnalyzeBranch(Block& block, FunctionId function_id, Location pc, Instruction inst,
+                       Opcode opcode);
 
     void AnalyzeBRA(Block& block, FunctionId function_id, Location pc, Instruction inst,
                     bool is_absolute);
@@ -126,8 +159,7 @@ private:
     AnalysisState AnalyzeEXIT(Block& block, FunctionId function_id, Location pc, Instruction inst);
 
     /// Return the branch target block id
-    [[nodiscard]] BlockId AddLabel(const Block& block, Stack stack, Location pc,
-                                   FunctionId function_id);
+    BlockId AddLabel(const Block& block, Stack stack, Location pc, FunctionId function_id);
 
     Environment& env;
     boost::container::small_vector<Function, 1> functions;
diff --git a/src/shader_recompiler/frontend/maxwell/program.cpp b/src/shader_recompiler/frontend/maxwell/program.cpp
index 67a98ba57..49d1f4bfb 100644
--- a/src/shader_recompiler/frontend/maxwell/program.cpp
+++ b/src/shader_recompiler/frontend/maxwell/program.cpp
@@ -8,40 +8,53 @@
 #include "shader_recompiler/frontend/maxwell/program.h"
 #include "shader_recompiler/frontend/maxwell/termination_code.h"
 #include "shader_recompiler/frontend/maxwell/translate/translate.h"
+#include "shader_recompiler/ir_opt/passes.h"
 
 namespace Shader::Maxwell {
+namespace {
+void TranslateCode(Environment& env, const Flow::Function& cfg_function, IR::Function& function,
+                   std::span<IR::Block*> block_map, IR::Block* block_memory) {
+    const size_t num_blocks{cfg_function.blocks.size()};
+    function.blocks.reserve(num_blocks);
 
-Program::Function::~Function() {
-    std::ranges::for_each(blocks, &std::destroy_at<IR::Block>);
-}
-
-Program::Program(Environment& env, const Flow::CFG& cfg) {
-    std::vector<IR::Block*> block_map;
-    functions.reserve(cfg.Functions().size());
+    for (const Flow::BlockId block_id : cfg_function.blocks) {
+        const Flow::Block& flow_block{cfg_function.blocks_data[block_id]};
 
-    for (const Flow::Function& cfg_function : cfg.Functions()) {
-        Function& function{functions.emplace_back()};
+        function.blocks.emplace_back(std::construct_at(block_memory, Translate(env, flow_block)));
+        block_map[flow_block.id] = function.blocks.back().get();
+        ++block_memory;
+    }
+}
 
-        const size_t num_blocks{cfg_function.blocks.size()};
-        IR::Block* block_memory{block_alloc_pool.allocate(num_blocks)};
-        function.blocks.reserve(num_blocks);
+void EmitTerminationInsts(const Flow::Function& cfg_function,
+                          std::span<IR::Block* const> block_map) {
+    for (const Flow::BlockId block_id : cfg_function.blocks) {
+        const Flow::Block& flow_block{cfg_function.blocks_data[block_id]};
+        EmitTerminationCode(flow_block, block_map);
+    }
+}
 
-        block_map.resize(cfg_function.blocks_data.size());
+void TranslateFunction(Environment& env, const Flow::Function& cfg_function, IR::Function& function,
+                       IR::Block* block_memory) {
+    std::vector<IR::Block*> block_map;
+    block_map.resize(cfg_function.blocks_data.size());
 
-        // Visit the instructions of all blocks
-        for (const Flow::BlockId block_id : cfg_function.blocks) {
-            const Flow::Block& flow_block{cfg_function.blocks_data[block_id]};
+    TranslateCode(env, cfg_function, function, block_map, block_memory);
+    EmitTerminationInsts(cfg_function, block_map);
+}
+} // Anonymous namespace
 
-            IR::Block* const block{std::construct_at(block_memory, Translate(env, flow_block))};
-            ++block_memory;
-            function.blocks.push_back(block);
-            block_map[flow_block.id] = block;
-        }
-        // Now that all blocks are defined, emit the termination instructions
-        for (const Flow::BlockId block_id : cfg_function.blocks) {
-            const Flow::Block& flow_block{cfg_function.blocks_data[block_id]};
-            EmitTerminationCode(flow_block, block_map);
-        }
+Program::Program(Environment& env, const Flow::CFG& cfg) {
+    functions.reserve(cfg.Functions().size());
+    for (const Flow::Function& cfg_function : cfg.Functions()) {
+        TranslateFunction(env, cfg_function, functions.emplace_back(),
+                          block_alloc_pool.allocate(cfg_function.blocks.size()));
+    }
+    std::ranges::for_each(functions, Optimization::SsaRewritePass);
+    for (IR::Function& function : functions) {
+        Optimization::Invoke(Optimization::DeadCodeEliminationPass, function);
+        Optimization::Invoke(Optimization::IdentityRemovalPass, function);
+        // Optimization::Invoke(Optimization::VerificationPass, function);
     }
 }
 
@@ -50,16 +63,16 @@ std::string DumpProgram(const Program& program) {
     std::map<const IR::Inst*, size_t> inst_to_index;
     std::map<const IR::Block*, size_t> block_to_index;
 
-    for (const Program::Function& function : program.functions) {
-        for (const IR::Block* const block : function.blocks) {
-            block_to_index.emplace(block, index);
+    for (const IR::Function& function : program.functions) {
+        for (const auto& block : function.blocks) {
+            block_to_index.emplace(block.get(), index);
             ++index;
         }
     }
     std::string ret;
-    for (const Program::Function& function : program.functions) {
+    for (const IR::Function& function : program.functions) {
         ret += fmt::format("Function\n");
-        for (const IR::Block* const block : function.blocks) {
+        for (const auto& block : function.blocks) {
             ret += IR::DumpBlock(*block, block_to_index, inst_to_index, index) + '\n';
         }
     }
diff --git a/src/shader_recompiler/frontend/maxwell/program.h b/src/shader_recompiler/frontend/maxwell/program.h
index 7814b2c01..36e678a9e 100644
--- a/src/shader_recompiler/frontend/maxwell/program.h
+++ b/src/shader_recompiler/frontend/maxwell/program.h
@@ -4,13 +4,16 @@
 
 #pragma once
 
+#include <memory>
 #include <string>
 #include <vector>
 
+#include <boost/container/small_vector.hpp>
 #include <boost/pool/pool_alloc.hpp>
 
 #include "shader_recompiler/environment.h"
 #include "shader_recompiler/frontend/ir/basic_block.h"
+#include "shader_recompiler/frontend/ir/function.h"
 #include "shader_recompiler/frontend/maxwell/control_flow.h"
 
 namespace Shader::Maxwell {
@@ -22,16 +25,10 @@ public:
     explicit Program(Environment& env, const Flow::CFG& cfg);
 
 private:
-    struct Function {
-        ~Function();
-
-        std::vector<IR::Block*> blocks;
-    };
-
     boost::pool_allocator<IR::Block, boost::default_user_allocator_new_delete,
                           boost::details::pool::null_mutex>
         block_alloc_pool;
-    std::vector<Function> functions;
+    boost::container::small_vector<IR::Function, 1> functions;
 };
 
 [[nodiscard]] std::string DumpProgram(const Program& program);
diff --git a/src/shader_recompiler/frontend/maxwell/termination_code.cpp b/src/shader_recompiler/frontend/maxwell/termination_code.cpp
index a4ea5c5e3..ed5137f20 100644
--- a/src/shader_recompiler/frontend/maxwell/termination_code.cpp
+++ b/src/shader_recompiler/frontend/maxwell/termination_code.cpp
@@ -47,12 +47,19 @@ static IR::U1 GetCond(IR::Condition cond, IR::IREmitter& ir) {
 
 static void EmitBranch(const Flow::Block& flow_block, std::span<IR::Block* const> block_map,
                        IR::IREmitter& ir) {
+    const auto add_immediate_predecessor = [&](Flow::BlockId label) {
+        block_map[label]->AddImmediatePredecessor(&ir.block);
+    };
     if (flow_block.cond == true) {
+        add_immediate_predecessor(flow_block.branch_true);
         return ir.Branch(block_map[flow_block.branch_true]);
     }
     if (flow_block.cond == false) {
+        add_immediate_predecessor(flow_block.branch_false);
         return ir.Branch(block_map[flow_block.branch_false]);
     }
+    add_immediate_predecessor(flow_block.branch_true);
+    add_immediate_predecessor(flow_block.branch_false);
     return ir.BranchConditional(GetCond(flow_block.cond, ir), block_map[flow_block.branch_true],
                                 block_map[flow_block.branch_false]);
 }
diff --git a/src/shader_recompiler/frontend/maxwell/termination_code.h b/src/shader_recompiler/frontend/maxwell/termination_code.h
index b0d667942..04e044534 100644
--- a/src/shader_recompiler/frontend/maxwell/termination_code.h
+++ b/src/shader_recompiler/frontend/maxwell/termination_code.h
@@ -11,6 +11,7 @@
 
 namespace Shader::Maxwell {
 
+/// Emit termination instructions and collect immediate predecessors
 void EmitTerminationCode(const Flow::Block& flow_block, std::span<IR::Block* const> block_map);
 
 } // namespace Shader::Maxwell
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/impl.h b/src/shader_recompiler/frontend/maxwell/translate/impl/impl.h
index bc607b002..8be7d6ff1 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/impl.h
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/impl.h
@@ -208,7 +208,7 @@ public:
     void P2R_reg(u64 insn);
     void P2R_cbuf(u64 insn);
     void P2R_imm(u64 insn);
-    void PBK(u64 insn);
+    void PBK();
     void PCNT(u64 insn);
     void PEXIT(u64 insn);
     void PIXLD(u64 insn);
@@ -252,7 +252,7 @@ public:
     void SHR_reg(u64 insn);
     void SHR_cbuf(u64 insn);
     void SHR_imm(u64 insn);
-    void SSY(u64 insn);
+    void SSY();
     void ST(u64 insn);
     void STG(u64 insn);
     void STL(u64 insn);
diff --git a/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp b/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
index c907c1ffb..0f52696d1 100644
--- a/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
+++ b/src/shader_recompiler/frontend/maxwell/translate/impl/not_implemented.cpp
@@ -762,7 +762,7 @@ void TranslatorVisitor::P2R_imm(u64) {
     ThrowNotImplemented(Opcode::P2R_imm);
 }
 
-void TranslatorVisitor::PBK(u64) {
+void TranslatorVisitor::PBK() {
     // PBK is a no-op
 }
 
@@ -938,8 +938,8 @@ void TranslatorVisitor::SHR_imm(u64) {
     ThrowNotImplemented(Opcode::SHR_imm);
 }
 
-void TranslatorVisitor::SSY(u64) {
-    ThrowNotImplemented(Opcode::SSY);
+void TranslatorVisitor::SSY() {
+    // SSY is a no-op
 }
 
 void TranslatorVisitor::ST(u64) {