4 files changed, 7 insertions, 4 deletions

diff --git a/drivers/gpu/nvgpu/common/linux/dma.c b/drivers/gpu/nvgpu/common/linux/dma.c
index ca657da2..b62c4593 100644
--- a/drivers/gpu/nvgpu/common/linux/dma.c
+++ b/drivers/gpu/nvgpu/common/linux/dma.c
@@ -28,9 +28,11 @@
 #include <nvgpu/vidmem.h>
 
 #include <nvgpu/linux/dma.h>
+#include <nvgpu/linux/vidmem.h>
 
 #include "gk20a/gk20a.h"
 #include "gk20a/platform_gk20a.h"
+
 #include "os_linux.h"
 
 /*
@@ -348,6 +350,7 @@ int nvgpu_dma_alloc_flags_vid_at(struct gk20a *g, unsigned long flags,
 
         mem->aligned_size = size;
         mem->aperture = APERTURE_VIDMEM;
+        mem->vidmem_alloc = (struct nvgpu_page_alloc *)(uintptr_t)addr;
         mem->allocator = vidmem_alloc;
         mem->priv.flags = flags;
 
diff --git a/drivers/gpu/nvgpu/common/linux/nvgpu_mem.c b/drivers/gpu/nvgpu/common/linux/nvgpu_mem.c
index e1f8a598..8740ac3d 100644
--- a/drivers/gpu/nvgpu/common/linux/nvgpu_mem.c
+++ b/drivers/gpu/nvgpu/common/linux/nvgpu_mem.c
@@ -25,6 +25,7 @@
 #include <nvgpu/vidmem.h>
 
 #include <nvgpu/linux/dma.h>
+#include <nvgpu/linux/vidmem.h>
 
 #include "os_linux.h"
 
@@ -355,7 +356,7 @@ u64 nvgpu_mem_get_addr(struct gk20a *g, struct nvgpu_mem *mem)
         /*
          * Otherwise get the vidmem address.
          */
-        alloc = nvgpu_vidmem_get_page_alloc(mem->priv.sgt->sgl);
+        alloc = mem->vidmem_alloc;
 
         /* This API should not be used with > 1 chunks */
         WARN_ON(alloc->nr_chunks != 1);
diff --git a/drivers/gpu/nvgpu/common/mm/vidmem.c b/drivers/gpu/nvgpu/common/mm/vidmem.c
index 18f46c03..d1c5a2e8 100644
--- a/drivers/gpu/nvgpu/common/mm/vidmem.c
+++ b/drivers/gpu/nvgpu/common/mm/vidmem.c
@@ -195,7 +195,7 @@ int nvgpu_vidmem_clear(struct gk20a *g, struct nvgpu_mem *mem)
         if (g->mm.vidmem.ce_ctx_id == (u32)~0)
                 return -EINVAL;
 
-        alloc = nvgpu_vidmem_get_page_alloc(mem->priv.sgt->sgl);
+        alloc = mem->vidmem_alloc;
 
         nvgpu_sgt_for_each_sgl(sgl, &alloc->sgt) {
                 if (gk20a_last_fence)
diff --git a/drivers/gpu/nvgpu/common/pramin.c b/drivers/gpu/nvgpu/common/pramin.c
index b7bc7439..b6166f51 100644
--- a/drivers/gpu/nvgpu/common/pramin.c
+++ b/drivers/gpu/nvgpu/common/pramin.c
@@ -23,7 +23,6 @@
 #include <nvgpu/pramin.h>
 #include <nvgpu/page_allocator.h>
 #include <nvgpu/enabled.h>
-#include <nvgpu/vidmem.h>
 
 #include "gk20a/gk20a.h"
 
@@ -55,7 +54,7 @@ void nvgpu_pramin_access_batched(struct gk20a *g, struct nvgpu_mem *mem,
         if (!g->regs && nvgpu_is_enabled(g, NVGPU_DRIVER_IS_DYING))
                 return;
 
-        alloc = nvgpu_vidmem_get_page_alloc(mem->priv.sgt->sgl);
+        alloc = mem->vidmem_alloc;
         sgt = &alloc->sgt;
 
         nvgpu_sgt_for_each_sgl(sgl, sgt) {