perf(libkernel): find direct mappings that share physical memory through an index

core/libs/prx/libkernel/DirectMemory/DirectMemory.cpp

@@ -348,15 +348,77 @@ struct DirectMapping {
 
 std::map<std::uintptr_t, DirectMapping> g_directMappings;
 
+#if defined(__linux__)
+struct PhysicalSpan {
+    std::uint64_t end;
+    std::vector<std::uintptr_t> bases;
+};
+
+std::map<std::uint64_t, PhysicalSpan> g_physicalSpans;
+
+void SplitPhysicalSpan(std::uint64_t at) {
+    auto span = g_physicalSpans.upper_bound(at);
+    if (span == g_physicalSpans.begin()) return;
+    --span;
+    if (span->first == at || span->second.end <= at) return;
+    g_physicalSpans.emplace_hint(std::next(span), at, PhysicalSpan{span->second.end, span->second.bases});
+    span->second.end = at;
+}
+#endif
+
+std::map<std::uintptr_t, DirectMapping>::iterator InsertMapping(std::uintptr_t base, DirectMapping mapping) {
+    const auto [inserted, added] = g_directMappings.emplace(base, std::move(mapping));
+    if (!added) throw std::runtime_error("direct mapping inserted over a live one");
+#if defined(__linux__)
+    const auto first = inserted->second.phys;
+    const auto last = first + (inserted->second.end - base);
+    SplitPhysicalSpan(first);
+    SplitPhysicalSpan(last);
+    auto span = g_physicalSpans.lower_bound(first);
+    for (auto cursor = first; cursor < last; ++span) {
+        if (span == g_physicalSpans.end() || span->first > cursor) span = g_physicalSpans.emplace_hint(span, cursor, PhysicalSpan{span == g_physicalSpans.end() ? last : std::min(last, span->first), {}});
+        auto& bases = span->second.bases;
+        bases.insert(std::upper_bound(bases.begin(), bases.end(), base), base);
+        cursor = span->second.end;
+    }
+#endif
+    return inserted;
+}
+
+std::map<std::uintptr_t, DirectMapping>::iterator RemoveMapping(std::map<std::uintptr_t, DirectMapping>::iterator mapping) {
+#if defined(__linux__)
+    const auto base = mapping->first;
+    const auto first = mapping->second.phys;
+    const auto last = first + (mapping->second.end - base);
+    for (auto span = g_physicalSpans.lower_bound(first); span != g_physicalSpans.end() && span->first < last;) {
+        auto& bases = span->second.bases;
+        bases.erase(std::lower_bound(bases.begin(), bases.end(), base));
+        span = bases.empty() ? g_physicalSpans.erase(span) : std::next(span);
+    }
+    for (auto span = g_physicalSpans.lower_bound(first); span != g_physicalSpans.end() && span->first <= last;) {
+        if (span != g_physicalSpans.begin()) {
+            const auto before = std::prev(span);
+            if (before->second.end == span->first && before->second.bases == span->second.bases) {
+                before->second.end = span->second.end;
+                span = g_physicalSpans.erase(span);
+                continue;
+            }
+        }
+        ++span;
+    }
+#endif
+    return g_directMappings.erase(mapping);
+}
+
 void EraseMappings(std::uintptr_t start, std::uintptr_t end) {
     auto it = g_directMappings.lower_bound(start);
     if (it != g_directMappings.begin() && std::prev(it)->second.end > start) --it;
     while (it != g_directMappings.end() && it->first < end) {
         const auto base = it->first;
         const auto mapping = it->second;
-        it = g_directMappings.erase(it);
-        if (base < start) g_directMappings.emplace(base, DirectMapping{start, mapping.phys, mapping.memoryType, mapping.backing});
-        if (mapping.end > end) it = g_directMappings.emplace(end, DirectMapping{mapping.end, mapping.phys + end - base, mapping.memoryType, mapping.backing}).first;
+        it = RemoveMapping(it);
+        if (base < start) InsertMapping(base, DirectMapping{start, mapping.phys, mapping.memoryType, mapping.backing});
+        if (mapping.end > end) it = InsertMapping(end, DirectMapping{mapping.end, mapping.phys + end - base, mapping.memoryType, mapping.backing});
     }
 }
 
@@ -369,14 +431,14 @@ void ErasePhysMappings(std::uint64_t first, std::uint64_t last) {
             ++it;
             continue;
         }
-        it = g_directMappings.erase(it);
+        it = RemoveMapping(it);
         if (mapping.phys < first) {
             const auto keep = first - mapping.phys;
-            g_directMappings.emplace(base, DirectMapping{base + keep, mapping.phys, mapping.memoryType, mapping.backing});
+            InsertMapping(base, DirectMapping{base + keep, mapping.phys, mapping.memoryType, mapping.backing});
         }
         if (mappingLast > last) {
             const auto skip = last - mapping.phys;
-            it = g_directMappings.emplace(base + skip, DirectMapping{mapping.end, last, mapping.memoryType, mapping.backing}).first;
+            it = InsertMapping(base + skip, DirectMapping{mapping.end, last, mapping.memoryType, mapping.backing});
         }
     }
 }
