3 files changed, 35 insertions, 5 deletions

diff --git a/drivers/gpu/nvgpu/gp10b/gr_gp10b.c b/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
index 0a60612d..576cdf45 100644
--- a/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
@@ -510,13 +510,14 @@ static int gr_gp10b_init_ctx_state(struct gk20a *g)
 
 int gr_gp10b_alloc_gr_ctx(struct gk20a *g,
                           struct gr_ctx_desc **gr_ctx, struct vm_gk20a *vm,
+                          u32 class,
                           u32 flags)
 {
         int err;
 
         gk20a_dbg_fn("");
 
-        err = gr_gk20a_alloc_gr_ctx(g, gr_ctx, vm, flags);
+        err = gr_gk20a_alloc_gr_ctx(g, gr_ctx, vm, class, flags);
         if (err)
                 return err;
 
@@ -566,9 +567,15 @@ int gr_gp10b_alloc_gr_ctx(struct gk20a *g,
                         goto fail_free_betacb;
                 }
 
-                (*gr_ctx)->t18x.preempt_mode = flags;
+                (*gr_ctx)->preempt_mode = flags;
         }
 
+        if (class == PASCAL_COMPUTE_A)
+                if (flags == NVGPU_GR_PREEMPTION_MODE_CILP)
+                        (*gr_ctx)->preempt_mode = NVGPU_GR_PREEMPTION_MODE_CILP;
+                else
+                        (*gr_ctx)->preempt_mode = NVGPU_GR_PREEMPTION_MODE_CTA;
+
         gk20a_dbg_fn("done");
 
         return err;
@@ -610,16 +617,24 @@ static void gr_gp10b_update_ctxsw_preemption_mode(struct gk20a *g,
         struct gr_ctx_desc *gr_ctx = ch_ctx->gr_ctx;
         u32 gfxp_preempt_option =
                 ctxsw_prog_main_image_graphics_preemption_options_control_gfxp_f();
+        u32 cilp_preempt_option =
+                ctxsw_prog_main_image_compute_preemption_options_control_cilp_f();
         int err;
 
         gk20a_dbg_fn("");
 
-        if (gr_ctx->t18x.preempt_mode == NVGPU_GR_PREEMPTION_MODE_GFXP) {
+        if (gr_ctx->preempt_mode == NVGPU_GR_PREEMPTION_MODE_GFXP) {
                 gk20a_dbg_info("GfxP: %x", gfxp_preempt_option);
                 gk20a_mem_wr32(ctx_ptr + ctxsw_prog_main_image_graphics_preemption_options_o(), 0,
                                 gfxp_preempt_option);
         }
 
+        if (gr_ctx->preempt_mode == NVGPU_GR_PREEMPTION_MODE_CILP) {
+                gk20a_dbg_info("CILP: %x", cilp_preempt_option);
+                gk20a_mem_wr32(ctx_ptr + ctxsw_prog_main_image_compute_preemption_options_o(), 0,
+                                cilp_preempt_option);
+        }
+
         if (gr_ctx->t18x.preempt_ctxsw_buffer.gpu_va) {
                 u32 addr;
                 u32 size;
diff --git a/drivers/gpu/nvgpu/gp10b/gr_gp10b.h b/drivers/gpu/nvgpu/gp10b/gr_gp10b.h
index 6bbda564..1b99cafb 100644
--- a/drivers/gpu/nvgpu/gp10b/gr_gp10b.h
+++ b/drivers/gpu/nvgpu/gp10b/gr_gp10b.h
@@ -39,14 +39,13 @@ struct gr_t18x {
 };
 
 struct gr_ctx_desc_t18x {
-        int preempt_mode;
         struct mem_desc preempt_ctxsw_buffer;
         struct mem_desc spill_ctxsw_buffer;
         struct mem_desc betacb_ctxsw_buffer;
         struct mem_desc pagepool_ctxsw_buffer;
 };
 
-#define NVGPU_GR_PREEMPTION_MODE_WFI            0
 #define NVGPU_GR_PREEMPTION_MODE_GFXP           1
+#define NVGPU_GR_PREEMPTION_MODE_CILP           3
 
 #endif
diff --git a/drivers/gpu/nvgpu/gp10b/hw_ctxsw_prog_gp10b.h b/drivers/gpu/nvgpu/gp10b/hw_ctxsw_prog_gp10b.h
index 0892f12e..7872c19c 100644
--- a/drivers/gpu/nvgpu/gp10b/hw_ctxsw_prog_gp10b.h
+++ b/drivers/gpu/nvgpu/gp10b/hw_ctxsw_prog_gp10b.h
@@ -254,4 +254,20 @@ static inline u32 ctxsw_prog_main_image_full_preemption_ptr_o(void)
 {
         return 0x00000068;
 }
+static inline u32 ctxsw_prog_main_image_compute_preemption_options_o(void)
+{
+        return 0x00000084;
+}
+static inline u32 ctxsw_prog_main_image_compute_preemption_options_control_f(u32 v)
+{
+        return (v & 0x3) << 0;
+}
+static inline u32 ctxsw_prog_main_image_compute_preemption_options_control_cta_f(void)
+{
+        return 0x1;
+}
+static inline u32 ctxsw_prog_main_image_compute_preemption_options_control_cilp_f(void)
+{
+        return 0x2;
+}
 #endif