diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.c
index 771995093cac77e5cddeac26ae179f044ac71b2c..c942f6b4be1ba2ebaa2976af2de5c88d2d9aed6d 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.c
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.c
@@ -2573,21 +2573,6 @@ static uint32_t amdgpu_vm_get_block_size(uint64_t vm_size)
 		return ((bits + 3) / 2);
 }
 
-/**
- * amdgpu_vm_set_fragment_size - adjust fragment size in PTE
- *
- * @adev: amdgpu_device pointer
- * @fragment_size_default: the default fragment size if it's set auto
- */
-void amdgpu_vm_set_fragment_size(struct amdgpu_device *adev,
-				 uint32_t fragment_size_default)
-{
-	if (amdgpu_vm_fragment_size == -1)
-		adev->vm_manager.fragment_size = fragment_size_default;
-	else
-		adev->vm_manager.fragment_size = amdgpu_vm_fragment_size;
-}
-
 /**
  * amdgpu_vm_adjust_size - adjust vm size, block size and fragment size
  *
@@ -2595,22 +2580,29 @@ void amdgpu_vm_set_fragment_size(struct amdgpu_device *adev,
  * @vm_size: the default vm size if it's set auto
  */
 void amdgpu_vm_adjust_size(struct amdgpu_device *adev, uint32_t vm_size,
-			   uint32_t fragment_size_default)
+			   uint32_t fragment_size_default, unsigned max_level)
 {
-	/* adjust vm size firstly */
-	if (amdgpu_vm_size != -1)
+	/* adjust vm size first, but only for two level setups for now */
+	if (amdgpu_vm_size != -1 && max_level == 1)
 		vm_size = amdgpu_vm_size;
 
 	adev->vm_manager.max_pfn = (uint64_t)vm_size << 18;
+	adev->vm_manager.num_level = max_level;
 
-	/* block size depends on vm size */
-	if (amdgpu_vm_block_size == -1)
+	/* block size depends on vm size and hw setup*/
+	if (adev->vm_manager.num_level > 1)
+		/* Use fixed block_size for multi level page tables */
+		adev->vm_manager.block_size = 9;
+	else if (amdgpu_vm_block_size == -1)
 		adev->vm_manager.block_size =
 			amdgpu_vm_get_block_size(vm_size);
 	else
 		adev->vm_manager.block_size = amdgpu_vm_block_size;
 
-	amdgpu_vm_set_fragment_size(adev, fragment_size_default);
+	if (amdgpu_vm_fragment_size == -1)
+		adev->vm_manager.fragment_size = fragment_size_default;
+	else
+		adev->vm_manager.fragment_size = amdgpu_vm_fragment_size;
 
 	DRM_INFO("vm size is %u GB, block size is %u-bit, fragment size is %u-bit\n",
 		 vm_size, adev->vm_manager.block_size,
diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.h b/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.h
index c80d45dd2bd31b454fbfac9728f9704b698288c2..54e540d5e8d7d85068040b8d453ee3212934ac19 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.h
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_vm.h
@@ -324,10 +324,8 @@ struct amdgpu_bo_va_mapping *amdgpu_vm_bo_lookup_mapping(struct amdgpu_vm *vm,
 							 uint64_t addr);
 void amdgpu_vm_bo_rmv(struct amdgpu_device *adev,
 		      struct amdgpu_bo_va *bo_va);
-void amdgpu_vm_set_fragment_size(struct amdgpu_device *adev,
-				 uint32_t fragment_size_default);
 void amdgpu_vm_adjust_size(struct amdgpu_device *adev, uint32_t vm_size,
-			   uint32_t fragment_size_default);
+			   uint32_t fragment_size_default, unsigned max_level);
 int amdgpu_vm_ioctl(struct drm_device *dev, void *data, struct drm_file *filp);
 bool amdgpu_vm_need_pipeline_sync(struct amdgpu_ring *ring,
 				  struct amdgpu_job *job);
diff --git a/drivers/gpu/drm/amd/amdgpu/gmc_v6_0.c b/drivers/gpu/drm/amd/amdgpu/gmc_v6_0.c
index 6098c773711ff248bbbb508c2259906de7d478c4..49224bf383248cbeee70400240e935a8b0b6795e 100644
--- a/drivers/gpu/drm/amd/amdgpu/gmc_v6_0.c
+++ b/drivers/gpu/drm/amd/amdgpu/gmc_v6_0.c
@@ -832,7 +832,7 @@ static int gmc_v6_0_sw_init(void *handle)
 	if (r)
 		return r;
 
-	amdgpu_vm_adjust_size(adev, 64, 9);
+	amdgpu_vm_adjust_size(adev, 64, 9, 1);
 
 	adev->mc.mc_mask = 0xffffffffffULL;
 
@@ -877,7 +877,6 @@ static int gmc_v6_0_sw_init(void *handle)
 	 * amdkfd will use VMIDs 8-15
 	 */
 	adev->vm_manager.id_mgr[0].num_ids = AMDGPU_NUM_OF_VMIDS;
-	adev->vm_manager.num_level = 1;
 	amdgpu_vm_manager_init(adev);
 
 	/* base offset of vram pages */
diff --git a/drivers/gpu/drm/amd/amdgpu/gmc_v7_0.c b/drivers/gpu/drm/amd/amdgpu/gmc_v7_0.c
index 8b460e9d44314822a6b9b4fd302351c847c64eed..c39cf8440afd1320ad9484ccad636e020855c027 100644
--- a/drivers/gpu/drm/amd/amdgpu/gmc_v7_0.c
+++ b/drivers/gpu/drm/amd/amdgpu/gmc_v7_0.c
@@ -971,7 +971,7 @@ static int gmc_v7_0_sw_init(void *handle)
 	 * Currently set to 4GB ((1 << 20) 4k pages).
 	 * Max GPUVM size for cayman and SI is 40 bits.
 	 */
-	amdgpu_vm_adjust_size(adev, 64, 9);
+	amdgpu_vm_adjust_size(adev, 64, 9, 1);
 
 	/* Set the internal MC address mask
 	 * This is the max address of the GPU's
@@ -1026,7 +1026,6 @@ static int gmc_v7_0_sw_init(void *handle)
 	 * amdkfd will use VMIDs 8-15
 	 */
 	adev->vm_manager.id_mgr[0].num_ids = AMDGPU_NUM_OF_VMIDS;
-	adev->vm_manager.num_level = 1;
 	amdgpu_vm_manager_init(adev);
 
 	/* base offset of vram pages */
diff --git a/drivers/gpu/drm/amd/amdgpu/gmc_v8_0.c b/drivers/gpu/drm/amd/amdgpu/gmc_v8_0.c
index 1fd7f9daab0a0fd7d4b5db7d87883d84dbd28635..421e751a04646f9aabc866aa5ed8f643c4dcc042 100644
--- a/drivers/gpu/drm/amd/amdgpu/gmc_v8_0.c
+++ b/drivers/gpu/drm/amd/amdgpu/gmc_v8_0.c
@@ -1068,7 +1068,7 @@ static int gmc_v8_0_sw_init(void *handle)
 	 * Currently set to 4GB ((1 << 20) 4k pages).
 	 * Max GPUVM size for cayman and SI is 40 bits.
 	 */
-	amdgpu_vm_adjust_size(adev, 64, 9);
+	amdgpu_vm_adjust_size(adev, 64, 9, 1);
 
 	/* Set the internal MC address mask
 	 * This is the max address of the GPU's
@@ -1123,7 +1123,6 @@ static int gmc_v8_0_sw_init(void *handle)
 	 * amdkfd will use VMIDs 8-15
 	 */
 	adev->vm_manager.id_mgr[0].num_ids = AMDGPU_NUM_OF_VMIDS;
-	adev->vm_manager.num_level = 1;
 	amdgpu_vm_manager_init(adev);
 
 	/* base offset of vram pages */
diff --git a/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c b/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
index 30eb625a991cc4887c67ae9c7d1ea84e2d25bd92..729e4d591293204fc0cef98ee291c3d73eb26769 100644
--- a/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
+++ b/drivers/gpu/drm/amd/amdgpu/gmc_v9_0.c
@@ -769,16 +769,11 @@ static int gmc_v9_0_sw_init(void *handle)
 	switch (adev->asic_type) {
 	case CHIP_RAVEN:
 		adev->mc.vram_type = AMDGPU_VRAM_TYPE_UNKNOWN;
-		if (adev->rev_id == 0x0 || adev->rev_id == 0x1) {
-			adev->vm_manager.max_pfn = 1ULL << 36;
-			adev->vm_manager.block_size = 9;
-			adev->vm_manager.num_level = 3;
-			amdgpu_vm_set_fragment_size(adev, 9);
-		} else {
+		if (adev->rev_id == 0x0 || adev->rev_id == 0x1)
+			amdgpu_vm_adjust_size(adev, 256 * 1024, 9, 3);
+		else
 			/* vm_size is 64GB for legacy 2-level page support */
-			amdgpu_vm_adjust_size(adev, 64, 9);
-			adev->vm_manager.num_level = 1;
-		}
+			amdgpu_vm_adjust_size(adev, 64, 9, 1);
 		break;
 	case CHIP_VEGA10:
 		/* XXX Don't know how to get VRAM type yet. */
@@ -788,19 +783,12 @@ static int gmc_v9_0_sw_init(void *handle)
 		 * vm size is 256TB (48bit), maximum size of Vega10,
 		 * block size 512 (9bit)
 		 */
-		adev->vm_manager.max_pfn = 1ULL << 36;
-		adev->vm_manager.block_size = 9;
-		adev->vm_manager.num_level = 3;
-		amdgpu_vm_set_fragment_size(adev, 9);
+		amdgpu_vm_adjust_size(adev, 256 * 1024, 9, 3);
 		break;
 	default:
 		break;
 	}
 
-	DRM_INFO("vm size is %llu GB, block size is %u-bit, fragment size is %u-bit\n",
-		 adev->vm_manager.max_pfn >> 18, adev->vm_manager.block_size,
-		 adev->vm_manager.fragment_size);
-
 	/* This interrupt is VMC page fault.*/
 	r = amdgpu_irq_add_id(adev, AMDGPU_IH_CLIENTID_VMC, 0,
 				&adev->mc.vm_fault);