4 files changed, 13 insertions, 4 deletions

diff --git a/drivers/gpu/nvgpu/gk20a/dbg_gpu_gk20a.c b/drivers/gpu/nvgpu/gk20a/dbg_gpu_gk20a.c
index 321cebb2..309fe75a 100644
--- a/drivers/gpu/nvgpu/gk20a/dbg_gpu_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/dbg_gpu_gk20a.c
@@ -891,8 +891,8 @@ static int nvgpu_dbg_gpu_ioctl_smpc_ctxsw_mode(struct dbg_session_gk20a *dbg_s,
                 goto clean_up;
         }
 
-        err = gr_gk20a_update_smpc_ctxsw_mode(g, ch_gk20a,
-                      args->mode == NVGPU_DBG_GPU_SMPC_CTXSW_MODE_CTXSW);
+        err = g->ops.gr.update_smpc_ctxsw_mode(g, ch_gk20a,
+                                args->mode == NVGPU_DBG_GPU_SMPC_CTXSW_MODE_CTXSW);
         if (err) {
                 gk20a_err(dev_from_gk20a(g),
                           "error (%d) during smpc ctxsw mode update\n", err);
@@ -927,8 +927,8 @@ static int nvgpu_dbg_gpu_ioctl_hwpm_ctxsw_mode(struct dbg_session_gk20a *dbg_s,
                 goto clean_up;
         }
 
-        err = gr_gk20a_update_hwpm_ctxsw_mode(g, ch_gk20a,
-                      args->mode == NVGPU_DBG_GPU_HWPM_CTXSW_MODE_CTXSW);
+        err = g->ops.gr.update_hwpm_ctxsw_mode(g, ch_gk20a,
+                                args->mode == NVGPU_DBG_GPU_HWPM_CTXSW_MODE_CTXSW);
         if (err)
                 gk20a_err(dev_from_gk20a(g),
                           "error (%d) during pm ctxsw mode update\n", err);
diff --git a/drivers/gpu/nvgpu/gk20a/gk20a.h b/drivers/gpu/nvgpu/gk20a/gk20a.h
index 9d8dc5f7..d9cc3d4f 100644
--- a/drivers/gpu/nvgpu/gk20a/gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/gk20a.h
@@ -172,6 +172,12 @@ struct gpu_ops {
                 void (*update_ctxsw_preemption_mode)(struct gk20a *g,
                                 struct channel_ctx_gk20a *ch_ctx,
                                 void *ctx_ptr);
+                int (*update_smpc_ctxsw_mode)(struct gk20a *g,
+                                struct channel_gk20a *c,
+                                bool enable);
+                int (*update_hwpm_ctxsw_mode)(struct gk20a *g,
+                                struct channel_gk20a *c,
+                                bool enable);
                 int (*dump_gr_regs)(struct gk20a *g,
                                 struct gk20a_debug_output *o);
                 int (*update_pc_sampling)(struct channel_gk20a *ch,
diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index b0b5571f..ada67edd 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -8308,4 +8308,6 @@ void gk20a_init_gr_ops(struct gpu_ops *gops)
         gops->gr.handle_sm_exception = gr_gk20a_handle_sm_exception;
         gops->gr.handle_tex_exception = gr_gk20a_handle_tex_exception;
         gops->gr.get_lrf_tex_ltc_dram_override = NULL;
+        gops->gr.update_smpc_ctxsw_mode = gr_gk20a_update_smpc_ctxsw_mode;
+        gops->gr.update_hwpm_ctxsw_mode = gr_gk20a_update_hwpm_ctxsw_mode;
 }
diff --git a/drivers/gpu/nvgpu/gk20a/mm_gk20a.h b/drivers/gpu/nvgpu/gk20a/mm_gk20a.h
index 52d6c4e5..5390536e 100644
--- a/drivers/gpu/nvgpu/gk20a/mm_gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/mm_gk20a.h
@@ -76,6 +76,7 @@ struct zcull_ctx_desc {
 struct pm_ctx_desc {
         struct mem_desc mem;
         u32 pm_mode;
+        bool ctx_was_enabled;   /* Used in the virtual case only */
 };
 
 struct gk20a;