diff --git a/drivers/gpu/nvgpu/common/gr/gr_falcon.c b/drivers/gpu/nvgpu/common/gr/gr_falcon.c
index 22c07c3cd..774a8a0c3 100644
--- a/drivers/gpu/nvgpu/common/gr/gr_falcon.c
+++ b/drivers/gpu/nvgpu/common/gr/gr_falcon.c
@@ -90,7 +90,6 @@ int nvgpu_gr_falcon_bind_fecs_elpg(struct gk20a *g)
 #ifdef CONFIG_NVGPU_LS_PMU
 	struct nvgpu_pmu *pmu = g->pmu;
 	struct mm_gk20a *mm = &g->mm;
-	struct vm_gk20a *vm = mm->pmu.vm;
 	int err = 0;
 	u32 size;
 	u32 data;
@@ -109,12 +108,10 @@ int nvgpu_gr_falcon_bind_fecs_elpg(struct gk20a *g)
 
 	nvgpu_log(g, gpu_dbg_gr, "FECS PG buffer size = %u", size);
 
-	if (nvgpu_pmu_pg_buf_get_cpu_va(g, pmu) == NULL) {
-		err = nvgpu_dma_alloc_map_sys(vm, size, nvgpu_pmu_pg_buf(g, pmu));
-		if (err != 0) {
-			nvgpu_err(g, "failed to allocate memory");
-			return -ENOMEM;
-		}
+	err = nvgpu_pmu_pg_buf_alloc(g, pmu, size);
+	if (err != 0) {
+		nvgpu_err(g, "failed to allocate pg_buf memory");
+		return err;
 	}
 
 	data = g->ops.gr.falcon.get_fecs_current_ctx_data(g,
diff --git a/drivers/gpu/nvgpu/common/pmu/pg/pmu_pg.c b/drivers/gpu/nvgpu/common/pmu/pg/pmu_pg.c
index b13b883ba..3f3b64948 100644
--- a/drivers/gpu/nvgpu/common/pmu/pg/pmu_pg.c
+++ b/drivers/gpu/nvgpu/common/pmu/pg/pmu_pg.c
@@ -1188,22 +1188,21 @@ u64 nvgpu_pmu_pg_buf_get_gpu_va(struct gk20a *g, struct nvgpu_pmu *pmu)
 	return pmu->pg->pg_buf.gpu_va;
 }
 
-struct nvgpu_mem *nvgpu_pmu_pg_buf(struct gk20a *g, struct nvgpu_pmu *pmu)
+int nvgpu_pmu_pg_buf_alloc(struct gk20a *g, struct nvgpu_pmu *pmu, u32 size)
 {
-	if (!is_pg_supported(g, pmu->pg)) {
-		return NULL;
+	struct mm_gk20a *mm = &g->mm;
+	struct vm_gk20a *vm = mm->pmu.vm;
+	int err = 0;
+
+	if (!nvgpu_mem_is_valid(&pmu->pg->pg_buf)) {
+		err = nvgpu_dma_alloc_map_sys(vm, size, &pmu->pg->pg_buf);
+		if (err != 0) {
+			nvgpu_err(g, "failed to allocate pg_buf");
+			return err;
+		}
 	}
 
-	return &pmu->pg->pg_buf;
-}
-
-void *nvgpu_pmu_pg_buf_get_cpu_va(struct gk20a *g, struct nvgpu_pmu *pmu)
-{
-	if (!is_pg_supported(g, pmu->pg)) {
-		return NULL;
-	}
-
-	return pmu->pg->pg_buf.cpu_va;
+	return err;
 }
 
 int nvgpu_pmu_restore_golden_img_state(struct gk20a *g)
diff --git a/drivers/gpu/nvgpu/include/nvgpu/pmu/pmu_pg.h b/drivers/gpu/nvgpu/include/nvgpu/pmu/pmu_pg.h
index 55402e774..fbb515d28 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/pmu/pmu_pg.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/pmu/pmu_pg.h
@@ -176,9 +176,7 @@ int nvgpu_pmu_elpg_statistics(struct gk20a *g, u32 pg_engine_id,
 		struct pmu_pg_stats_data *pg_stat_data);
 void nvgpu_pmu_save_zbc(struct gk20a *g, u32 entries);
 bool nvgpu_pmu_is_lpwr_feature_supported(struct gk20a *g, u32 feature_id);
-
+int nvgpu_pmu_pg_buf_alloc(struct gk20a *g, struct nvgpu_pmu *pmu, u32 size);
 u64 nvgpu_pmu_pg_buf_get_gpu_va(struct gk20a *g, struct nvgpu_pmu *pmu);
-struct nvgpu_mem *nvgpu_pmu_pg_buf(struct gk20a *g, struct nvgpu_pmu *pmu);
-void *nvgpu_pmu_pg_buf_get_cpu_va(struct gk20a *g, struct nvgpu_pmu *pmu);
 
 #endif /* NVGPU_PMU_PG_H */