From 0b6f8890a9c8714175d7e332b15bf161de0610e3 Mon Sep 17 00:00:00 2001
From: chang-l <liuc@nvidia.com>
Date: Mon, 18 Nov 2024 15:59:59 -0800
Subject: [PATCH] Sync stream for scatter_op

---
 cpp/src/wholememory_ops/scatter_op_impl_mapped.cu | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cpp/src/wholememory_ops/scatter_op_impl_mapped.cu b/cpp/src/wholememory_ops/scatter_op_impl_mapped.cu
index 77f570f90..2d2bef81e 100644
--- a/cpp/src/wholememory_ops/scatter_op_impl_mapped.cu
+++ b/cpp/src/wholememory_ops/scatter_op_impl_mapped.cu
@@ -18,6 +18,7 @@
 #include <wholememory/env_func_ptrs.h>
 #include <wholememory/wholememory.h>
 
+#include "cuda_macros.hpp"
 #include "wholememory_ops/functions/gather_scatter_func.h"
 
 namespace wholememory_ops {
@@ -41,6 +42,7 @@ wholememory_error_code_t wholememory_scatter_mapped(
                       wholememory_desc,
                       stream,
                       scatter_sms);
+  WM_CUDA_CHECK(cudaStreamSynchronize(stream));
 }
 
 }  // namespace wholememory_ops