diff --git a/drivers/gpu/nvgpu/os/linux/driver_common.c b/drivers/gpu/nvgpu/os/linux/driver_common.c
index b81cc90c2..4d345aaaf 100644
--- a/drivers/gpu/nvgpu/os/linux/driver_common.c
+++ b/drivers/gpu/nvgpu/os/linux/driver_common.c
@@ -98,8 +98,7 @@ static void nvgpu_init_gr_vars(struct gk20a *g)
 	gk20a_init_gr(g);
 
 	nvgpu_log_info(g, "total ram pages : %lu", totalram_pages);
-	g->gr.max_comptag_mem = totalram_pages
-				 >> (10 - (PAGE_SHIFT - 10));
+	g->gr.max_comptag_mem = totalram_size_in_mb;
 
 	/* Deduct the part taken by the running system */
 	if (platform->comptag_mem_deduct &&
diff --git a/drivers/gpu/nvgpu/os/linux/os_linux.h b/drivers/gpu/nvgpu/os/linux/os_linux.h
index 13c20a81c..5f35db09d 100644
--- a/drivers/gpu/nvgpu/os/linux/os_linux.h
+++ b/drivers/gpu/nvgpu/os/linux/os_linux.h
@@ -165,4 +165,6 @@ static inline struct device *dev_from_gk20a(struct gk20a *g)
 
 #define INTERFACE_NAME "nvhost%s-gpu"
 
+#define totalram_size_in_mb (totalram_pages >> (10 - (PAGE_SHIFT - 10)))
+
 #endif
diff --git a/drivers/gpu/nvgpu/os/linux/vgpu/vgpu_linux.c b/drivers/gpu/nvgpu/os/linux/vgpu/vgpu_linux.c
index 831913161..33c067c6a 100644
--- a/drivers/gpu/nvgpu/os/linux/vgpu/vgpu_linux.c
+++ b/drivers/gpu/nvgpu/os/linux/vgpu/vgpu_linux.c
@@ -439,8 +439,7 @@ int vgpu_probe(struct platform_device *pdev)
 	gk20a_init_gr(gk20a);
 
 	nvgpu_log_info(gk20a, "total ram pages : %lu", totalram_pages);
-	gk20a->gr.max_comptag_mem = totalram_pages
-				 >> (10 - (PAGE_SHIFT - 10));
+	gk20a->gr.max_comptag_mem = totalram_size_in_mb;
 
 	nvgpu_ref_init(&gk20a->refcount);