move ntrace patch to third-party (exclude ext_tag) (NVIDIA#4)

Summary: Pull Request resolved: facebookresearch#4 Introduce ntrace into NCCL to trace important IB events per collective communication. It requires a separate module ntrace_rt to enable in Makefile based build. For now, ENABLE_NTRACE is always off when building from Makefile. Co-author: Xianghuai Zhang <[email protected]> Reviewed By: kingchc Differential Revision: D38637718 fbshipit-source-id: d689a31f2525bd423c9e87671b9fc8ddabb2835a
minsii · Oct 7, 2022 · 2ea9277 · 2ea9277
1 parent 7dacc35
commit 2ea9277
Show file tree

Hide file tree

Showing 4 changed files with 111 additions and 7 deletions.
diff --git a/src/include/ibvwrap.h b/src/include/ibvwrap.h
@@ -1049,6 +1049,10 @@ typedef enum ibv_return_enum
     IBV_SUCCESS = 0,                   //!< The operation was successful
 } ibv_return_t;
 
+// Explicitly include NTRACE after verbs type definition
+// to pass these types to ntrace_rt.h.
+#include "ntrace_profiler.h"
+
 ncclResult_t wrap_ibv_symbols(void);
 ncclResult_t wrap_ibv_fork_init(void);
 ncclResult_t wrap_ibv_get_device_list(struct ibv_device ***ret, int *num_devices);
