From 2cd63dac4228121522fa966bf7786be27e1e0165 Mon Sep 17 00:00:00 2001
From: akolliasAMD <99202231+akolliasAMD@users.noreply.github.com>
Date: Tue, 30 Aug 2022 16:18:37 -0600
Subject: [PATCH] added stream synch after hipMemset (#609)

[ROCm/rccl commit: 06bce9d0c9706d6915617a7f94a1d6a5a54f9610]
---
 projects/rccl/src/include/alloc.h                             | 1 +
 projects/rccl/test/common/PtrUnion.cpp                        | 1 +
 projects/rccl/tools/MultiRank/rccl-allreduce-multirank.cc     | 2 +-
 projects/rccl/tools/MultiRank/rccl-reducescatter-multirank.cc | 2 +-
 projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp         | 2 +-
 5 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/projects/rccl/src/include/alloc.h b/projects/rccl/src/include/alloc.h
index 94c6b65627..e9e07f7ae2 100644
--- a/projects/rccl/src/include/alloc.h
+++ b/projects/rccl/src/include/alloc.h
@@ -93,6 +93,7 @@ static ncclResult_t ncclCudaCallocDebug(const char *filefunc, int line, T** ptr,
   CUDACHECK(hipStreamDestroy(stream));
 #else
   CUDACHECK(hipMemset(*ptr, 0, nelem*sizeof(T)));
+  CUDACHECK(hipStreamSynchronize(NULL));
 #endif
   INFO(NCCL_ALLOC, "%s:%d Cuda Alloc Size %ld pointer %p", filefunc, line, nelem*sizeof(T), *ptr);
   int dev;
diff --git a/projects/rccl/test/common/PtrUnion.cpp b/projects/rccl/test/common/PtrUnion.cpp
index 708b86d606..2340d97c7c 100644
--- a/projects/rccl/test/common/PtrUnion.cpp
+++ b/projects/rccl/test/common/PtrUnion.cpp
@@ -99,6 +99,7 @@ namespace RcclUnitTesting
       ERROR("Unable to call hipMemset\n");
       return TEST_FAIL;
     }
+    hipStreamSynchronize(NULL);
     return TEST_SUCCESS;
   }
 
diff --git a/projects/rccl/tools/MultiRank/rccl-allreduce-multirank.cc b/projects/rccl/tools/MultiRank/rccl-allreduce-multirank.cc
index b4d36635a6..6334f3c465 100644
--- a/projects/rccl/tools/MultiRank/rccl-allreduce-multirank.cc
+++ b/projects/rccl/tools/MultiRank/rccl-allreduce-multirank.cc
@@ -172,7 +172,7 @@ int main(int argc, char* argv[])
       HIPCHECK(hipMalloc(recvbuff+i, size*sizeof(float)));
       h_recvbuff[i] = (float*) malloc (size *sizeof(float));
       HIPCHECK(hipMemset(recvbuff[i], 0, size*sizeof(float)));
-
+      HIPCHECK(hipStreamSynchronize(NULL));
       HIPCHECK(hipStreamCreate(s+i));
   }
 
diff --git a/projects/rccl/tools/MultiRank/rccl-reducescatter-multirank.cc b/projects/rccl/tools/MultiRank/rccl-reducescatter-multirank.cc
index 13aeeac612..0b1fe5b524 100644
--- a/projects/rccl/tools/MultiRank/rccl-reducescatter-multirank.cc
+++ b/projects/rccl/tools/MultiRank/rccl-reducescatter-multirank.cc
@@ -174,7 +174,7 @@ int main(int argc, char* argv[])
       HIPCHECK(hipMalloc(recvbuff+i, recvsize*sizeof(float)));
       h_recvbuff[i] = (float*) malloc (recvsize *sizeof(float));
       HIPCHECK(hipMemset(recvbuff[i], 0, recvsize*sizeof(float)));
-
+      HIPCHECK(hipStreamSynchronize(NULL));
       HIPCHECK(hipStreamCreate(s+i));
   }
 
diff --git a/projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp b/projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp
index 8c566275cc..bb64c466c1 100644
--- a/projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp
+++ b/projects/rccl/tools/rccl-prim-test/rccl_prim_test.cpp
@@ -616,7 +616,7 @@ int main(int argc,char* argv[])
       HIPCHECK(hipStreamSynchronize(stream[i]));
       HIPCHECK(hipMemset(d_profiling_data[i], 0, sizeof(struct profiling_data_t)));
     }
-
+    HIPCHECK(hipStreamSynchronize(NULL));
     auto start = std::chrono::high_resolution_clock::now();
     for (int j = 0; j < iters; j ++) {
       for (int i = 0; i < nGpu; i ++) {