diff --git a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
index 42e6125c0..398e2b420 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
+++ b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
@@ -1262,15 +1262,11 @@ void vgpu_gr_init_cyclestats(struct gk20a *g)
 {
 #if defined(CONFIG_NVGPU_CYCLESTATS)
 	bool snapshots_supported = true;
-	u32 max_css_buffer_size;
 
 	/* cyclestats not supported on vgpu */
 	nvgpu_set_enabled(g, NVGPU_SUPPORT_CYCLE_STATS, false);
 
-	max_css_buffer_size = vgpu_css_get_buffer_size(g);
-
-	/* snapshots not supported if the buffer size is 0 */
-	if (max_css_buffer_size == 0) {
+	if (vgpu_css_init(g) != 0) {
 		snapshots_supported = false;
 	}
 
diff --git a/drivers/gpu/nvgpu/common/vgpu/perf/cyclestats_snapshot_vgpu.c b/drivers/gpu/nvgpu/common/vgpu/perf/cyclestats_snapshot_vgpu.c
index 2be0e62e4..a0ec94c4c 100644
--- a/drivers/gpu/nvgpu/common/vgpu/perf/cyclestats_snapshot_vgpu.c
+++ b/drivers/gpu/nvgpu/common/vgpu/perf/cyclestats_snapshot_vgpu.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2019, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -32,19 +32,13 @@
 #include "cyclestats_snapshot_vgpu.h"
 #include "common/vgpu/ivc/comm_vgpu.h"
 
-static struct tegra_hv_ivm_cookie *css_cookie;
-
-static int vgpu_css_reserve_mempool(struct gk20a *g,
-					struct tegra_hv_ivm_cookie **cookie_p)
+int vgpu_css_init(struct gk20a *g)
 {
+	struct vgpu_priv_data *priv = vgpu_get_priv_data(g);
 	struct tegra_hv_ivm_cookie *cookie;
 	u32 mempool;
 	int err;
 
-	if (cookie_p == NULL) {
-		return -EINVAL;
-	}
-
 	err = nvgpu_dt_read_u32_index(g, "mempool-css", 1, &mempool);
 	if (err) {
 		nvgpu_err(g, "dt missing mempool-css");
@@ -58,40 +52,28 @@ static int vgpu_css_reserve_mempool(struct gk20a *g,
 		return -EINVAL;
 	}
 
-	*cookie_p = cookie;
+	priv->css_cookie = cookie;
 
 	return 0;
 }
 
 u32 vgpu_css_get_buffer_size(struct gk20a *g)
 {
-	struct tegra_hv_ivm_cookie *cookie;
-	u32 size;
-	int err;
+	struct vgpu_priv_data *priv = vgpu_get_priv_data(g);
 
 	nvgpu_log_fn(g, " ");
 
-	if (css_cookie) {
-		size = (u32)vgpu_ivm_get_size(css_cookie);
-		nvgpu_log_info(g, "buffer size = 0x%08x", size);
-		return size;
+	if (NULL == priv->css_cookie) {
+		return 0U;
 	}
 
-	err = vgpu_css_reserve_mempool(g, &cookie);
-	if (0 != err) {
-		return 0;
-	}
-
-	size = vgpu_ivm_get_size(cookie);
-
-	vgpu_ivm_mempool_unreserve(cookie);
-	nvgpu_log_info(g, "buffer size = 0x%08x", size);
-	return size;
+	return vgpu_ivm_get_size(priv->css_cookie);
 }
 
 static int vgpu_css_init_snapshot_buffer(struct gk20a *g)
 {
 	struct gk20a_cs_snapshot *data = g->cs_data;
+	struct vgpu_priv_data *priv = vgpu_get_priv_data(g);
 	void *buf = NULL;
 	int err;
 	u64 size;
@@ -102,12 +84,11 @@ static int vgpu_css_init_snapshot_buffer(struct gk20a *g)
 		return 0;
 	}
 
-	err = vgpu_css_reserve_mempool(g, &css_cookie);
-	if (0 != err) {
-		return err;
+	if (NULL == priv->css_cookie) {
+		return -EINVAL;
 	}
 
-	size = vgpu_ivm_get_size(css_cookie);
+	size = vgpu_ivm_get_size(priv->css_cookie);
 	/* Make sure buffer size is large enough */
 	if (size < CSS_MIN_HW_SNAPSHOT_SIZE) {
 		nvgpu_info(g, "mempool size 0x%llx too small", size);
@@ -115,7 +96,7 @@ static int vgpu_css_init_snapshot_buffer(struct gk20a *g)
 		goto fail;
 	}
 
-	buf = vgpu_ivm_mempool_map(css_cookie);
+	buf = vgpu_ivm_mempool_map(priv->css_cookie);
 	if (!buf) {
 		nvgpu_info(g, "vgpu_ivm_mempool_map failed");
 		err = -EINVAL;
@@ -129,25 +110,21 @@ static int vgpu_css_init_snapshot_buffer(struct gk20a *g)
 	(void) memset(data->hw_snapshot, 0xff, size);
 	return 0;
 fail:
-	vgpu_ivm_mempool_unreserve(css_cookie);
-	css_cookie = NULL;
 	return err;
 }
 
 void vgpu_css_release_snapshot_buffer(struct gk20a *g)
 {
 	struct gk20a_cs_snapshot *data = g->cs_data;
+	struct vgpu_priv_data *priv = vgpu_get_priv_data(g);
 
 	if (!data->hw_snapshot) {
 		return;
 	}
 
-	vgpu_ivm_mempool_unmap(css_cookie, data->hw_snapshot);
+	vgpu_ivm_mempool_unmap(priv->css_cookie, data->hw_snapshot);
 	data->hw_snapshot = NULL;
 
-	vgpu_ivm_mempool_unreserve(css_cookie);
-	css_cookie = NULL;
-
 	nvgpu_log_info(g, "cyclestats(vgpu): buffer for snapshots released\n");
 }
 
diff --git a/drivers/gpu/nvgpu/common/vgpu/perf/cyclestats_snapshot_vgpu.h b/drivers/gpu/nvgpu/common/vgpu/perf/cyclestats_snapshot_vgpu.h
index 0436db827..24b86bed8 100644
--- a/drivers/gpu/nvgpu/common/vgpu/perf/cyclestats_snapshot_vgpu.h
+++ b/drivers/gpu/nvgpu/common/vgpu/perf/cyclestats_snapshot_vgpu.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2019, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -29,6 +29,7 @@ struct gk20a;
 struct nvgpu_channel;
 struct gk20a_cs_snapshot_client;
 
+int vgpu_css_init(struct gk20a *g);
 void vgpu_css_release_snapshot_buffer(struct gk20a *g);
 int vgpu_css_flush_snapshots(struct nvgpu_channel *ch,
 			u32 *pending, bool *hw_overflow);
diff --git a/drivers/gpu/nvgpu/include/nvgpu/vgpu/vgpu.h b/drivers/gpu/nvgpu/include/nvgpu/vgpu/vgpu.h
index 144504c35..fbb7874b9 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/vgpu/vgpu.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/vgpu/vgpu.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2014-2019, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2014-2020, NVIDIA CORPORATION.  All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a
  * copy of this software and associated documentation files (the "Software"),
@@ -42,6 +42,7 @@ struct vgpu_priv_data {
 	u32 num_freqs;
 	unsigned long *freqs;
 	struct nvgpu_mutex vgpu_clk_get_freq_lock;
+	struct tegra_hv_ivm_cookie *css_cookie;
 };
 
 struct vgpu_priv_data *vgpu_get_priv_data(struct gk20a *g);