2 files changed, 8 insertions, 5 deletions

diff --git a/drivers/gpu/nvgpu/common/linux/driver_common.c b/drivers/gpu/nvgpu/common/linux/driver_common.c
index 0d29223f..80e7698b 100644
--- a/drivers/gpu/nvgpu/common/linux/driver_common.c
+++ b/drivers/gpu/nvgpu/common/linux/driver_common.c
@@ -20,6 +20,7 @@
 #include <nvgpu/nvgpu_common.h>
 #include <nvgpu/soc.h>
 #include <nvgpu/bug.h>
+#include <nvgpu/enabled.h>
 
 #include "gk20a/gk20a_scale.h"
 #include "gk20a/gk20a.h"
@@ -133,6 +134,9 @@ static void nvgpu_init_mm_vars(struct gk20a *g)
         g->mm.disable_bigpage = platform->disable_bigpage;
         g->mm.vidmem_is_vidmem = platform->vidmem_is_vidmem;
 
+        __nvgpu_set_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES,
+                            platform->unify_address_spaces);
+
         nvgpu_mutex_init(&g->mm.tlb_lock);
         nvgpu_mutex_init(&g->mm.priv_lock);
 }
diff --git a/drivers/gpu/nvgpu/common/mm/vm.c b/drivers/gpu/nvgpu/common/mm/vm.c
index bdc8554c..904284ec 100644
--- a/drivers/gpu/nvgpu/common/mm/vm.c
+++ b/drivers/gpu/nvgpu/common/mm/vm.c
@@ -22,10 +22,10 @@
 #include <nvgpu/list.h>
 #include <nvgpu/rbtree.h>
 #include <nvgpu/semaphore.h>
+#include <nvgpu/enabled.h>
 
 #include "gk20a/gk20a.h"
 #include "gk20a/mm_gk20a.h"
-#include "gk20a/platform_gk20a.h"
 
 int vm_aspace_id(struct vm_gk20a *vm)
 {
@@ -255,7 +255,6 @@ int nvgpu_init_vm(struct mm_gk20a *mm,
         u64 user_lp_vma_start, user_lp_vma_limit;
         u64 kernel_vma_start, kernel_vma_limit;
         struct gk20a *g = mm->g;
-        struct gk20a_platform *p = gk20a_get_platform(g->dev);
 
         if (WARN_ON(kernel_reserved + low_hole > aperture_size))
                 return -ENOMEM;
@@ -275,7 +274,7 @@ int nvgpu_init_vm(struct mm_gk20a *mm,
         vm->vma[gmmu_page_size_small]  = &vm->user;
         vm->vma[gmmu_page_size_big]    = &vm->user;
         vm->vma[gmmu_page_size_kernel] = &vm->kernel;
-        if (!p->unify_address_spaces)
+        if (!nvgpu_is_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES))
                 vm->vma[gmmu_page_size_big] = &vm->user_lp;
 
         vm->va_start  = low_hole;
@@ -293,7 +292,7 @@ int nvgpu_init_vm(struct mm_gk20a *mm,
 
         /* Setup vma limits. */
         if (kernel_reserved + low_hole < aperture_size) {
-                if (p->unify_address_spaces) {
+                if (nvgpu_is_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES)) {
                         user_vma_start = low_hole;
                         user_vma_limit = vm->va_limit - kernel_reserved;
                         user_lp_vma_start = user_vma_limit;
@@ -346,7 +345,7 @@ int nvgpu_init_vm(struct mm_gk20a *mm,
          * Determine if big pages are possible in this VM. If a split address
          * space is used then check the user_lp vma instead of the user vma.
          */
-        if (p->unify_address_spaces)
+        if (nvgpu_is_enabled(g, NVGPU_MM_UNIFY_ADDRESS_SPACES))
                 vm->big_pages = nvgpu_big_pages_possible(vm, user_vma_start,
                                          user_vma_limit - user_vma_start);
         else