drm/amd/display: move vmid determination logic to a module

Currently vmid is decided internally inside dc. With the introduction
of new asics we are required to coordinate vmid use with external
components.

This change converts vmid logic to a DAL module allowing vmid to be
passed in as a parameter to DC.

Signed-off-by: Dmytro Laktyushkin <[email protected]>
Reviewed-by: Charlene Liu <[email protected]>
Acked-by: Bhawanpreet Lakha <[email protected]>
Signed-off-by: Alex Deucher <[email protected]>
diff --git a/drivers/gpu/drm/amd/display/dc/Makefile b/drivers/gpu/drm/amd/display/dc/Makefile
index 00eaa69..55ce5b6 100644
--- a/drivers/gpu/drm/amd/display/dc/Makefile
+++ b/drivers/gpu/drm/amd/display/dc/Makefile
@@ -50,8 +50,11 @@
 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
+
+ifdef CONFIG_DRM_AMD_DC_DCN2_0
+DISPLAY_CORE += dc_vm_helper.o
+endif
 
 AMD_DISPLAY_CORE = $(addprefix $(AMDDALPATH)/dc/core/,$(DISPLAY_CORE))
 
diff --git a/drivers/gpu/drm/amd/display/dc/core/dc.c b/drivers/gpu/drm/amd/display/dc/core/dc.c
index 8bc8ac3..a29bb9b 100644
--- a/drivers/gpu/drm/amd/display/dc/core/dc.c
+++ b/drivers/gpu/drm/amd/display/dc/core/dc.c
@@ -808,26 +808,6 @@ void dc_destroy(struct dc **dc)
 	*dc = NULL;
 }
 
-#ifdef CONFIG_DRM_AMD_DC_DCN2_0
-bool dc_init_memory_hub(struct dc *dc, struct dc_addr_space_config *config)
-{
-	// Memory hub init isn't done as part of dc_create because in windows, dal/dc is
-	// constructed before the vm config is setup in kmd so there's no way
-	// they can give it to us at boot/dc_create
-	bool vmSupported;
-
-	// Call HWSS to setup HUBBUB for address config
-	dc->hwss.init_dchub(dc->hwseq, dc, config);
-
-	// Pre-init system aperture start/end for all HUBP instances (if not gating?)
-	// or cache system aperture if using power gating
-	memcpy(&dc->vm_config, config, sizeof(struct dc_addr_space_config));
-
-	vmSupported = (dc->ctx->asic_id.chip_family == FAMILY_NV) ? true : false;
-	return vmSupported;
-}
-
-#endif
 static void enable_timing_multisync(
 		struct dc *dc,
 		struct dc_state *ctx)
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
index 6ce87b6..a96d8de 100644
--- a/drivers/gpu/drm/amd/display/dc/core/dc_vm_helper.c
+++ b/drivers/gpu/drm/amd/display/dc/core/dc_vm_helper.c
@@ -24,8 +24,9 @@
  */
 
 #include "vm_helper.h"
+#include "dc.h"
 
-static void mark_vmid_used(struct vm_helper *vm_helper, unsigned int pos, uint8_t hubp_idx)
+void vm_helper_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];
 
@@ -33,91 +34,43 @@ static void mark_vmid_used(struct vm_helper *vm_helper, unsigned int pos, uint8_
 	vmids.vmid_usage[1] = 1 << pos;
 }
 
-static void add_ptb_to_table(struct vm_helper *vm_helper, unsigned int vmid, uint64_t ptb)
+int dc_setup_system_context(struct dc *dc, struct dc_phy_addr_space_config *pa_config)
 {
-	vm_helper->ptb_assigned_to_vmid[vmid] = ptb;
-	vm_helper->num_vmids_available--;
+	int num_vmids = 0;
+
+	/* Call HWSS to setup HUBBUB for address config */
+	if (dc->hwss.init_sys_ctx) {
+		num_vmids = dc->hwss.init_sys_ctx(dc->hwseq, dc, pa_config);
+
+		/* Pre-init system aperture start/end for all HUBP instances (if not gating?)
+		 * or cache system aperture if using power gating
+		 */
+		memcpy(&dc->vm_pa_config, pa_config, sizeof(struct dc_phy_addr_space_config));
+		dc->vm_pa_config.valid = true;
+	}
+
+	return num_vmids;
 }
 
