From cc00fb431b5d492a09da13af6ca76c1230e9b52d Mon Sep 17 00:00:00 2001
From: Balint Joo <bjoo@nvidia.com>
Date: Fri, 24 May 2024 15:56:34 +0000
Subject: [PATCH] Made async alloc work with default memory space

---
 benchmarks/async_alloc/Makefile        | 2 ++
 benchmarks/async_alloc/async_alloc.cpp | 6 ++++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/benchmarks/async_alloc/Makefile b/benchmarks/async_alloc/Makefile
index 5d5157dff5f..008fb1c0fa1 100644
--- a/benchmarks/async_alloc/Makefile
+++ b/benchmarks/async_alloc/Makefile
@@ -1,3 +1,5 @@
+#KOKKOS_DEVICES=OpenMP
+#KOKKOS_ARCH="ZEN2"
 KOKKOS_DEVICES=Cuda
 KOKKOS_CUDA_OPTIONS=enable_lambda
 KOKKOS_ARCH = "ZEN2,Ada89"
diff --git a/benchmarks/async_alloc/async_alloc.cpp b/benchmarks/async_alloc/async_alloc.cpp
index 7b71d1479e7..308f3de9bf9 100644
--- a/benchmarks/async_alloc/async_alloc.cpp
+++ b/benchmarks/async_alloc/async_alloc.cpp
@@ -5,6 +5,8 @@
 // Kokkos Headers
 #include "Kokkos_Core.hpp"
 
+using MemorySpace = Kokkos::DefaultExecutionSpace::memory_space;
+
 // Store the individual timings
 std::vector<std::pair<size_t, double>> inner_loop_times;
 
@@ -24,7 +26,7 @@ std::pair<double, double> test(bool up) {
 
   Kokkos::Timer first_alloc_timer;
   {  // Prime the pump - first long alloc -- Time it.
-    Kokkos::View<float *, Kokkos::CudaSpace> dummy("unlabeled", 64);
+    Kokkos::View<float *, MemorySpace> dummy("unlabeled", 64);
   }
   double first_alloc_time = first_alloc_timer.seconds();
 
@@ -36,7 +38,7 @@ std::pair<double, double> test(bool up) {
   for (size_t num : sizes) {
     inner_loop_timer.reset();
     for (int i = 0; i < iters; i++) {
-      Kokkos::View<float *, Kokkos::CudaSpace> a("unlabeled", num);
+      Kokkos::View<float *, MemorySpace> a("unlabeled", num);
     }
     double inner_loop_time = inner_loop_timer.seconds();