From 48a23243a4de0d295230b731694c4b80829f4484 Mon Sep 17 00:00:00 2001
From: Changho Hwang <changhohwang@microsoft.com>
Date: Wed, 22 Mar 2023 12:06:35 +0000
Subject: [PATCH] Dealloc more resources

---
 src/ib.cc   |  2 ++
 src/init.cc | 25 +++++++++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/ib.cc b/src/ib.cc
index 8eef4182..1c7c0c71 100644
--- a/src/ib.cc
+++ b/src/ib.cc
@@ -142,6 +142,8 @@ mscclppResult_t mscclppIbContextDestroy(struct mscclppIbContext *ctx)
     }
     ibv_destroy_cq(ctx->qps[i].cq);
     free(ctx->qps[i].wcs);
+    free(ctx->qps[i].sges);
+    free(ctx->qps[i].wrs);
   }
   if (ctx->pd != NULL) {
     ibv_dealloc_pd(ctx->pd);
diff --git a/src/init.cc b/src/init.cc
index d2d88c83..f2c72560 100644
--- a/src/init.cc
+++ b/src/init.cc
@@ -145,6 +145,31 @@ mscclppResult_t mscclppCommDestroy(mscclppComm_t comm){
   if (comm == NULL)
     return mscclppSuccess;
 
+  for (int i = 0; i < comm->nConns; ++i) {
+    struct mscclppConn *conn = &comm->conns[i];
+    if (conn->cpuProxyFlagGdrDesc) {
+      // IB
+      MSCCLPPCHECK(mscclppGdrCudaFree(conn->cpuProxyFlagGdrDesc));
+    } else if (conn->devConn->proxyFlag) {
+      // P2P
+      MSCCLPPCHECK(mscclppCudaFree(conn->devConn->proxyFlag));
+    }
+  }
+
+  for (int i = 0; i < MSCCLPP_PROXY_MAX_NUM; ++i) {
+    struct mscclppProxyState *proxyState = comm->proxyState[i];
+    if (proxyState) {
+      MSCCLPPCHECK(mscclppGdrCudaFree(proxyState->triggerFifo.desc));
+      MSCCLPPCHECK(mscclppGdrCudaFree(proxyState->fifoHead.desc));
+      MSCCLPPCHECK(mscclppGdrCudaFree(proxyState->fifoTail.desc));
+      free(proxyState);
+    }
+  }
+
+  if (comm->stream != NULL) {
+    CUDACHECK(cudaStreamDestroy(comm->stream));
+  }
+
   for (int i = 0; i < MSCCLPP_IB_MAX_DEVS; ++i) {
     if (comm->ibContext[i]) {
       MSCCLPPCHECK(mscclppIbContextDestroy(comm->ibContext[i]));