Convert listen_comm function pointers to virtual methods

Replace accept() and close() function pointers in nccl_net_ofi_listen_comm
with virtual methods. This eliminates manual function pointer assignment
and provides compile-time verification that all implementations provide
these methods.

Before: listen_comm->accept(listen_comm, &recv_comm);
After:  listen_comm->accept(&recv_comm);

Signed-off-by: Bibrak Qamar Chandio <bibracha@amazon.com>

Author: bibrakc

include/gin/nccl_ofi_gin.h

@@ -29,18 +29,18 @@ inline nccl_ofi_device_copy &get_device_copy()
 class nccl_ofi_gin_listen_comm {
 private:
 	nccl_net_ofi_ep_t *ep;
-	nccl_net_ofi_listen_comm_t *l_comm;
+	nccl_net_ofi_listen_comm *l_comm;
 
 public:
 	nccl_ofi_gin_listen_comm(int dev_arg, nccl_net_ofi_ep_t *ep_arg,
-				 nccl_net_ofi_listen_comm_t *l_comm_arg)
+				 nccl_net_ofi_listen_comm *l_comm_arg)
 	    : ep(ep_arg), l_comm(l_comm_arg)
 	{
 	}
 
 	~nccl_ofi_gin_listen_comm()
 	{
-		int ret = l_comm->close(l_comm);
+		int ret = l_comm->close();
 		if (ret != 0) {
 			NCCL_OFI_WARN("GIN: Unable to close net listen comm: %d", ret);
 		}

include/nccl_ofi.h

@@ -111,12 +111,11 @@ class nccl_net_ofi_ep_t;
 class nccl_net_ofi_plugin_t;
 
 struct nccl_net_ofi_comm;
-struct nccl_net_ofi_listen_comm;
+class nccl_net_ofi_listen_comm;
 struct nccl_net_ofi_send_comm;
 struct nccl_net_ofi_recv_comm;
 
 typedef struct nccl_net_ofi_comm nccl_net_ofi_comm_t;
-typedef struct nccl_net_ofi_listen_comm nccl_net_ofi_listen_comm_t;
 typedef struct nccl_net_ofi_send_comm nccl_net_ofi_send_comm_t;
 typedef struct nccl_net_ofi_recv_comm nccl_net_ofi_recv_comm_t;
 
@@ -682,7 +681,7 @@ class nccl_net_ofi_ep_t {
 	 * handle is set to COMM_CREATE_START.
 	 */
 	virtual int listen(nccl_net_ofi_conn_handle_t *handle,
-			   nccl_net_ofi_listen_comm_t **listen_comm) = 0;
+			   nccl_net_ofi_listen_comm **listen_comm) = 0;
 
 	/* Create a connection to a process that has called
 	 * listen().
@@ -864,12 +863,14 @@ struct nccl_net_ofi_comm {
 /**
  * Listen Communicator - Communicator for a listen/accept pairing
  */
-struct nccl_net_ofi_listen_comm {
+class nccl_net_ofi_listen_comm {
+public:
 	nccl_net_ofi_comm_t base;
 
-	int (*accept)(nccl_net_ofi_listen_comm_t *listen_comm,
-			       nccl_net_ofi_recv_comm_t **recv_comm);
-	int (*close)(nccl_net_ofi_listen_comm_t *listen_comm);
+	virtual ~nccl_net_ofi_listen_comm() = default;
+
+	virtual int accept(nccl_net_ofi_recv_comm_t **recv_comm) = 0;
+	virtual int close() = 0;
 };
 
 struct nccl_net_ofi_send_comm {

include/nccl_ofi_rdma.h

@@ -684,11 +684,10 @@ typedef struct nccl_net_ofi_rdma_recv_comm {
 	uint64_t remote_mr_key[MAX_NUM_RAILS];
 } nccl_net_ofi_rdma_recv_comm_t;
 
-typedef struct nccl_net_ofi_rdma_listen_comm {
-	/* This base listen communicator must be the first member of
-	 * this struct. This allows casting between pointers of this
-	 * struct and its base struct. */
-	nccl_net_ofi_listen_comm_t base;
+class nccl_net_ofi_rdma_listen_comm : public nccl_net_ofi_listen_comm {
+public:
+	int accept(nccl_net_ofi_recv_comm_t **recv_comm) override;
+	int close() override;
 
 	/* Associated listener from connection manager */
 	nccl_ofi_cm_listener *listener;
@@ -698,7 +697,7 @@ typedef struct nccl_net_ofi_rdma_listen_comm {
 
 	/* Stage of connection establishment on listen side */
 	nccl_ofi_comm_stage_t stage;
-} nccl_net_ofi_rdma_listen_comm_t;
+};
 
 
 class nccl_net_ofi_rdma_domain_rail_t {
@@ -1035,7 +1034,7 @@ class nccl_net_ofi_rdma_ep_t : public nccl_net_ofi_ep_t {
 	int cleanup_resources() override;
 
 	int listen(nccl_net_ofi_conn_handle_t *handle,
-		   nccl_net_ofi_listen_comm_t **listen_comm) override;
+		   nccl_net_ofi_listen_comm **listen_comm) override;
 
 	/**
 	 * @brief	Execute the connect functionality from listen/connect/accept

include/nccl_ofi_sendrecv.h

@@ -50,19 +50,19 @@ class nccl_net_ofi_sendrecv_mr_handle_t : public nccl_net_ofi_mr_handle_t {
 	ofi_mr_ptr mr;
 };
 
-typedef struct nccl_net_ofi_sendrecv_listen_comm {
-	/* This base listen communicator must be the first member of
-	 * this struct. This allows casting between pointers of this
-	 * struct and its base struct. */
-	nccl_net_ofi_listen_comm_t base;
+class nccl_net_ofi_sendrecv_listen_comm : public nccl_net_ofi_listen_comm {
+public:
+	int accept(nccl_net_ofi_recv_comm_t **recv_comm) override;
+	int close() override;
 
 	struct fid_ep *local_ep;
 	fi_addr_t local_ep_addr;
 	/* Saves temporary state when creating receive communicator object */
 	save_comm_state_t state;
 
 	nccl_ofi_cm_listener *listener;
-} nccl_net_ofi_sendrecv_listen_comm_t;
+};
+
 
 typedef struct nccl_net_ofi_sendrecv_send_comm {
 	/* This base send communicator must be the first member of this
@@ -183,7 +183,7 @@ class nccl_net_ofi_sendrecv_ep_t : public nccl_net_ofi_ep_t {
 	nccl_net_ofi_sendrecv_ep_t(nccl_net_ofi_sendrecv_domain_t *domain_arg);
 
 	int listen(nccl_net_ofi_conn_handle_t *handle,
-		   nccl_net_ofi_listen_comm_t **listen_comm) override;
+		   nccl_net_ofi_listen_comm **listen_comm) override;
 
 	int connect(nccl_net_ofi_conn_handle_t *handle,
 		    nccl_net_ofi_send_comm_t **send_comm,

src/gin/nccl_ofi_gin.cpp

@@ -160,7 +160,7 @@ int nccl_ofi_gin_listen_comm::connect(nccl_net_ofi_conn_handle_t *handles[], int
 			}
 		}
 		if (r_comm == nullptr) {
-			ret = l_comm->accept(l_comm, &r_comm);
+			ret = l_comm->accept(&r_comm);
 			if (ret != 0) {
 				NCCL_OFI_WARN("Error in bootstrap ring accept: %d", ret);
 				return ret;

src/gin/nccl_ofi_gin_api.cpp

@@ -149,7 +149,7 @@ static ncclResult_t nccl_ofi_gin_listen(void *ctx, int dev, void *handle, void *
 
 		nccl_net_ofi_ep_t *ep = device->get_ep(0, static_cast<long>(comm_id));
 
-		nccl_net_ofi_listen_comm_t *l_comm = nullptr;
+		nccl_net_ofi_listen_comm *l_comm = nullptr;
 		int ret = ep->listen(static_cast<nccl_net_ofi_conn_handle_t *>(handle), &l_comm);
 		if (ret != 0) {
 			NCCL_OFI_WARN("GIN: error listening on device %i.", dev);

src/nccl_ofi_api.cpp

@@ -147,8 +147,8 @@ ncclResult_t nccl_net_ofi_listen(int dev_id, void *handle, void **lComm,
 	int ret = 0;
 	nccl_net_ofi_device_t *device = nullptr;
 	nccl_net_ofi_ep_t *ep = nullptr;
-	nccl_net_ofi_listen_comm_t **listen_comm =
-		reinterpret_cast<nccl_net_ofi_listen_comm_t **>(lComm);
+	nccl_net_ofi_listen_comm **listen_comm =
+		reinterpret_cast<nccl_net_ofi_listen_comm **>(lComm);
 
 	/* Validate plugin */
 	if (OFI_UNLIKELY(plugin == nullptr)) {
@@ -308,13 +308,13 @@ ncclResult_t nccl_net_ofi_accept(void *lComm, void **rComm)
 	}
 
 	/* Invoke listen communicator accept() function */
-	nccl_net_ofi_listen_comm_t *listen_comm =
-		reinterpret_cast<nccl_net_ofi_listen_comm_t *>(lComm);
+	nccl_net_ofi_listen_comm *listen_comm =
+		reinterpret_cast<nccl_net_ofi_listen_comm *>(lComm);
 	nccl_net_ofi_recv_comm_t **recv_comm =
 		reinterpret_cast<nccl_net_ofi_recv_comm_t **>(rComm);
 	int ret = 0;
 	try {
-		ret = listen_comm->accept(listen_comm, recv_comm);
+		ret = listen_comm->accept(recv_comm);
 	}
 	catch (const std::exception &e) {
 		NCCL_OFI_WARN("Caught exception in plugin accept: %s", e.what());
@@ -642,8 +642,8 @@ ncclResult_t nccl_net_ofi_closeRecv(void *rComm)
 
 ncclResult_t nccl_net_ofi_closeListen(void *lComm)
 {
-	nccl_net_ofi_listen_comm_t *listen_comm =
-		(nccl_net_ofi_listen_comm_t *)lComm;
+	nccl_net_ofi_listen_comm *listen_comm =
+		(nccl_net_ofi_listen_comm *)lComm;
 
 	/* neuron has a cleanup race between the atexit handler and *
 	 * calling close on all the communicators, so be more silent
@@ -659,7 +659,7 @@ ncclResult_t nccl_net_ofi_closeListen(void *lComm)
 		return check_return(ncclInternalError);
 	}
 
-	int ret = listen_comm->close(listen_comm);
+	int ret = listen_comm->close();
 	return nccl_net_ofi_retval_translate_impl(ret);
 }