diff --git a/core/libs/prx/libSceAgcDriver/tests/SpirvModule.cpp b/core/libs/prx/libSceAgcDriver/tests/SpirvModule.cpp index cf6e8a94f..d067e9720 100644 --- a/core/libs/prx/libSceAgcDriver/tests/SpirvModule.cpp +++ b/core/libs/prx/libSceAgcDriver/tests/SpirvModule.cpp @@ -2,6 +2,7 @@ #include "SpirvBackend/SpirvSpecialization.hpp" #include #include +#include #include #include #include @@ -285,8 +286,21 @@ void testFoldingLeavesNoDeadBlocksOrConstants() { checkFoldedSwitch(foldedSwitchModule(32u, {15u, 16u, 1u}), 32u, "the OpTypeInt width and the LocalSize literals equal the dead ids"); } +void testIdsBeyondTheBoundAreRejected() { + auto words = foldedSwitchModule(13u, {1u, 1u, 1u}); + words[3] = 12u; + bool rejected = false; + try { + static_cast(ShaderRecompiler::SpecializeSpirv(words)); + } catch (const std::exception&) { + rejected = true; + } + check(rejected, "a module whose ids reach its bound was specialized"); +} + int main() { testFoldingLeavesNoDeadBlocksOrConstants(); + testIdsBeyondTheBoundAreRejected(); testStructExtractsAreNotShuffled(); testRepeatedTypesShareOneId(); testOperandShapesDoNotCollide(); diff --git a/core/shader/recompiler/SpirvBackend/src/SpirvSpecialization.cpp b/core/shader/recompiler/SpirvBackend/src/SpirvSpecialization.cpp index d9cc99512..441a136b1 100644 --- a/core/shader/recompiler/SpirvBackend/src/SpirvSpecialization.cpp +++ b/core/shader/recompiler/SpirvBackend/src/SpirvSpecialization.cpp @@ -428,11 +428,105 @@ struct ScalarType { bool boolean; }; +template +class IdTable { +public: + void Reset(std::size_t bound) { + items.assign(bound, T{}); + present.assign(bound, 0u); + count = 0; + } + bool contains(std::uint32_t id) const { return id < present.size() && present[id] != 0u; } + const T* find(std::uint32_t id) const { return contains(id) ? &items[id] : nullptr; } + const T& at(std::uint32_t id) const { + if (!contains(id)) throw std::out_of_range("prepared SPIR-V id has no entry"); + return items[id]; + } + T& at(std::uint32_t id) { + if (!contains(id)) throw std::out_of_range("prepared SPIR-V id has no entry"); + return items[id]; + } + bool emplace(std::uint32_t id, T value) { + check(id); + if (present[id] != 0u) return false; + items[id] = std::move(value); + present[id] = 1u; + ++count; + return true; + } + T& operator[](std::uint32_t id) { + check(id); + if (present[id] == 0u) { + items[id] = T{}; + present[id] = 1u; + ++count; + } + return items[id]; + } + void erase(std::uint32_t id) { + if (!contains(id)) return; + present[id] = 0u; + items[id] = T{}; + --count; + } + std::size_t size() const { return count; } + bool empty() const { return count == 0; } + class Iterator { + public: + Iterator(IdTable& table, std::uint32_t id) : table(table), id(id) { skip(); } + std::pair operator*() const { return {id, table.items[id]}; } + Iterator& operator++() { + ++id; + skip(); + return *this; + } + bool operator!=(const Iterator& other) const { return id != other.id; } + + private: + void skip() { + while (id < table.present.size() && table.present[id] == 0u) ++id; + } + IdTable& table; + std::uint32_t id; + }; + Iterator begin() { return Iterator(*this, 0u); } + Iterator end() { return Iterator(*this, static_cast(present.size())); } + +private: + void check(std::uint32_t id) const { + if (id >= present.size()) throw std::runtime_error("prepared SPIR-V id exceeds the module bound"); + } + std::vector items; + std::vector present; + std::size_t count = 0; +}; + +class IdSet { +public: + void Reset(std::size_t bound) { present.assign(bound, 0u); } + bool contains(std::uint32_t id) const { return id < present.size() && present[id] != 0u; } + bool insert(std::uint32_t id) { + if (id >= present.size()) throw std::runtime_error("prepared SPIR-V id exceeds the module bound"); + if (present[id] != 0u) return false; + present[id] = 1u; + return true; + } + +private: + std::vector present; +}; + class Specialization { public: explicit Specialization(std::span words) { if (words.size() < 5u || words[0] != spv::MagicNumber) throw std::runtime_error("invalid prepared SPIR-V header"); header.assign(words.begin(), words.begin() + 5); + bound = header[3]; + types.Reset(bound); + values.Reset(bound); + resultTypes.Reset(bound); + vectorTypes.Reset(bound); + removed.Reset(bound); for (std::size_t cursor = 5; cursor < words.size();) { const auto count = words[cursor] >> 16u; if (count == 0u || count > words.size() - cursor) throw std::runtime_error("truncated prepared SPIR-V instruction"); @@ -479,15 +573,15 @@ public: private: std::optional value(std::uint32_t id) const { - const auto found = values.find(id); - return found != values.end() ? std::optional(found->second) : std::nullopt; + const auto* found = values.find(id); + return found != nullptr ? std::optional(*found) : std::nullopt; } std::optional evaluate(const Instruction& instruction) const { const auto op = Opcode(instruction); if (instruction.size() < 4u) return std::nullopt; - const auto type = types.find(instruction[1]); - if (type == types.end() || type->second.width > 32u) return std::nullopt; + const auto* type = types.find(instruction[1]); + if (type == nullptr || type->width > 32u) return std::nullopt; const auto left = value(instruction[3]); if (!left) return std::nullopt; if (op == spv::OpCopyObject) return left; @@ -505,8 +599,8 @@ private: case spv::OpBitwiseAnd: return *left & *right; case spv::OpBitwiseOr: return *left | *right; case spv::OpBitwiseXor: return *left ^ *right; - case spv::OpShiftRightLogical: return *right < type->second.width ? std::optional(*left >> *right) : std::nullopt; - case spv::OpShiftLeftLogical: return *right < type->second.width ? std::optional(*left << *right) : std::nullopt; + case spv::OpShiftRightLogical: return *right < type->width ? std::optional(*left >> *right) : std::nullopt; + case spv::OpShiftLeftLogical: return *right < type->width ? std::optional(*left << *right) : std::nullopt; case spv::OpIEqual: return *left == *right; case spv::OpINotEqual: return *left != *right; case spv::OpULessThan: return *left < *right; @@ -532,7 +626,8 @@ private: bool fold() { bool changed = false; bool function = false; - std::map> extracts; + IdTable> extracts; + extracts.Reset(bound); for (auto& instruction : instructions) { if (instruction.empty()) continue; auto op = Opcode(instruction); @@ -554,14 +649,14 @@ private: op = Opcode(instruction); if (op == spv::OpCompositeExtract && instruction.size() == 5u) extracts.emplace(instruction[2], std::pair{instruction[3], instruction[4]}); if (op == spv::OpCompositeConstruct && instruction.size() == 7u) { - const auto first = extracts.find(instruction[3]); - bool shuffle = first != extracts.end() && vectorTypes.contains(instruction[1]) && resultTypes.contains(first->second.first) && vectorTypes.contains(resultTypes.at(first->second.first)); + const auto* first = extracts.find(instruction[3]); + bool shuffle = first != nullptr && vectorTypes.contains(instruction[1]) && resultTypes.contains(first->first) && vectorTypes.contains(resultTypes.at(first->first)); for (std::size_t index = 4; shuffle && index < instruction.size(); ++index) { - const auto found = extracts.find(instruction[index]); - shuffle = found != extracts.end() && found->second.first == first->second.first; + const auto* found = extracts.find(instruction[index]); + shuffle = found != nullptr && found->first == first->first; } if (shuffle) { - const auto source = first->second.first; + const auto source = first->first; bool identity = resultTypes.contains(source) && resultTypes.at(source) == instruction[1]; Instruction replacement = Make(spv::OpVectorShuffle, {instruction[1], instruction[2], source, source}); for (std::size_t index = 3; index < instruction.size(); ++index) { @@ -587,7 +682,8 @@ private: } bool prune() { - std::map> blocks; + IdTable> blocks; + blocks.Reset(bound); std::vector entries; std::uint32_t block = 0; bool entry = false; @@ -622,7 +718,8 @@ private: } return true; }; - std::set loopHeaders; + IdSet loopHeaders; + loopHeaders.Reset(bound); std::map continueHeaders; for (const auto& [label, indices] : blocks) { for (const auto index : indices) { @@ -631,7 +728,8 @@ private: continueHeaders.emplace(instructions[index].at(2), label); } } - std::map> edges; + IdTable> edges; + edges.Reset(bound); for (const auto& [label, indices] : blocks) { auto& terminal = instructions[indices.back()]; auto op = Opcode(terminal); @@ -668,18 +766,19 @@ private: else if (op == spv::OpBranchConditional) successors = {terminal.at(2), terminal.at(3)}; else if (op == spv::OpSwitch) { successors.push_back(terminal.at(2)); - const auto selectorType = resultTypes.find(terminal.at(1)); - const auto stride = selectorType != resultTypes.end() && types.contains(selectorType->second) && types.at(selectorType->second).width == 64u ? 3u : 2u; + const auto* selectorType = resultTypes.find(terminal.at(1)); + const auto stride = selectorType != nullptr && types.contains(*selectorType) && types.at(*selectorType).width == 64u ? 3u : 2u; for (std::size_t index = 3; index + stride <= terminal.size(); index += stride) successors.push_back(terminal[index + stride - 1u]); } } - std::set live; + IdSet live; + live.Reset(bound); auto pending = entries; while (!pending.empty()) { const auto label = pending.back(); pending.pop_back(); if (!blocks.contains(label)) throw std::runtime_error("prepared branch references a missing block"); - if (!live.insert(label).second) continue; + if (!live.insert(label)) continue; const auto& successors = edges.at(label); pending.insert(pending.end(), successors.begin(), successors.end()); } @@ -730,14 +829,15 @@ private: return changed; } - std::set referencedIds() const { - std::set referenced; + IdSet referencedIds() const { + IdSet referenced; + referenced.Reset(bound); for (const auto& instruction : instructions) { if (instruction.empty()) continue; std::uint32_t switchLiteralWords = 1u; if (Opcode(instruction) == spv::OpSwitch && instruction.size() > 1u) { - const auto selectorType = resultTypes.find(instruction[1]); - if (selectorType != resultTypes.end() && types.contains(selectorType->second) && types.at(selectorType->second).width == 64u) switchLiteralWords = 2u; + const auto* selectorType = resultTypes.find(instruction[1]); + if (selectorType != nullptr && types.contains(*selectorType) && types.at(*selectorType).width == 64u) switchLiteralWords = 2u; } VisitIdOperands(instruction, switchLiteralWords, [&](std::size_t index) { referenced.insert(instruction[index]); }); } @@ -800,7 +900,8 @@ private: } bool propagateCopies() { - std::map copies; + IdTable copies; + copies.Reset(bound); for (const auto& instruction : instructions) { if (Opcode(instruction) != spv::OpCopyObject) continue; if (instruction.size() != 4u) throw std::runtime_error("invalid prepared copy instruction"); @@ -845,7 +946,8 @@ private: } void removeDeadComputations() { - std::map definitions; + IdTable definitions; + definitions.Reset(bound); std::vector pending; bool function = false; for (std::size_t index = 0; index < instructions.size(); ++index) { @@ -871,13 +973,14 @@ private: const auto resultIndex = hasResult ? (hasType ? 2u : 1u) : 0u; for (std::size_t operand = 1; operand < instruction.size(); ++operand) if (operand != resultIndex) pending.push_back(instruction[operand]); } - std::set live; + IdSet live; + live.Reset(bound); while (!pending.empty()) { const auto id = pending.back(); pending.pop_back(); - const auto found = definitions.find(id); - if (found == definitions.end() || !live.insert(id).second) continue; - const auto& instruction = instructions[found->second]; + const auto* found = definitions.find(id); + if (found == nullptr || !live.insert(id)) continue; + const auto& instruction = instructions[*found]; if (!VisitInputs(instruction, [&](std::size_t operand) { pending.push_back(instruction[operand]); })) throw std::runtime_error("missing prepared pure instruction operands"); } for (const auto& [id, index] : definitions) { @@ -890,11 +993,12 @@ private: std::vector header; std::vector instructions; std::vector constants; - std::map types; - std::map values; - std::map resultTypes; - std::set vectorTypes; - std::set removed; + std::uint32_t bound = 0; + IdTable types; + IdTable values; + IdTable resultTypes; + IdSet vectorTypes; + IdSet removed; }; }