3 files changed, 27 insertions, 8 deletions

diff --git a/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c b/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
index 4929f4d1..7f2f5a65 100644
--- a/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/fifo_gv11b.c
@@ -185,7 +185,6 @@ static u32 gv11b_userd_gp_get(struct gk20a *g, struct channel_gk20a *c)
 
         return gk20a_mem_rd32(g, userd_mem,
                         offset + ram_userd_gp_get_w());
-
 }
 
 static void gv11b_userd_gp_put(struct gk20a *g, struct channel_gk20a *c)
@@ -199,17 +198,13 @@ static void gv11b_userd_gp_put(struct gk20a *g, struct channel_gk20a *c)
         smp_mb();
 
         gv11b_ring_channel_doorbell(c);
-
 }
 
 static void channel_gv11b_unbind(struct channel_gk20a *ch)
 {
         gk20a_dbg_fn("");
 
-        gv11b_free_subctx_header(ch);
-
         channel_gk20a_unbind(ch);
-
 }
 
 static u32 gv11b_fifo_get_num_fifos(struct gk20a *g)
@@ -231,4 +226,5 @@ void gv11b_init_fifo(struct gpu_ops *gops)
         gops->fifo.setup_ramfc = channel_gv11b_setup_ramfc;
         gops->fifo.unbind_channel = channel_gv11b_unbind;
         gops->fifo.eng_runlist_base_size = fifo_eng_runlist_base__size_1_v;
+        gops->fifo.free_channel_ctx_header = gv11b_free_subctx_header;
 }
diff --git a/drivers/gpu/nvgpu/gv11b/gr_gv11b.c b/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
index 10b1aebb..dfb46701 100644
--- a/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/gr_gv11b.c
@@ -1853,8 +1853,6 @@ static int gr_gv11b_commit_inst(struct channel_gk20a *c, u64 gpu_va)
         /* point this address to engine_wfi_ptr */
         gk20a_mem_wr32(c->g, &c->inst_block, ram_in_engine_wfi_target_w(),
                 ram_in_engine_cs_wfi_v() |
-                ram_in_engine_wfi_target_f(
-                        ram_in_engine_wfi_target_sys_mem_ncoh_v()) |
                 ram_in_engine_wfi_mode_f(ram_in_engine_wfi_mode_virtual_v()) |
                 ram_in_engine_wfi_ptr_lo_f(addr_lo));
 
@@ -1919,6 +1917,29 @@ static int gr_gv11b_commit_global_timeslice(struct gk20a *g,
         return 0;
 }
 
+static void gv11b_restore_context_header(struct gk20a *g,
+                                        struct mem_desc *ctxheader)
+{
+        u32 va_lo, va_hi;
+        struct gr_gk20a *gr = &g->gr;
+
+        va_hi = gk20a_mem_rd(g, ctxheader,
+                        ctxsw_prog_main_image_context_buffer_ptr_hi_o());
+        va_lo = gk20a_mem_rd(g, ctxheader,
+                        ctxsw_prog_main_image_context_buffer_ptr_o());
+        gk20a_mem_wr_n(g, ctxheader, 0,
+                       gr->ctx_vars.local_golden_image,
+                       gr->ctx_vars.golden_image_size);
+        gk20a_mem_wr(g, ctxheader,
+                        ctxsw_prog_main_image_context_buffer_ptr_hi_o(), va_hi);
+        gk20a_mem_wr(g, ctxheader,
+                        ctxsw_prog_main_image_context_buffer_ptr_o(), va_lo);
+        gk20a_mem_wr(g, ctxheader,
+                        ctxsw_prog_main_image_num_restore_ops_o(), 0);
+        gk20a_mem_wr(g, ctxheader,
+                        ctxsw_prog_main_image_num_save_ops_o(), 0);
+}
+
 void gv11b_init_gr(struct gpu_ops *gops)
 {
         gp10b_init_gr(gops);
@@ -1971,5 +1992,6 @@ void gv11b_init_gr(struct gpu_ops *gops)
         gops->gr.program_sm_id_numbering =
                         gr_gv11b_program_sm_id_numbering;
         gops->gr.commit_inst = gr_gv11b_commit_inst;
+        gops->gr.restore_context_header = gv11b_restore_context_header;
 
 }
diff --git a/drivers/gpu/nvgpu/gv11b/subctx_gv11b.c b/drivers/gpu/nvgpu/gv11b/subctx_gv11b.c
index 3acc53f6..b0d0a192 100644
--- a/drivers/gpu/nvgpu/gv11b/subctx_gv11b.c
+++ b/drivers/gpu/nvgpu/gv11b/subctx_gv11b.c
@@ -44,6 +44,7 @@ int gv11b_alloc_subctx_header(struct channel_gk20a *c)
 {
         struct ctx_header_desc *ctx = &c->ch_ctx.ctx_header;
         struct gk20a *g = c->g;
+        struct gr_gk20a *gr = &g->gr;
         int ret = 0;
 
         gk20a_dbg_fn("");
@@ -51,7 +52,7 @@ int gv11b_alloc_subctx_header(struct channel_gk20a *c)
         if (ctx->mem.gpu_va == 0) {
                 ret = gk20a_gmmu_alloc_attr_sys(g,
                                 DMA_ATTR_NO_KERNEL_MAPPING,
-                                ctxsw_prog_fecs_header_v(),
+                                gr->ctx_vars.golden_image_size,
                                 &ctx->mem);
                 if (ret) {
                         gk20a_err(dev_from_gk20a(g),