diff --git a/drivers/gpu/nvgpu/common/gr/ctx.c b/drivers/gpu/nvgpu/common/gr/ctx.c
index 7ab13625a..2c48e849a 100644
--- a/drivers/gpu/nvgpu/common/gr/ctx.c
+++ b/drivers/gpu/nvgpu/common/gr/ctx.c
@@ -612,3 +612,25 @@ u32 nvgpu_gr_ctx_get_ctx_id(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx)
 
 	return gr_ctx->ctx_id;
 }
+
+int nvgpu_gr_ctx_zcull_setup(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
+	bool set_zcull_ptr)
+{
+	nvgpu_log_fn(g, " ");
+
+	if (gr_ctx->zcull_ctx.gpu_va == 0ULL &&
+	    g->ops.gr.ctxsw_prog.is_zcull_mode_separate_buffer(
+			gr_ctx->zcull_ctx.ctx_sw_mode)) {
+		return -EINVAL;
+	}
+
+	g->ops.gr.ctxsw_prog.set_zcull(g, &gr_ctx->mem,
+		gr_ctx->zcull_ctx.ctx_sw_mode);
+
+	if (set_zcull_ptr) {
+		g->ops.gr.ctxsw_prog.set_zcull_ptr(g, &gr_ctx->mem,
+			gr_ctx->zcull_ctx.gpu_va);
+	}
+
+	return 0;
+}
diff --git a/drivers/gpu/nvgpu/common/gr/subctx.c b/drivers/gpu/nvgpu/common/gr/subctx.c
index 02bb42294..f9c69bdaa 100644
--- a/drivers/gpu/nvgpu/common/gr/subctx.c
+++ b/drivers/gpu/nvgpu/common/gr/subctx.c
@@ -107,3 +107,13 @@ void nvgpu_gr_subctx_load_ctx_header(struct gk20a *g,
 
 	g->ops.gr.ctxsw_prog.set_type_per_veid_header(g, ctxheader);
 }
+
+void nvgpu_gr_subctx_zcull_setup(struct gk20a *g, struct nvgpu_gr_subctx *subctx,
+		struct nvgpu_gr_ctx *gr_ctx)
+{
+
+	nvgpu_log_fn(g, " ");
+
+	g->ops.gr.ctxsw_prog.set_zcull_ptr(g, &subctx->ctx_header,
+		gr_ctx->zcull_ctx.gpu_va);
+}
diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index 81756a60d..5035d5b50 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -634,19 +634,10 @@ int gr_gk20a_fecs_ctx_bind_channel(struct gk20a *g,
 static int gr_gk20a_ctx_zcull_setup(struct gk20a *g, struct channel_gk20a *c,
 		struct nvgpu_gr_ctx *gr_ctx)
 {
-	struct nvgpu_mem *mem = NULL;
 	int ret = 0;
 
 	nvgpu_log_fn(g, " ");
 
-	mem = &gr_ctx->mem;
-
-	if (gr_ctx->zcull_ctx.gpu_va == 0ULL &&
-	    g->ops.gr.ctxsw_prog.is_zcull_mode_separate_buffer(
-			gr_ctx->zcull_ctx.ctx_sw_mode)) {
-		return -EINVAL;
-	}
-
 	ret = gk20a_disable_channel_tsg(g, c);
 	if (ret != 0) {
 		nvgpu_err(g, "failed to disable channel/TSG");
@@ -659,14 +650,13 @@ static int gr_gk20a_ctx_zcull_setup(struct gk20a *g, struct channel_gk20a *c,
 		return ret;
 	}
 
-	g->ops.gr.ctxsw_prog.set_zcull(g, mem, gr_ctx->zcull_ctx.ctx_sw_mode);
-
 	if (c->subctx != NULL) {
-		g->ops.gr.ctxsw_prog.set_zcull_ptr(g, &c->subctx->ctx_header,
-			gr_ctx->zcull_ctx.gpu_va);
+		ret = nvgpu_gr_ctx_zcull_setup(g, gr_ctx, false);
+		if (ret == 0) {
+			nvgpu_gr_subctx_zcull_setup(g, c->subctx, gr_ctx);
+		}
 	} else {
-		g->ops.gr.ctxsw_prog.set_zcull_ptr(g, mem,
-			gr_ctx->zcull_ctx.gpu_va);
+		ret = nvgpu_gr_ctx_zcull_setup(g, gr_ctx, true);
 	}
 
 	gk20a_enable_channel_tsg(g, c);
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
index 265181ca9..8b82c6a7a 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
@@ -179,4 +179,6 @@ void nvgpu_gr_ctx_patch_write(struct gk20a *g,
 
 u32 nvgpu_gr_ctx_get_ctx_id(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx);
 
+int nvgpu_gr_ctx_zcull_setup(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
+	bool set_zcull_ptr);
 #endif /* NVGPU_INCLUDE_GR_CTX_H */
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gr/subctx.h b/drivers/gpu/nvgpu/include/nvgpu/gr/subctx.h
index 0e5017dbc..8400ddffd 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gr/subctx.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gr/subctx.h
@@ -43,4 +43,7 @@ void nvgpu_gr_subctx_load_ctx_header(struct gk20a *g,
 	struct nvgpu_gr_subctx *subctx,
 	struct nvgpu_gr_ctx *gr_ctx, u64 gpu_va);
 
+void nvgpu_gr_subctx_zcull_setup(struct gk20a *g, struct nvgpu_gr_subctx *subctx,
+		struct nvgpu_gr_ctx *gr_ctx);
+
 #endif /* NVGPU_GR_SUBCTX_H */