From e48f942068b535f9bc8e06cf2eecc0f0fd100344 Mon Sep 17 00:00:00 2001
From: Changho Hwang <changhohwang@microsoft.com>
Date: Wed, 6 Sep 2023 12:30:04 +0000
Subject: [PATCH] Gather size 4

---
 src/ib.cc                | 4 ++--
 test/mp_unit/ib_tests.cu | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/ib.cc b/src/ib.cc
index 50fddc96..9fd16082 100644
--- a/src/ib.cc
+++ b/src/ib.cc
@@ -72,8 +72,8 @@ IbQp::IbQp(ibv_context* ctx, ibv_pd* pd, int port, int maxCqSize, int maxCqPollN
   qpInitAttr.qp_type = IBV_QPT_RC;
   qpInitAttr.cap.max_send_wr = maxSendWr;
   qpInitAttr.cap.max_recv_wr = maxRecvWr;
-  qpInitAttr.cap.max_send_sge = 1;
-  qpInitAttr.cap.max_recv_sge = 1;
+  qpInitAttr.cap.max_send_sge = maxNumSgesPerWr;
+  qpInitAttr.cap.max_recv_sge = maxNumSgesPerWr;
   qpInitAttr.cap.max_inline_data = 0;
 
   struct ibv_qp* _qp = ibv_create_qp(pd, &qpInitAttr);
diff --git a/test/mp_unit/ib_tests.cu b/test/mp_unit/ib_tests.cu
index 2a7d3841..d4c5539a 100644
--- a/test/mp_unit/ib_tests.cu
+++ b/test/mp_unit/ib_tests.cu
@@ -36,7 +36,7 @@ void IbPeerToPeerTest::SetUp() {
   bootstrap->initialize(id);
 
   ibCtx = std::make_shared<mscclpp::IbCtx>(ibDevName);
-  qp = ibCtx->createQp(1024, 1, 8192, 0, 64, 1);
+  qp = ibCtx->createQp(1024, 1, 8192, 0, 64, 4);
 
   int remoteRank = (gEnv->rank == 0) ? 1 : 0;
 
@@ -315,7 +315,7 @@ TEST_F(IbPeerToPeerTest, SendGather) {
 
   mscclpp::Timer timeout(3);
 
-  const int numDataSrcs = 1;
+  const int numDataSrcs = 4;
   const int nelemPerMr = 1024;
 
   // Gather send from rank 0 to 1