diff --git a/drivers/gpu/nvgpu/common/gr/ctx.c b/drivers/gpu/nvgpu/common/gr/ctx.c
index c5faf15f9..c44c6f731 100644
--- a/drivers/gpu/nvgpu/common/gr/ctx.c
+++ b/drivers/gpu/nvgpu/common/gr/ctx.c
@@ -576,39 +576,6 @@ void nvgpu_gr_ctx_patch_write(struct gk20a *g,
 	}
 }
 
-void nvgpu_gr_ctx_reset_patch_count(struct gk20a *g,
-	struct nvgpu_gr_ctx *gr_ctx)
-{
-	u32 tmp;
-
-	tmp = g->ops.gr.ctxsw_prog.get_patch_count(g, &gr_ctx->mem);
-	if (tmp == 0U) {
-		gr_ctx->patch_ctx.data_count = 0;
-	}
-}
-
-void nvgpu_gr_ctx_set_patch_ctx(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
-	bool set_patch_addr)
-{
-	g->ops.gr.ctxsw_prog.set_patch_count(g, &gr_ctx->mem,
-		gr_ctx->patch_ctx.data_count);
-	if (set_patch_addr) {
-		g->ops.gr.ctxsw_prog.set_patch_addr(g, &gr_ctx->mem,
-			gr_ctx->patch_ctx.mem.gpu_va);
-	}
-}
-
-void nvgpu_gr_ctx_init_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
-	u32 graphics_preempt_mode)
-{
-	gr_ctx->graphics_preempt_mode = graphics_preempt_mode;
-}
-
-u32 nvgpu_gr_ctx_get_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx)
-{
-	return gr_ctx->graphics_preempt_mode;
-}
-
 void nvgpu_gr_ctx_init_compute_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
 	u32 compute_preempt_mode)
 {
@@ -623,9 +590,19 @@ u32 nvgpu_gr_ctx_get_compute_preemption_mode(struct nvgpu_gr_ctx *gr_ctx)
 bool nvgpu_gr_ctx_check_valid_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
 	u32 graphics_preempt_mode, u32 compute_preempt_mode)
 {
+#ifdef CONFIG_NVGPU_GRAPHICS
 	if ((graphics_preempt_mode == 0U) && (compute_preempt_mode == 0U)) {
 		return false;
 	}
+#else
+	if (graphics_preempt_mode != 0U) {
+		return false;
+	}
+
+	if (compute_preempt_mode == 0U) {
+		return false;
+	}
+#endif
 
 #ifndef CONFIG_NVGPU_CILP
 	if (compute_preempt_mode > NVGPU_PREEMPTION_MODE_COMPUTE_CTA) {
@@ -640,11 +617,13 @@ bool nvgpu_gr_ctx_check_valid_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
 	}
 #endif
 
+#ifdef CONFIG_NVGPU_GRAPHICS
 	/* Do not allow lower preemption modes than current ones */
 	if ((graphics_preempt_mode != 0U) &&
 	    (graphics_preempt_mode < gr_ctx->graphics_preempt_mode)) {
 		return false;
 	}
+#endif
 
 	if ((compute_preempt_mode != 0U) &&
 	    (compute_preempt_mode < gr_ctx->compute_preempt_mode)) {
@@ -689,6 +668,17 @@ u32 nvgpu_gr_ctx_get_tsgid(struct nvgpu_gr_ctx *gr_ctx)
 }
 
 #ifdef CONFIG_NVGPU_GRAPHICS
+void nvgpu_gr_ctx_init_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
+	u32 graphics_preempt_mode)
+{
+	gr_ctx->graphics_preempt_mode = graphics_preempt_mode;
+}
+
+u32 nvgpu_gr_ctx_get_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx)
+{
+	return gr_ctx->graphics_preempt_mode;
+}
+
 void nvgpu_gr_ctx_set_zcull_ctx(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
 	u32 mode, u64 gpu_va)
 {
@@ -926,6 +916,28 @@ void nvgpu_gr_ctx_set_cilp_preempt_pending(struct nvgpu_gr_ctx *gr_ctx,
 #endif
 
 #ifdef CONFIG_NVGPU_DEBUGGER
+void nvgpu_gr_ctx_reset_patch_count(struct gk20a *g,
+	struct nvgpu_gr_ctx *gr_ctx)
+{
+	u32 tmp;
+
+	tmp = g->ops.gr.ctxsw_prog.get_patch_count(g, &gr_ctx->mem);
+	if (tmp == 0U) {
+		gr_ctx->patch_ctx.data_count = 0;
+	}
+}
+
+void nvgpu_gr_ctx_set_patch_ctx(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
+	bool set_patch_addr)
+{
+	g->ops.gr.ctxsw_prog.set_patch_count(g, &gr_ctx->mem,
+		gr_ctx->patch_ctx.data_count);
+	if (set_patch_addr) {
+		g->ops.gr.ctxsw_prog.set_patch_addr(g, &gr_ctx->mem,
+			gr_ctx->patch_ctx.mem.gpu_va);
+	}
+}
+
 int nvgpu_gr_ctx_alloc_pm_ctx(struct gk20a *g,
 	struct nvgpu_gr_ctx *gr_ctx,
 	struct nvgpu_gr_ctx_desc *gr_ctx_desc,
diff --git a/drivers/gpu/nvgpu/common/gr/obj_ctx.c b/drivers/gpu/nvgpu/common/gr/obj_ctx.c
index 1cb2f1e81..d65d1724d 100644
--- a/drivers/gpu/nvgpu/common/gr/obj_ctx.c
+++ b/drivers/gpu/nvgpu/common/gr/obj_ctx.c
@@ -114,6 +114,7 @@ static int nvgpu_gr_obj_ctx_init_ctxsw_preemption_mode(struct gk20a *g,
 	return 0;
 }
 
+#ifdef CONFIG_NVGPU_GRAPHICS
 static int nvgpu_gr_obj_ctx_set_graphics_preemption_mode(struct gk20a *g,
 	struct nvgpu_gr_config *config, struct nvgpu_gr_ctx_desc *gr_ctx_desc,
 	struct nvgpu_gr_ctx *gr_ctx, struct vm_gk20a *vm,
@@ -123,7 +124,6 @@ static int nvgpu_gr_obj_ctx_set_graphics_preemption_mode(struct gk20a *g,
 
 	/* set preemption modes */
 	switch (graphics_preempt_mode) {
-#ifdef CONFIG_NVGPU_GRAPHICS
 	case NVGPU_PREEMPTION_MODE_GRAPHICS_GFXP:
 		{
 		u32 rtv_cb_size;
@@ -164,7 +164,6 @@ static int nvgpu_gr_obj_ctx_set_graphics_preemption_mode(struct gk20a *g,
 			graphics_preempt_mode);
 		break;
 		}
-#endif /* CONFIG_NVGPU_GRAPHICS */
 	case NVGPU_PREEMPTION_MODE_GRAPHICS_WFI:
 		nvgpu_gr_ctx_init_graphics_preemption_mode(gr_ctx,
 			graphics_preempt_mode);
@@ -176,11 +175,10 @@ static int nvgpu_gr_obj_ctx_set_graphics_preemption_mode(struct gk20a *g,
 		break;
 	}
 
-#ifdef CONFIG_NVGPU_GRAPHICS
 fail:
-#endif
 	return err;
 }
+#endif
 
 static int nvgpu_gr_obj_ctx_set_compute_preemption_mode(struct gk20a *g,
 	struct nvgpu_gr_ctx *gr_ctx, u32 class_num, u32 compute_preempt_mode)
@@ -239,12 +237,14 @@ int nvgpu_gr_obj_ctx_set_ctxsw_preemption_mode(struct gk20a *g,
 		goto fail;
 	}
 
+#ifdef CONFIG_NVGPU_GRAPHICS
 	err = nvgpu_gr_obj_ctx_set_graphics_preemption_mode(g, config,
 				gr_ctx_desc, gr_ctx, vm, graphics_preempt_mode);
 
 	if (err != 0) {
 		goto fail;
 	}
+#endif
 
 	err = nvgpu_gr_obj_ctx_set_compute_preemption_mode(g, gr_ctx,
 					class_num, compute_preempt_mode);
diff --git a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
index 08fcf07b6..2d50f1eb8 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/gr/ctx.h
@@ -152,16 +152,6 @@ void nvgpu_gr_ctx_patch_write(struct gk20a *g,
 	struct nvgpu_gr_ctx *gr_ctx,
 	u32 addr, u32 data, bool patch);
 
-void nvgpu_gr_ctx_reset_patch_count(struct gk20a *g,
-	struct nvgpu_gr_ctx *gr_ctx);
-void nvgpu_gr_ctx_set_patch_ctx(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
-	bool set_patch_addr);
-
-void nvgpu_gr_ctx_init_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
-	u32 graphics_preempt_mode);
-
-u32 nvgpu_gr_ctx_get_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx);
-
 void nvgpu_gr_ctx_init_compute_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
 	u32 compute_preempt_mode);
 u32 nvgpu_gr_ctx_get_compute_preemption_mode(struct nvgpu_gr_ctx *gr_ctx);
@@ -213,6 +203,11 @@ struct nvgpu_mem *nvgpu_gr_ctx_get_gfxp_rtvcb_ctxsw_buffer(
 void nvgpu_gr_ctx_set_preemption_buffer_va(struct gk20a *g,
 	struct nvgpu_gr_ctx *gr_ctx);
 
+void nvgpu_gr_ctx_init_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx,
+	u32 graphics_preempt_mode);
+
+u32 nvgpu_gr_ctx_get_graphics_preemption_mode(struct nvgpu_gr_ctx *gr_ctx);
+
 bool nvgpu_gr_ctx_desc_force_preemption_gfxp(
 		struct nvgpu_gr_ctx_desc *gr_ctx_desc);
 
@@ -236,6 +231,11 @@ int nvgpu_gr_ctx_alloc_pm_ctx(struct gk20a *g,
 void nvgpu_gr_ctx_free_pm_ctx(struct gk20a *g, struct vm_gk20a *vm,
 	struct nvgpu_gr_ctx *gr_ctx);
 
+void nvgpu_gr_ctx_reset_patch_count(struct gk20a *g,
+	struct nvgpu_gr_ctx *gr_ctx);
+void nvgpu_gr_ctx_set_patch_ctx(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx,
+	bool set_patch_addr);
+
 u32 nvgpu_gr_ctx_get_ctx_id(struct gk20a *g, struct nvgpu_gr_ctx *gr_ctx);
 u32 nvgpu_gr_ctx_read_ctx_id(struct nvgpu_gr_ctx *gr_ctx);