From 1cd2c294fafd873dbc4e707a12e7ec1bea18b1d4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Orell=20B=C3=BChler?= <orell@orellbuehler.dev>
Date: Sun, 19 May 2024 19:50:46 +0200
Subject: [PATCH] allow total memory hiding

---
 src/hook.c | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/hook.c b/src/hook.c
index 45965e2..6876cdb 100644
--- a/src/hook.c
+++ b/src/hook.c
@@ -41,6 +41,7 @@
 #include "utlist.h"
 
 #define ENV_NVSHARE_ENABLE_SINGLE_OVERSUB  "NVSHARE_ENABLE_SINGLE_OVERSUB"
+#define ENV_NVSHARE_HIDE_MEMINFO_RESERVE_MIB "NVSHARE_HIDE_MEMINFO_RESERVE_MIB"
 
 #define MEMINFO_RESERVE_MIB 1536           /* MiB */
 #define KERN_SYNC_DURATION_BIG 10          /* seconds */
@@ -84,6 +85,7 @@ int pending_kernel_window = 1;
 pthread_mutex_t kcount_mutex;
 
 int enable_single_oversub = 0;
+int hide_meminfo_reserve_mib = 0;
 int nvml_ok = 1;
 
 /* Representation of a CUDA memory allocation */
@@ -330,6 +332,11 @@ static void initialize_libnvshare(void)
 		log_warn("Enabling GPU memory oversubscription for this"
 		         " application");
 	}
+	value = getenv(ENV_NVSHARE_HIDE_MEMINFO_RESERVE_MIB);
+	if (value != NULL) {
+		hide_meminfo_reserve_mib = 1;
+		log_warn("Hiding %d MiB from CUDA memory info", MEMINFO_RESERVE_MIB);
+	}
 
 	bootstrap_cuda();
 }
@@ -782,6 +789,10 @@ CUresult cuMemGetInfo(size_t *free, size_t *total)
 	reserve_mib = (MEMINFO_RESERVE_MIB) MiB;
 	*free = *total - (size_t) reserve_mib;
 
+	if (hide_meminfo_reserve_mib == 1) {
+		*total -= reserve_mib;
+	}
+
 	log_debug("nvshare's cuMemGetInfo returning free=%.2f MiB,"
 		  " total=%.2f MiB", toMiB(*free), toMiB(*total));
 	return result;