Tabulate based cuda memory summary

DESCRIPTION

@@ -50,7 +50,9 @@ Suggests:
     palmerpenguins,
     mvtnorm,
     numDeriv,
-    katex
+    katex,
+    tabulate,
+    prettyunits
 VignetteBuilder: knitr
 Collate:
     'R7.R'

R/cuda.R

@@ -154,8 +154,111 @@ cuda_memory_stats <- function(device = cuda_current_device()) {
 #' @rdname cuda_memory_stats
 #' @export
 cuda_memory_summary <- function(device = cuda_current_device()) {
+  if (rlang::is_installed("tabulate")) {
+    tabulate_memory_summary(device)
+  } else {
+    cli::cli_warn(
+      .frequency = "once", .frequency_id = "tabulate",
+      "Install the {.var tabulate} package for a better display."
+    )
+    result <- cuda_memory_stats(device)
+    print(result)  
+  }
+}
+
+tabulate_memory_summary <- function(device = cuda_current_device()) {
+  rlang::check_installed("prettyunits")
   result <- cuda_memory_stats(device)
-  print(result)
+
+  title <- tabulate::tabulate_table()
+  title %>% 
+    tabulate::table_add_row(paste0("CUDA memory summary, device ID ", device))
+
+  OOM <- tabulate::tabulate_table()
+  OOM %>% 
+    tabulate::table_add_row(c(
+      paste0("CUDA OOMs: ", result$num_ooms),
+      paste0("cudaMalloc retries: ", result$num_alloc_retries) 
+    )) %>% 
+    tabulate::format_hide_border()
+
+  title %>% 
+    tabulate::table_add_row(OOM)
+
+  metrics <- tabulate::tabulate_table()
+  metrics %>% 
+    tabulate::table_add_row(c("Metric", "Cur Usage", "Peak Usage", "Tot Alloc", "Tot Freed"))
+
+  i <- 2
+
+  add_metric_row <- function(table, title, metrics, format) {
+    table %>% 
+      tabulate::table_add_row(c(
+        title, 
+        format(metrics$current),
+        format(metrics$peak),
+        format(metrics$allocated),
+        format(metrics$freed)
+      ))
+    if ((i+1) %% 3 != 0) {
+      table[i,] %>% tabulate::format_hide_border_top()
+      table[i, 1] %>% tabulate::format_font_align("right")
+    } else {
+      table[i,1] %>% tabulate::format_font_align("left")
+    }
+
+    i <<- i+1
+    invisible(table)
+  }
+
+  add_metric <- function(table, title, metrics, format = prettyunits::pretty_bytes) {
+    table %>% 
+      add_metric_row(title, metrics$all, format) %>% 
+      add_metric_row("from large pool", metrics$large_pool, format) %>% 
+      add_metric_row("from small pool", metrics$small_pool, format)
+    invisible(table)
+  }
+
+  format_int <- function(x) {
+    format(x, big.mark = ",", scientific = FALSE)
+  }
+
+  metrics %>% 
+    add_metric("Allocated memory", result$allocated_bytes) %>% 
+    add_metric("Active memory", result$active_bytes) %>% 
+    add_metric("GPU reserved memory", result$reserved_bytes) %>% 
+    add_metric("Allocations", result$allocation, format_int) %>% 
+    add_metric("Active allocs", result$active, format_int) %>% 
+    add_metric("GPU reserved segments", result$segment, format_int)
+
+  metrics %>% 
+    tabulate::table_add_row(c(
+      "Oversize allocations",
+      sapply(result$oversize_allocations, format_int)
+    ))
+
+  metrics %>% 
+    tabulate::table_add_row(c(
+      "Oversize GPU segments",
+      sapply(result$oversize_segments, format_int)
+    ))
+
+  space <- "."
+
+  metrics[1,] %>% tabulate::format_hide_border_top()
+  metrics[,1] %>% tabulate::format_border_left(value = space)
+  metrics[,1] %>% tabulate::format_corner(value = "-")
+  metrics[,5] %>% tabulate::format_border_right(value = space)
+  metrics[,5] %>% tabulate::format_corner(value="-")
+
+  title %>% 
+    tabulate::table_add_row(metrics)
+
+  title[1,1] %>% tabulate::format_font_align("center")
+  title[2,1] %>% tabulate::format_font_align("center")
+  print(title)
+
+  invisible(result)
 }
 
 #' @export