2 files changed, 5 insertions, 1 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);