From 6e77c99cc80628b58743d116ddd316bbb9579bf2 Mon Sep 17 00:00:00 2001
From: Chang Liu <liuc@nvidia.com>
Date: Thu, 21 Nov 2024 14:55:00 -0800
Subject: [PATCH] [Bugfix] Sync stream for scatter_op (#235)

This is to address the issue found by: https://github.com/rapidsai/wholegraph/pull/229

Authors:
  - Chang Liu (https://github.com/chang-l)

Approvers:
  - https://github.com/linhu-nv
  - Brad Rees (https://github.com/BradReesWork)

URL: https://github.com/rapidsai/wholegraph/pull/235
---
 cpp/src/wholememory_ops/scatter_op_impl_mapped.cu | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cpp/src/wholememory_ops/scatter_op_impl_mapped.cu b/cpp/src/wholememory_ops/scatter_op_impl_mapped.cu
index 77f570f90..dbe68b120 100644
--- a/cpp/src/wholememory_ops/scatter_op_impl_mapped.cu
+++ b/cpp/src/wholememory_ops/scatter_op_impl_mapped.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -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