2 files changed, 10 insertions, 9 deletions

diff --git a/drivers/gpu/nvgpu/gp10b/fifo_gp10b.c b/drivers/gpu/nvgpu/gp10b/fifo_gp10b.c
index 89b5527d..aa38dc54 100644
--- a/drivers/gpu/nvgpu/gp10b/fifo_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/fifo_gp10b.c
@@ -156,10 +156,10 @@ static int gp10b_fifo_resetup_ramfc(struct channel_gk20a *c)
 
         if (new_syncpt && new_syncpt != old_syncpt) {
                 /* disable channel */
-                c->g->ops.fifo.disable_channel(c);
+                gk20a_disable_channel_tsg(c->g, c);
 
                 /* preempt the channel */
-                WARN_ON(c->g->ops.fifo.preempt_channel(c->g, c->hw_chid));
+                WARN_ON(gk20a_fifo_preempt(c->g, c));
 
                 v = pbdma_allowed_syncpoints_0_valid_f(1);
 
@@ -173,9 +173,7 @@ static int gp10b_fifo_resetup_ramfc(struct channel_gk20a *c)
         }
 
         /* enable channel */
-        gk20a_writel(c->g, ccsr_channel_r(c->hw_chid),
-                gk20a_readl(c->g, ccsr_channel_r(c->hw_chid)) |
-                ccsr_channel_enable_set_true_f());
+        gk20a_enable_channel_tsg(c->g, c);
 
         gk20a_dbg_fn("done");
 
diff --git a/drivers/gpu/nvgpu/gp10b/gr_gp10b.c b/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
index 205d2c97..5db65175 100644
--- a/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
+++ b/drivers/gpu/nvgpu/gp10b/gr_gp10b.c
@@ -2061,18 +2061,21 @@ static int gr_gp10b_set_preemption_mode(struct channel_gk20a *ch,
         if (gk20a_mem_begin(g, mem))
                 return -ENOMEM;
 
-        g->ops.fifo.disable_channel(ch);
-        err = g->ops.fifo.preempt_channel(g, ch->hw_chid);
+        err = gk20a_disable_channel_tsg(g, ch);
         if (err)
                 goto unmap_ctx;
 
+        err = gk20a_fifo_preempt(g, ch);
+        if (err)
+                goto enable_ch;
+
         if (g->ops.gr.update_ctxsw_preemption_mode) {
                 g->ops.gr.update_ctxsw_preemption_mode(ch->g, ch_ctx, mem);
                 g->ops.gr.commit_global_cb_manager(g, ch, true);
         }
 
-        g->ops.fifo.enable_channel(ch);
-
+enable_ch:
+        gk20a_enable_channel_tsg(g, ch);
 unmap_ctx:
         gk20a_mem_end(g, mem);