diff --git a/drivers/gpu/nvgpu/common/gr/gr.c b/drivers/gpu/nvgpu/common/gr/gr.c
index 126fd3f14..0da1e19ae 100644
--- a/drivers/gpu/nvgpu/common/gr/gr.c
+++ b/drivers/gpu/nvgpu/common/gr/gr.c
@@ -399,6 +399,13 @@ static int gr_init_setup_sw(struct gk20a *g)
 		goto clean_up;
 	}
 
+	err = nvgpu_gr_hwpm_map_init(g, &g->gr.hwpm_map,
+		g->gr.ctx_vars.pm_ctxsw_image_size);
+	if (err != 0) {
+		nvgpu_err(g, "hwpm_map init failed");
+		goto clean_up;
+	}
+
 	err = nvgpu_gr_config_init_map_tiles(g, gr->config);
 	if (err != 0) {
 		goto clean_up;
diff --git a/drivers/gpu/nvgpu/common/gr/gr_falcon.c b/drivers/gpu/nvgpu/common/gr/gr_falcon.c
index 3918763a4..60c000042 100644
--- a/drivers/gpu/nvgpu/common/gr/gr_falcon.c
+++ b/drivers/gpu/nvgpu/common/gr/gr_falcon.c
@@ -171,15 +171,6 @@ int nvgpu_gr_falcon_init_ctx_state(struct gk20a *g)
 		if (err != 0) {
 			goto out;
 		}
-
-		if (g->gr.ctx_vars.pm_ctxsw_image_size != 0U) {
-			err = nvgpu_gr_hwpm_map_init(g, &g->gr.hwpm_map,
-				g->gr.ctx_vars.pm_ctxsw_image_size);
-			if (err != 0) {
-				nvgpu_err(g, "hwpm_map init failed");
-				goto out;
-			}
-		}
 	}
 
 out:
