2 files changed, 10 insertions, 0 deletions

diff --git a/drivers/gpu/nvgpu/vgpu/gp10b/vgpu_gr_gp10b.c b/drivers/gpu/nvgpu/vgpu/gp10b/vgpu_gr_gp10b.c
index 21fec478..749d6a15 100644
--- a/drivers/gpu/nvgpu/vgpu/gp10b/vgpu_gr_gp10b.c
+++ b/drivers/gpu/nvgpu/vgpu/gp10b/vgpu_gr_gp10b.c
@@ -56,6 +56,7 @@ static int vgpu_gr_gp10b_alloc_gr_ctx(struct gk20a *g,
         struct gr_ctx_desc *gr_ctx;
         u32 graphics_preempt_mode = 0;
         u32 compute_preempt_mode = 0;
+        struct vgpu_priv_data *priv = vgpu_get_priv_data(g);
         int err;
 
         gk20a_dbg_fn("");
@@ -71,6 +72,14 @@ static int vgpu_gr_gp10b_alloc_gr_ctx(struct gk20a *g,
         if (flags & NVGPU_ALLOC_OBJ_FLAGS_CILP)
                 compute_preempt_mode = NVGPU_COMPUTE_PREEMPTION_MODE_CILP;
 
+        if (priv->constants.force_preempt_mode && !graphics_preempt_mode &&
+                !compute_preempt_mode) {
+                graphics_preempt_mode = PASCAL_A == class ?
+                                        NVGPU_GRAPHICS_PREEMPTION_MODE_GFXP : 0;
+                compute_preempt_mode = PASCAL_COMPUTE_A == class ?
+                                        NVGPU_COMPUTE_PREEMPTION_MODE_CTA : 0;
+        }
+
         if (graphics_preempt_mode || compute_preempt_mode) {
                 if (g->ops.gr.set_ctxsw_preemption_mode) {
                         err = g->ops.gr.set_ctxsw_preemption_mode(g, gr_ctx, vm,
diff --git a/include/linux/tegra_vgpu.h b/include/linux/tegra_vgpu.h
index 3e3bbf58..c0964083 100644
--- a/include/linux/tegra_vgpu.h
+++ b/include/linux/tegra_vgpu.h
@@ -442,6 +442,7 @@ struct tegra_vgpu_constants_params {
          */
         u16 gpc_tpc_mask[TEGRA_VGPU_MAX_GPC_COUNT];
         u32 hwpm_ctx_size;
+        u8 force_preempt_mode;
 };
 
 struct tegra_vgpu_channel_cyclestats_snapshot_params {