diff --git a/drivers/gpu/drm/amd/display/dc/Makefile b/drivers/gpu/drm/amd/display/dc/Makefile
index 6da4e4f844b2..f581d5394caa 100644
--- a/drivers/gpu/drm/amd/display/dc/Makefile
+++ b/drivers/gpu/drm/amd/display/dc/Makefile
@@ -41,8 +41,8 @@ AMD_DC = $(addsuffix /Makefile, $(addprefix $(FULL_AMD_DISPLAY_PATH)/dc/,$(DC_LI
 include $(AMD_DC)
 
 DISPLAY_CORE = dc.o dc_link.o dc_resource.o dc_hw_sequencer.o dc_sink.o \
-dc_surface.o dc_link_hwss.o dc_link_dp.o dc_link_ddc.o dc_debug.o dc_stream.o \
-dc_vm_helper.o
+dc_surface.o dc_link_hwss.o dc_link_dp.o dc_link_ddc.o dc_debug.o dc_stream.o
+
 
 AMD_DISPLAY_CORE = $(addprefix $(AMDDALPATH)/dc/core/,$(DISPLAY_CORE))
 
diff --git a/drivers/gpu/drm/amd/display/dc/core/dc_vm_helper.c b/drivers/gpu/drm/amd/display/dc/core/dc_vm_helper.c
deleted file mode 100644
index 6ce87b682a32..000000000000
--- a/drivers/gpu/drm/amd/display/dc/core/dc_vm_helper.c
+++ /dev/null
@@ -1,123 +0,0 @@
-/*
- * Copyright 2018 Advanced Micro Devices, Inc.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- *
- * Authors: AMD
- *
- */
-
-#include "vm_helper.h"
-
-static void mark_vmid_used(struct vm_helper *vm_helper, unsigned int pos, uint8_t hubp_idx)
-{
-	struct vmid_usage vmids = vm_helper->hubp_vmid_usage[hubp_idx];
-
-	vmids.vmid_usage[0] = vmids.vmid_usage[1];
-	vmids.vmid_usage[1] = 1 << pos;
-}
-
-static void add_ptb_to_table(struct vm_helper *vm_helper, unsigned int vmid, uint64_t ptb)
-{
-	vm_helper->ptb_assigned_to_vmid[vmid] = ptb;
-	vm_helper->num_vmids_available--;
-}
-
-static void clear_entry_from_vmid_table(struct vm_helper *vm_helper, unsigned int vmid)
-{
-	vm_helper->ptb_assigned_to_vmid[vmid] = 0;
-	vm_helper->num_vmids_available++;
-}
-
-static void evict_vmids(struct vm_helper *vm_helper)
-{
-	int i;
-	uint16_t ord = 0;
-
-	for (i = 0; i < vm_helper->num_vmid; i++)
-		ord |= vm_helper->hubp_vmid_usage[i].vmid_usage[0] | vm_helper->hubp_vmid_usage[i].vmid_usage[1];
-
-	// At this point any positions with value 0 are unused vmids, evict them
-	for (i = 1; i < vm_helper->num_vmid; i++) {
-		if (ord & (1u << i))
-			clear_entry_from_vmid_table(vm_helper, i);
-	}
-}
-
-// Return value of -1 indicates vmid table unitialized or ptb dne in the table
-static int get_existing_vmid_for_ptb(struct vm_helper *vm_helper, uint64_t ptb)
-{
-	int i;
-
-	for (i = 0; i < vm_helper->num_vmid; i++) {
-		if (vm_helper->ptb_assigned_to_vmid[i] == ptb)
-			return i;
-	}
-
-	return -1;
-}
-
-// Expected to be called only when there's an available vmid
-static int get_next_available_vmid(struct vm_helper *vm_helper)
-{
-	int i;
-
-	for (i = 1; i < vm_helper->num_vmid; i++) {
-		if (vm_helper->ptb_assigned_to_vmid[i] == 0)
-			return i;
-	}
-
-	return -1;
-}
-
-uint8_t get_vmid_for_ptb(struct vm_helper *vm_helper, int64_t ptb, uint8_t hubp_idx)
-{
-	unsigned int vmid = 0;
-	int vmid_exists = -1;
-
-	// Physical address gets vmid 0
-	if (ptb == 0)
-		return 0;
-
-	vmid_exists = get_existing_vmid_for_ptb(vm_helper, ptb);
-
-	if (vmid_exists != -1) {
-		mark_vmid_used(vm_helper, vmid_exists, hubp_idx);
-		vmid = vmid_exists;
-	} else {
-		if (vm_helper->num_vmids_available == 0)
-			evict_vmids(vm_helper);
-
-		vmid = get_next_available_vmid(vm_helper);
-		mark_vmid_used(vm_helper, vmid, hubp_idx);
-		add_ptb_to_table(vm_helper, vmid, ptb);
-	}
-
-	return vmid;
-}
-
-void init_vm_helper(struct vm_helper *vm_helper, unsigned int num_vmid, unsigned int num_hubp)
-{
-	vm_helper->num_vmid = num_vmid;
-	vm_helper->num_hubp = num_hubp;
-	vm_helper->num_vmids_available = num_vmid - 1;
-
-	memset(vm_helper->hubp_vmid_usage, 0, sizeof(vm_helper->hubp_vmid_usage[0]) * MAX_HUBP);
-	memset(vm_helper->ptb_assigned_to_vmid, 0, sizeof(vm_helper->ptb_assigned_to_vmid[0]) * MAX_VMID);
-}
diff --git a/drivers/gpu/drm/amd/display/dc/dc.h b/drivers/gpu/drm/amd/display/dc/dc.h
index c4bd9216dd61..ffd1fe5df99e 100644
--- a/drivers/gpu/drm/amd/display/dc/dc.h
+++ b/drivers/gpu/drm/amd/display/dc/dc.h
@@ -458,6 +458,7 @@ struct dc_callback_init {
 };
 
 struct dc *dc_create(const struct dc_init_data *init_params);
+int dc_get_vmid_use_vector(struct dc *dc);
 void dc_init_callbacks(struct dc *dc,
 		const struct dc_callback_init *init_params);
 void dc_destroy(struct dc **dc);
diff --git a/drivers/gpu/drm/amd/display/dc/dc_hw_types.h b/drivers/gpu/drm/amd/display/dc/dc_hw_types.h
index c91b8aad78c9..479c5f8352f9 100644
--- a/drivers/gpu/drm/amd/display/dc/dc_hw_types.h
+++ b/drivers/gpu/drm/amd/display/dc/dc_hw_types.h
@@ -99,6 +99,8 @@ struct dc_plane_address {
 	};
 
 	union large_integer page_table_base;
+
+	uint8_t vmid;
 };
 
 struct dc_size {
diff --git a/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.c b/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.c
index aea2b63db137..4cb7fe8b3903 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.c
+++ b/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.c
@@ -317,8 +317,7 @@ void hubp1_program_pixel_format(
 bool hubp1_program_surface_flip_and_addr(
 	struct hubp *hubp,
 	const struct dc_plane_address *address,
-	bool flip_immediate,
-	uint8_t vmid)
+	bool flip_immediate)
 {
 	struct dcn10_hubp *hubp1 = TO_DCN10_HUBP(hubp);
 
diff --git a/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hw_sequencer.c b/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hw_sequencer.c
index d9edf7f84cfa..be245d4fe5c2 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hw_sequencer.c
+++ b/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hw_sequencer.c
@@ -1236,8 +1236,7 @@ static void dcn10_update_plane_addr(const struct dc *dc, struct pipe_ctx *pipe_c
 	pipe_ctx->plane_res.hubp->funcs->hubp_program_surface_flip_and_addr(
 			pipe_ctx->plane_res.hubp,
 			&plane_state->address,
-			plane_state->flip_immediate,
-			0);
+			plane_state->flip_immediate);
 
 	plane_state->status.requested_address = plane_state->address;
 
diff --git a/drivers/gpu/drm/amd/display/dc/inc/hw/hubp.h b/drivers/gpu/drm/amd/display/dc/inc/hw/hubp.h
index 5420ad2da96f..e4909d21b9f2 100644
--- a/drivers/gpu/drm/amd/display/dc/inc/hw/hubp.h
+++ b/drivers/gpu/drm/amd/display/dc/inc/hw/hubp.h
@@ -80,8 +80,7 @@ struct hubp_funcs {
 	bool (*hubp_program_surface_flip_and_addr)(
 		struct hubp *hubp,
 		const struct dc_plane_address *address,
-		bool flip_immediate,
-		uint8_t vmid);
+		bool flip_immediate);
 
 	void (*hubp_program_pte_vm)(
 		struct hubp *hubp,
diff --git a/drivers/gpu/drm/amd/display/dc/inc/hw/vmid.h b/drivers/gpu/drm/amd/display/dc/inc/hw/vmid.h
index 037beb0a2a27..76de0e4284e0 100644
--- a/drivers/gpu/drm/amd/display/dc/inc/hw/vmid.h
+++ b/drivers/gpu/drm/amd/display/dc/inc/hw/vmid.h
@@ -44,6 +44,7 @@ struct dcn_vmid_page_table_config {
 	uint64_t	page_table_end_addr;
 	enum dcn_hubbub_page_table_depth	depth;
 	enum dcn_hubbub_page_table_block_size	block_size;
+	uint64_t	page_table_base_addr;
 };
 
 #endif /* DAL_DC_INC_HW_VMID_H_ */
diff --git a/drivers/gpu/drm/amd/display/dc/inc/vm_helper.h b/drivers/gpu/drm/amd/display/dc/inc/vm_helper.h
index 193407f76a80..8bfcef0a3675 100644
--- a/drivers/gpu/drm/amd/display/dc/inc/vm_helper.h
+++ b/drivers/gpu/drm/amd/display/dc/inc/vm_helper.h
@@ -28,29 +28,21 @@
 
 #include "dc_types.h"
 
-#define MAX_VMID 16
 #define MAX_HUBP 6
 
 struct vmid_usage {
-	uint16_t vmid_usage[2];
+	int vmid_usage[2];
 };
 
 struct vm_helper {
 	unsigned int num_vmid;
-	unsigned int num_hubp;
-	unsigned int num_vmids_available;
-	uint64_t ptb_assigned_to_vmid[MAX_VMID];
 	struct vmid_usage hubp_vmid_usage[MAX_HUBP];
 };
 
-uint8_t get_vmid_for_ptb(
-		struct vm_helper *vm_helper,
-		int64_t ptb,
-		uint8_t pipe_idx);
+void vm_helper_mark_vmid_used(struct vm_helper *vm_helper, unsigned int pos, uint8_t hubp_idx);
 
-void init_vm_helper(
+void vm_helper_init(
 	struct vm_helper *vm_helper,
-	unsigned int num_vmid,
-	unsigned int num_hubp);
+	unsigned int num_vmid);
 
 #endif /* DC_INC_VM_HELPER_H_ */