ofi: manage ofi resource lifecycle with smart ptrs

To ensure that Libfabric resources managed by transport datatypes are
automatically cleaned up during stack unwinding, this wraps Libfabric
resources in unique_ptr with custom destructors that call "fi_close".
Adds and refactors existing "ofiutils" helper functions to create and
initialize the unique_ptr Libfabric resources. Also adds a ofi_result
struct as a custom return type to return both a Libfabric resource
unique_ptr and a return code in a single object.

I only use the unique_ptr wrapper for classes that own the associated
Libfabric resource and the connection management endpoint class'
reference to a domain. I continue to use raw pointers for other classes
that just use the resource owned by someone else.

To support the unique_ptr requirement to not be copyable, this updates
the RDMA transport protocol rail datatypes that own Libfabric resources
to delete the copy constructor and copy assignment operator.

Signed-off-by: Aviv Benchorin <[email protected]>

Author: AvivBenchorin

include/Makefile.am

@@ -38,6 +38,7 @@ noinst_HEADERS = \
 	nccl_ofi_system.h \
 	nccl_ofi_topo.h \
 	nccl_ofi_tracepoint.h \
+	ofi/resource_wrapper.h \
 	platform-aws.h \
 	internal/tuner/nccl_defaults.h \
 	stats/histogram.h \

include/cm/nccl_ofi_cm_resources.h

@@ -15,6 +15,7 @@
 #include "cm/nccl_ofi_cm_reqs.h"
 
 #include "nccl_ofi_freelist.h"
