From cd1a71ad5681dd51925da5fd18a12ed60277a8c6 Mon Sep 17 00:00:00 2001
From: "romain.biessy" <romain.biessy@codeplay.com>
Date: Mon, 29 Jul 2024 18:08:02 +0200
Subject: [PATCH] Wait before freeing USM pointers

---
 tests/unit_tests/sparse_blas/source/sparse_spmm_usm.cpp | 2 ++
 tests/unit_tests/sparse_blas/source/sparse_spmv_usm.cpp | 2 ++
 tests/unit_tests/sparse_blas/source/sparse_spsv_usm.cpp | 2 ++
 3 files changed, 6 insertions(+)

diff --git a/tests/unit_tests/sparse_blas/source/sparse_spmm_usm.cpp b/tests/unit_tests/sparse_blas/source/sparse_spmm_usm.cpp
index 766915537..5778430a6 100644
--- a/tests/unit_tests/sparse_blas/source/sparse_spmm_usm.cpp
+++ b/tests/unit_tests/sparse_blas/source/sparse_spmm_usm.cpp
@@ -157,6 +157,8 @@ int test_spmm(sycl::device *dev, sparse_matrix_format_t format, intType nrows_A,
                                       a_host.data(), reset_nnz, static_cast<std::size_t>(nrows_A));
             }
             if (reset_nnz > nnz) {
+                // Wait before freeing usm pointers
+                ev_spmm.wait_and_throw();
                 ia_usm_uptr = malloc_device_uptr<intType>(main_queue, ia_host.size());
                 ja_usm_uptr = malloc_device_uptr<intType>(main_queue, ja_host.size());
                 a_usm_uptr = malloc_device_uptr<fpType>(main_queue, a_host.size());
diff --git a/tests/unit_tests/sparse_blas/source/sparse_spmv_usm.cpp b/tests/unit_tests/sparse_blas/source/sparse_spmv_usm.cpp
index fdeb57913..ded92a770 100644
--- a/tests/unit_tests/sparse_blas/source/sparse_spmv_usm.cpp
+++ b/tests/unit_tests/sparse_blas/source/sparse_spmv_usm.cpp
@@ -149,6 +149,8 @@ int test_spmv(sycl::device *dev, sparse_matrix_format_t format, intType nrows_A,
                                       a_host.data(), reset_nnz, static_cast<std::size_t>(nrows_A));
             }
             if (reset_nnz > nnz) {
+                // Wait before freeing usm pointers
+                ev_spmv.wait_and_throw();
                 ia_usm_uptr = malloc_device_uptr<intType>(main_queue, ia_host.size());
                 ja_usm_uptr = malloc_device_uptr<intType>(main_queue, ja_host.size());
                 a_usm_uptr = malloc_device_uptr<fpType>(main_queue, a_host.size());
diff --git a/tests/unit_tests/sparse_blas/source/sparse_spsv_usm.cpp b/tests/unit_tests/sparse_blas/source/sparse_spsv_usm.cpp
index 1bc025851..3b58db914 100644
--- a/tests/unit_tests/sparse_blas/source/sparse_spsv_usm.cpp
+++ b/tests/unit_tests/sparse_blas/source/sparse_spsv_usm.cpp
@@ -145,6 +145,8 @@ int test_spsv(sycl::device *dev, sparse_matrix_format_t format, intType m, doubl
                                       a_host.data(), reset_nnz, mu);
             }
             if (reset_nnz > nnz) {
+                // Wait before freeing usm pointers
+                ev_spsv.wait_and_throw();
                 ia_usm_uptr = malloc_device_uptr<intType>(main_queue, ia_host.size());
                 ja_usm_uptr = malloc_device_uptr<intType>(main_queue, ja_host.size());
                 a_usm_uptr = malloc_device_uptr<fpType>(main_queue, a_host.size());