diff --git a/drivers/gpu/nvgpu/gk20a/gk20a.h b/drivers/gpu/nvgpu/gk20a/gk20a.h
index 1157390af..ba7c7e975 100644
--- a/drivers/gpu/nvgpu/gk20a/gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/gk20a.h
@@ -334,6 +334,7 @@ struct gpu_ops {
 					  void *inst_ptr, int size);
 		u32 (*get_big_page_sizes)(void);
 		u32 (*get_physical_addr_bits)(struct gk20a *g);
+		int (*init_mm_setup_hw)(struct gk20a *g);
 		int (*init_bar2_vm)(struct gk20a *g);
 		int (*init_bar2_mm_hw_setup)(struct gk20a *g);
 		const struct gk20a_mmu_level *
diff --git a/drivers/gpu/nvgpu/gk20a/mm_gk20a.c b/drivers/gpu/nvgpu/gk20a/mm_gk20a.c
index 79bfb6875..4b39df42d 100644
--- a/drivers/gpu/nvgpu/gk20a/mm_gk20a.c
+++ b/drivers/gpu/nvgpu/gk20a/mm_gk20a.c
@@ -391,7 +391,7 @@ int gk20a_init_mm_setup_sw(struct gk20a *g)
 }
 
 /* make sure gk20a_init_mm_support is called before */
-static int gk20a_init_mm_setup_hw(struct gk20a *g)
+int gk20a_init_mm_setup_hw(struct gk20a *g)
 {
 	struct mm_gk20a *mm = &g->mm;
 	struct inst_desc *inst_block = &mm->bar1.inst_block;
@@ -435,9 +435,8 @@ int gk20a_init_mm_support(struct gk20a *g)
 	if (err)
 		return err;
 
-	err = gk20a_init_mm_setup_hw(g);
-	if (err)
-		return err;
+	if (g->ops.mm.init_mm_setup_hw)
+		err = g->ops.mm.init_mm_setup_hw(g);
 
 	return err;
 }
@@ -3116,5 +3115,6 @@ void gk20a_init_mm(struct gpu_ops *gops)
 	gops->mm.get_physical_addr_bits = gk20a_mm_get_physical_addr_bits;
 	gops->mm.get_mmu_levels = gk20a_mm_get_mmu_levels;
 	gops->mm.init_pdb = gk20a_mm_init_pdb;
+	gops->mm.init_mm_setup_hw = gk20a_init_mm_setup_hw;
 }
 
diff --git a/drivers/gpu/nvgpu/gk20a/mm_gk20a.h b/drivers/gpu/nvgpu/gk20a/mm_gk20a.h
index 8f0f736ba..b87780241 100644
--- a/drivers/gpu/nvgpu/gk20a/mm_gk20a.h
+++ b/drivers/gpu/nvgpu/gk20a/mm_gk20a.h
@@ -319,6 +319,7 @@ struct channel_gk20a;
 
 int gk20a_init_mm_support(struct gk20a *g);
 int gk20a_init_mm_setup_sw(struct gk20a *g);
+int gk20a_init_mm_setup_hw(struct gk20a *g);
 
 int gk20a_mm_fb_flush(struct gk20a *g);
 void gk20a_mm_l2_flush(struct gk20a *g, bool invalidate);
diff --git a/drivers/gpu/nvgpu/gm20b/mm_gm20b.c b/drivers/gpu/nvgpu/gm20b/mm_gm20b.c
index f85a17188..6b9c6da8b 100644
--- a/drivers/gpu/nvgpu/gm20b/mm_gm20b.c
+++ b/drivers/gpu/nvgpu/gm20b/mm_gm20b.c
@@ -126,4 +126,5 @@ void gm20b_init_mm(struct gpu_ops *gops)
 	gops->mm.get_physical_addr_bits = gk20a_mm_get_physical_addr_bits;
 	gops->mm.get_mmu_levels = gk20a_mm_get_mmu_levels;
 	gops->mm.init_pdb = gk20a_mm_init_pdb;
+	gops->mm.init_mm_setup_hw = gk20a_init_mm_setup_hw;
 }