5 files changed, 37 insertions, 3 deletions

diff --git a/drivers/gpu/nvgpu/common/fb/fb_gp106.c b/drivers/gpu/nvgpu/common/fb/fb_gp106.c
index c8a32cd2..6dacceb5 100644
--- a/drivers/gpu/nvgpu/common/fb/fb_gp106.c
+++ b/drivers/gpu/nvgpu/common/fb/fb_gp106.c
@@ -52,3 +52,18 @@ void gp106_fb_reset(struct gk20a *g)
         val &= ~fb_mmu_priv_level_mask_write_violation_m();
         gk20a_writel(g, fb_mmu_priv_level_mask_r(), val);
 }
+
+size_t gp106_fb_get_vidmem_size(struct gk20a *g)
+{
+        u32 range = gk20a_readl(g, fb_mmu_local_memory_range_r());
+        u32 mag = fb_mmu_local_memory_range_lower_mag_v(range);
+        u32 scale = fb_mmu_local_memory_range_lower_scale_v(range);
+        u32 ecc = fb_mmu_local_memory_range_ecc_mode_v(range);
+        size_t bytes = ((size_t)mag << scale) * SZ_1M;
+
+        if (ecc) {
+                bytes = bytes / 16U * 15U;
+        }
+
+        return bytes;
+}
diff --git a/drivers/gpu/nvgpu/common/fb/fb_gp106.h b/drivers/gpu/nvgpu/common/fb/fb_gp106.h
index d5ee87f4..bb4ccd6b 100644
--- a/drivers/gpu/nvgpu/common/fb/fb_gp106.h
+++ b/drivers/gpu/nvgpu/common/fb/fb_gp106.h
@@ -25,4 +25,6 @@
 struct gpu_ops;
 
 void gp106_fb_reset(struct gk20a *g);
+size_t gp106_fb_get_vidmem_size(struct gk20a *g);
+
 #endif
diff --git a/drivers/gpu/nvgpu/common/fb/fb_gv100.c b/drivers/gpu/nvgpu/common/fb/fb_gv100.c
index 848d6efb..4cae5312 100644
--- a/drivers/gpu/nvgpu/common/fb/fb_gv100.c
+++ b/drivers/gpu/nvgpu/common/fb/fb_gv100.c
@@ -281,3 +281,18 @@ int gv100_fb_enable_nvlink(struct gk20a *g)
 
         return 0;
 }
+
+size_t gv100_fb_get_vidmem_size(struct gk20a *g)
+{
+        u32 range = gk20a_readl(g, fb_mmu_local_memory_range_r());
+        u32 mag = fb_mmu_local_memory_range_lower_mag_v(range);
+        u32 scale = fb_mmu_local_memory_range_lower_scale_v(range);
+        u32 ecc = fb_mmu_local_memory_range_ecc_mode_v(range);
+        size_t bytes = ((size_t)mag << scale) * SZ_1M;
+
+        if (ecc) {
+                bytes = bytes / 16U * 15U;
+        }
+
+        return bytes;
+}
diff --git a/drivers/gpu/nvgpu/common/fb/fb_gv100.h b/drivers/gpu/nvgpu/common/fb/fb_gv100.h
index 195baccf..d47fded2 100644
--- a/drivers/gpu/nvgpu/common/fb/fb_gv100.h
+++ b/drivers/gpu/nvgpu/common/fb/fb_gv100.h
@@ -33,4 +33,6 @@ void gv100_fb_disable_hub_intr(struct gk20a *g);
 int gv100_fb_memory_unlock(struct gk20a *g);
 int gv100_fb_init_nvlink(struct gk20a *g);
 int gv100_fb_enable_nvlink(struct gk20a *g);
+size_t gv100_fb_get_vidmem_size(struct gk20a *g);
+
 #endif
diff --git a/drivers/gpu/nvgpu/common/mm/vidmem.c b/drivers/gpu/nvgpu/common/mm/vidmem.c
index 1ace2333..3f5e0fbb 100644
--- a/drivers/gpu/nvgpu/common/mm/vidmem.c
+++ b/drivers/gpu/nvgpu/common/mm/vidmem.c
@@ -41,7 +41,7 @@ void nvgpu_vidmem_destroy(struct gk20a *g)
 {
         struct nvgpu_timeout timeout;
 
-        if (!g->ops.mm.get_vidmem_size)
+        if (!g->ops.fb.get_vidmem_size)
                 return;
 
         nvgpu_timeout_init(g, &timeout, 100, NVGPU_TIMER_RETRY_TIMER);
@@ -293,8 +293,8 @@ int nvgpu_vidmem_init(struct mm_gk20a *mm)
         static struct nvgpu_alloc_carveout bootstrap_co =
                 NVGPU_CARVEOUT("bootstrap-region", 0, 0);
 
-        size = g->ops.mm.get_vidmem_size ?
-                        g->ops.mm.get_vidmem_size(g) : 0;
+        size = g->ops.fb.get_vidmem_size ?
+                        g->ops.fb.get_vidmem_size(g) : 0;
         if (!size)
                 return 0;