From 4fb06f94d19b3ee245f15a61aeb8e332018cc9db Mon Sep 17 00:00:00 2001
From: kjayapra-amd <karthik.jayaprakash@amd.com>
Date: Tue, 7 May 2024 18:32:20 -0400
Subject: [PATCH] SWDEV-463872 - Request the same va_addr on the second device,
 if not fail.

Change-Id: I791625b6434a3980b35b8d6d6d3b6070fa159614
---
 rocclr/device/pal/palresource.cpp | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/rocclr/device/pal/palresource.cpp b/rocclr/device/pal/palresource.cpp
index dd801157bc..7f145fb010 100644
--- a/rocclr/device/pal/palresource.cpp
+++ b/rocclr/device/pal/palresource.cpp
@@ -1340,11 +1340,10 @@ bool Resource::create(MemoryType memType, CreateParams* params, bool forceLinear
     createInfo.flags.busAddressable = true;
   } else if (memoryType() == VaRange) {
     createInfo.flags.virtualAlloc = true;
+    createInfo.vaRange = Pal::VaRange::Svm;
     if (params->owner_->getSvmPtr() != nullptr) {
-      createInfo.vaRange = Pal::VaRange::Svm;
       createInfo.flags.useReservedGpuVa = true;
       createInfo.pReservedGpuVaOwner = params->svmBase_->iMem();
-      desc_.SVMRes_ = true;
     }
   }
 
@@ -1367,7 +1366,16 @@ bool Resource::create(MemoryType memType, CreateParams* params, bool forceLinear
     mapCount_++;
   }
   if (memoryType() == VaRange) {
-    params->owner_->setSvmPtr(reinterpret_cast<void*>(memRef_->iMem()->Desc().gpuVirtAddr));
+    void* gpu_virt_addr = reinterpret_cast<void*>(memRef_->iMem()->Desc().gpuVirtAddr);
+
+    // if svm_ptr was not set, set the gpu_virt_addr
+    if (params->owner_->getSvmPtr() == nullptr) {
+      params->owner_->setSvmPtr(gpu_virt_addr);
+    } else if (params->owner_->getSvmPtr() != gpu_virt_addr) {
+      LogPrintfError("Cannot get different ptrs in va_range, svm_ptr: %p new_svm_ptr:%p \n",
+                     params->owner_->getSvmPtr(), gpu_virt_addr);
+      return false;
+    }
   }
   return true;
 }
@@ -2222,7 +2230,10 @@ bool ResourceCache::addGpuMemory(Resource::Descriptor* desc, GpuMemoryReference*
   size_t size = ref->iMem()->Desc().size;
 
   // Check if runtime can free suballocation
-  if ((desc->type_ == Resource::Local) && !desc->SVMRes_) {
+  if (desc->type_ == Resource::VaRange) {
+    // We do no sub allocate VA Range.
+    result = true;
+  } else if ((desc->type_ == Resource::Local) && !desc->SVMRes_) {
     result = mem_sub_alloc_local_.Free(&lockCacheOps_, ref, offset);
   } else if ((desc->type_ == Resource::Local) && desc->SVMRes_) {
     result = mem_sub_alloc_coarse_.Free(&lockCacheOps_, ref, offset);
@@ -2279,7 +2290,10 @@ GpuMemoryReference* ResourceCache::findGpuMemory(Resource::Descriptor* desc, Pal
   GpuMemoryReference* ref = nullptr;
 
   // Check if the runtime can suballocate memory
-  if ((desc->type_ == Resource::Local) && !desc->SVMRes_) {
+  if (desc->type_ == Resource::VaRange) {
+    // Do not use suballocator for VA_Range.
+    return nullptr;
+  } else if ((desc->type_ == Resource::Local) && !desc->SVMRes_) {
     ref = mem_sub_alloc_local_.Allocate(size, alignment, reserved_va, offset);
   } else if ((desc->type_ == Resource::Local) && desc->SVMRes_) {
     ref = mem_sub_alloc_coarse_.Allocate(size, alignment, reserved_va, offset);