-static void clear_entry_from_vmid_table(struct vm_helper *vm_helper, unsigned int vmid)
+void dc_setup_vm_context(struct dc *dc, struct dc_virtual_addr_space_config *va_config, int vmid)
 {
-	vm_helper->ptb_assigned_to_vmid[vmid] = 0;
-	vm_helper->num_vmids_available++;
+	dc->hwss.init_vm_ctx(dc->hwseq, dc, va_config, vmid);
 }
 
-static void evict_vmids(struct vm_helper *vm_helper)
+int dc_get_vmid_use_vector(struct dc *dc)
 {
 	int i;
-	uint16_t ord = 0;
+	int in_use = 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);
-	}
+	for (i = 0; i < dc->vm_helper->num_vmid; i++)
+		in_use |= dc->vm_helper->hubp_vmid_usage[i].vmid_usage[0]
+			| dc->vm_helper->hubp_vmid_usage[i].vmid_usage[1];
+	return in_use;
 }
 
-// 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)
+void vm_helper_init(struct vm_helper *vm_helper, unsigned int num_vmid)
 {
 	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 07a264f..5ebfeee 100644
--- a/drivers/gpu/drm/amd/display/dc/dc.h
+++ b/drivers/gpu/drm/amd/display/dc/dc.h
@@ -403,21 +403,17 @@ struct dc_phy_addr_space_config {
 		uint64_t page_table_end_addr;
 		uint64_t page_table_base_addr;
 	} gart_config;
+
+	bool valid;
 };
 
 struct dc_virtual_addr_space_config {
+	uint64_t	page_table_base_addr;
 	uint64_t	page_table_start_addr;
 	uint64_t	page_table_end_addr;
 	uint32_t	page_table_block_size_in_bytes;
 	uint8_t		page_table_depth; // 1 = 1 level, 2 = 2 level, etc.  0 = invalid
 };
