7 files changed, 29 insertions, 20 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) {
diff --git a/drivers/gpu/nvgpu/include/nvgpu/linux/vidmem.h b/drivers/gpu/nvgpu/include/nvgpu/linux/vidmem.h
index 76bbb05b..ec02faec 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/linux/vidmem.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/linux/vidmem.h
@@ -24,13 +24,18 @@ struct dma_buf;
 struct gk20a;
 
 #ifdef CONFIG_GK20A_VIDMEM
+
 struct gk20a *nvgpu_vidmem_buf_owner(struct dma_buf *dmabuf);
 int nvgpu_vidmem_export_linux(struct gk20a *g, size_t bytes);
 
+void nvgpu_vidmem_set_page_alloc(struct scatterlist *sgl, u64 addr);
+struct nvgpu_page_alloc *nvgpu_vidmem_get_page_alloc(struct scatterlist *sgl);
+
 int nvgpu_vidmem_buf_access_memory(struct gk20a *g, struct dma_buf *dmabuf,
                 void *buffer, u64 offset, u64 size, u32 cmd);
 
 #else /* !CONFIG_GK20A_VIDMEM */
+
 static inline struct gk20a *nvgpu_vidmem_buf_owner(struct dma_buf *dmabuf)
 {
         return NULL;
@@ -41,16 +46,28 @@ static inline int nvgpu_vidmem_export_linux(struct gk20a *g, size_t bytes)
         return -ENOSYS;
 }
 
+static inline void nvgpu_vidmem_set_page_alloc(struct scatterlist *sgl,
+                                               u64 addr)
+{
+}
+
+static inline struct nvgpu_page_alloc *nvgpu_vidmem_get_page_alloc(
+        struct scatterlist *sgl)
+{
+        return NULL;
+}
+
 static inline int nvgpu_vidmem_buf_access_memory(struct gk20a *g,
-                                             struct dma_buf *dmabuf,
-                                             void *buffer, u64 offset,
-                                             u64 size, u32 cmd)
+                                                 struct dma_buf *dmabuf,
+                                                 void *buffer, u64 offset,
+                                                 u64 size, u32 cmd)
 {
         return -ENOSYS;
 }
 
 #endif
 
+
 struct nvgpu_vidmem_linux {
         struct dma_buf  *dmabuf;
         void            *dmabuf_priv;
diff --git a/drivers/gpu/nvgpu/include/nvgpu/nvgpu_mem.h b/drivers/gpu/nvgpu/include/nvgpu/nvgpu_mem.h
index 23a1bad7..537409a8 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/nvgpu_mem.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/nvgpu_mem.h
@@ -39,6 +39,7 @@ struct nvgpu_sgt;
 struct gk20a;
 struct nvgpu_allocator;
 struct nvgpu_gmmu_attrs;
+struct nvgpu_page_alloc;
 
 #define NVGPU_MEM_DMA_ERROR             (~0ULL)
 
@@ -162,6 +163,7 @@ struct nvgpu_mem {
         /*
          * Fields only populated for vidmem allocations.
          */
+        struct nvgpu_page_alloc                 *vidmem_alloc;
         struct nvgpu_allocator                  *allocator;
         struct nvgpu_list_node                   clear_list_entry;
 
diff --git a/drivers/gpu/nvgpu/include/nvgpu/vidmem.h b/drivers/gpu/nvgpu/include/nvgpu/vidmem.h
index b89c710d..9e9f8301 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/vidmem.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/vidmem.h
@@ -27,7 +27,6 @@
 #include <nvgpu/errno.h>
 #include <nvgpu/nvgpu_mem.h>
 
-struct scatterlist;
 struct work_struct;
 
 struct gk20a;
@@ -74,8 +73,6 @@ struct nvgpu_vidmem_buf *nvgpu_vidmem_user_alloc(struct gk20a *g, size_t bytes);
 
 void nvgpu_vidmem_buf_free(struct gk20a *g, struct nvgpu_vidmem_buf *buf);
 
-struct nvgpu_page_alloc *nvgpu_vidmem_get_page_alloc(struct scatterlist *sgl);
-void nvgpu_vidmem_set_page_alloc(struct scatterlist *sgl, u64 addr);
 bool nvgpu_addr_is_vidmem_page_alloc(u64 addr);
 int nvgpu_vidmem_get_space(struct gk20a *g, u64 *space);
 
@@ -93,16 +90,6 @@ int nvgpu_vidmem_clear(struct gk20a *g, struct nvgpu_mem *mem);
  * When VIDMEM support is not present this interface is used.
  */
 
-static inline struct nvgpu_page_alloc *
-nvgpu_vidmem_get_page_alloc(struct scatterlist *sgl)
-{
-        return NULL;
-}
-
-static inline void nvgpu_vidmem_set_page_alloc(struct scatterlist *sgl, u64 addr)
-{
-}
-
 static inline bool nvgpu_addr_is_vidmem_page_alloc(u64 addr)
 {
         return false;