From c4cc26076330a9812f3848bcf4c98adaab58c718 Mon Sep 17 00:00:00 2001
From: Cliff Burdick <30670611+cliffburdick@users.noreply.github.com>
Date: Tue, 19 Jul 2022 12:46:08 -0700
Subject: [PATCH] Using async allocation in median() (#220)

---
 include/matx_reduce.h | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/include/matx_reduce.h b/include/matx_reduce.h
index 6663d03f4..2f06b8e19 100644
--- a/include/matx_reduce.h
+++ b/include/matx_reduce.h
@@ -1149,7 +1149,11 @@ void inline median(TensorType &dest,
   constexpr int RANK_IN = TensorInType::Rank();
   static_assert(RANK_IN <= 2 && (RANK_IN == TensorType::Rank() + 1));
 
-  auto tmp_sort = make_tensor<T>(in.Shape());
+  T *tmp_alloc;
+  matxAlloc(reinterpret_cast<void **>(&tmp_alloc),
+                sizeof(T) * TotalSize(in), MATX_ASYNC_DEVICE_MEMORY,
+                stream);
+  auto tmp_sort = make_tensor<T, non_owning>(tmp_alloc, in.Shape());
 
   // If the rank is 0 we're finding the median of a vector
   if constexpr (RANK_IN == 1) {
@@ -1187,6 +1191,8 @@ void inline median(TensorType &dest,
       (dest = (sv + sv2) / 2.0f).run(stream);
     }
   }
+
+  matxFree(tmp_alloc);
 #endif  
 }