Merge pull request #408 from oneandonlydean/perf/agc-flat-queue-registers

perf(agc): keep a queue's registers in flat arrays instead of a std::map

core/libs/prx/libSceAgcDriver/Execution/include/QueueState.hpp

@@ -1,7 +1,13 @@
 #ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_QUEUESTATE_HPP
 #define CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_QUEUESTATE_HPP
 
+#include <bit>
+#include <cstddef>
 #include <cstdint>
+#include <initializer_list>
+#include <iterator>
+#include <stdexcept>
+#include <utility>
 #include <map>
 #include <array>
 #include <optional>
@@ -10,7 +16,128 @@
 
 namespace AgcDriver {
 
-using Registers = std::map<std::uint32_t, std::uint32_t>;
+class Registers {
+public:
+    using key_type = std::uint32_t;
+    using mapped_type = std::uint32_t;
+    using value_type = std::pair<const std::uint32_t, std::uint32_t>;
+
+    class const_iterator {
+    public:
+        using iterator_category = std::forward_iterator_tag;
+        using value_type = Registers::value_type;
+        using difference_type = std::ptrdiff_t;
+        using pointer = const value_type*;
+        using reference = value_type;
+        struct Arrow {
+            value_type entry;
+            const value_type* operator->() const { return &entry; }
+        };
+        const_iterator() = default;
+        const_iterator(const Registers* owner, std::size_t index) : owner(owner), index(index) {}
+        value_type operator*() const { return {static_cast<std::uint32_t>(index), owner->values[index]}; }
+        Arrow operator->() const { return {**this}; }
+        const_iterator& operator++() {
+            index = owner->nextPresent(index + 1);
+            return *this;
+        }
+        const_iterator operator++(int) {
+            auto previous = *this;
+            ++*this;
+            return previous;
+        }
+        bool operator==(const const_iterator& other) const { return index == other.index; }
+        bool operator!=(const const_iterator& other) const { return index != other.index; }
+
+    private:
+        const Registers* owner = nullptr;
+        std::size_t index = 0;
+    };
+    using iterator = const_iterator;
+
+    Registers() = default;
+    Registers(std::initializer_list<std::pair<std::uint32_t, std::uint32_t>> entries) {
+        for (const auto& [offset, value] : entries) emplace(offset, value);
+    }
+
+    const_iterator begin() const { return {this, nextPresent(0)}; }
+    const_iterator end() const { return {this, End}; }
+    const_iterator find(std::uint32_t offset) const { return contains(offset) ? const_iterator{this, offset} : end(); }
+    const_iterator lower_bound(std::uint32_t offset) const { return {this, nextPresent(offset)}; }
+    const_iterator upper_bound(std::uint32_t offset) const { return {this, nextPresent(static_cast<std::size_t>(offset) + 1)}; }
+    bool contains(std::uint32_t offset) const { return offset < values.size() && ((present[offset / 64] >> (offset % 64)) & 1u) != 0; }
+    std::size_t count(std::uint32_t offset) const { return contains(offset) ? 1 : 0; }
+    std::size_t size() const { return entries; }
+    bool empty() const { return entries == 0; }
+    void clear() {
+        values.clear();
+        present.clear();
+        entries = 0;
+    }
+    std::pair<const_iterator, bool> emplace(std::uint32_t offset, std::uint32_t value) {
+        if (contains(offset)) return {const_iterator{this, offset}, false};
+        mark(offset) = value;
+        return {const_iterator{this, offset}, true};
+    }
+    std::pair<const_iterator, bool> insert_or_assign(std::uint32_t offset, std::uint32_t value) {
+        if (contains(offset)) {
+            values[offset] = value;
+            return {const_iterator{this, offset}, false};
+        }
+        mark(offset) = value;
+        return {const_iterator{this, offset}, true};
+    }
+    std::uint32_t& operator[](std::uint32_t offset) {
+        if (contains(offset)) return values[offset];
+        return mark(offset) = 0;
+    }
+    std::uint32_t& at(std::uint32_t offset) {
+        if (!contains(offset)) throw std::out_of_range("register is not set");
+        return values[offset];
+    }
+    const std::uint32_t& at(std::uint32_t offset) const {
+        if (!contains(offset)) throw std::out_of_range("register is not set");
+        return values[offset];
+    }
+    std::size_t erase(std::uint32_t offset) {
+        if (!contains(offset)) return 0;
+        present[offset / 64] &= ~(std::uint64_t{1} << (offset % 64));
+        --entries;
+        return 1;
+    }
+    bool operator==(const Registers& other) const {
+        auto a = begin();
+        auto b = other.begin();
+        for (; a != end() && b != other.end(); ++a, ++b) {
+            if ((*a).first != (*b).first || (*a).second != (*b).second) return false;
+        }
+        return a == end() && b == other.end();
+    }
+
+private:
+    static constexpr std::size_t End = ~std::size_t{0};
+    std::uint32_t& mark(std::uint32_t offset) {
+        if (offset >= values.size()) {
+            const auto words = static_cast<std::size_t>(offset) / 64 + 1;
+            values.resize(words * 64, 0);
+            present.resize(words, 0);
+        }
+        present[offset / 64] |= std::uint64_t{1} << (offset % 64);
+        ++entries;
+        return values[offset];
+    }
+    std::size_t nextPresent(std::size_t index) const {
+        for (std::size_t word = index / 64; word < present.size(); ++word) {
+            auto bits = present[word];
+            if (word == index / 64) bits &= ~std::uint64_t{0} << (index % 64);
+            if (bits != 0) return word * 64 + static_cast<std::size_t>(std::countr_zero(bits));
+        }
+        return End;
+    }
+    std::vector<std::uint32_t> values;
+    std::vector<std::uint64_t> present;
+    std::size_t entries = 0;
+};
 
 inline Registers InitialContextRegisters() {
     Registers result{

core/libs/prx/libSceAgcDriver/tests/Pm4.cpp

@@ -99,6 +99,29 @@ void testRegisters() {
     expectFailure([&] { AgcDriver::Pm4::Validate(makePacket(0x9f, {0, 0, 0x80000000, 0}), 0x20); }, "compute");
 }
 
+void testRegisterFile() {
+    AgcDriver::Registers registers{{0x300, 3}, {0x10, 1}, {0x41, 2}};
+    check(registers.size() == 3 && !registers.contains(0x11) && registers.at(0x41) == 2, "register file lookup");
+    check(registers.find(0x12) == registers.end() && registers.find(0x10)->second == 1, "register file find");
+    check(!registers.emplace(0x10, 9).second && registers.at(0x10) == 1, "register file emplace replaced a value");
+    check(registers.insert_or_assign(0x10, 7).second == false && registers.at(0x10) == 7, "register file assignment");
+    check(registers.lower_bound(0x11)->first == 0x41 && registers.upper_bound(0x41)->first == 0x300 && registers.lower_bound(0x301) == registers.end(), "register file bounds");
+    std::vector<std::pair<std::uint32_t, std::uint32_t>> order;
+    for (const auto& [offset, value] : registers) order.emplace_back(offset, value);
+    check(order == std::vector<std::pair<std::uint32_t, std::uint32_t>>{{0x10, 7}, {0x41, 2}, {0x300, 3}}, "register file order");
+    auto copy = registers;
+    copy[0x7000] = 5;
+    check(copy.size() == 4 && registers.size() == 3 && !registers.contains(0x7000) && !(copy == registers), "register file copy");
+    check(copy.erase(0x7000) == 1 && copy.erase(0x7000) == 0 && copy == registers, "register file erase");
+    bool threw = false;
+    try {
+        static_cast<void>(registers.at(0x42));
+    } catch (const std::out_of_range&) {
+        threw = true;
+    }
+    check(threw, "register file read an unset register");
+}
+
 void testContextAndBases() {
     AgcDriver::QueueState state;
     execute(state, makePacket(0x69, {0x10, 17}));
@@ -514,6 +537,7 @@ int main(int argc, char** argv) {
         testCatalog();
         testWriteChangedKeepsUntouchedBytes();
         testRegisters();
+        testRegisterFile();
         testContextAndBases();
         testIndexedDraw();
         testAutoDraw();