From c865fcd4bb3d04b710f1ba17d0d412be0fcac5b6 Mon Sep 17 00:00:00 2001
From: Brandon Potter <brandon.potter@amd.com>
Date: Tue, 6 Aug 2024 12:08:35 -0700
Subject: [PATCH] Fix memcpy_wave bug

---
 src/util.hpp | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/src/util.hpp b/src/util.hpp
index 0d4dad806a..67fdac6786 100644
--- a/src/util.hpp
+++ b/src/util.hpp
@@ -302,19 +302,34 @@ __device__ __forceinline__ void memcpy_wg(void* dst, void* src, size_t size) {
 }
 
 __device__ __forceinline__ void memcpy_wave(void* dst, void* src, size_t size) {
-  uint8_t* dst_bytes{static_cast<uint8_t*>(dst)};
-  uint8_t* src_bytes{static_cast<uint8_t*>(src)};
+  int wave_tid = get_flat_block_id() % WF_SIZE;
+  int wave_size{wave_SZ()};
 
   int cpy_size{};
-  int thread_id{get_flat_block_id()};
+  uint8_t* dst_bytes{nullptr};
+  uint8_t* dst_def{nullptr};
+  uint8_t* src_bytes{nullptr};
+  uint8_t* src_def{nullptr};
+
+  dst_def = reinterpret_cast<uint8_t*>(dst);
+  src_def = reinterpret_cast<uint8_t*>(src);
+  dst_bytes = dst_def;
+  src_bytes = src_def;
+
   for (int j{8}; j > 1; j >>= 1) {
     cpy_size = size / j;
-    for (int i{thread_id}; i < cpy_size; i += WF_SIZE) {
-      store_asm(src_bytes, dst_bytes, j);
+    for (int i{wave_tid}; i < cpy_size; i += wave_size) {
+      dst_bytes = dst_def;
+      src_bytes = src_def;
+
       src_bytes += i * j;
       dst_bytes += i * j;
-      size -= cpy_size * j;
+
+      store_asm(src_bytes, dst_bytes, j);
     }
+    size -= cpy_size * j;
+    dst_def += cpy_size * j;
+    src_def += cpy_size * j;
   }
 
   if (size == 1) {