@@ -1077,6 +1081,7 @@ ncclResult_t wrap_ibv_create_cq(struct ibv_cq **ret, struct ibv_context *context
 ncclResult_t wrap_ibv_destroy_cq(struct ibv_cq *cq);
 static inline ncclResult_t wrap_ibv_poll_cq(struct ibv_cq *cq, int num_entries, struct ibv_wc *wc, int* num_done) {
   int done = cq->context->ops.poll_cq(cq, num_entries, wc); /*returns the number of wcs or 0 on success, a negative number otherwise*/
+  NTRACE_PROFILING_RECORD(ibv_poll_cq, cq, num_entries, wc, done);
   if (done < 0) {
     WARN("Call to ibv_poll_cq() returned %d", done);
     return ncclSystemError;
@@ -1092,6 +1097,7 @@ static inline int ibv_post_send(struct ibv_qp *qp, struct ibv_send_wr *wr, struc
 }
 
 static inline ncclResult_t wrap_ibv_post_send(struct ibv_qp *qp, struct ibv_send_wr *wr, struct ibv_send_wr **bad_wr) {
+  NTRACE_PROFILING_RECORD(ibv_post_send, qp, wr, bad_wr, NULL);
   int ret = qp->context->ops.post_send(qp, wr, bad_wr); /*returns 0 on success, or the value of errno on failure (which indicates the failure reason)*/
   if (ret != IBV_SUCCESS) {
     WARN("ibv_post_send() failed with error %s, Bad WR %p, First WR %p", strerror(ret), wr, *bad_wr);
@@ -1101,6 +1107,7 @@ static inline ncclResult_t wrap_ibv_post_send(struct ibv_qp *qp, struct ibv_send
 }
 
 static inline ncclResult_t wrap_ibv_post_recv(struct ibv_qp *qp, struct ibv_recv_wr *wr, struct ibv_recv_wr **bad_wr) {
+  NTRACE_PROFILING_RECORD(ibv_post_recv, qp, wr, bad_wr, NULL);
   int ret = qp->context->ops.post_recv(qp, wr, bad_wr); /*returns 0 on success, or the value of errno on failure (which indicates the failure reason)*/
   if (ret != IBV_SUCCESS) {
     WARN("ibv_post_recv() failed with error %s", strerror(ret));

diff --git a/src/include/ntrace_profiler.h b/src/include/ntrace_profiler.h
@@ -0,0 +1,26 @@
+/* (c) Facebook, Inc. and its affiliates. Confidential and proprietary. */
+
+#ifndef NCCL_NTRACE_PROFILER_H_
+#define NCCL_NTRACE_PROFILER_H_
+
+#ifdef ENABLE_NTRACE
+#include "ntrace_rt.h"
+
+#define NTRACE_PROFILING_RECORD(profile_state, ...) \
+  do {                                              \
+    ntrace_log_##profile_state(__VA_ARGS__);        \
+  } while (0)
+
+static inline void ntraceProfilingDump(void) {
+  ntrace_dump();
+}
+
+#else
+#define NTRACE_PROFILING_RECORD(profile_state, ...) \
+  do { /* no op */                                  \
+  } while (0)
+
+static inline void ntraceProfilingDump(void){/* no op */};
+#endif
+
+#endif /* end of NCCL_NTRACE_PROFILER_H_ */
diff --git a/src/misc/ibvwrap.cc b/src/misc/ibvwrap.cc
@@ -193,6 +193,44 @@ ncclResult_t wrap_ibv_symbols(void) {
   } \
   return ncclSuccess;
 
+// _NO_RETURN version of the above check macros to enable ntrace post event recording.
+// Compare to the original version, it sets ncclret without return.
+#define IBV_PTR_CHECK_NO_RETURN(name_internal, call, retval, error_retval, name, ncclret) \
+  if (name_internal == NULL) {             \
+     WARN("lib wrapper not initialized."); \
+     ncclret = ncclInternalError;          \
+  } else {                                 \
+    retval = call;                         \
+    if (retval == error_retval) {          \
+      WARN("Call to " name " failed");     \
+      ncclret = ncclSystemError;           \
+    }                                      \
+  }
+
+#define IBV_INT_CHECK_RET_ERRNO_NO_RETURN(name_internal, call, success_retval, name, ncclret) \
+  if (name_internal == NULL) {             \
+     WARN("lib wrapper not initialized."); \
+     ncclret = ncclInternalError;          \
+  } else {                                 \
+    int ret = call;                        \
+    if (ret != success_retval) {           \
+      WARN("Call to " name " failed with error %s", strerror(ret)); \
+      ncclret = ncclSystemError;           \
+    }                                      \
+  }
+
+#define IBV_INT_CHECK_NO_RETURN(name_internal, call, error_retval, name, ncclret) \
+  if (name_internal == NULL) {             \
+     WARN("lib wrapper not initialized."); \
+     ncclret = ncclInternalError;          \
+  } else {                                 \
+    int ret = call;                        \
+    if (ret == error_retval) {             \
+      WARN("Call to " name " failed");     \
+      ncclret = ncclSystemError;           \
+    }                                      \
+  }
+
 #define IBV_PASSTHRU(name_internal, call) \
   if (name_internal == NULL) { \
      WARN("lib wrapper not initialized."); \
@@ -224,11 +262,17 @@ const char *wrap_ibv_get_device_name(struct ibv_device *device) {
 }
 
 ncclResult_t wrap_ibv_open_device(struct ibv_context **ret, struct ibv_device *device) { /*returns 0 on success, -1 on failure*/
-  IBV_PTR_CHECK(ibv_internal_open_device, ibv_internal_open_device(device), *ret, NULL, "ibv_open_device");
+  ncclResult_t ncclret = ncclSuccess;
+  IBV_PTR_CHECK_NO_RETURN(ibv_internal_open_device, ibv_internal_open_device(device), *ret, NULL, "ibv_open_device", ncclret);
+  NTRACE_PROFILING_RECORD(ibv_open_device, *ret, device);
+  return ncclret;
 }
 
 ncclResult_t wrap_ibv_close_device(struct ibv_context *context) { /*returns 0 on success, -1 on failure*/
-  IBV_INT_CHECK(ibv_internal_close_device, ibv_internal_close_device(context), -1, "ibv_close_device");
+  ncclResult_t ncclret = ncclSuccess;
+  NTRACE_PROFILING_RECORD(ibv_close_device, context);
+  IBV_INT_CHECK_NO_RETURN(ibv_internal_close_device, ibv_internal_close_device(context), -1, "ibv_close_device", ncclret);
+  return ncclret;
 }
 
 ncclResult_t wrap_ibv_get_async_event(struct ibv_context *context, struct ibv_async_event *event) { /*returns 0 on success, and -1 on error*/
@@ -244,11 +288,17 @@ ncclResult_t wrap_ibv_query_device(struct ibv_context *context, struct ibv_devic
 }
 
 ncclResult_t wrap_ibv_query_port(struct ibv_context *context, uint8_t port_num, struct ibv_port_attr *port_attr) { /*returns 0 on success, or the value of errno on failure (which indicates the failure reason)*/
-  IBV_INT_CHECK_RET_ERRNO(ibv_internal_query_port, ibv_internal_query_port(context, port_num, port_attr), 0, "ibv_query_port");
+  ncclResult_t ncclret = ncclSuccess;
+  IBV_INT_CHECK_RET_ERRNO_NO_RETURN(ibv_internal_query_port, ibv_internal_query_port(context, port_num, port_attr), 0, "ibv_query_port", ncclret);
+  NTRACE_PROFILING_RECORD(ibv_query_port, context, port_num, *port_attr);
+  return ncclret;
 }
 
 ncclResult_t wrap_ibv_query_gid(struct ibv_context *context, uint8_t port_num, int index, union ibv_gid *gid) {
-  IBV_INT_CHECK_RET_ERRNO(ibv_internal_query_gid, ibv_internal_query_gid(context, port_num, index, gid), 0, "ibv_query_gid");
+   ncclResult_t ncclret = ncclSuccess;
+   IBV_INT_CHECK_RET_ERRNO_NO_RETURN(ibv_internal_query_gid, ibv_internal_query_gid(context, port_num, index, gid), 0, "ibv_query_gid", ncclret);
+   NTRACE_PROFILING_RECORD(ibv_query_gid, context, port_num, index, *gid);
+   return ncclret;
 }
 
 ncclResult_t wrap_ibv_query_qp(struct ibv_qp *qp, struct ibv_qp_attr *attr, int attr_mask, struct ibv_qp_init_attr *init_attr) {
@@ -308,15 +358,24 @@ ncclResult_t wrap_ibv_destroy_cq(struct ibv_cq *cq) {
 }
 
 ncclResult_t wrap_ibv_destroy_qp(struct ibv_qp *qp) {
-  IBV_INT_CHECK_RET_ERRNO(ibv_internal_destroy_qp, ibv_internal_destroy_qp(qp), 0, "ibv_destroy_qp");
+  ncclResult_t ncclret = ncclSuccess;
+  NTRACE_PROFILING_RECORD(ibv_destroy_qp, qp);
+  IBV_INT_CHECK_RET_ERRNO_NO_RETURN(ibv_internal_destroy_qp, ibv_internal_destroy_qp(qp), 0, "ibv_destroy_qp", ncclret);
+  return ncclret;
 }
 
 ncclResult_t wrap_ibv_create_qp(struct ibv_qp **ret, struct ibv_pd *pd, struct ibv_qp_init_attr *qp_init_attr) {
-  IBV_PTR_CHECK(ibv_internal_create_qp, ibv_internal_create_qp(pd, qp_init_attr), *ret, NULL, "ibv_create_qp");
+  ncclResult_t ncclret = ncclSuccess;
+  IBV_PTR_CHECK_NO_RETURN(ibv_internal_create_qp, ibv_internal_create_qp(pd, qp_init_attr), *ret, NULL, "ibv_create_qp", ncclret);
+  NTRACE_PROFILING_RECORD(ibv_create_qp, *ret, qp_init_attr);
+  return ncclret;
 }
 
 ncclResult_t wrap_ibv_modify_qp(struct ibv_qp *qp, struct ibv_qp_attr *attr, int attr_mask) { /*returns 0 on success, or the value of errno on failure (which indicates the failure reason)*/
-  IBV_INT_CHECK_RET_ERRNO(ibv_internal_modify_qp, ibv_internal_modify_qp(qp, attr, attr_mask), 0, "ibv_modify_qp");
+  ncclResult_t ncclret = ncclSuccess;
+  NTRACE_PROFILING_RECORD(ibv_modify_qp, qp, attr, attr_mask);
+  IBV_INT_CHECK_RET_ERRNO_NO_RETURN(ibv_internal_modify_qp, ibv_internal_modify_qp(qp, attr, attr_mask), 0, "ibv_modify_qp", ncclret);
+  return ncclret;
 }
 
 ncclResult_t wrap_ibv_event_type_str(char **ret, enum ibv_event_type event) {

diff --git a/src/proxy.cc b/src/proxy.cc
@@ -10,6 +10,17 @@
 #include "socket.h"
 #include "shm.h"
 #include "profiler.h"
+
+// when build NCCL with ntrace_rt.h, we ensure ntrace_rt.h is always
+// referenced via ibvwrap.h after verbs type definition, to obtain type
+// definition while avoiding cross-reference issues.
+#ifdef ENABLE_NTRACE
+#include "ibvwrap.h"
+#else
+// for reference to no-op ntraceProfilingDump
+#include "ntrace_profiler.h"
+#endif
+
 #define ENABLE_TIMER 0
 #include "timer.h"
 
@@ -733,6 +744,7 @@ ncclResult_t ncclProxyProgressDestroy(struct ncclComm* comm) {
   }
 
   ncclProfilingDump();
+  ntraceProfilingDump();
   TIME_PRINT("Proxy");
   return ncclSuccess;
 }