Using async allocation in median() (#220)

NVIDIA · Jul 19, 2022 · c4cc260 · c4cc260
1 parent 7c11d77
commit c4cc260
Showing 1 changed file with 7 additions and 1 deletion.
diff --git 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  
 }