From 5580cb75743c92ec4f098d74089f6afc6fff5770 Mon Sep 17 00:00:00 2001
From: deeksha-amd <deeksha.goplani@amd.com>
Date: Thu, 1 May 2025 04:31:11 +0530
Subject: [PATCH] Added new tests for improving the code coverage (#1656)

Signed-off-by: Deeksha Goplani <deeksha.goplani@amd.com>

[ROCm/rccl commit: 248683846526b1363837b70d9b7316658dd20538]
---
 projects/rccl/test/StandaloneTests.cpp | 71 ++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/projects/rccl/test/StandaloneTests.cpp b/projects/rccl/test/StandaloneTests.cpp
index 35d4814c3f..f51b5683aa 100644
--- a/projects/rccl/test/StandaloneTests.cpp
+++ b/projects/rccl/test/StandaloneTests.cpp
@@ -299,4 +299,75 @@ namespace RcclUnitTesting
       }
     }
   }
+  /**
+   * \brief Verify the device associated with communicator in both single and multi-device scenarios
+   * ******************************************************************************************/
+  TEST(Standalone, CommCuDevice_Check)
+  {
+    int numDevices;
+    HIPCALL(hipGetDeviceCount(&numDevices));
+    if (numDevices < 1) {
+      GTEST_SKIP() << "No devices available.";
+    }
+
+    // Test single comm initialization
+    ncclComm_t comm;
+    ncclUniqueId id;
+    NCCLCHECK(ncclGetUniqueId(&id));
+    HIPCALL(hipSetDevice(0));
+    NCCLCHECK(ncclCommInitRank(&comm, 1, id, 0));
+
+    // Verify device assignment
+    int device;
+    NCCLCHECK(ncclCommCuDevice(comm, &device));
+    ASSERT_EQ(device, 0);
+    NCCLCHECK(ncclCommDestroy(comm));
+
+    // Test multi-device scenario if available
+    if (numDevices > 1) {
+      std::vector<ncclComm_t> comms(numDevices);
+
+      // Initialize all communicators at once
+      NCCLCHECK(ncclCommInitAll(comms.data(), numDevices, nullptr));
+
+      // Verify device assignments
+      for (int i = 0; i < numDevices; i++) {
+        int assignedDevice;
+        NCCLCHECK(ncclCommCuDevice(comms[i], &assignedDevice));
+        ASSERT_EQ(assignedDevice, i);
+      }
+
+      // Clean up
+      for (int i = 0; i < numDevices; i++) {
+        NCCLCHECK(ncclCommDestroy(comms[i]));
+      }
+    }
+  }
+
+  /**
+   * \brief verifies that ncclCommUserRank correctly fails when provided with an invalid (null) communicator handle
+   * ******************************************************************************************/
+  TEST(Standalone, SplitComms_RankCheck_Basic_Failure) {
+    // Check for multi-gpu
+    int numDevices;
+    HIPCALL(hipGetDeviceCount(&numDevices));
+    if (numDevices < 2) {
+        GTEST_SKIP() << "This test requires at least 2 devices.";
+    }
+
+    // Initialize the original comms
+    std::vector<ncclComm_t> comms(numDevices);
+    NCCLCHECK(ncclCommInitAll(comms.data(), numDevices, nullptr));
+
+    // Create an invalid comm handle that will cause a failure
+    ncclComm_t invalidComm = nullptr;
+
+    // This NCCL_CHECK will fail because we're trying to query rank from a null communicator
+    int rank;
+    NCCLCHECK(ncclCommUserRank(invalidComm, &rank));
+
+    // Clean up comms
+    for (auto& comm : comms)
+      NCCLCHECK(ncclCommDestroy(comm));
+  }
 }