memory: Fixes for direct memory allocation. (#615)

* memory: Size direct memory based on requested flexible memory.

* memory: Guard against OrbisProcParam without an OrbisKernelMemParam.

* memory: Account for alignment in direct memory suitability checks and add more debugging.

src/core/libraries/kernel/memory_management.cpp

@@ -15,7 +15,8 @@ namespace Libraries::Kernel {
 
 u64 PS4_SYSV_ABI sceKernelGetDirectMemorySize() {
     LOG_WARNING(Kernel_Vmm, "called");
-    return SCE_KERNEL_MAIN_DMEM_SIZE;
+    const auto* memory = Core::Memory::Instance();
+    return memory->GetTotalDirectSize();
 }
 
 int PS4_SYSV_ABI sceKernelAllocateDirectMemory(s64 searchStart, s64 searchEnd, u64 len,
@@ -52,8 +53,8 @@ int PS4_SYSV_ABI sceKernelAllocateDirectMemory(s64 searchStart, s64 searchEnd, u
 
 s32 PS4_SYSV_ABI sceKernelAllocateMainDirectMemory(size_t len, size_t alignment, int memoryType,
                                                    s64* physAddrOut) {
-    return sceKernelAllocateDirectMemory(0, SCE_KERNEL_MAIN_DMEM_SIZE, len, alignment, memoryType,
-                                         physAddrOut);
+    const auto searchEnd = static_cast<s64>(sceKernelGetDirectMemorySize());
+    return sceKernelAllocateDirectMemory(0, searchEnd, len, alignment, memoryType, physAddrOut);
 }
 
 s32 PS4_SYSV_ABI sceKernelCheckedReleaseDirectMemory(u64 start, size_t len) {
@@ -78,7 +79,7 @@ s32 PS4_SYSV_ABI sceKernelAvailableDirectMemorySize(u64 searchStart, u64 searchE
     if (physAddrOut == nullptr || sizeOut == nullptr) {
         return ORBIS_KERNEL_ERROR_EINVAL;
     }
-    if (searchEnd > SCE_KERNEL_MAIN_DMEM_SIZE) {
+    if (searchEnd > sceKernelGetDirectMemorySize()) {
         return ORBIS_KERNEL_ERROR_EINVAL;
     }
     if (searchEnd <= searchStart) {

src/core/libraries/kernel/memory_management.h

@@ -6,7 +6,7 @@
 #include "common/bit_field.h"
 #include "common/types.h"
 
-constexpr u64 SCE_KERNEL_MAIN_DMEM_SIZE = 4608_MB; // ~ 4.5GB
+constexpr u64 SCE_KERNEL_MAIN_DMEM_SIZE = 5056_MB; // ~ 5GB
 
 namespace Libraries::Kernel {
 

src/core/linker.cpp

@@ -68,11 +68,19 @@ void Linker::Execute() {
     }
 
     // Configure used flexible memory size.
-    // if (auto* mem_param = GetProcParam()->mem_param) {
-    //      if (u64* flexible_size = mem_param->flexible_memory_size) {
-    //          memory->SetTotalFlexibleSize(*flexible_size);
-    //      }
-    //  }
+    if (const auto* proc_param = GetProcParam()) {
+        if (proc_param->size >=
+            offsetof(OrbisProcParam, mem_param) + sizeof(OrbisKernelMemParam*)) {
+            if (const auto* mem_param = proc_param->mem_param) {
+                if (mem_param->size >=
+                    offsetof(OrbisKernelMemParam, flexible_memory_size) + sizeof(u64*)) {
+                    if (const auto* flexible_size = mem_param->flexible_memory_size) {
+                        memory->SetupMemoryRegions(*flexible_size);
+                    }
+                }
+            }
+        }
+    }
 
     // Init primary thread.
     Common::SetCurrentThreadName("GAME_MainThread");

src/core/memory.cpp

@@ -11,9 +11,11 @@
 
 namespace Core {
 
+constexpr u64 SCE_DEFAULT_FLEXIBLE_MEMORY_SIZE = 448_MB;
+
 MemoryManager::MemoryManager() {
-    // Insert an area that covers direct memory physical block.
-    dmem_map.emplace(0, DirectMemoryArea{0, SCE_KERNEL_MAIN_DMEM_SIZE});
+    // Set up the direct and flexible memory regions.
+    SetupMemoryRegions(SCE_DEFAULT_FLEXIBLE_MEMORY_SIZE);
 
     // Insert a virtual memory area that covers the entire area we manage.
     const VAddr system_managed_base = impl.SystemManagedVirtualBase();
@@ -35,6 +37,19 @@ MemoryManager::MemoryManager() {
 
 MemoryManager::~MemoryManager() = default;
 
+void MemoryManager::SetupMemoryRegions(u64 flexible_size) {
+    total_flexible_size = flexible_size;
+    total_direct_size = SCE_KERNEL_MAIN_DMEM_SIZE - flexible_size;
+
+    // Insert an area that covers direct memory physical block.
+    // Note that this should never be called after direct memory allocations have been made.
+    dmem_map.clear();
+    dmem_map.emplace(0, DirectMemoryArea{0, total_direct_size});
+
+    LOG_INFO(Kernel_Vmm, "Configured memory regions: flexible size = {:#x}, direct size = {:#x}",
+             total_flexible_size, total_direct_size);
+}
+
 PAddr MemoryManager::Allocate(PAddr search_start, PAddr search_end, size_t size, u64 alignment,
                               int memory_type) {
     std::scoped_lock lk{mutex};
@@ -42,12 +57,17 @@ PAddr MemoryManager::Allocate(PAddr search_start, PAddr search_end, size_t size,
     auto dmem_area = FindDmemArea(search_start);
 
     const auto is_suitable = [&] {
-        return dmem_area->second.is_free && dmem_area->second.size >= size;
+        const auto aligned_base = alignment > 0 ? Common::AlignUp(dmem_area->second.base, alignment)
+                                                : dmem_area->second.base;
+        const auto alignment_size = aligned_base - dmem_area->second.base;
+        const auto remaining_size =
+            dmem_area->second.size >= alignment_size ? dmem_area->second.size - alignment_size : 0;
+        return dmem_area->second.is_free && remaining_size >= size;
     };
     while (!is_suitable() && dmem_area->second.GetEnd() <= search_end) {
         dmem_area++;
     }
-    ASSERT_MSG(is_suitable(), "Unable to find free direct memory area");
+    ASSERT_MSG(is_suitable(), "Unable to find free direct memory area: size = {:#x}", size);
 
     // Align free position
     PAddr free_addr = dmem_area->second.base;
@@ -333,14 +353,19 @@ int MemoryManager::DirectQueryAvailable(PAddr search_start, PAddr search_end, si
             continue;
         }
 
-        if (dmem_area->second.size > max_size) {
-            paddr = dmem_area->second.base;
-            max_size = dmem_area->second.size;
+        const auto aligned_base = alignment > 0 ? Common::AlignUp(dmem_area->second.base, alignment)
+                                                : dmem_area->second.base;
+        const auto alignment_size = aligned_base - dmem_area->second.base;
+        const auto remaining_size =
+            dmem_area->second.size >= alignment_size ? dmem_area->second.size - alignment_size : 0;
+        if (remaining_size > max_size) {
+            paddr = aligned_base;
+            max_size = remaining_size;
         }
         dmem_area++;
     }
 
-    *phys_addr_out = alignment > 0 ? Common::AlignUp(paddr, alignment) : paddr;
+    *phys_addr_out = paddr;
     *size_out = max_size;
     return ORBIS_OK;
 }
@@ -418,7 +443,8 @@ MemoryManager::DMemHandle MemoryManager::CarveDmemArea(PAddr addr, size_t size)
 
     const PAddr start_in_area = addr - area.base;
     const PAddr end_in_vma = start_in_area + size;
-    ASSERT_MSG(end_in_vma <= area.size, "Mapping cannot fit inside free region");
+    ASSERT_MSG(end_in_vma <= area.size, "Mapping cannot fit inside free region: size = {:#x}",
+               size);
 
     if (end_in_vma != area.size) {
         // Split VMA at the end of the allocated region

src/core/memory.h

@@ -130,8 +130,8 @@ public:
         rasterizer = rasterizer_;
     }
 
-    void SetTotalFlexibleSize(u64 size) {
-        total_flexible_size = size;
+    u64 GetTotalDirectSize() const {
+        return total_direct_size;
     }
 
     u64 GetAvailableFlexibleSize() const {
@@ -142,6 +142,8 @@ public:
         return impl.SystemReservedVirtualBase();
     }
 
+    void SetupMemoryRegions(u64 flexible_size);
+
     PAddr Allocate(PAddr search_start, PAddr search_end, size_t size, u64 alignment,
                    int memory_type);
 
@@ -217,7 +219,8 @@ private:
     DMemMap dmem_map;
     VMAMap vma_map;
     std::recursive_mutex mutex;
-    size_t total_flexible_size = 448_MB;
+    size_t total_direct_size{};
+    size_t total_flexible_size{};
     size_t flexible_usage{};
     Vulkan::Rasterizer* rasterizer{};
 };