-
-struct dc_addr_space_config {
-	struct dc_phy_addr_space_config		pa_config;
-	struct dc_virtual_addr_space_config	va_config;
-	uint32_t	valid:1;
-};
-
 #endif
 
 struct dc_bounding_box_overrides {
@@ -449,7 +445,7 @@ struct dc {
 #endif
 	struct dc_context *ctx;
 #ifdef CONFIG_DRM_AMD_DC_DCN2_0
-	struct dc_addr_space_config vm_config;
+	struct dc_phy_addr_space_config vm_pa_config;
 #endif
 
 	uint8_t link_count;
@@ -539,8 +535,11 @@ struct dc_callback_init {
 };
 
 struct dc *dc_create(const struct dc_init_data *init_params);
+int dc_get_vmid_use_vector(struct dc *dc);
 #ifdef CONFIG_DRM_AMD_DC_DCN2_0
-bool dc_init_memory_hub(struct dc *dc, struct dc_addr_space_config *config);
+void dc_setup_vm_context(struct dc *dc, struct dc_virtual_addr_space_config *va_config, int vmid);
+/* Returns the number of vmids supported */
+int dc_setup_system_context(struct dc *dc, struct dc_phy_addr_space_config *pa_config);
 #endif
 void dc_init_callbacks(struct dc *dc,
 		const struct dc_callback_init *init_params);
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 6b1431d..22db568 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 a48d314..934bacc 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.c
+++ b/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.c
@@ -339,8 +339,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_hubp.h b/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.h
index 91dc42e..31c8fdd 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.h
+++ b/drivers/gpu/drm/amd/display/dc/dcn10/dcn10_hubp.h
@@ -720,8 +720,7 @@ void hubp1_dcc_control(struct hubp *hubp,
 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);
 
 #endif
 bool hubp1_is_flip_pending(struct 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 8ad4997..d8c54f96 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
@@ -1297,8 +1297,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/dcn20/dcn20_hubbub.c b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubbub.c
index 724f1c5..ece6e13 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubbub.c
+++ b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubbub.c
@@ -47,6 +47,11 @@
 #define FN(reg_name, field_name) \
 	hubbub1->shifts->field_name, hubbub1->masks->field_name
 
+#ifdef NUM_VMID
+#undef NUM_VMID
+#endif
+#define NUM_VMID 16
+
 bool hubbub2_dcc_support_swizzle(
 		enum swizzle_mode_values swizzle,
 		unsigned int bytes_per_element,
@@ -294,15 +299,6 @@ bool hubbub2_get_dcc_compression_cap(struct hubbub *hubbub,
 	return true;
 }
 
-void hubbub2_setup_vmid_ptb(struct hubbub *hubbub,
-		uint64_t ptb,
-		uint8_t vmid)
-{
-	struct dcn20_hubbub *hubbub1 = TO_DCN20_HUBBUB(hubbub);
-
-	dcn20_vmid_set_ptb(&hubbub1->vmid[vmid], ptb);
-}
-
 static enum dcn_hubbub_page_table_depth page_table_depth_to_hw(unsigned int page_table_depth)
 {
 	enum dcn_hubbub_page_table_depth depth = 0;
@@ -347,49 +343,53 @@ static enum dcn_hubbub_page_table_block_size page_table_block_size_to_hw(unsigne
 	return block_size;
 }
 
-void hubbub2_init_dchub(struct hubbub *hubbub,
-		struct hubbub_addr_config *config)
+void hubbub2_init_vm_ctx(struct hubbub *hubbub,
+		struct dcn_hubbub_virt_addr_config *va_config,
+		int vmid)
 {
-	int i;
 	struct dcn20_hubbub *hubbub1 = TO_DCN20_HUBBUB(hubbub);
-	struct dcn_vmid_page_table_config phys_config;
 	struct dcn_vmid_page_table_config virt_config;
 
-	REG_SET(DCN_VM_FB_LOCATION_BASE, 0,
-			FB_BASE, config->pa_config.system_aperture.fb_base);
-	REG_SET(DCN_VM_FB_LOCATION_TOP, 0,
-			FB_TOP, config->pa_config.system_aperture.fb_top);
-	REG_SET(DCN_VM_FB_OFFSET, 0,
-			FB_OFFSET, config->pa_config.system_aperture.fb_offset);
-	REG_SET(DCN_VM_AGP_BOT, 0,
-			AGP_BOT, config->pa_config.system_aperture.agp_bot);
-	REG_SET(DCN_VM_AGP_TOP, 0,
-			AGP_TOP, config->pa_config.system_aperture.agp_top);
-	REG_SET(DCN_VM_AGP_BASE, 0,
-			AGP_BASE, config->pa_config.system_aperture.agp_base);
+	virt_config.page_table_start_addr = va_config->page_table_start_addr >> 12;
+	virt_config.page_table_end_addr = va_config->page_table_end_addr >> 12;
+	virt_config.depth = page_table_depth_to_hw(va_config->page_table_depth);
+	virt_config.block_size = page_table_block_size_to_hw(va_config->page_table_block_size);
+	virt_config.page_table_base_addr = va_config->page_table_base_addr;
 
-	if (config->pa_config.gart_config.page_table_start_addr != config->pa_config.gart_config.page_table_end_addr) {
+	dcn20_vmid_setup(&hubbub1->vmid[vmid], &virt_config);
+}
+
+int hubbub2_init_dchub_sys_ctx(struct hubbub *hubbub,
+		struct dcn_hubbub_phys_addr_config *pa_config)
+{
+	struct dcn20_hubbub *hubbub1 = TO_DCN20_HUBBUB(hubbub);
+	struct dcn_vmid_page_table_config phys_config;
+
+	REG_SET(DCN_VM_FB_LOCATION_BASE, 0,
+			FB_BASE, pa_config->system_aperture.fb_base);
+	REG_SET(DCN_VM_FB_LOCATION_TOP, 0,
+			FB_TOP, pa_config->system_aperture.fb_top);
+	REG_SET(DCN_VM_FB_OFFSET, 0,
+			FB_OFFSET, pa_config->system_aperture.fb_offset);
+	REG_SET(DCN_VM_AGP_BOT, 0,
+			AGP_BOT, pa_config->system_aperture.agp_bot);
+	REG_SET(DCN_VM_AGP_TOP, 0,
+			AGP_TOP, pa_config->system_aperture.agp_top);
+	REG_SET(DCN_VM_AGP_BASE, 0,
+			AGP_BASE, pa_config->system_aperture.agp_base);
+
+	if (pa_config->gart_config.page_table_start_addr != pa_config->gart_config.page_table_end_addr) {
 		phys_config.depth = 1;
 		phys_config.block_size = 4096;
-		phys_config.page_table_start_addr = config->pa_config.gart_config.page_table_start_addr >> 12;
-		phys_config.page_table_end_addr = config->pa_config.gart_config.page_table_end_addr >> 12;
+		phys_config.page_table_start_addr = pa_config->gart_config.page_table_start_addr >> 12;
+		phys_config.page_table_end_addr = pa_config->gart_config.page_table_end_addr >> 12;
+		phys_config.page_table_base_addr = pa_config->gart_config.page_table_base_addr;
 
 		// Init VMID 0 based on PA config
 		dcn20_vmid_setup(&hubbub1->vmid[0], &phys_config);
-		dcn20_vmid_set_ptb(&hubbub1->vmid[0], config->pa_config.gart_config.page_table_base_addr);
 	}
 
-	if (config->va_config.page_table_start_addr != config->va_config.page_table_end_addr) {
-		// Init VMID 1-15 based on VA config
-		for (i = 1; i < 16; i++) {
-			virt_config.page_table_start_addr = config->va_config.page_table_start_addr >> 12;
-			virt_config.page_table_end_addr = config->va_config.page_table_end_addr >> 12;
-			virt_config.depth = page_table_depth_to_hw(config->va_config.page_table_depth);
-			virt_config.block_size = page_table_block_size_to_hw(config->va_config.page_table_block_size);
-
-			dcn20_vmid_setup(&hubbub1->vmid[i], &virt_config);
-		}
-	}
+	return NUM_VMID;
 }
 
 void hubbub2_update_dchub(struct hubbub *hubbub,
@@ -564,8 +564,8 @@ static void hubbub2_program_watermarks(
 
 static const struct hubbub_funcs hubbub2_funcs = {
 	.update_dchub = hubbub2_update_dchub,
-	.init_dchub = hubbub2_init_dchub,
-	.setup_vmid_ptb = hubbub2_setup_vmid_ptb,
+	.init_dchub_sys_ctx = hubbub2_init_dchub_sys_ctx,
+	.init_vm_ctx = hubbub2_init_vm_ctx,
 	.dcc_support_swizzle = hubbub2_dcc_support_swizzle,
 	.dcc_support_pixel_format = hubbub2_dcc_support_pixel_format,
 	.get_dcc_compression_cap = hubbub2_get_dcc_compression_cap,
diff --git a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubbub.h b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubbub.h
index 6ca4a5a9f..a7b6ca2 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubbub.h
+++ b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubbub.h
@@ -89,9 +89,11 @@ bool hubbub2_initialize_vmids(struct hubbub *hubbub,
 		const struct dc_dcc_surface_param *input,
 		struct dc_surface_dcc_cap *output);
 
-void hubbub2_init_dchub(struct hubbub *hubbub,
-		struct hubbub_addr_config *config);
-
+int hubbub2_init_dchub_sys_ctx(struct hubbub *hubbub,
+		struct dcn_hubbub_phys_addr_config *pa_config);
+void hubbub2_init_vm_ctx(struct hubbub *hubbub,
+		struct dcn_hubbub_virt_addr_config *va_config,
+		int vmid);
 void hubbub2_update_dchub(struct hubbub *hubbub,
 		struct dchub_init_data *dh_data);
 
@@ -102,8 +104,4 @@ void hubbub2_get_dchub_ref_freq(struct hubbub *hubbub,
 void hubbub2_wm_read_state(struct hubbub *hubbub,
 		struct dcn_hubbub_wm *wm);
 
-void hubbub2_setup_vmid_ptb(struct hubbub *hubbub,
-		uint64_t ptb,
-		uint8_t vmid);
-
 #endif
diff --git a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubp.c b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubp.c
index dcba2c5..d3f7dd3 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubp.c
+++ b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubp.c
@@ -447,8 +447,7 @@ bool hubp2_dmdata_status_done(struct hubp *hubp)
 bool hubp2_program_surface_flip_and_addr(
 	struct hubp *hubp,
 	const struct dc_plane_address *address,
-	bool flip_immediate,
-	uint8_t vmid)
+	bool flip_immediate)
 {
 	struct dcn20_hubp *hubp2 = TO_DCN20_HUBP(hubp);
 
@@ -458,7 +457,7 @@ bool hubp2_program_surface_flip_and_addr(
 
 	// Program VMID reg
 	REG_UPDATE(VMID_SETTINGS_0,
-			VMID, vmid);
+			VMID, address->vmid);
 
 	if (address->type == PLN_ADDR_TYPE_GRPH_STEREO) {
 		REG_UPDATE(DCSURF_FLIP_CONTROL, SURFACE_FLIP_MODE_FOR_STEREOSYNC, 0x1);
diff --git a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubp.h b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubp.h
index f790ab9..d5acc34 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubp.h
+++ b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hubp.h
@@ -260,8 +260,7 @@ void hubp2_program_deadline(
 bool hubp2_program_surface_flip_and_addr(
 	struct hubp *hubp,
 	const struct dc_plane_address *address,
-	bool flip_immediate,
-	uint8_t vmid);
+	bool flip_immediate);
 
 void hubp2_program_surface_config(
 	struct hubp *hubp,
diff --git a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hwseq.c b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hwseq.c
index d55b15f..20bc3b9f 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hwseq.c
+++ b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_hwseq.c
@@ -1122,15 +1122,13 @@ void dcn20_enable_plane(
 		print_rq_dlg_ttu(dc, pipe_ctx);
 	}
 */
-	if (dc->vm_config.valid) {
+	if (dc->vm_pa_config.valid) {
 		struct vm_system_aperture_param apt;
 
 		apt.sys_default.quad_part = 0;
-		apt.sys_high.quad_part = 0;
-		apt.sys_low.quad_part = 0;
 
-		apt.sys_high.quad_part = dc->vm_config.pa_config.system_aperture.start_addr;
-		apt.sys_low.quad_part = dc->vm_config.pa_config.system_aperture.end_addr;
+		apt.sys_high.quad_part = dc->vm_pa_config.system_aperture.start_addr;
+		apt.sys_low.quad_part = dc->vm_pa_config.system_aperture.end_addr;
 
 		// Program system aperture settings
 		pipe_ctx->plane_res.hubp->funcs->hubp_set_vm_system_aperture_settings(pipe_ctx->plane_res.hubp, &apt);
@@ -1565,26 +1563,43 @@ void dcn20_disable_stream(struct pipe_ctx *pipe_ctx, int option)
 	dce110_disable_stream(pipe_ctx, option);
 }
 
-static void dcn20_init_dchub(struct dce_hwseq *hws, struct dc *dc, struct dc_addr_space_config *config)
+static void dcn20_init_vm_ctx(
+		struct dce_hwseq *hws,
+		struct dc *dc,
+		struct dc_virtual_addr_space_config *va_config,
+		int vmid)
 {
-	struct hubbub_addr_config hubbub_config;
+	struct dcn_hubbub_virt_addr_config config;
 
-	hubbub_config.pa_config.system_aperture.fb_top = config->pa_config.system_aperture.fb_top;
-	hubbub_config.pa_config.system_aperture.fb_offset = config->pa_config.system_aperture.fb_offset;
-	hubbub_config.pa_config.system_aperture.fb_base = config->pa_config.system_aperture.fb_base;
-	hubbub_config.pa_config.system_aperture.agp_top = config->pa_config.system_aperture.agp_top;
-	hubbub_config.pa_config.system_aperture.agp_bot = config->pa_config.system_aperture.agp_bot;
-	hubbub_config.pa_config.system_aperture.agp_base = config->pa_config.system_aperture.agp_base;
-	hubbub_config.pa_config.gart_config.page_table_start_addr = config->pa_config.gart_config.page_table_start_addr;
-	hubbub_config.pa_config.gart_config.page_table_end_addr = config->pa_config.gart_config.page_table_end_addr;
-	hubbub_config.pa_config.gart_config.page_table_base_addr = config->pa_config.gart_config.page_table_base_addr;
+	if (vmid == 0) {
+		ASSERT(0); /* VMID cannot be 0 for vm context */
+		return;
+	}
 
-	hubbub_config.va_config.page_table_start_addr = config->va_config.page_table_start_addr;
-	hubbub_config.va_config.page_table_end_addr = config->va_config.page_table_end_addr;
-	hubbub_config.va_config.page_table_block_size = config->va_config.page_table_block_size_in_bytes;
-	hubbub_config.va_config.page_table_depth = config->va_config.page_table_depth;
+	config.page_table_start_addr = va_config->page_table_start_addr;
+	config.page_table_end_addr = va_config->page_table_end_addr;
+	config.page_table_block_size = va_config->page_table_block_size_in_bytes;
+	config.page_table_depth = va_config->page_table_depth;
+	config.page_table_base_addr = va_config->page_table_base_addr;
 
-	dc->res_pool->hubbub->funcs->init_dchub(dc->res_pool->hubbub, &hubbub_config);
+	dc->res_pool->hubbub->funcs->init_vm_ctx(dc->res_pool->hubbub, &config, vmid);
+}
+
+static int dcn20_init_sys_ctx(struct dce_hwseq *hws, struct dc *dc, struct dc_phy_addr_space_config *pa_config)
+{
+	struct dcn_hubbub_phys_addr_config config;
+
+	config.system_aperture.fb_top = pa_config->system_aperture.fb_top;
+	config.system_aperture.fb_offset = pa_config->system_aperture.fb_offset;
+	config.system_aperture.fb_base = pa_config->system_aperture.fb_base;
+	config.system_aperture.agp_top = pa_config->system_aperture.agp_top;
+	config.system_aperture.agp_bot = pa_config->system_aperture.agp_bot;
+	config.system_aperture.agp_base = pa_config->system_aperture.agp_base;
+	config.gart_config.page_table_start_addr = pa_config->gart_config.page_table_start_addr;
+	config.gart_config.page_table_end_addr = pa_config->gart_config.page_table_end_addr;
+	config.gart_config.page_table_base_addr = pa_config->gart_config.page_table_base_addr;
+
+	return dc->res_pool->hubbub->funcs->init_dchub_sys_ctx(dc->res_pool->hubbub, &config);
 }
 
 static bool patch_address_for_sbs_tb_stereo(
@@ -1619,30 +1634,19 @@ static void dcn20_update_plane_addr(const struct dc *dc, struct pipe_ctx *pipe_c
 	bool addr_patched = false;
 	PHYSICAL_ADDRESS_LOC addr;
 	struct dc_plane_state *plane_state = pipe_ctx->plane_state;
-	uint8_t vmid;
 
 	if (plane_state == NULL)
 		return;
 
 	addr_patched = patch_address_for_sbs_tb_stereo(pipe_ctx, &addr);
 
-	// Call Helper to assign correct VMID to this PTB
-	vmid = get_vmid_for_ptb(dc->vm_helper,
-			plane_state->address.page_table_base.quad_part,
-			pipe_ctx->pipe_idx);
-
-	// Call hubbub to program PTB of VMID only if its VA
-	// PA PTB is a one-time setup at init
-	if (vmid > 0 && dc->res_pool->hubbub->funcs->setup_vmid_ptb)
-		dc->res_pool->hubbub->funcs->setup_vmid_ptb(dc->res_pool->hubbub,
-				plane_state->address.page_table_base.quad_part,
-				vmid);
+	// Call Helper to track VMID use
+	vm_helper_mark_vmid_used(dc->vm_helper, plane_state->address.vmid, pipe_ctx->plane_res.hubp->inst);
 
 	pipe_ctx->plane_res.hubp->funcs->hubp_program_surface_flip_and_addr(
 			pipe_ctx->plane_res.hubp,
 			&plane_state->address,
-			plane_state->flip_immediate,
-			vmid);
+			plane_state->flip_immediate);
 
 	plane_state->status.requested_address = plane_state->address;
 
@@ -1991,7 +1995,8 @@ void dcn20_hw_sequencer_construct(struct dc *dc)
 	dc->hwss.blank_pixel_data = dcn20_blank_pixel_data;
 	dc->hwss.dmdata_status_done = dcn20_dmdata_status_done;
 	dc->hwss.disable_stream = dcn20_disable_stream;
-	dc->hwss.init_dchub = dcn20_init_dchub;
+	dc->hwss.init_sys_ctx = dcn20_init_sys_ctx;
+	dc->hwss.init_vm_ctx = dcn20_init_vm_ctx;
 	dc->hwss.disable_stream_gating = dcn20_disable_stream_gating;
 	dc->hwss.enable_stream_gating = dcn20_enable_stream_gating;
 	dc->hwss.setup_vupdate_interrupt = dcn20_setup_vupdate_interrupt;
diff --git a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_resource.c b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_resource.c
index 26af860..4cc66c2 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_resource.c
+++ b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_resource.c
@@ -2899,7 +2899,7 @@ static bool construct(
 
 	// Init the vm_helper
 	if (dc->vm_helper)
-		init_vm_helper(dc->vm_helper, 16, pool->base.pipe_count);
+		vm_helper_init(dc->vm_helper, 16);
 
 	/*************************************************
 	 *  Create resources                             *
diff --git a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_vmid.c b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_vmid.c
index 50953c0..27679ef 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_vmid.c
+++ b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_vmid.c
@@ -51,12 +51,9 @@ void dcn20_vmid_setup(struct dcn20_vmid *vmid, const struct dcn_vmid_page_table_
 	REG_SET_2(CNTL, 0,
 			VM_CONTEXT0_PAGE_TABLE_DEPTH, config->depth,
 			VM_CONTEXT0_PAGE_TABLE_BLOCK_SIZE, config->block_size);
-}
 
-void dcn20_vmid_set_ptb(struct dcn20_vmid *vmid, uint64_t base)
-{
 	REG_SET(PAGE_TABLE_BASE_ADDR_HI32, 0,
-			VM_CONTEXT0_PAGE_DIRECTORY_ENTRY_HI32, (base >> 32) & 0xFFFFFFFF);
+			VM_CONTEXT0_PAGE_DIRECTORY_ENTRY_HI32, (config->page_table_base_addr >> 32) & 0xFFFFFFFF);
 	REG_SET(PAGE_TABLE_BASE_ADDR_LO32, 0,
-			VM_CONTEXT0_PAGE_DIRECTORY_ENTRY_LO32, base & 0xFFFFFFFF);
+			VM_CONTEXT0_PAGE_DIRECTORY_ENTRY_LO32, config->page_table_base_addr & 0xFFFFFFFF);
 }
diff --git a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_vmid.h b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_vmid.h
index 9c2f701..02fafb0 100644
--- a/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_vmid.h
+++ b/drivers/gpu/drm/amd/display/dc/dcn20/dcn20_vmid.h
@@ -86,6 +86,5 @@ struct dcn20_vmid {
 };
 
 void dcn20_vmid_setup(struct dcn20_vmid *vmid, const struct dcn_vmid_page_table_config *config);
-void dcn20_vmid_set_ptb(struct dcn20_vmid *vmid, uint64_t base);
 
 #endif /* DAL_DC_DCN20_DCN20_VMID_H_ */
diff --git a/drivers/gpu/drm/amd/display/dc/inc/hw/dchubbub.h b/drivers/gpu/drm/amd/display/dc/inc/hw/dchubbub.h
index d56fd7d8..959f5b6 100644
--- a/drivers/gpu/drm/amd/display/dc/inc/hw/dchubbub.h
+++ b/drivers/gpu/drm/amd/display/dc/inc/hw/dchubbub.h
@@ -87,6 +87,7 @@ struct dcn_hubbub_virt_addr_config {
 	uint64_t				page_table_end_addr;
 	enum dcn_hubbub_page_table_block_size	page_table_block_size;
 	enum dcn_hubbub_page_table_depth	page_table_depth;
+	uint64_t				page_table_base_addr;
 };
 
 struct hubbub_addr_config {
@@ -105,18 +106,13 @@ struct hubbub_funcs {
 			struct dchub_init_data *dh_data);
 
 #ifdef CONFIG_DRM_AMD_DC_DCN2_0
-	void (*init_dchub)(
+	int (*init_dchub_sys_ctx)(
 			struct hubbub *hubbub,
-			struct hubbub_addr_config *config);
-	void (*setup_vmid_ptb)(
+			struct dcn_hubbub_phys_addr_config *pa_config);
+	void (*init_vm_ctx)(
 			struct hubbub *hubbub,
-			uint64_t ptb,
-			uint8_t vmid);
-
-	void (*set_ptb)(
-			struct hubbub *hubbub,
-			uint8_t vmid,
-			uint64_t base_addr);
+			struct dcn_hubbub_virt_addr_config *va_config,
+			int vmid);
 
 #endif
 	bool (*get_dcc_compression_cap)(struct hubbub *hubbub,
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 3424778..51bff87 100644
--- a/drivers/gpu/drm/amd/display/dc/inc/hw/hubp.h
+++ b/drivers/gpu/drm/amd/display/dc/inc/hw/hubp.h
@@ -83,8 +83,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 037beb0..76de0e4 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/hw_sequencer.h b/drivers/gpu/drm/amd/display/dc/inc/hw_sequencer.h
index 13b113d..4ffe42c 100644
--- a/drivers/gpu/drm/amd/display/dc/inc/hw_sequencer.h
+++ b/drivers/gpu/drm/amd/display/dc/inc/hw_sequencer.h
@@ -75,7 +75,8 @@ struct resource_pool;
 struct resource_context;
 struct stream_resource;
 #ifdef CONFIG_DRM_AMD_DC_DCN2_0
-struct dc_addr_space_config;
+struct dc_phy_addr_space_config;
+struct dc_virtual_addr_space_config;
 #endif
 
 struct hw_sequencer_funcs {
@@ -132,11 +133,15 @@ struct hw_sequencer_funcs {
 		struct dchub_init_data *dh_data);
 
 #ifdef CONFIG_DRM_AMD_DC_DCN2_0
-	void (*init_dchub)(
-		struct dce_hwseq *hws,
-		struct dc *dc,
-		struct dc_addr_space_config *dh_data);
-
+	int (*init_sys_ctx)(
+			struct dce_hwseq *hws,
+			struct dc *dc,
+			struct dc_phy_addr_space_config *pa_config);
+	void (*init_vm_ctx)(
+			struct dce_hwseq *hws,
+			struct dc *dc,
+			struct dc_virtual_addr_space_config *va_config,
+			int vmid);
 #endif
 	void (*update_mpcc)(
 		struct dc *dc,
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 193407f..8bfcef0 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_ */