nvhost_as_gpu: share Switch 2 device ranges between a process's address spaces

Each address space allocated its own SMMU range for the same memory
outside the heap. The caches key memory by device address, so these were
separate copies, and presentation used whichever range was recorded
last: Fire Emblem: Fortune's Weave drew into one alias while VI showed
another, and the screen stayed black.

Key these ranges by (ASID, CPU address, size) and count references. The
first address space allocates the range, the last release frees it, and
a closed session's entries are dropped so a reused ASID can't inherit
them. EDEN_NX2_PRIVATE_DEVICE_RANGES=1 restores a range per address
space.

Assisted-by: Claude Opus 5.5 (debugging, implementation and drafting)
Signed-off-by: duck912 <350528965+Duck912@users.noreply.github.com>

4 files changed, 90 insertions, 11 deletions

diff --git a/src/core/hle/service/nvdrv/core/container.cpp b/src/core/hle/service/nvdrv/core/container.cpp
index 9edce03f64..22db243e3f 100644
--- a/src/core/hle/service/nvdrv/core/container.cpp
+++ b/src/core/hle/service/nvdrv/core/container.cpp
@@ -11,6 +11,7 @@
 #include "core/hle/service/nvdrv/core/heap_mapper.h"
 #include "core/hle/service/nvdrv/core/nvmap.h"
 #include "core/hle/service/nvdrv/core/syncpoint_manager.h"
+#include "core/hle/service/nvdrv/devices/nx2_stub_device.h"
 #include "core/memory.h"
 #include "video_core/host1x/host1x.h"
 
@@ -130,6 +131,7 @@ void Container::CloseSession(SessionId session_id) {
         session.has_preallocated_area = false;
     }
     session.is_active = false;
+    Devices::Nx2GpuMemoryRegistry::Get().ForgetDeviceRanges(session.asid.id);
     smmu.UnregisterProcess(impl->sessions[session_id.id].asid);
     impl->id_pool.emplace_front(session_id.id);
 }
diff --git a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp
index e66c72cd26..9d696c027e 100644
--- a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp
+++ b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp
@@ -5,6 +5,7 @@
 // SPDX-FileCopyrightText: 2021 Skyline Team and Contributors
 // SPDX-License-Identifier: GPL-3.0-or-later
 
+#include <cstdlib>
 #include <cstring>
 #include <utility>
 
