2 files changed, 7 insertions, 2 deletions

diff --git a/drivers/gpu/nvgpu/vgpu/gr_vgpu.c b/drivers/gpu/nvgpu/vgpu/gr_vgpu.c
index c7996fd9..1c72ca76 100644
--- a/drivers/gpu/nvgpu/vgpu/gr_vgpu.c
+++ b/drivers/gpu/nvgpu/vgpu/gr_vgpu.c
@@ -1026,7 +1026,7 @@ int vgpu_gr_update_smpc_ctxsw_mode(struct gk20a *g,
 }
 
 int vgpu_gr_update_hwpm_ctxsw_mode(struct gk20a *g,
-        struct channel_gk20a *ch, bool enable)
+        struct channel_gk20a *ch, u64 gpu_va, bool enable)
 {
         struct tsg_gk20a *tsg;
         struct nvgpu_gr_ctx *ch_ctx;
@@ -1041,6 +1041,11 @@ int vgpu_gr_update_hwpm_ctxsw_mode(struct gk20a *g,
         if (!tsg)
                 return -EINVAL;
 
+        if (gpu_va) {
+                nvgpu_err(g, "gpu_va suppose to be allocated by this function.");
+                return -EINVAL;
+        }
+
         ch_ctx = &tsg->gr_ctx;
         pm_ctx = &ch_ctx->pm_ctx;
 
diff --git a/drivers/gpu/nvgpu/vgpu/gr_vgpu.h b/drivers/gpu/nvgpu/vgpu/gr_vgpu.h
index 1f55823c..d6b25b97 100644
--- a/drivers/gpu/nvgpu/vgpu/gr_vgpu.h
+++ b/drivers/gpu/nvgpu/vgpu/gr_vgpu.h
@@ -58,7 +58,7 @@ int vgpu_gr_set_sm_debug_mode(struct gk20a *g,
 int vgpu_gr_update_smpc_ctxsw_mode(struct gk20a *g,
         struct channel_gk20a *ch, bool enable);
 int vgpu_gr_update_hwpm_ctxsw_mode(struct gk20a *g,
-        struct channel_gk20a *ch, bool enable);
+        struct channel_gk20a *ch, u64 gpu_va, bool enable);
 int vgpu_gr_clear_sm_error_state(struct gk20a *g,
                 struct channel_gk20a *ch, u32 sm_id);
 int vgpu_gr_suspend_contexts(struct gk20a *g,