@@ -396,10 +458,16 @@ void ValidatePhysicalRange(std::uint64_t phys, std::size_t len) {
 void WatchMapping(std::uintptr_t address, std::size_t len) {
     GuestWriteWatch::GuestWriteWatchRegister_nid_postfix(reinterpret_cast<const void*>(address), len);
     const auto end = address + len;
+    std::vector<std::uintptr_t> sharing;
     for (auto added = g_directMappings.lower_bound(address); added != g_directMappings.end() && added->first < end; ++added) {
         const auto addedPhys = added->second.phys;
         const auto addedPhysEnd = addedPhys + (added->second.end - added->first);
-        for (const auto& [base, other] : g_directMappings) {
+        sharing.clear();
+        for (auto span = g_physicalSpans.lower_bound(addedPhys); span != g_physicalSpans.end() && span->first < addedPhysEnd; ++span) sharing.insert(sharing.end(), span->second.bases.begin(), span->second.bases.end());
+        std::sort(sharing.begin(), sharing.end());
+        sharing.erase(std::unique(sharing.begin(), sharing.end()), sharing.end());
+        for (const auto base : sharing) {
+            const auto& other = g_directMappings.at(base);
             if ((base >= address && base < end) || other.backing != added->second.backing) continue;
             const auto first = std::max(addedPhys, other.phys);
             const auto last = std::min(addedPhysEnd, other.phys + (other.end - base));
@@ -481,7 +549,7 @@ void AddMapping(std::uintptr_t address, std::size_t len, std::uint64_t phys, int
             bytes += PS5_PAGE_SIZE;
         }
         page.backing->Map(address + offset, bytes, page.offset, nativeProt);
-        g_directMappings.emplace(address + offset, DirectMapping{address + offset + bytes, phys + offset, page.backing->MemoryType(), page.backing});
+        InsertMapping(address + offset, DirectMapping{address + offset + bytes, phys + offset, page.backing->MemoryType(), page.backing});
         offset += bytes;
     }
 #if defined(__linux__)
@@ -945,13 +1013,13 @@ int DoMtypeprotect(const void* addr, size_t len, int type, int prot) {
         while (it != g_directMappings.end() && it->first < end) {
             const auto base = it->first;
             const auto mapping = it->second;
-            it = g_directMappings.erase(it);
-            if (base < first) g_directMappings.emplace(base, DirectMapping{first, mapping.phys, mapping.memoryType, mapping.backing});
+            it = RemoveMapping(it);
+            if (base < first) InsertMapping(base, DirectMapping{first, mapping.phys, mapping.memoryType, mapping.backing});
             const auto low = std::max(base, first);
             const auto high = std::min(mapping.end, end);
-            g_directMappings.emplace(low, DirectMapping{high, mapping.phys + low - base, type, mapping.backing});
+            InsertMapping(low, DirectMapping{high, mapping.phys + low - base, type, mapping.backing});
             physical.emplace_back(mapping.phys + low - base, high - low);
-            if (mapping.end > end) it = g_directMappings.emplace(end, DirectMapping{mapping.end, mapping.phys + end - base, mapping.memoryType, mapping.backing}).first;
+            if (mapping.end > end) it = InsertMapping(end, DirectMapping{mapping.end, mapping.phys + end - base, mapping.memoryType, mapping.backing});
         }
     }
     for (const auto& [phys, bytes] : physical) DirectMemoryRetype(static_cast<int64_t>(phys), bytes, type);

core/libs/tests/GuestMemory.cpp

@@ -31,6 +31,7 @@
 #include <fstream>
 #include <string>
 #include <sys/mman.h>
+#include <time.h>
 #include <unistd.h>
 #endif
 #ifdef _WIN32
@@ -1604,6 +1605,90 @@ static void CheckDirectMemoryWriteWatch() {
     Require(sceKernelMunmap(flexible, page * 3) == 0);
     Require(sceKernelReleaseDirectMemory(phys, page * 3) == 0);
 }
+
+static void CheckDirectMemoryWriteWatchFollowsRemovedMappings() {
+    if (!GuestWriteWatch::GuestWriteWatchAvailable_nid_postfix()) {
+        std::puts("write watch unavailable: removed direct mappings not tested");
+        return;
+    }
+    constexpr std::size_t page = 0x4000;
+    std::int64_t phys = 0;
+    Require(sceKernelAllocateDirectMemory(0, 0x7fffffffffll, page * 4, 0, 0, &phys) == 0);
+    void* whole = nullptr;
+    Require(sceKernelMapDirectMemory(&whole, page * 4, 3, 0, phys, 0) == 0);
+    const auto address = reinterpret_cast<std::uintptr_t>(whole);
+    const auto watched = [](const void* pointer, std::size_t offset) { return GuestWriteWatch::GuestWriteWatchCovers_nid_postfix(reinterpret_cast<std::uintptr_t>(pointer) + offset, page); };
+    void* pair = nullptr;
+    Require(sceKernelMapDirectMemory(&pair, page * 2, 3, 0, phys + page, 0) == 0);
+    Require(watched(whole, 0) && !watched(whole, page) && !watched(whole, page * 2) && watched(whole, page * 3));
+    Require(!watched(pair, 0) && !watched(pair, page));
+    Require(sceKernelMunmap(pair, page) == 0);
+    void* single = nullptr;
+    Require(sceKernelMapDirectMemory(&single, page, 3, 0, phys + page, 0) == 0);
+    Require(!watched(single, 0) && watched(whole, 0) && watched(whole, page * 3));
+    Require(sceKernelMunmap(single, page) == 0);
+    Require(sceKernelMunmap(static_cast<unsigned char*>(pair) + page, page) == 0);
+    void* replaced = whole;
+    Require(sceKernelMapDirectMemory(&replaced, page, 3, 0x10, phys + page * 3, 0) == 0 && replaced == whole);
+    Require(!watched(whole, 0) && !watched(whole, page * 3));
+    std::uintptr_t start = 0;
+    std::uintptr_t end = 0;
+    std::uint64_t offset = 0;
+    int type = 0;
+    Require(QueryDirectMapping(address, &start, &end, &offset, &type) && start == address && end == address + page && offset == static_cast<std::uint64_t>(phys) + page * 3);
+    void* alone = nullptr;
+    Require(sceKernelMapDirectMemory(&alone, page, 3, 0, phys, 0) == 0);
+    Require(watched(alone, 0));
+    void* again = nullptr;
+    Require(sceKernelMapDirectMemory(&again, page, 3, 0, phys, 0) == 0);
+    Require(!watched(alone, 0) && !watched(again, 0));
+    Require(sceKernelMunmap(again, page) == 0);
+    Require(sceKernelMunmap(alone, page) == 0);
+    Require(sceKernelMunmap(whole, page * 4) == 0);
+    Require(sceKernelReleaseDirectMemory(phys, page * 4) == 0);
+}
+
+static double MapAcrossBackingsSeconds(std::int64_t phys, std::size_t backings) {
+    constexpr std::size_t page = 0x4000;
+    const auto threadSeconds = [] {
+        timespec now{};
+        Require(clock_gettime(CLOCK_THREAD_CPUTIME_ID, &now) == 0);
+        return static_cast<double>(now.tv_sec) + static_cast<double>(now.tv_nsec) * 1e-9;
+    };
+    auto best = std::numeric_limits<double>::max();
+    for (int attempt = 0; attempt < 3; ++attempt) {
+        void* mapped = nullptr;
+        const auto before = threadSeconds();
+        Require(sceKernelMapDirectMemory(&mapped, page * backings, 3, 0, phys, 0) == 0);
+        best = std::min(best, threadSeconds() - before);
+        const auto last = reinterpret_cast<std::uintptr_t>(mapped) + page * (backings - 1);
+        std::uintptr_t start = 0;
+        std::uintptr_t end = 0;
+        std::uint64_t offset = 0;
+        int type = 0;
+        Require(QueryDirectMapping(last, &start, &end, &offset, &type) && start == last && end == last + page);
+        Require(sceKernelMunmap(mapped, page * backings) == 0);
+    }
+    return best;
+}
+
+static void CheckMappingCostFollowsTheMappingCount() {
+    constexpr std::size_t page = 0x4000;
+    constexpr std::size_t few = 1024;
+    constexpr std::size_t many = few * 16;
+    std::int64_t phys = 0;
+    Require(sceKernelAllocateDirectMemory(0, 0x7fffffffffll, page * many, 0, 0, &phys) == 0);
+    Require(sceKernelReleaseDirectMemory(phys, page * many) == 0);
+    for (std::size_t index = 0; index < many; ++index) {
+        const auto wanted = phys + static_cast<std::int64_t>(page * index);
+        std::int64_t piece = 0;
+        Require(sceKernelAllocateDirectMemory(wanted, wanted + static_cast<std::int64_t>(page), page, 0, 0, &piece) == 0 && piece == wanted);
+    }
+    const auto small = MapAcrossBackingsSeconds(phys, few);
+    const auto large = MapAcrossBackingsSeconds(phys, many);
+    Require(large <= small * 48);
+    Require(sceKernelReleaseDirectMemory(phys, page * many) == 0);
+}
 #endif
 
 static void CheckVirtualQuerySplitFlexibleRanges() {
@@ -1789,6 +1874,8 @@ int main() {
     CheckWriteWatchArmedCollect();
     CheckWriteWatchFreshCollect();
     CheckDirectMemoryWriteWatch();
+    CheckDirectMemoryWriteWatchFollowsRemovedMappings();
+    CheckMappingCostFollowsTheMappingCount();
     CheckDirectMemoryBackingNeedsNoFilesystem();
     CheckDirectMemorySharedBacking();
 #endif