2 files changed, 6 insertions, 4 deletions

diff --git a/drivers/gpu/nvgpu/vgpu/gr_vgpu.c b/drivers/gpu/nvgpu/vgpu/gr_vgpu.c
index c7960c1f..dd515f41 100644
--- a/drivers/gpu/nvgpu/vgpu/gr_vgpu.c
+++ b/drivers/gpu/nvgpu/vgpu/gr_vgpu.c
@@ -526,7 +526,7 @@ static int vgpu_gr_alloc_obj_ctx(struct channel_gk20a  *c,
         } else {
                 if (!tsg->tsg_gr_ctx) {
                         tsg->vm = c->vm;
-                        gk20a_vm_get(tsg->vm);
+                        nvgpu_vm_get(tsg->vm);
                         err = g->ops.gr.alloc_gr_ctx(g, &tsg->tsg_gr_ctx,
                                                 c->vm,
                                                 args->class_num,
@@ -536,7 +536,7 @@ static int vgpu_gr_alloc_obj_ctx(struct channel_gk20a  *c,
                         if (err) {
                                 nvgpu_err(g,
                                         "fail to allocate TSG gr ctx buffer, err=%d", err);
-                                gk20a_vm_put(tsg->vm);
+                                nvgpu_vm_put(tsg->vm);
                                 tsg->vm = NULL;
                                 goto out;
                         }
diff --git a/drivers/gpu/nvgpu/vgpu/mm_vgpu.c b/drivers/gpu/nvgpu/vgpu/mm_vgpu.c
index 7e42c198..cfa9e428 100644
--- a/drivers/gpu/nvgpu/vgpu/mm_vgpu.c
+++ b/drivers/gpu/nvgpu/vgpu/mm_vgpu.c
@@ -22,6 +22,8 @@
 #include "vgpu/vgpu.h"
 #include "gk20a/mm_gk20a.h"
 
+#include "common/linux/vm_priv.h"
+
 static int vgpu_init_mm_setup_sw(struct gk20a *g)
 {
         struct mm_gk20a *mm = &g->mm;
@@ -216,7 +218,7 @@ static void vgpu_vm_remove_support(struct vm_gk20a *vm)
         nvgpu_rbtree_enum_start(0, &node, vm->mapped_buffers);
         while (node) {
                 mapped_buffer = mapped_buffer_from_rbtree_node(node);
-                gk20a_vm_unmap_locked(mapped_buffer, NULL);
+                nvgpu_vm_unmap_locked(mapped_buffer, NULL);
                 nvgpu_rbtree_enum_start(0, &node, vm->mapped_buffers);
         }
 
@@ -454,7 +456,7 @@ static int vgpu_vm_bind_channel(struct gk20a_as_share *as_share,
         }
 
         if (ch->vm)
-                gk20a_vm_get(ch->vm);
+                nvgpu_vm_get(ch->vm);
 
         return err;
 }