diff --git a/drivers/gpu/nvgpu/common/gr/hwpm_map.c b/drivers/gpu/nvgpu/common/gr/hwpm_map.c
index c7c400aa8..e550f97aa 100644
--- a/drivers/gpu/nvgpu/common/gr/hwpm_map.c
+++ b/drivers/gpu/nvgpu/common/gr/hwpm_map.c
@@ -42,6 +42,10 @@ int nvgpu_gr_hwpm_map_init(struct gk20a *g, struct nvgpu_gr_hwpm_map **hwpm_map,
 {
 	struct nvgpu_gr_hwpm_map *tmp_map;
 
+	if (size == 0U) {
+		return -EINVAL;
+	}
+
 	tmp_map = nvgpu_kzalloc(g, sizeof(*tmp_map));
 	if (tmp_map == NULL) {
 		return -ENOMEM;
@@ -65,6 +69,11 @@ void nvgpu_gr_hwpm_map_deinit(struct gk20a *g,
 	nvgpu_kfree(g, hwpm_map);
 }
 
+u32 nvgpu_gr_hwpm_map_get_size(struct nvgpu_gr_hwpm_map *hwpm_map)
+{
+	return hwpm_map->pm_ctxsw_image_size;
+}
+
 static int map_cmp(const void *a, const void *b)
 {
 	const struct ctxsw_buf_offset_map_entry *e1;
diff --git a/drivers/gpu/nvgpu/common/vgpu/gr/ctx_vgpu.c b/drivers/gpu/nvgpu/common/vgpu/gr/ctx_vgpu.c
index a55525c78..351efd509 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gr/ctx_vgpu.c
+++ b/drivers/gpu/nvgpu/common/vgpu/gr/ctx_vgpu.c
@@ -31,6 +31,7 @@
 #include <nvgpu/gk20a.h>
 #include <nvgpu/gr/global_ctx.h>
 #include <nvgpu/gr/ctx.h>
+#include <nvgpu/gr/hwpm_map.h>
 
 #include "common/gr/ctx_priv.h"
 
@@ -173,7 +174,7 @@ int vgpu_gr_alloc_pm_ctx(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
 	}
 
 	pm_ctx->mem.gpu_va = nvgpu_vm_alloc_va(vm,
-			g->gr.ctx_vars.pm_ctxsw_image_size,
+			nvgpu_gr_hwpm_map_get_size(g->gr.hwpm_map),
 			GMMU_PAGE_SIZE_KERNEL);
 
 	if (!pm_ctx->mem.gpu_va) {
@@ -181,7 +182,7 @@ int vgpu_gr_alloc_pm_ctx(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
 		return -ENOMEM;
 	}
 
-	pm_ctx->mem.size = g->gr.ctx_vars.pm_ctxsw_image_size;
+	pm_ctx->mem.size = nvgpu_gr_hwpm_map_get_size(g->gr.hwpm_map);
 	return 0;
 }
 
diff --git a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
index e2a60ea5d..450abbafb 100644
--- a/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
+++ b/drivers/gpu/nvgpu/common/vgpu/gr/gr_vgpu.c
@@ -40,6 +40,7 @@
 #include <nvgpu/gr/zbc.h>
 #include <nvgpu/gr/zcull.h>
 #include <nvgpu/gr/fecs_trace.h>
+#include <nvgpu/gr/hwpm_map.h>
 #include <nvgpu/cyclestats_snapshot.h>
 #include <nvgpu/power_features/pg.h>
 
@@ -710,6 +711,13 @@ static int vgpu_gr_init_gr_setup_sw(struct gk20a *g)
 		goto clean_up;
 	}
 
+	err = nvgpu_gr_hwpm_map_init(g, &g->gr.hwpm_map,
+		g->gr.ctx_vars.pm_ctxsw_image_size);
+	if (err != 0) {
+		nvgpu_err(g, "hwpm_map init failed");
+		goto clean_up;
+	}
+
 	err = vgpu_gr_init_gr_zcull(g, gr);
 	if (err) {
 		goto clean_up;
@@ -920,8 +928,6 @@ int vgpu_gr_update_hwpm_ctxsw_mode(struct gk20a *g,
 				"failed to allocate pm ctxt buffer");
 			return err;
 		}
-		nvgpu_gr_ctx_get_pm_ctx_mem(gr_ctx)->size =
-			g->gr.ctx_vars.pm_ctxsw_image_size;
 	}
 
 	msg.cmd = TEGRA_VGPU_CMD_CHANNEL_SET_HWPM_CTXSW_MODE;
diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index e42daa562..717c46a97 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -228,7 +228,7 @@ int gr_gk20a_update_hwpm_ctxsw_mode(struct gk20a *g,
 	if (mode != NVGPU_GR_CTX_HWPM_CTXSW_MODE_NO_CTXSW) {
 		nvgpu_gr_ctx_set_size(g->gr.gr_ctx_desc,
 			NVGPU_GR_CTX_PM_CTX,
-			g->gr.ctx_vars.pm_ctxsw_image_size);
+			nvgpu_gr_hwpm_map_get_size(g->gr.hwpm_map));
 
 		ret = nvgpu_gr_ctx_alloc_pm_ctx(g, gr_ctx,
 			g->gr.gr_ctx_desc, c->vm,
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gr/hwpm_map.h b/drivers/gpu/nvgpu/include/nvgpu/gr/hwpm_map.h
index 2bedba341..b2de73bf3 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gr/hwpm_map.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gr/hwpm_map.h
@@ -42,6 +42,8 @@ int nvgpu_gr_hwpm_map_init(struct gk20a *g, struct nvgpu_gr_hwpm_map **hwpm_map,
 void nvgpu_gr_hwpm_map_deinit(struct gk20a *g,
 	struct nvgpu_gr_hwpm_map *hwpm_map);
 
+u32 nvgpu_gr_hwpm_map_get_size(struct nvgpu_gr_hwpm_map *hwpm_map);
+
 int nvgpu_gr_hwmp_map_find_priv_offset(struct gk20a *g,
 	struct nvgpu_gr_hwpm_map *hwpm_map,
 	u32 addr, u32 *priv_offset);