+#include "ofi/resource_wrapper.h"
 
 namespace nccl_ofi_cm {
 
@@ -28,8 +29,20 @@ class endpoint
 {
 public:
 	/* Memory registration handle */
-	struct mr_handle_t {
-		fid_mr *mr;
+	class mr_handle_t {
+	public:
+		/* Default constructor */
+		mr_handle_t() = default;
+
+		/* Move constructor and assignment */
+		mr_handle_t(mr_handle_t&&) = default;
+		mr_handle_t& operator=(mr_handle_t&&) = default;
+
+		/* Delete copy operations since smart pointers are non-copyable */
+		mr_handle_t(const mr_handle_t&) = delete;
+		mr_handle_t& operator=(const mr_handle_t&) = delete;
+
+		ofi_mr_ptr mr;
 		uint64_t mr_key;
 		endpoint &ep;
 	};
@@ -41,6 +54,14 @@ class endpoint
 	 *      OFI domain against which to construct this ep
 	 */
 	endpoint(nccl_net_ofi_domain_t &domain);
+	
+	/* Move constructor and assignment */
+	endpoint(endpoint&&) = default;
+	endpoint& operator=(endpoint&&) = default;
+
+	/* Delete copy operations since smart pointers are non-copyable */
+	endpoint(const endpoint&) = delete;
+	endpoint& operator=(const endpoint&) = delete;
 
 	/**
 	 * Destructor. Closes OFI endpoint if not already closed, as well as
@@ -64,21 +85,21 @@ class endpoint
 	 *
 	 * @param req: used for the context of the operation
 	 */
-	int send(nccl_ofi_cm_conn_msg &conn_msg, size_t size, mr_handle_t mr_handle,
+	int send(nccl_ofi_cm_conn_msg &conn_msg, size_t size, mr_handle_t &mr_handle,
 		 fi_addr_t dest_addr, nccl_ofi_cm_req &req);
 
 	/**
 	 * Post a recv to the endpoint, with given parameters
 	 *
 	 * @param req: used for the context of the operation
 	 */
-	int recv(nccl_ofi_cm_conn_msg &conn_msg, size_t size, mr_handle_t mr_handle,
+	int recv(nccl_ofi_cm_conn_msg &conn_msg, size_t size, mr_handle_t &mr_handle,
 		 nccl_ofi_cm_req &req);
 
 	/**
 	 * Close associated ofi_ep, while leaving other resources open
 	 */
-	int close_ofi_ep();
+	void close_ofi_ep();
 
 	/* Menory registration/deregistration. Note: these functions are static
 	   to be usable with the freelist interface */
@@ -87,12 +108,12 @@ class endpoint
 	static int dereg_mr(void *handle_ptr);
 private:
 	/* Input to CM */
-	fid_domain *ofi_domain;
+	ofi_domain_ptr &ofi_domain;
 	nccl_ofi_idpool_t &mr_key_pool;
 
 	/* Created by CM */
-	fid_ep *ofi_ep;
-	fid_av *av;
+	ofi_av_ptr av;
+	ofi_ep_ptr ofi_ep;
 };
 
 /**

include/nccl_ofi.h

@@ -19,6 +19,7 @@
 #include "nccl_ofi_topo.h"
 #include "nccl_ofi_idpool.h"
 #include "nccl_ofi_mr.h"
+#include "ofi/resource_wrapper.h"
 
 /*
  * NCCL_NET_HANDLE_MAXSIZE is a limited resource (and defined in NCCL).
@@ -148,7 +149,8 @@ struct nccl_net_ofi_req {
 	int (*test)(nccl_net_ofi_req_t *req, int *done, int *size);
 };
 
-struct nccl_net_ofi_mr_handle_t {
+class nccl_net_ofi_mr_handle_t {
+public:
 	/**
 	 * @brief	Default constructor
 	 */
@@ -458,15 +460,15 @@ class nccl_net_ofi_domain_t {
 	 * depending on the transport; in that case, this will be the domain object
 	 * associated with the "leader NIC".
 	 */
-	virtual struct fid_domain *get_ofi_domain_for_cm() = 0;
+	virtual ofi_domain_ptr &get_ofi_domain_for_cm() = 0;
 
 	/**
 	 * Retrieve an fid_cq object associated with this domain to be used for 
 	 * connection management. There may be more than one fid_cq per domain, depending
 	 * on the transport; in that case, this will be the cq object associated with the
 	 * "leader NIC".
 	 */
-	virtual struct fid_cq *get_ofi_cq_for_cm() = 0;
+	virtual ofi_cq_ptr &get_ofi_cq_for_cm() = 0;
 
 	/* Create a new endpoint
 	 *

include/nccl_ofi_ofiutils.h

@@ -8,6 +8,7 @@
 #include <rdma/fabric.h>
 
 #include "nccl_ofi_param.h"
+#include "ofi/resource_wrapper.h"
 
 /*
  * Memeory util functions to ensure that the compiler does not optimize
@@ -28,22 +29,68 @@ int nccl_ofi_ofiutils_get_providers(const char *prov_include,
 				    unsigned int *num_prov_infos);
 
 
-/*
- * @brief	Allocates and initialises libfabric endpoint and AV.
+/**
+ * @brief	Release libfabric endpoint and address vector
+ */
+void nccl_ofi_ofiutils_ep_release(ofi_ep_ptr& ep, ofi_av_ptr& av, int dev_id);
+
+/**
+ * @brief	Create and initialize libfabric fabric
  *
- * @param cq:	Completion queue to which the new endpoint will be bound
- * @return	Endpoint ep
- * @return	Address vector av
+ * @param info:		Fabric info for fabric creation
+ * @return		Result containing error code and fabric pointer
  */
-int nccl_ofi_ofiutils_init_connection(struct fi_info *info, struct fid_domain *domain,
-				      struct fid_ep **ep,   struct fid_av **av,
-				      struct fid_cq *cq);
+ofi_fabric_result nccl_ofi_ofiutils_fabric_create(struct fi_info *info);
 
-/*
- * @brief	Release libfabric endpoint and address vector
+/**
+ * @brief	Create and initialize libfabric domain
+ *
+ * @param fabric:	Fabric handle
+ * @param info:		Fabric info for domain creation
+ * @return		Result containing error code and domain pointer
+ */
+ofi_domain_result nccl_ofi_ofiutils_domain_create(ofi_fabric_ptr& fabric, struct fi_info *info);
+
+/**
+ * @brief	Create and initialize libfabric endpoint
+ *
+ * @param info:		Fabric info for endpoint creation
+ * @param domain:	Fabric domain
+ * @param av:		Address vector to which the new endpoint will be bound
+ * @param cq:		Completion queue to which the new endpoint will be bound
+ * @return		Result containing error code and endpoint pointer
+ */
+ofi_ep_result nccl_ofi_ofiutils_ep_create(struct fi_info *info, ofi_domain_ptr &domain,
+					  ofi_av_ptr &av, ofi_cq_ptr &cq);
+
+/**
+ * @brief	Create and initialize libfabric address vector
+ *
+ * @param domain:	Domain handle
+ * @return		Result containing error code and address vector pointer
+ */
+ofi_av_result nccl_ofi_ofiutils_av_create(ofi_domain_ptr &domain);
+
+/**
+ * @brief	Create and initialize libfabric completion queue
+ *
+ * @param domain:	Domain handle
+ * @param cq_attr:	CQ attributes
+ * @return		Result containing error code and completion queue pointer
+ */
+ofi_cq_result nccl_ofi_ofiutils_cq_create(ofi_domain_ptr &domain, struct fi_cq_attr *cq_attr);
+
+/**
+ * @brief	Register memory region with libfabric using fi_mr_regattr
+ *
+ * @param domain:	Domain handle
+ * @param mr_attr:	Memory region attributes structure
+ * @param flags:	Registration flags
+ * @return		Result containing error code and memory region pointer
  */
-void nccl_ofi_ofiutils_ep_release(struct fid_ep *ep, struct fid_av *av,
-				  int dev_id);
+ofi_mr_result nccl_ofi_ofiutils_mr_regattr(ofi_domain_ptr &domain, 
+					   struct fi_mr_attr *mr_attr, 
+					   uint64_t flags);
 
 /*
  * @brief	Free libfabric NIC info list.