shader: Add support for forward declarations

11 files changed, 79 insertions, 68 deletions

diff --git a/externals/sirit b/externals/sirit
-Subproject c374bfd9fdff02a0cff85d005488967b1b0f675
+Subproject f819ade0efe925a782090dea9e1bf300fedffb3
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.cpp b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
index e29e448c7..0895414b4 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.cpp
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.cpp
@@ -64,31 +64,49 @@ EmitSPIRV::EmitSPIRV(IR::Program& program) {
     std::system("spirv-cross shader.spv");
 }
 
+template <auto method, typename... Args>
+static void SetDefinition(EmitSPIRV& emit, EmitContext& ctx, IR::Inst* inst, Args... args) {
+    const Id forward_id{inst->Definition<Id>()};
+    const bool has_forward_id{Sirit::ValidId(forward_id)};
+    Id current_id{};
+    if (has_forward_id) {
+        current_id = ctx.ExchangeCurrentId(forward_id);
+    }
+    const Id new_id{(emit.*method)(ctx, std::forward<Args>(args)...)};
+    if (has_forward_id) {
+        ctx.ExchangeCurrentId(current_id);
+    } else {
+        inst->SetDefinition<Id>(new_id);
+    }
+}
+
 template <auto method>
 static void Invoke(EmitSPIRV& emit, EmitContext& ctx, IR::Inst* inst) {
     using M = decltype(method);
     using std::is_invocable_r_v;
     if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&>) {
-        ctx.Define(inst, (emit.*method)(ctx));
+        SetDefinition<method>(emit, ctx, inst);
     } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, Id>) {
-        ctx.Define(inst, (emit.*method)(ctx, ctx.Def(inst->Arg(0))));
+        SetDefinition<method>(emit, ctx, inst, ctx.Def(inst->Arg(0)));
     } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, Id, Id>) {
-        ctx.Define(inst, (emit.*method)(ctx, ctx.Def(inst->Arg(0)), ctx.Def(inst->Arg(1))));
+        SetDefinition<method>(emit, ctx, inst, ctx.Def(inst->Arg(0)), ctx.Def(inst->Arg(1)));
     } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, Id, Id, Id>) {
-        ctx.Define(inst, (emit.*method)(ctx, ctx.Def(inst->Arg(0)), ctx.Def(inst->Arg(1)),
-                                        ctx.Def(inst->Arg(2))));
+        SetDefinition<method>(emit, ctx, inst, ctx.Def(inst->Arg(0)), ctx.Def(inst->Arg(1)),
+                              ctx.Def(inst->Arg(2)));
+    } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, IR::Inst*>) {
+        SetDefinition<method>(emit, ctx, inst, inst);
     } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, IR::Inst*, Id, Id>) {
-        ctx.Define(inst, (emit.*method)(ctx, inst, ctx.Def(inst->Arg(0)), ctx.Def(inst->Arg(1))));
+        SetDefinition<method>(emit, ctx, inst, inst, ctx.Def(inst->Arg(0)), ctx.Def(inst->Arg(1)));
     } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, IR::Inst*, Id, Id, Id>) {
-        ctx.Define(inst, (emit.*method)(ctx, inst, ctx.Def(inst->Arg(0)), ctx.Def(inst->Arg(1)),
-                                        ctx.Def(inst->Arg(2))));
+        SetDefinition<method>(emit, ctx, inst, inst, ctx.Def(inst->Arg(0)), ctx.Def(inst->Arg(1)),
+                              ctx.Def(inst->Arg(2)));
     } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, Id, u32>) {
-        ctx.Define(inst, (emit.*method)(ctx, ctx.Def(inst->Arg(0)), inst->Arg(1).U32()));
+        SetDefinition<method>(emit, ctx, inst, ctx.Def(inst->Arg(0)), inst->Arg(1).U32());
     } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, const IR::Value&>) {
-        ctx.Define(inst, (emit.*method)(ctx, inst->Arg(0)));
+        SetDefinition<method>(emit, ctx, inst, inst->Arg(0));
     } else if constexpr (is_invocable_r_v<Id, M, EmitSPIRV&, EmitContext&, const IR::Value&,
                                           const IR::Value&>) {
-        ctx.Define(inst, (emit.*method)(ctx, inst->Arg(0), inst->Arg(1)));
+        SetDefinition<method>(emit, ctx, inst, inst->Arg(0), inst->Arg(1));
     } else if constexpr (is_invocable_r_v<void, M, EmitSPIRV&, EmitContext&, IR::Inst*>) {
         (emit.*method)(ctx, inst);
     } else if constexpr (is_invocable_r_v<void, M, EmitSPIRV&, EmitContext&>) {
@@ -122,11 +140,28 @@ static Id TypeId(const EmitContext& ctx, IR::Type type) {
 
 Id EmitSPIRV::EmitPhi(EmitContext& ctx, IR::Inst* inst) {
     const size_t num_args{inst->NumArgs()};
-    boost::container::small_vector<Id, 64> operands;
+    boost::container::small_vector<Id, 32> operands;
     operands.reserve(num_args * 2);
     for (size_t index = 0; index < num_args; ++index) {
+        // Phi nodes can have forward declarations, if an argument is not defined provide a forward
+        // declaration of it. Invoke will take care of giving it the right definition when it's
+        // actually defined.
+        const IR::Value arg{inst->Arg(index)};
+        Id def{};
+        if (arg.IsImmediate()) {
+            // Let the context handle immediate definitions, as it already knows how
+            def = ctx.Def(arg);
+        } else {
+            IR::Inst* const arg_inst{arg.Inst()};
+            def = arg_inst->Definition<Id>();
+            if (!Sirit::ValidId(def)) {
+                // If it hasn't been defined, get a forward declaration
+                def = ctx.ForwardDeclarationId();
+                arg_inst->SetDefinition<Id>(def);
+            }
+        }
         IR::Block* const phi_block{inst->PhiBlock(index)};
-        operands.push_back(ctx.Def(inst->Arg(index)));
+        operands.push_back(def);
         operands.push_back(ctx.BlockLabel(phi_block));
     }
     const Id result_type{TypeId(ctx, inst->Arg(0).Type())};
diff --git a/src/shader_recompiler/backend/spirv/emit_spirv.h b/src/shader_recompiler/backend/spirv/emit_spirv.h
index 6b09757d1..7d76377b5 100644
--- a/src/shader_recompiler/backend/spirv/emit_spirv.h
+++ b/src/shader_recompiler/backend/spirv/emit_spirv.h
@@ -6,8 +6,6 @@
 
 #include <sirit/sirit.h>
 
-#include <boost/container/flat_map.hpp>
-
 #include "common/common_types.h"
 #include "shader_recompiler/frontend/ir/microinstruction.h"
 #include "shader_recompiler/frontend/ir/program.h"
@@ -16,37 +14,6 @@ namespace Shader::Backend::SPIRV {
 
 using Sirit::Id;
 
-class DefMap {
-public:
-    void Define(IR::Inst* inst, Id def_id) {
-        const InstInfo info{.use_count{inst->UseCount()}, .def_id{def_id}};
-        const auto it{map.insert(map.end(), std::make_pair(inst, info))};
-        if (it == map.end()) {
-            throw LogicError("Defining already defined instruction");
-        }
-    }
-
-    [[nodiscard]] Id Consume(IR::Inst* inst) {
-        const auto it{map.find(inst)};
-        if (it == map.end()) {
-            throw LogicError("Consuming undefined instruction");
-        }
-        const Id def_id{it->second.def_id};
-        if (--it->second.use_count == 0) {
-            map.erase(it);
-        }
-        return def_id;
-    }
-
-private:
-    struct InstInfo {
-        int use_count;
-        Id def_id;
-    };
-
-    boost::container::flat_map<IR::Inst*, InstInfo> map;
-};
-
 class VectorTypes {
 public:
     void Define(Sirit::Module& sirit_ctx, Id base_type, std::string_view name) {
@@ -76,7 +43,7 @@ public:
 
     [[nodiscard]] Id Def(const IR::Value& value) {
         if (!value.IsImmediate()) {
-            return def_map.Consume(value.Inst());
+            return value.Inst()->Definition<Id>();
         }
         switch (value.Type()) {
         case IR::Type::U1:
@@ -90,10 +57,6 @@ public:
         }
     }
 
-    void Define(IR::Inst* inst, Id def_id) {
-        def_map.Define(inst, def_id);
-    }
-
     [[nodiscard]] Id BlockLabel(IR::Block* block) const {
         const auto it{std::ranges::lower_bound(block_label_map, block, {},
                                                &std::pair<IR::Block*, Id>::first)};
@@ -117,7 +80,6 @@ public:
     Id local_invocation_id{};
 
 private:
-    DefMap def_map;
     std::vector<std::pair<IR::Block*, Id>> block_label_map;
 };
 
diff --git a/src/shader_recompiler/frontend/ir/basic_block.cpp b/src/shader_recompiler/frontend/ir/basic_block.cpp
index c97626712..5ae91dd7d 100644
--- a/src/shader_recompiler/frontend/ir/basic_block.cpp
+++ b/src/shader_recompiler/frontend/ir/basic_block.cpp
@@ -26,7 +26,7 @@ void Block::AppendNewInst(Opcode op, std::initializer_list<Value> args) {
 }
 
 Block::iterator Block::PrependNewInst(iterator insertion_point, Opcode op,
-                                      std::initializer_list<Value> args, u64 flags) {
+                                      std::initializer_list<Value> args, u32 flags) {
     Inst* const inst{inst_pool->Create(op, flags)};
     const auto result_it{instructions.insert(insertion_point, *inst)};
 
diff --git a/src/shader_recompiler/frontend/ir/basic_block.h b/src/shader_recompiler/frontend/ir/basic_block.h
index 3205705e7..778b32e43 100644
--- a/src/shader_recompiler/frontend/ir/basic_block.h
+++ b/src/shader_recompiler/frontend/ir/basic_block.h
@@ -42,7 +42,7 @@ public:
 
     /// Prepends a new instruction to this basic block before the insertion point.
     iterator PrependNewInst(iterator insertion_point, Opcode op,
-                            std::initializer_list<Value> args = {}, u64 flags = 0);
+                            std::initializer_list<Value> args = {}, u32 flags = 0);
 
     /// Set the branches to jump to when all instructions have executed.
     void SetBranches(Condition cond, Block* branch_true, Block* branch_false);
diff --git a/src/shader_recompiler/frontend/ir/ir_emitter.h b/src/shader_recompiler/frontend/ir/ir_emitter.h
index 4decb46bc..24b012a39 100644
--- a/src/shader_recompiler/frontend/ir/ir_emitter.h
+++ b/src/shader_recompiler/frontend/ir/ir_emitter.h
@@ -178,7 +178,7 @@ private:
     }
 
     template <typename T>
-    requires(sizeof(T) <= sizeof(u64) && std::is_trivially_copyable_v<T>) struct Flags {
+    requires(sizeof(T) <= sizeof(u32) && std::is_trivially_copyable_v<T>) struct Flags {
         Flags() = default;
         Flags(T proxy_) : proxy{proxy_} {}
 
@@ -187,7 +187,7 @@ private:
 
     template <typename T = Value, typename FlagType, typename... Args>
     T Inst(Opcode op, Flags<FlagType> flags, Args... args) {
-        u64 raw_flags{};
+        u32 raw_flags{};
         std::memcpy(&raw_flags, &flags.proxy, sizeof(flags.proxy));
         auto it{block->PrependNewInst(insertion_point, op, {Value{args}...}, raw_flags)};
         return T{Value{&*it}};
diff --git a/src/shader_recompiler/frontend/ir/microinstruction.cpp b/src/shader_recompiler/frontend/ir/microinstruction.cpp
index 9279b9692..ee76db9ad 100644
--- a/src/shader_recompiler/frontend/ir/microinstruction.cpp
+++ b/src/shader_recompiler/frontend/ir/microinstruction.cpp
@@ -31,7 +31,7 @@ static void RemovePseudoInstruction(IR::Inst*& inst, IR::Opcode expected_opcode)
     inst = nullptr;
 }
 
-Inst::Inst(IR::Opcode op_, u64 flags_) noexcept : op{op_}, flags{flags_} {
+Inst::Inst(IR::Opcode op_, u32 flags_) noexcept : op{op_}, flags{flags_} {
     if (op == Opcode::Phi) {
         std::construct_at(&phi_args);
     } else {
diff --git a/src/shader_recompiler/frontend/ir/microinstruction.h b/src/shader_recompiler/frontend/ir/microinstruction.h
index ddf0f90a9..5b244fa0b 100644
--- a/src/shader_recompiler/frontend/ir/microinstruction.h
+++ b/src/shader_recompiler/frontend/ir/microinstruction.h
@@ -12,6 +12,7 @@
 
 #include <boost/intrusive/list.hpp>
 
+#include "common/bit_cast.h"
 #include "common/common_types.h"
 #include "shader_recompiler/frontend/ir/opcodes.h"
 #include "shader_recompiler/frontend/ir/type.h"
@@ -25,7 +26,7 @@ constexpr size_t MAX_ARG_COUNT = 4;
 
 class Inst : public boost::intrusive::list_base_hook<> {
 public:
-    explicit Inst(Opcode op_, u64 flags_) noexcept;
+    explicit Inst(Opcode op_, u32 flags_) noexcept;
     ~Inst();
 
     Inst& operator=(const Inst&) = delete;
@@ -86,13 +87,25 @@ public:
     void ReplaceUsesWith(Value replacement);
 
     template <typename FlagsType>
-    requires(sizeof(FlagsType) <= sizeof(u64) && std::is_trivially_copyable_v<FlagsType>)
+    requires(sizeof(FlagsType) <= sizeof(u32) && std::is_trivially_copyable_v<FlagsType>)
         [[nodiscard]] FlagsType Flags() const noexcept {
         FlagsType ret;
         std::memcpy(&ret, &flags, sizeof(ret));
         return ret;
     }
 
+    /// Intrusively store the host definition of this instruction.
+    template <typename DefinitionType>
+    void SetDefinition(DefinitionType def) {
+        definition = Common::BitCast<u32>(def);
+    }
+
+    /// Return the intrusively stored host definition of this instruction.
+    template <typename DefinitionType>
+    [[nodiscard]] DefinitionType Definition() const noexcept {
+        return Common::BitCast<DefinitionType>(definition);
+    }
+
 private:
     struct NonTriviallyDummy {
         NonTriviallyDummy() noexcept {}
@@ -103,7 +116,8 @@ private:
 
     IR::Opcode op{};
     int use_count{};
-    u64 flags{};
+    u32 flags{};
+    u32 definition{};
     union {
         NonTriviallyDummy dummy{};
         std::array<Value, MAX_ARG_COUNT> args;
diff --git a/src/shader_recompiler/frontend/ir/modifiers.h b/src/shader_recompiler/frontend/ir/modifiers.h
index 28bb9e798..c288eede0 100644
--- a/src/shader_recompiler/frontend/ir/modifiers.h
+++ b/src/shader_recompiler/frontend/ir/modifiers.h
@@ -6,13 +6,13 @@
 
 namespace Shader::IR {
 
-enum class FmzMode {
+enum class FmzMode : u8 {
     None, // Denorms are not flushed, NAN is propagated (nouveau)
     FTZ,  // Flush denorms to zero, NAN is propagated (D3D11, NVN, GL, VK)
     FMZ,  // Flush denorms to zero, x * 0 == 0 (D3D9)
 };
 
-enum class FpRounding {
+enum class FpRounding : u8 {
     RN, // Round to nearest even,
     RM, // Round towards negative infinity
     RP, // Round towards positive infinity
@@ -21,8 +21,8 @@ enum class FpRounding {
 
 struct FpControl {
     bool no_contraction{false};
-    FpRounding rounding : 8 = FpRounding::RN;
-    FmzMode fmz_mode : 8 = FmzMode::FTZ;
+    FpRounding rounding{FpRounding::RN};
+    FmzMode fmz_mode{FmzMode::FTZ};
 };
-static_assert(sizeof(FpControl) <= sizeof(u64));
+static_assert(sizeof(FpControl) <= sizeof(u32));
 } // namespace Shader::IR
diff --git a/src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp b/src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp
index 34393e1d5..08fd364bb 100644
--- a/src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp
+++ b/src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp
@@ -161,8 +161,8 @@ std::optional<StorageBufferAddr> Track(const IR::Value& value, const Bias* bias)
             return std::nullopt;
         }
         const StorageBufferAddr storage_buffer{
-            .index = index.U32(),
-            .offset = offset.U32(),
+            .index{index.U32()},
+            .offset{offset.U32()},
         };
         if (bias && !MeetsBias(storage_buffer, *bias)) {
             // We have to blacklist some addresses in case we wrongly point to them
diff --git a/src/shader_recompiler/main.cpp b/src/shader_recompiler/main.cpp
index 29f65966c..3b110af61 100644
--- a/src/shader_recompiler/main.cpp
+++ b/src/shader_recompiler/main.cpp
@@ -76,5 +76,5 @@ int main() {
     fmt::print(stdout, "{}\n", cfg.Dot());
     IR::Program program{TranslateProgram(*inst_pool, *block_pool, env, cfg)};
     fmt::print(stdout, "{}\n", IR::DumpProgram(program));
-    // Backend::SPIRV::EmitSPIRV spirv{program};
+    Backend::SPIRV::EmitSPIRV spirv{program};
 }