1 files changed, 6 insertions, 0 deletions

diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index 94f66e83..54c96680 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -670,6 +670,12 @@ void gr_gk20a_ctx_patch_write_end(struct gk20a *g,
                                         struct channel_ctx_gk20a *ch_ctx)
 {
         gk20a_mem_end(g, &ch_ctx->patch_ctx.mem);
+        /* Write context count to context image if it is mapped */
+        if (ch_ctx->gr_ctx->mem.cpu_va) {
+                gk20a_mem_wr(g, &ch_ctx->gr_ctx->mem,
+                             ctxsw_prog_main_image_patch_count_o(),
+                             ch_ctx->patch_ctx.data_count);
+        }
 }
 
 void gr_gk20a_ctx_patch_write(struct gk20a *g,