diff --git a/drivers/gpu/nvgpu/common/gr/ctx.c b/drivers/gpu/nvgpu/common/gr/ctx.c
index 565e26a5e..233a8d493 100644
--- a/drivers/gpu/nvgpu/common/gr/ctx.c
+++ b/drivers/gpu/nvgpu/common/gr/ctx.c
@@ -650,3 +650,26 @@ int nvgpu_gr_ctx_zcull_setup(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
 
 	return 0;
 }
+
+int nvgpu_gr_ctx_set_smpc_mode(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
+	bool enable)
+{
+	int err;
+
+	if (!nvgpu_mem_is_valid(&gr_ctx->mem)) {
+		nvgpu_err(g, "no graphics context allocated");
+		return -EFAULT;
+	}
+
+	/* Channel gr_ctx buffer is gpu cacheable.
+	   Flush and invalidate before cpu update. */
+	err = g->ops.mm.l2_flush(g, true);
+	if (err != 0) {
+		nvgpu_err(g, "l2_flush failed");
+		return err;
+	}
+
+	g->ops.gr.ctxsw_prog.set_pm_smpc_mode(g, &gr_ctx->mem, enable);
+
+	return err;
+}
diff --git a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
index 5dcf2b29b..184b1bc25 100644
--- a/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/gr_gk20a.c
@@ -1391,8 +1391,6 @@ int gr_gk20a_update_smpc_ctxsw_mode(struct gk20a *g,
 				    bool enable_smpc_ctxsw)
 {
 	struct tsg_gk20a *tsg;
-	struct nvgpu_gr_ctx *gr_ctx = NULL;
-	struct nvgpu_mem *mem = NULL;
 	int ret;
 
 	nvgpu_log_fn(g, " ");
@@ -1402,13 +1400,6 @@ int gr_gk20a_update_smpc_ctxsw_mode(struct gk20a *g,
 		return -EINVAL;
 	}
 
-	gr_ctx = tsg->gr_ctx;
-	mem = &gr_ctx->mem;
-	if (!nvgpu_mem_is_valid(mem)) {
-		nvgpu_err(g, "no graphics context allocated");
-		return -EFAULT;
-	}
-
 	ret = gk20a_disable_channel_tsg(g, c);
 	if (ret != 0) {
 		nvgpu_err(g, "failed to disable channel/TSG");
@@ -1421,15 +1412,7 @@ int gr_gk20a_update_smpc_ctxsw_mode(struct gk20a *g,
 		goto out;
 	}
 
-	/* Channel gr_ctx buffer is gpu cacheable.
-	   Flush and invalidate before cpu update. */
-	ret = g->ops.mm.l2_flush(g, true);
-	if (ret != 0) {
-		nvgpu_err(g, "l2_flush failed");
-		goto out;
-	}
-
-	g->ops.gr.ctxsw_prog.set_pm_smpc_mode(g, mem, enable_smpc_ctxsw);
+	ret = nvgpu_gr_ctx_set_smpc_mode(g, tsg->gr_ctx, enable_smpc_ctxsw);
 
 out:
 	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 63ff91915..49c95f246 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
@@ -182,4 +182,7 @@ u32 nvgpu_gr_ctx_get_ctx_id(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx);
 int nvgpu_gr_ctx_init_zcull(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);
+
+int nvgpu_gr_ctx_set_smpc_mode(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
+	bool enable);
 #endif /* NVGPU_INCLUDE_GR_CTX_H */