From 6470a57dac510645ba7ede5bd236f19a52bca031 Mon Sep 17 00:00:00 2001
From: Rahul Manocha <rmanocha@amd.com>
Date: Thu, 1 Aug 2024 16:38:58 -0700
Subject: [PATCH] SWDEV-474617 - Add Peer Device access check for memcpy2D
 tests and fix for SVM Memory test.

Change-Id: Idff7c02087f99c59f456e0027bbe0bc19cf67059


[ROCm/hip-tests commit: 8b1599604141e9e9df70410e1e93eb3df6024cbf]
---
 .../hipSVMTestFineGrainMemoryConsistency.cpp  | 19 +++++++++++--------
 .../unit/memory/memcpy2d_tests_common.hh      | 14 ++++++++------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/projects/hip-tests/catch/unit/memory/hipSVMTestFineGrainMemoryConsistency.cpp b/projects/hip-tests/catch/unit/memory/hipSVMTestFineGrainMemoryConsistency.cpp
index a4eb2d60d6..cd5dd8fa6a 100644
--- a/projects/hip-tests/catch/unit/memory/hipSVMTestFineGrainMemoryConsistency.cpp
+++ b/projects/hip-tests/catch/unit/memory/hipSVMTestFineGrainMemoryConsistency.cpp
@@ -234,16 +234,19 @@ void launch_kernels_and_verify(std::vector<hipStream_t> &streams, unsigned int n
 *  - HIP_VERSION >= 5.7
 */
 TEST_CASE("test_svm_fine_grain_memory_consistency") {
-  int pcieAtomic = 0;
-  HIP_CHECK(hipDeviceGetAttribute(&pcieAtomic, hipDeviceAttributeHostNativeAtomicSupported, 0));
-  if (!pcieAtomic) {
-    fprintf(stderr, "Device doesn't support pcie atomic, Skipped\n");
-    REQUIRE(true);
-    return;
-  }
-  const int num_elements = 2167;
   int num_devices = 0;
   HIP_CHECK(hipGetDeviceCount(&num_devices));
+
+  for(int id = 0; id < num_devices; id++) {
+    int pcieAtomic = 0;
+    HIP_CHECK(hipDeviceGetAttribute(&pcieAtomic, hipDeviceAttributeHostNativeAtomicSupported, id));
+    if (!pcieAtomic) {
+      fprintf(stderr, "Device doesn't support pcie atomic, Skipped\n");
+      REQUIRE(true);
+      return;
+    }
+  }
+  const int num_elements = 2167;
   std::vector<hipStream_t> streams(num_devices);
 
   for (int d = 0; d < num_devices; d++) {
diff --git a/projects/hip-tests/catch/unit/memory/memcpy2d_tests_common.hh b/projects/hip-tests/catch/unit/memory/memcpy2d_tests_common.hh
index 54564dd9dd..1c1ce86cb1 100644
--- a/projects/hip-tests/catch/unit/memory/memcpy2d_tests_common.hh
+++ b/projects/hip-tests/catch/unit/memory/memcpy2d_tests_common.hh
@@ -72,12 +72,7 @@ void Memcpy2DDeviceToDeviceShell(F memcpy_func, const hipStream_t kernel_stream
   const size_t src_cols_mult = GENERATE(1, 2);
 
   INFO("Src device: " << src_device << ", Dst device: " << dst_device);
-
-  HIP_CHECK(hipSetDevice(src_device));
-  if constexpr (enable_peer_access) {
-    if (src_device == dst_device) {
-      return;
-    }
+  if (device_count > 1) {
     int can_access_peer = 0;
     HIP_CHECK(hipDeviceCanAccessPeer(&can_access_peer, src_device, dst_device));
     if (!can_access_peer) {
@@ -86,6 +81,13 @@ void Memcpy2DDeviceToDeviceShell(F memcpy_func, const hipStream_t kernel_stream
       HipTest::HIP_SKIP_TEST(msg.c_str());
       return;
     }
+  }
+
+  HIP_CHECK(hipSetDevice(src_device));
+  if constexpr (enable_peer_access) {
+    if (src_device == dst_device) {
+      return;
+    }
     HIP_CHECK(hipDeviceEnablePeerAccess(dst_device, 0));
   }