From a44a5b896de63c83093ec086c4c5c5cfdf7e63ea Mon Sep 17 00:00:00 2001
From: kjayapra-amd <karthik.jayaprakash@amd.com>
Date: Fri, 29 Mar 2024 14:32:53 -0400
Subject: [PATCH] SWDEV-413997 - Fixing test cases where setaccess is missed
 and validation data structure not updated.

Change-Id: I65d8e77b7ce9b4378b56cfa85c50ed23af69dd90


[ROCm/hip-tests commit: 7d0f8e43b7c415a3ac055238b2957f859713bb36]
---
 .../catch/unit/virtualMemoryManagement/hipMemMap.cc    | 10 +++++++++-
 .../unit/virtualMemoryManagement/hipMemSetGetAccess.cc |  6 ++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/projects/hip-tests/catch/unit/virtualMemoryManagement/hipMemMap.cc b/projects/hip-tests/catch/unit/virtualMemoryManagement/hipMemMap.cc
index e856993e6b..310976028d 100644
--- a/projects/hip-tests/catch/unit/virtualMemoryManagement/hipMemMap.cc
+++ b/projects/hip-tests/catch/unit/virtualMemoryManagement/hipMemMap.cc
@@ -224,8 +224,11 @@ TEST_CASE("Unit_hipMemMap_PhysicalMemory_Map2MultVMMs") {
   for (int buf = 0; buf < num_buf; buf++) {
     HIP_CHECK(hipMemMap(ptrA[buf], size_mem, 0, handle, 0));
   }
+  // Set access for all the buffers.
+  for (int buf = 0; buf < num_buf; buf++) {
+    HIP_CHECK(hipMemSetAccess(ptrA[buf], size_mem, &accessDesc, 1));
+  }
   // Copy data to VMM via ptrA[0]
-  HIP_CHECK(hipMemSetAccess(ptrA[0], size_mem, &accessDesc, 1));
   HIP_CHECK(hipMemcpyHtoD(ptrA[0], A_h.data(), buffer_size));
   // Validate the data contained in VMM using ptrA[0], ptrA[1],
   // ......, ptrA[num_buf-1]
@@ -234,6 +237,11 @@ TEST_CASE("Unit_hipMemMap_PhysicalMemory_Map2MultVMMs") {
     HIP_CHECK(hipMemcpyDtoH(B_h.data(), ptrA[buf], buffer_size));
     REQUIRE(true == std::equal(B_h.begin(), B_h.end(), A_h.data()));
   }
+
+  for (int buf = 0; buf < num_buf; buf++) {
+    HIP_CHECK(hipMemUnmap(ptrA[buf], size_mem));
+  }
+
   // Release resources
   HIP_CHECK(hipMemRelease(handle));
   for (int buf = 0; buf < num_buf; buf++) {
diff --git a/projects/hip-tests/catch/unit/virtualMemoryManagement/hipMemSetGetAccess.cc b/projects/hip-tests/catch/unit/virtualMemoryManagement/hipMemSetGetAccess.cc
index 8533341c88..def3883670 100644
--- a/projects/hip-tests/catch/unit/virtualMemoryManagement/hipMemSetGetAccess.cc
+++ b/projects/hip-tests/catch/unit/virtualMemoryManagement/hipMemSetGetAccess.cc
@@ -361,11 +361,17 @@ TEST_CASE("Unit_hipMemSetAccess_FuncTstOnMultDev") {
     accessDesc.flags = hipMemAccessFlagsProtReadWrite;
     // Make the address accessible to GPU deviceId
     std::vector<int> A_h(N), B_h(N);
+    // Set Device, for kernel launch to be launched in the right device.
     HIP_CHECK(hipMemSetAccess(ptrA, size_mem, &accessDesc, 1));
     for (int idx = 0; idx < N; idx++) {
       A_h[idx] = idx;
     }
     HIP_CHECK(hipMemcpyHtoD(ptrA, A_h.data(), buffer_size));
+        // Set the A_h to verify with square kernel.
+    for (int idx = 0; idx < N; idx++) {
+      A_h[idx] = idx * idx;
+    }
+    HIP_CHECK(hipSetDevice(deviceId));
     // Launch square kernel
     hipLaunchKernelGGL(square_kernel, dim3(N / THREADS_PER_BLOCK), dim3(THREADS_PER_BLOCK), 0, 0,
                        static_cast<int*>(ptrA));