From 741ce10538537778e1369153c0a3fe2afb2c370d Mon Sep 17 00:00:00 2001
From: Matthew Martineau <mmartineau@nvidia.com>
Date: Fri, 14 Jun 2024 05:23:13 -0700
Subject: [PATCH] Ensure transformVector adds to mem pool and async frees.
 Needs some generalisation.

---
 src/distributed/distributed_manager.cu | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/distributed/distributed_manager.cu b/src/distributed/distributed_manager.cu
index d1485c9e..066e7bfb 100644
--- a/src/distributed/distributed_manager.cu
+++ b/src/distributed/distributed_manager.cu
@@ -5272,6 +5272,7 @@ void DistributedManager<TemplateConfig<AMGX_device, t_vecPrec, t_matPrec, t_indP
 template <AMGX_VecPrecision t_vecPrec, AMGX_MatPrecision t_matPrec, AMGX_IndPrecision t_indPrec>
 void DistributedManager<TemplateConfig<AMGX_device, t_vecPrec, t_matPrec, t_indPrec> >::transformVector(VVector_v &v)
 {
+    amgx::memory::setAsyncFreeFlag(true);
     if (this->neighbors.size() == 0) { return; }
     else if (this->renumbering.size() == 0)
     {
@@ -5309,6 +5310,11 @@ void DistributedManager<TemplateConfig<AMGX_device, t_vecPrec, t_matPrec, t_indP
 
     cudaCheckError();
     v.set_transformed();
+    amgx::thrust::global_thread_handle::joinDevicePools();
+    amgx::memory::setAsyncFreeFlag(true);
+
+    // XXX Is this needed?
+    amgx::thrust::global_thread_handle::cudaFreeWait();
 }
 
 template <AMGX_VecPrecision t_vecPrec, AMGX_MatPrecision t_matPrec, AMGX_IndPrecision t_indPrec>