1 files changed, 88 insertions, 32 deletions

diff --git a/src/core/hle/service/nvdrv/core/nvmap.cpp b/src/core/hle/service/nvdrv/core/nvmap.cpp
index 0ca05257e..1b59c6b15 100644
--- a/src/core/hle/service/nvdrv/core/nvmap.cpp
+++ b/src/core/hle/service/nvdrv/core/nvmap.cpp
@@ -2,14 +2,19 @@
 // SPDX-FileCopyrightText: 2022 Skyline Team and Contributors
 // SPDX-License-Identifier: GPL-3.0-or-later
 
+#include <functional>
+
 #include "common/alignment.h"
 #include "common/assert.h"
 #include "common/logging/log.h"
+#include "core/hle/service/nvdrv/core/container.h"
+#include "core/hle/service/nvdrv/core/heap_mapper.h"
 #include "core/hle/service/nvdrv/core/nvmap.h"
 #include "core/memory.h"
 #include "video_core/host1x/host1x.h"
 
 using Core::Memory::YUZU_PAGESIZE;
+constexpr size_t BIG_PAGE_SIZE = YUZU_PAGESIZE * 16;
 
 namespace Service::Nvidia::NvCore {
 NvMap::Handle::Handle(u64 size_, Id id_)
@@ -17,9 +22,9 @@ NvMap::Handle::Handle(u64 size_, Id id_)
     flags.raw = 0;
 }
 
