Added tests for coll_reg (#1700)

Changes to coll_reg

Co-authored-by: Welling <awelling@ctr2-alola-login-01.amd.com>

CMakeLists.txt

@@ -40,6 +40,7 @@ option(TRACE                                   "Enable additional tracing"
 option(FAULT_INJECTION                         "Enable fault injection"                        ON)
 option(FORCE_REDUCE_PIPELINING                 "Force reduce pipelining"                       OFF)
 option(DISABLE_CHEAP_THREADFENCE               "Compile-time killswitch for simpler fence"     OFF)
+option(RCCL_EXPOSE_STATIC                      "Expose internal static functions for testing"  OFF)
 
 # Default GPU architectures to build
 #==================================================================================================
@@ -371,6 +372,9 @@ if(ROCTX)
   endif()
 endif()
 
+if(RCCL_EXPOSE_STATIC)
+  add_definitions(-DRCCL_EXPOSE_STATIC)
+endif()
 # Determine version from makefiles/version.mk and fill in templates
 #==================================================================================================
 ## parse version from Makefile NCCL_MAJOR, NCCL_MINOR, NCCL_PATCH must exist
@@ -741,6 +745,7 @@ if (BUILD_TESTS)
     # Create a list of files which needs to be modified to remove static
     set(TEST_NONSTATIC_SOURCE_FILES
       ${HIPIFY_SRC_DIR}/misc/alt_rsmi.cc
+      ${HIPIFY_SRC_DIR}/register/coll_reg.cc
       ${HIPIFY_SRC_DIR}/transport/shm.cc
       ${HIPIFY_SRC_DIR}/transport/p2p.cc
     )

test/CMakeLists.txt

@@ -142,6 +142,7 @@ if(BUILD_TESTS)
       AllocTests.cpp
       ArgCheckTests.cpp
       IpcsocketTests.cpp
+      CollRegTests.cpp
       ShmTests.cpp
       P2pTests.cpp
       BitOpsTests.cpp

test/CollRegTests.cpp

@@ -0,0 +1,102 @@
+/*************************************************************************
+ * Copyright (c) 2023 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include <gtest/gtest.h>
+#include <rccl/rccl.h>
+
+#include "TestBed.hpp"
+
+#include "CollRegUtils.hpp"
+
+namespace RcclUnitTesting
+{
+  
+TEST(CollReg, RegisterCheckP2P_ConnectedWithP2PFlag) {
+  ncclComm comm = {};
+  ncclConnector conn = {};
+  conn.connected = 1;
+  conn.conn.flags = NCCL_P2P_WRITE;
+
+  bool needReg = false;
+  EXPECT_EQ(registerCheckP2PConnection(&comm, &conn, nullptr, 1, &needReg), ncclSuccess);
+  EXPECT_TRUE(needReg);
+}
+
+TEST(CollReg, RegisterCheckP2P_ConnectedWithNoP2PFlag) {
+  ncclComm comm = {};
+  ncclConnector conn = {};
+  conn.connected = 1;
+  conn.conn.flags = 0;
+
+  bool needReg = true;
+  EXPECT_EQ(registerCheckP2PConnection(&comm, &conn, nullptr, 1, &needReg), ncclSuccess);
+  EXPECT_FALSE(needReg);
+}
+
+TEST(CollReg, RegisterCheckP2P_NotConnected_CanConnectTrue) {
+  // Save original canConnect
+  ScopedCanConnectOverride _(&(*ncclTransports[0]), mockCanConnect);
+  auto* originalCanConnect = ncclTransports[0]->canConnect;
+  ncclTransports[0]->canConnect = mockCanConnect;
+
+  ncclComm comm = {};
+  comm.rank = 0;
+
+  ncclPeerInfo peers[2] = {};
+  comm.peerInfo = peers;
+
+  // Set myInfo and peerInfo
+  peers[0].hostHash = 123;
+  peers[0].pidHash = 999;
+  peers[0].cudaDev = 0;
+  peers[0].busId = 0x1;
+  peers[0].comm = &comm;
+
+  peers[1].hostHash = 123;
+  peers[1].pidHash = 888;
+  peers[1].cudaDev = 1;
+  peers[1].busId = 0x2;
+  peers[1].comm = &comm;
+
+  ncclConnector conn = {};
+  conn.connected = 0;
+
+  ncclTopoGraph graph = {};
+
+  bool needReg = false;
+
+  EXPECT_EQ(registerCheckP2PConnection(&comm, &conn, &graph, 1, &needReg), ncclSuccess);
+  EXPECT_TRUE(needReg);  // Because mockCanConnect returns true
+
+  // Restore original canConnect
+  ncclTransports[0]->canConnect = originalCanConnect;
+}
+
+TEST(CollReg, RegisterCheckP2P_NotConnected_CanConnectFalse) {
+  ScopedCanConnectOverride _(&(*ncclTransports[0]), mockCanConnectFalse); 
+
+  ncclComm comm = {};
+  comm.rank = 0;
+  ncclPeerInfo peers[2] = {};
+  comm.peerInfo = peers;
+
+  // Set up myInfo and peerInfo to block P2P
+  peers[0].hostHash = 123; peers[0].pidHash = 999;
+  //peers[0].dev = 0;
+  peers[1].hostHash = 456; peers[1].pidHash = 888;  // different host
+  //peers[1].dev = 1;
+
+  ncclConnector conn = {};
+  conn.connected = 0;
+
+  ncclTopoGraph graph = {};
+
+  bool needReg = true;
+  EXPECT_EQ(registerCheckP2PConnection(&comm, &conn, &graph, 1, &needReg), ncclSuccess);
+  EXPECT_FALSE(needReg);  // should be false for different host
+}
+
+}

test/common/CollRegUtils.hpp

@@ -0,0 +1,52 @@
+/*************************************************************************
+ * Copyright (c) 2023 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * See LICENSE.txt for license information
+ ************************************************************************/
+
+#include <gtest/gtest.h>
+#include <rccl/rccl.h>
+#include <comm.h>
+#include <transport.h>
+#include "TestBed.hpp"
+
+ncclResult_t registerCheckP2PConnection(struct ncclComm* comm, struct ncclConnector* conn, struct ncclTopoGraph* graph, int peer, bool* needReg);
+
+// Global mock transport definition (outside any namespace)
+static ncclResult_t mockCanConnect(int* result, struct ncclComm*, struct ncclTopoGraph*,
+  struct ncclPeerInfo*, struct ncclPeerInfo*) {
+*result = 1;
+return ncclSuccess;
+}
+
+static ncclResult_t mockCanConnectFalse(int* result, struct ncclComm*, struct ncclTopoGraph*,
+  struct ncclPeerInfo*, struct ncclPeerInfo*) {
+*result = 0;  // Simulate failure to connect
+return ncclSuccess;
+}
+
+static struct ncclTransport mockTransport = {
+.canConnect = mockCanConnect,
+};
+
+struct ncclTransport* ncclTransports[] = {
+&mockTransport,
+nullptr
+};
+
+struct ScopedCanConnectOverride {
+  struct ncclTransport* transport;
+  ncclResult_t (*originalFn)(int*, struct ncclComm*, struct ncclTopoGraph*,
+                             struct ncclPeerInfo*, struct ncclPeerInfo*);
+
+  ScopedCanConnectOverride(struct ncclTransport* t,
+                           ncclResult_t (*mockFn)(int*, struct ncclComm*, struct ncclTopoGraph*,
+                                                  struct ncclPeerInfo*, struct ncclPeerInfo*))
+      : transport(t), originalFn(t->canConnect) {
+    t->canConnect = mockFn;
+  }
+
+  ~ScopedCanConnectOverride() {
+    transport->canConnect = originalFn;
+  }
+};