From ab2eeeca4c8ce90f21a1aa73af6732e5d301b213 Mon Sep 17 00:00:00 2001
From: "Milczarek, Slawomir" <slawomir.milczarek@intel.com>
Date: Mon, 22 Apr 2024 10:42:12 +0000
Subject: [PATCH] feature: Add TBX support for USM shared memory with USM host
 ptr

Related-To: NEO-11158

Signed-off-by: Milczarek, Slawomir <slawomir.milczarek@intel.com>
---
 .../os_agnostic_memory_manager.cpp            | 25 ++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

diff --git a/shared/source/memory_manager/os_agnostic_memory_manager.cpp b/shared/source/memory_manager/os_agnostic_memory_manager.cpp
index c14c54a655..7c22acbc07 100644
--- a/shared/source/memory_manager/os_agnostic_memory_manager.cpp
+++ b/shared/source/memory_manager/os_agnostic_memory_manager.cpp
@@ -136,7 +136,30 @@ GraphicsAllocation *OsAgnosticMemoryManager::allocateGraphicsMemoryWithAlignment
 }
 
 GraphicsAllocation *OsAgnosticMemoryManager::allocateUSMHostGraphicsMemory(const AllocationData &allocationData) {
-    return allocateGraphicsMemoryWithHostPtr(allocationData);
+    AllocationData allocData = allocationData;
+    if (allocData.type == AllocationType::svmCpu) {
+        auto &rootDeviceEnvironment = *executionEnvironment.rootDeviceEnvironments[allocData.rootDeviceIndex];
+        auto &productHelper = rootDeviceEnvironment.getHelper<ProductHelper>();
+        allocData.alignment = alignUpNonZero<size_t>(allocationData.alignment, productHelper.getSvmCpuAlignment());
+        allocData.size = alignUp(allocationData.size, allocData.alignment);
+    }
+    auto memoryAllocation = allocateGraphicsMemoryWithHostPtr(allocData);
+    if (memoryAllocation && allocData.type == AllocationType::svmCpu) {
+        void *gpuPtr = reserveCpuAddressRange(allocData.size, allocData.rootDeviceIndex);
+        if (nullptr == gpuPtr) {
+            cleanGraphicsMemoryCreatedFromHostPtr(memoryAllocation);
+            delete memoryAllocation;
+            return nullptr;
+        }
+        memoryAllocation->setReservedAddressRange(gpuPtr, allocData.size);
+        gpuPtr = alignUp(gpuPtr, allocData.alignment);
+
+        auto cpuPtr = const_cast<void *>(allocData.hostPtr);
+        auto gmmHelper = getGmmHelper(allocData.rootDeviceIndex);
+        auto canonizedGpuAddress = gmmHelper->canonize(reinterpret_cast<uint64_t>(gpuPtr));
+        memoryAllocation->setCpuPtrAndGpuAddress(cpuPtr, canonizedGpuAddress);
+    }
+    return memoryAllocation;
 }
 
 GraphicsAllocation *OsAgnosticMemoryManager::allocateGraphicsMemoryForNonSvmHostPtr(const AllocationData &allocationData) {