-NvResult NvMap::Handle::Alloc(Flags pFlags, u32 pAlign, u8 pKind, u64 pAddress) {
+NvResult NvMap::Handle::Alloc(Flags pFlags, u32 pAlign, u8 pKind, u64 pAddress,
+                              NvCore::SessionId pSessionId) {
     std::scoped_lock lock(mutex);
-
     // Handles cannot be allocated twice
     if (allocated) {
         return NvResult::AccessDenied;
@@ -28,6 +33,7 @@ NvResult NvMap::Handle::Alloc(Flags pFlags, u32 pAlign, u8 pKind, u64 pAddress)
     flags = pFlags;
     kind = pKind;
     align = pAlign < YUZU_PAGESIZE ? YUZU_PAGESIZE : pAlign;
+    session_id = pSessionId;
 
     // This flag is only applicable for handles with an address passed
     if (pAddress) {
@@ -63,7 +69,7 @@ NvResult NvMap::Handle::Duplicate(bool internal_session) {
     return NvResult::Success;
 }
 
-NvMap::NvMap(Tegra::Host1x::Host1x& host1x_) : host1x{host1x_} {}
+NvMap::NvMap(Container& core_, Tegra::Host1x::Host1x& host1x_) : host1x{host1x_}, core{core_} {}
 
 void NvMap::AddHandle(std::shared_ptr<Handle> handle_description) {
     std::scoped_lock lock(handles_lock);
@@ -78,12 +84,30 @@ void NvMap::UnmapHandle(Handle& handle_description) {
         handle_description.unmap_queue_entry.reset();
     }
 
+    // Free and unmap the handle from Host1x GMMU
+    if (handle_description.pin_virt_address) {
+        host1x.GMMU().Unmap(static_cast<GPUVAddr>(handle_description.pin_virt_address),
+                            handle_description.aligned_size);
+        host1x.Allocator().Free(handle_description.pin_virt_address,
+                                static_cast<u32>(handle_description.aligned_size));
+        handle_description.pin_virt_address = 0;
+    }
+
     // Free and unmap the handle from the SMMU
-    host1x.MemoryManager().Unmap(static_cast<GPUVAddr>(handle_description.pin_virt_address),
-                                 handle_description.aligned_size);
-    host1x.Allocator().Free(handle_description.pin_virt_address,
-                            static_cast<u32>(handle_description.aligned_size));
-    handle_description.pin_virt_address = 0;
+    const size_t map_size = handle_description.aligned_size;
+    if (!handle_description.in_heap) {
+        auto& smmu = host1x.MemoryManager();
+        size_t aligned_up = Common::AlignUp(map_size, BIG_PAGE_SIZE);
+        smmu.Unmap(handle_description.d_address, map_size);
+        smmu.Free(handle_description.d_address, static_cast<size_t>(aligned_up));
+        handle_description.d_address = 0;
+        return;
+    }
+    const VAddr vaddress = handle_description.address;
+    auto* session = core.GetSession(handle_description.session_id);
+    session->mapper->Unmap(vaddress, map_size);
+    handle_description.d_address = 0;
+    handle_description.in_heap = false;
 }
 
 bool NvMap::TryRemoveHandle(const Handle& handle_description) {
@@ -124,22 +148,33 @@ std::shared_ptr<NvMap::Handle> NvMap::GetHandle(Handle::Id handle) {
     }
 }
 
-VAddr NvMap::GetHandleAddress(Handle::Id handle) {
+DAddr NvMap::GetHandleAddress(Handle::Id handle) {
     std::scoped_lock lock(handles_lock);
     try {
-        return handles.at(handle)->address;
+        return handles.at(handle)->d_address;
     } catch (std::out_of_range&) {
         return 0;
     }
 }
 
-u32 NvMap::PinHandle(NvMap::Handle::Id handle) {
+DAddr NvMap::PinHandle(NvMap::Handle::Id handle, bool low_area_pin) {
     auto handle_description{GetHandle(handle)};
     if (!handle_description) [[unlikely]] {
         return 0;
     }
 
     std::scoped_lock lock(handle_description->mutex);
+    const auto map_low_area = [&] {
+        if (handle_description->pin_virt_address == 0) {
+            auto& gmmu_allocator = host1x.Allocator();
+            auto& gmmu = host1x.GMMU();
+            u32 address =
+                gmmu_allocator.Allocate(static_cast<u32>(handle_description->aligned_size));
+            gmmu.Map(static_cast<GPUVAddr>(address), handle_description->d_address,
+                     handle_description->aligned_size);
+            handle_description->pin_virt_address = address;
+        }
+    };
     if (!handle_description->pins) {
         // If we're in the unmap queue we can just remove ourselves and return since we're already
         // mapped
@@ -151,37 +186,58 @@ u32 NvMap::PinHandle(NvMap::Handle::Id handle) {
                 unmap_queue.erase(*handle_description->unmap_queue_entry);
                 handle_description->unmap_queue_entry.reset();
 
+                if (low_area_pin) {
+                    map_low_area();
+                    handle_description->pins++;
+                    return static_cast<DAddr>(handle_description->pin_virt_address);
+                }
+
                 handle_description->pins++;
-                return handle_description->pin_virt_address;
+                return handle_description->d_address;
             }
         }
 
+        using namespace std::placeholders;
         // If not then allocate some space and map it
-        u32 address{};
-        auto& smmu_allocator = host1x.Allocator();
-        auto& smmu_memory_manager = host1x.MemoryManager();
-        while ((address = smmu_allocator.Allocate(
-                    static_cast<u32>(handle_description->aligned_size))) == 0) {
-            // Free handles until the allocation succeeds
-            std::scoped_lock queueLock(unmap_queue_lock);
-            if (auto freeHandleDesc{unmap_queue.front()}) {
-                // Handles in the unmap queue are guaranteed not to be pinned so don't bother
-                // checking if they are before unmapping
-                std::scoped_lock freeLock(freeHandleDesc->mutex);
-                if (handle_description->pin_virt_address)
-                    UnmapHandle(*freeHandleDesc);
-            } else {
-                LOG_CRITICAL(Service_NVDRV, "Ran out of SMMU address space!");
+        DAddr address{};
+        auto& smmu = host1x.MemoryManager();
+        auto* session = core.GetSession(handle_description->session_id);
+        const VAddr vaddress = handle_description->address;
+        const size_t map_size = handle_description->aligned_size;
+        if (session->has_preallocated_area && session->mapper->IsInBounds(vaddress, map_size)) {
+            handle_description->d_address = session->mapper->Map(vaddress, map_size);
+            handle_description->in_heap = true;
+        } else {
+            size_t aligned_up = Common::AlignUp(map_size, BIG_PAGE_SIZE);
+            while ((address = smmu.Allocate(aligned_up)) == 0) {
+                // Free handles until the allocation succeeds
+                std::scoped_lock queueLock(unmap_queue_lock);
+                if (auto freeHandleDesc{unmap_queue.front()}) {
+                    // Handles in the unmap queue are guaranteed not to be pinned so don't bother
+                    // checking if they are before unmapping
+                    std::scoped_lock freeLock(freeHandleDesc->mutex);
+                    if (handle_description->d_address)
+                        UnmapHandle(*freeHandleDesc);
+                } else {
+                    LOG_CRITICAL(Service_NVDRV, "Ran out of SMMU address space!");
+                }
             }
+
+            handle_description->d_address = address;
+            smmu.Map(address, vaddress, map_size, session->asid, true);
+            handle_description->in_heap = false;
         }
+    }
 
-        smmu_memory_manager.Map(static_cast<GPUVAddr>(address), handle_description->address,
-                                handle_description->aligned_size);
-        handle_description->pin_virt_address = address;
+    if (low_area_pin) {
+        map_low_area();
     }
 
     handle_description->pins++;
-    return handle_description->pin_virt_address;
+    if (low_area_pin) {
+        return static_cast<DAddr>(handle_description->pin_virt_address);
+    }
+    return handle_description->d_address;
 }
 
 void NvMap::UnpinHandle(Handle::Id handle) {
@@ -232,7 +288,7 @@ std::optional<NvMap::FreeInfo> NvMap::FreeHandle(Handle::Id handle, bool interna
                 LOG_WARNING(Service_NVDRV, "User duplicate count imbalance detected!");
             } else if (handle_description->dupes == 0) {
                 // Force unmap the handle
-                if (handle_description->pin_virt_address) {
+                if (handle_description->d_address) {
                     std::scoped_lock queueLock(unmap_queue_lock);
                     UnmapHandle(*handle_description);
                 }