@@ -35,6 +36,13 @@ nvhost_as_gpu::nvhost_as_gpu(Core::System& system_, Module& module_, NvCore::Con
 
 nvhost_as_gpu::~nvhost_as_gpu() {
     Nx2GpuMemoryRegistry::Get().RemoveListener(this);
+    // Leave shared ranges to the process's other address spaces; like this address space's own
+    // ranges, one nobody else maps stays allocated.
+    for (const auto& [key, range] : nx2_device_ranges) {
+        if (range.shared) {
+            Nx2GpuMemoryRegistry::Get().ReleaseDeviceRange(range.asid, key.first, key.second);
+        }
+    }
 }
 
 NvResult nvhost_as_gpu::Ioctl1(DeviceFD fd, Ioctl command, std::span<const u8> input,
@@ -134,25 +142,37 @@ void nvhost_as_gpu::Nx2Map(u64 gpu_address, u64 cpu_address, u64 size) {
     }
     // The same CPU range is often mapped twice (at its own address and in a sparse region), so
     // reuse its SMMU mapping.
-    DAddr device_address{};
-    bool heap = false;
+    Nx2DeviceRange range{};
     if (const auto it = nx2_device_ranges.find({cpu_address, size}); it != nx2_device_ranges.end()) {
-        device_address = it->second.address;
-        heap = it->second.heap;
+        range = it->second;
     } else if (session->has_preallocated_area && session->mapper->IsInBounds(cpu_address, size)) {
-        device_address = session->mapper->Map(cpu_address, size);
-        heap = true;
+        range = {session->mapper->Map(cpu_address, size), true};
     } else {
-        auto& smmu = system.Host1x().MemoryManager();
-        device_address = smmu.Allocate(Common::AlignUp(size, 0x20000));
-        if (device_address == 0) {
+        const auto allocate = [&]() -> u64 {
+            auto& smmu = system.Host1x().MemoryManager();
+            const DAddr address = smmu.Allocate(Common::AlignUp(size, 0x20000));
+            if (address != 0) {
+                smmu.Map(address, cpu_address, size, session->asid, true);
+            }
+            return address;
+        };
+        // EDEN_NX2_PRIVATE_DEVICE_RANGES=1 (A/B aid): a range per address space.
+        static const bool private_ranges = std::getenv("EDEN_NX2_PRIVATE_DEVICE_RANGES") != nullptr;
+        if (private_ranges) {
+            range = {allocate(), false};
+        } else {
+            range = {Nx2GpuMemoryRegistry::Get().AcquireDeviceRange(session->asid.id, cpu_address,
+                                                                    size, allocate),
+                     false, true, session->asid.id};
+        }
+        if (range.address == 0) {
             LOG_CRITICAL(Service_NVDRV, "Switch 2 GPU memory {:#x}+{:#x}: out of SMMU space",
                          cpu_address, size);
             return;
         }
-        smmu.Map(device_address, cpu_address, size, session->asid, true);
     }
-    nx2_device_ranges[{cpu_address, size}] = {device_address, heap};
+    const DAddr device_address = range.address;
+    nx2_device_ranges[{cpu_address, size}] = range;
     Nx2GpuMemoryRegistry::Get().RecordDeviceRange(cpu_address, size, device_address);
     gmmu->Map(gpu_address, device_address, size, Tegra::PTEKind::PITCH, false);
     LOG_DEBUG(Service_NVDRV, "Switch 2 GPU mapping {:#x}+{:#x} -> CPU {:#x}", gpu_address, size,
@@ -177,6 +197,10 @@ void nvhost_as_gpu::Nx2Release(u64 cpu_address, u64 size) {
     }
     const Nx2DeviceRange range = it->second;
     nx2_device_ranges.erase(it);
+    if (range.shared &&
+        !Nx2GpuMemoryRegistry::Get().ReleaseDeviceRange(range.asid, cpu_address, size)) {
+        return; // another address space of the process still maps it
+    }
     auto* session = container.GetSession(nx2_session);
     auto* const mapper = range.heap && session != nullptr ? session->mapper.get() : nullptr;
     auto& smmu = system.Host1x().MemoryManager();
diff --git a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h
index 39b63ea369..37b6d0d4c4 100644
--- a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h
+++ b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h
@@ -238,6 +238,8 @@ private:
     struct Nx2DeviceRange {
         DAddr address;
         bool heap; ///< Mapped through the session's heap mapper (else an SMMU allocation)
+        bool shared{}; ///< A reference to the process's range in Nx2GpuMemoryRegistry
+        size_t asid{};
     };
     std::map<std::pair<u64, u64>, Nx2DeviceRange> nx2_device_ranges;
     void Nx2Release(u64 cpu_address, u64 size);
diff --git a/src/core/hle/service/nvdrv/devices/nx2_stub_device.h b/src/core/hle/service/nvdrv/devices/nx2_stub_device.h
index 869e17caf6..8dee9ccfb7 100644
--- a/src/core/hle/service/nvdrv/devices/nx2_stub_device.h
+++ b/src/core/hle/service/nvdrv/devices/nx2_stub_device.h
@@ -5,6 +5,7 @@
 
 #include <algorithm>
 #include <atomic>
+#include <compare>
 #include <functional>
 #include <map>
 #include <mutex>
@@ -166,12 +167,50 @@ public:
         {
             std::scoped_lock lk{device_mutex};
             device_ranges.clear();
+            shared_ranges.clear();
         }
         std::scoped_lock lk{mutex};
         mappings.clear();
         sparse_regions.clear();
     }
 
+    // Device (SMMU) range of a mapping outside the heap, shared by every address space of one
+    // process: the caches key memory by device address, so a range per address space made separate
+    // copies of the same memory, and presentation could read one the game never drew into. The
+    // first address space to map the range allocates it with `allocate` (0: out of space); every
+    // call takes a reference.
+    u64 AcquireDeviceRange(size_t asid, u64 cpu_address, u64 size,
+                           const std::function<u64()>& allocate) {
+        std::scoped_lock lk{device_mutex};
+        const SharedRangeKey key{asid, cpu_address, size};
+        if (const auto it = shared_ranges.find(key); it != shared_ranges.end()) {
+            ++it->second.references;
+            return it->second.device_address;
+        }
+        const u64 device_address = allocate();
+        if (device_address != 0) {
+            shared_ranges.emplace(key, SharedRange{device_address, 1});
+        }
+        return device_address;
+    }
+
+    // Drops a reference. True for the last one: the caller frees the range.
+    bool ReleaseDeviceRange(size_t asid, u64 cpu_address, u64 size) {
+        std::scoped_lock lk{device_mutex};
+        const auto it = shared_ranges.find(SharedRangeKey{asid, cpu_address, size});
+        if (it == shared_ranges.end() || --it->second.references != 0) {
+            return false;
+        }
+        shared_ranges.erase(it);
+        return true;
+    }
+
+    // The process's session closed and its ASID can be reused: stop sharing its ranges.
+    void ForgetDeviceRanges(size_t asid) {
+        std::scoped_lock lk{device_mutex};
+        std::erase_if(shared_ranges, [asid](const auto& entry) { return entry.first.asid == asid; });
+    }
+
     // Device (SMMU) address of a CPU range an address space mapped; used to present images.
     void RecordDeviceRange(u64 cpu_address, u64 size, u64 device_address) {
         std::scoped_lock lk{device_mutex};
@@ -197,6 +236,18 @@ private:
     std::mutex device_mutex;
     std::map<u64, std::pair<u64, u64>> device_ranges; // cpu start -> (size, device address)
 
+    struct SharedRangeKey {
+        size_t asid;
+        u64 cpu_address;
+        u64 size;
+        auto operator<=>(const SharedRangeKey&) const = default;
+    };
+    struct SharedRange {
+        u64 device_address;
+        u32 references;
+    };
+    std::map<SharedRangeKey, SharedRange> shared_ranges;
+
     struct Mapping {
         u64 gpu_address;
         u64 cpu_address;