4 files changed, 98 insertions, 1 deletions

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 cb4913b07..c1eea861d 100644
--- a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp
+++ b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.cpp
@@ -26,6 +26,8 @@ u32 nvhost_as_gpu::ioctl(Ioctl command, const std::vector<u8>& input, std::vecto
         return BindChannel(input, output);
     case IoctlCommand::IocGetVaRegionsCommand:
         return GetVARegions(input, output);
+    case IoctlCommand::IocUnmapBufferCommand:
+        return UnmapBuffer(input, output);
     }
 
     if (static_cast<IoctlCommand>(command.cmd.Value()) == IoctlCommand::IocRemapCommand)
@@ -125,6 +127,37 @@ u32 nvhost_as_gpu::MapBufferEx(const std::vector<u8>& input, std::vector<u8>& ou
         params.offset = gpu.memory_manager->MapBufferEx(object->addr, object->size);
     }
 
+    // Create a new mapping entry for this operation.
+    ASSERT_MSG(buffer_mappings.find(params.offset) == buffer_mappings.end(),
+               "Offset is already mapped");
+
+    BufferMapping mapping{};
+    mapping.nvmap_handle = params.nvmap_handle;
+    mapping.offset = params.offset;
+    mapping.size = object->size;
+
+    buffer_mappings[params.offset] = mapping;
+
+    std::memcpy(output.data(), &params, output.size());
+    return 0;
+}
+
+u32 nvhost_as_gpu::UnmapBuffer(const std::vector<u8>& input, std::vector<u8>& output) {
+    IoctlUnmapBuffer params{};
+    std::memcpy(&params, input.data(), input.size());
+
+    NGLOG_DEBUG(Service_NVDRV, "called, offset=0x{:X}", params.offset);
+
+    auto& gpu = Core::System::GetInstance().GPU();
+
+    auto itr = buffer_mappings.find(params.offset);
+
+    ASSERT_MSG(itr != buffer_mappings.end(), "Tried to unmap invalid mapping");
+
+    params.offset = gpu.memory_manager->UnmapBuffer(params.offset, itr->second.size);
+
+    buffer_mappings.erase(itr->second.offset);
+
     std::memcpy(output.data(), &params, output.size());
     return 0;
 }
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 f2dd0c3b3..d4c4b4db3 100644
--- a/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h
+++ b/src/core/hle/service/nvdrv/devices/nvhost_as_gpu.h
@@ -5,6 +5,7 @@
 #pragma once
 
 #include <memory>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 #include "common/common_types.h"
@@ -30,6 +31,7 @@ private:
         IocMapBufferExCommand = 0xC0284106,
         IocBindChannelCommand = 0x40044101,
         IocGetVaRegionsCommand = 0xC0404108,
+        IocUnmapBufferCommand = 0xC0084105,
     };
 
     struct IoctlInitalizeEx {
@@ -76,6 +78,11 @@ private:
     };
     static_assert(sizeof(IoctlMapBufferEx) == 40, "IoctlMapBufferEx is incorrect size");
 
+    struct IoctlUnmapBuffer {
+        u64_le offset;
+    };
+    static_assert(sizeof(IoctlUnmapBuffer) == 8, "IoctlUnmapBuffer is incorrect size");
+
     struct IoctlBindChannel {
         u32_le fd;
     };
@@ -98,12 +105,22 @@ private:
     static_assert(sizeof(IoctlGetVaRegions) == 16 + sizeof(IoctlVaRegion) * 2,
                   "IoctlGetVaRegions is incorrect size");
 
+    struct BufferMapping {
+        u64 offset;
+        u64 size;
+        u32 nvmap_handle;
+    };
+
+    /// Map containing the nvmap object mappings in GPU memory.
+    std::unordered_map<u64, BufferMapping> buffer_mappings;
+
     u32 channel{};
 
     u32 InitalizeEx(const std::vector<u8>& input, std::vector<u8>& output);
     u32 AllocateSpace(const std::vector<u8>& input, std::vector<u8>& output);
     u32 Remap(const std::vector<u8>& input, std::vector<u8>& output);
     u32 MapBufferEx(const std::vector<u8>& input, std::vector<u8>& output);
+    u32 UnmapBuffer(const std::vector<u8>& input, std::vector<u8>& output);
     u32 BindChannel(const std::vector<u8>& input, std::vector<u8>& output);
     u32 GetVARegions(const std::vector<u8>& input, std::vector<u8>& output);
 
