From 703fe8410e63607c7e9152498586428ebe012aa1 Mon Sep 17 00:00:00 2001
From: Dmugetsu <diegoavilafacebook75@gmail.com>
Date: Fri, 7 Mar 2025 15:13:57 -0600
Subject: [PATCH] Unified Memory Allocation WIP

---
 src/core/libraries/kernel/memory.cpp |  6 +++---
 src/core/libraries/kernel/memory.h   | 13 +++++++++----
 src/core/memory.cpp                  |  9 +++++++++
 src/core/memory.h                    | 18 ++++++++----------
 4 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/src/core/libraries/kernel/memory.cpp b/src/core/libraries/kernel/memory.cpp
index dc317237c..ccdf2c05a 100644
--- a/src/core/libraries/kernel/memory.cpp
+++ b/src/core/libraries/kernel/memory.cpp
@@ -22,7 +22,7 @@ namespace Libraries::Kernel {
 u64 PS4_SYSV_ABI sceKernelGetDirectMemorySize() {
     LOG_WARNING(Kernel_Vmm, "called");
     const auto* memory = Core::Memory::Instance();
-    return memory->GetTotalDirectSize();
+    return memory->GetTotalUnifiedMemorySize();
 }
 
 int PS4_SYSV_ABI sceKernelAllocateDirectMemory(s64 searchStart, s64 searchEnd, u64 len,
@@ -276,7 +276,7 @@ int PS4_SYSV_ABI sceKernelDirectMemoryQuery(u64 offset, int flags, OrbisQueryInf
 
 s32 PS4_SYSV_ABI sceKernelAvailableFlexibleMemorySize(size_t* out_size) {
     auto* memory = Core::Memory::Instance();
-    *out_size = memory->GetAvailableFlexibleSize();
+    *out_size = memory->GetTotalUnifiedMemorySize();
     LOG_INFO(Kernel_Vmm, "called size = {:#x}", *out_size);
     return ORBIS_OK;
 }
@@ -511,7 +511,7 @@ s32 PS4_SYSV_ABI sceKernelConfiguredFlexibleMemorySize(u64* sizeOut) {
     }
 
     auto* memory = Core::Memory::Instance();
-    *sizeOut = memory->GetTotalFlexibleSize();
+    *sizeOut = memory->GetTotalUnifiedMemorySize();
     return ORBIS_OK;
 }
 
diff --git a/src/core/libraries/kernel/memory.h b/src/core/libraries/kernel/memory.h
index 1b47f93a8..348e1d8ed 100644
--- a/src/core/libraries/kernel/memory.h
+++ b/src/core/libraries/kernel/memory.h
@@ -5,12 +5,17 @@
 
 #include "common/bit_field.h"
 #include "common/types.h"
+#include "core/linker.h"
+#include "core/memory.h"
+#include "src/common/memory_patcher.h"  
 
-constexpr u64 SCE_KERNEL_TOTAL_MEM = 5248_MB * 4;
-constexpr u64 SCE_KERNEL_TOTAL_MEM_PRO = 5888_MB * 4;
+constexpr bool USE_EXPANDED_MEMORY = true;
 
-constexpr u64 SCE_FLEXIBLE_MEMORY_BASE = 64_MB;
-constexpr u64 SCE_FLEXIBLE_MEMORY_SIZE = 512_MB;
+constexpr u64 SCE_KERNEL_TOTAL_MEM = (USE_EXPANDED_MEMORY ? 5248_MB * 1.2 : 5248_MB);
+constexpr u64 SCE_KERNEL_TOTAL_MEM_PRO = (USE_EXPANDED_MEMORY ? 5888_MB * 1.2 : 5888_MB);
+
+constexpr u64 SCE_FLEXIBLE_MEMORY_BASE = (USE_EXPANDED_MEMORY ? 64_MB * 1.2 : 64_MB);
+constexpr u64 SCE_FLEXIBLE_MEMORY_SIZE = (USE_EXPANDED_MEMORY ? 512_MB * 1.2 : 512_MB);
 
 namespace Core::Loader {
 class SymbolsResolver;
diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index 98d587e00..282179da9 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -37,14 +37,23 @@ MemoryManager::~MemoryManager() = default;
 void MemoryManager::SetupMemoryRegions(u64 flexible_size, bool use_extended_mem1,
                                        bool use_extended_mem2) {
     const bool is_neo = ::Libraries::Kernel::sceKernelIsNeoMode();
+
     auto total_size = is_neo ? SCE_KERNEL_TOTAL_MEM_PRO : SCE_KERNEL_TOTAL_MEM;
+    if (NeedsExtraMemory()) {
+        total_size = SCE_KERNEL_TOTAL_MEM * 1.2;
+    }
+
     if (!use_extended_mem1 && is_neo) {
         total_size -= 256_MB;
     }
     if (!use_extended_mem2 && !is_neo) {
         total_size -= 128_MB;
     }
+
     total_flexible_size = flexible_size - SCE_FLEXIBLE_MEMORY_BASE;
+    if (NeedsExtraMemory()) {
+        total_flexible_size = flexible_size - SCE_FLEXIBLE_MEMORY_BASE * 1.2;
+    }
     total_direct_size = total_size - flexible_size;
 
     // Insert an area that covers direct memory physical block.
diff --git a/src/core/memory.h b/src/core/memory.h
index 812ca5518..f89645237 100644
--- a/src/core/memory.h
+++ b/src/core/memory.h
@@ -36,6 +36,7 @@ enum class MemoryProt : u32 {
     GpuRead = 16,
     GpuWrite = 32,
     GpuReadWrite = 48,
+    UnifiedReadWrite = 50,
 };
 DECLARE_ENUM_FLAG_OPERATORS(MemoryProt)
 
@@ -144,16 +145,12 @@ public:
         return impl;
     }
 
-    u64 GetTotalDirectSize() const {
-        return total_direct_size;
+    u64 GetTotalUnifiedMemorySize() const {
+        return total_direct_size + total_flexible_size;
     }
 
-    u64 GetTotalFlexibleSize() const {
-        return total_flexible_size;
-    }
-
-    u64 GetAvailableFlexibleSize() const {
-        return total_flexible_size - flexible_usage;
+    u64 GetAvailableUnifiedMemorySize() const {
+        return GetTotalUnifiedMemorySize() - flexible_usage;
     }
 
     VAddr SystemReservedVirtualBase() noexcept {
@@ -168,8 +165,9 @@ public:
     }
     bool NeedsExtraMemory() {
         static const std::unordered_set<std::string> extra_memory_games = {
-            // here should be games with needs of extra memory as EXAMPLE
-            "CUSA03173"};
+            "CUSA28615", // Example game that needs extra memory
+            "CUSA03173"  // Add other games here if needed
+        };
 
         return extra_memory_games.find(MemoryPatcher::g_game_serial) != extra_memory_games.end();
     }