From cd0a2eba774e1bcdad886f3b5dcb953121237d38 Mon Sep 17 00:00:00 2001
From: Balint Joo <bjoo@nvidia.com>
Date: Fri, 24 May 2024 20:10:14 +0000
Subject: [PATCH] Made maximum contingent on precision - 1GiB for 32 bit - 16
 GiB for 64 bit - removed comparison of NULL with Ox0 - reapplied clang-format

---
 benchmarks/async_alloc/async_alloc.cpp | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/benchmarks/async_alloc/async_alloc.cpp b/benchmarks/async_alloc/async_alloc.cpp
index 308f3de9bf9..075702891d9 100644
--- a/benchmarks/async_alloc/async_alloc.cpp
+++ b/benchmarks/async_alloc/async_alloc.cpp
@@ -15,7 +15,12 @@ std::vector<std::pair<size_t, double>> inner_loop_times;
 std::pair<double, double> test(bool up) {
   int iters      = 50;
   size_t minimum = 8 / sizeof(float);  // 64K
-  size_t maximum = (16l * 1024l * 1024l * 1024l) / sizeof(float);
+
+  size_t gb = 1024 * 1024 * 1024 / sizeof(float);  // number of floats per GiB
+  size_t maximum = gb;  // on 32 bit, we make 1GiB the max
+
+  // On 64-bit we make 16 GiB the max
+  if constexpr (sizeof(size_t) == 8) maximum *= 16;
 
   std::vector<size_t> sizes;
   if (up) {
@@ -64,7 +69,7 @@ int main(int argc, char *argv[]) {
   char *env_string = getenv("KOKKOS_CUDA_MEMPOOL_SIZE");
   std::cout << "Async Malloc Benchmark: KOKKOS_CUDA_MEMPOOL_SIZE is ";
 
-  if (env_string == nullptr || env_string == 0x0)
+  if (env_string == nullptr)
     std::cout << "not set,";
   else
     std::cout << " " << env_string << ",";