diff --git a/src/core/hle/service/nvdrv/devices/nvmap.cpp b/src/core/hle/service/nvdrv/devices/nvmap.cpp
index 8d883209f..d66fb3a9c 100644
--- a/src/core/hle/service/nvdrv/devices/nvmap.cpp
+++ b/src/core/hle/service/nvdrv/devices/nvmap.cpp
@@ -30,6 +30,8 @@ u32 nvmap::ioctl(Ioctl command, const std::vector<u8>& input, std::vector<u8>& o
         return IocFromId(input, output);
     case IoctlCommand::Param:
         return IocParam(input, output);
+    case IoctlCommand::Free:
+        return IocFree(input, output);
     }
 
     UNIMPLEMENTED_MSG("Unimplemented ioctl");
@@ -45,6 +47,7 @@ u32 nvmap::IocCreate(const std::vector<u8>& input, std::vector<u8>& output) {
     object->id = next_id++;
     object->size = params.size;
     object->status = Object::Status::Created;
+    object->refcount = 1;
 
     u32 handle = next_handle++;
     handles[handle] = std::move(object);
@@ -101,6 +104,8 @@ u32 nvmap::IocFromId(const std::vector<u8>& input, std::vector<u8>& output) {
                             [&](const auto& entry) { return entry.second->id == params.id; });
     ASSERT(itr != handles.end());
 
+    itr->second->refcount++;
+
     // Return the existing handle instead of creating a new one.
     params.handle = itr->first;
 
@@ -142,4 +147,34 @@ u32 nvmap::IocParam(const std::vector<u8>& input, std::vector<u8>& output) {
     return 0;
 }
 
+u32 nvmap::IocFree(const std::vector<u8>& input, std::vector<u8>& output) {
+    enum FreeFlags {
+        Freed = 0,
+        NotFreedYet = 1,
+    };
+
+    IocFreeParams params;
+    std::memcpy(&params, input.data(), sizeof(params));
+
+    NGLOG_WARNING(Service_NVDRV, "(STUBBED) called");
+
+    auto itr = handles.find(params.handle);
+    ASSERT(itr != handles.end());
+
+    itr->second->refcount--;
+
+    params.refcount = itr->second->refcount;
+    params.size = itr->second->size;
+
+    if (itr->second->refcount == 0)
+        params.flags = Freed;
+    else
+        params.flags = NotFreedYet;
+
+    handles.erase(params.handle);
+
+    std::memcpy(output.data(), &params, sizeof(params));
+    return 0;
+}
+
 } // namespace Service::Nvidia::Devices
diff --git a/src/core/hle/service/nvdrv/devices/nvmap.h b/src/core/hle/service/nvdrv/devices/nvmap.h
index 431eb3773..5a3044167 100644
--- a/src/core/hle/service/nvdrv/devices/nvmap.h
+++ b/src/core/hle/service/nvdrv/devices/nvmap.h
@@ -34,6 +34,7 @@ public:
         u8 kind;
         VAddr addr;
         Status status;
+        u32 refcount;
     };
 
     std::shared_ptr<Object> GetObject(u32 handle) const {
@@ -59,7 +60,8 @@ private:
         FromId = 0xC0080103,
         Alloc = 0xC0200104,
         Param = 0xC00C0109,
-        GetId = 0xC008010E
+        GetId = 0xC008010E,
+        Free = 0xC0180105,
     };
 
     struct IocCreateParams {
@@ -102,11 +104,21 @@ private:
         u32_le value;
     };
 
+    struct IocFreeParams {
+        u32_le handle;
+        INSERT_PADDING_BYTES(4);
+        u64_le refcount;
+        u32_le size;
+        u32_le flags;
+    };
+    static_assert(sizeof(IocFreeParams) == 24, "IocFreeParams has wrong size");
+
     u32 IocCreate(const std::vector<u8>& input, std::vector<u8>& output);
     u32 IocAlloc(const std::vector<u8>& input, std::vector<u8>& output);
     u32 IocGetId(const std::vector<u8>& input, std::vector<u8>& output);
     u32 IocFromId(const std::vector<u8>& input, std::vector<u8>& output);
     u32 IocParam(const std::vector<u8>& input, std::vector<u8>& output);
+    u32 IocFree(const std::vector<u8>& input, std::vector<u8>& output);
 };
 
 } // namespace Service::Nvidia::Devices