diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu.h b/drivers/gpu/drm/amd/amdgpu/amdgpu.h
index ef195f293000..98b0cf8bd83d 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu.h
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu.h
@@ -288,6 +288,9 @@ struct amdgpu_ip_block_version {
 	const struct amd_ip_funcs *funcs;
 };
 
+#define HW_REV(_Major, _Minor, _Rev) \
+	((((uint32_t) (_Major)) << 16) | ((uint32_t) (_Minor) << 8) | ((uint32_t) (_Rev)))
+
 struct amdgpu_ip_block {
 	struct amdgpu_ip_block_status status;
 	const struct amdgpu_ip_block_version *version;
@@ -627,6 +630,11 @@ struct amdgpu_fw_vram_usage {
 	u64 size;
 	struct amdgpu_bo *reserved_bo;
 	void *va;
+
+	/* Offset on the top of VRAM, used as c2p write buffer.
+	*/
+	u64 mem_train_fb_loc;
+	bool mem_train_support;
 };
 
 /*
diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_atombios.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_atombios.c
index 1c9d40f97a9b..72232fccf61a 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_atombios.c
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_atombios.c
@@ -2038,6 +2038,11 @@ int amdgpu_atombios_init(struct amdgpu_device *adev)
 	if (adev->is_atom_fw) {
 		amdgpu_atomfirmware_scratch_regs_init(adev);
 		amdgpu_atomfirmware_allocate_fb_scratch(adev);
+		ret = amdgpu_atomfirmware_get_mem_train_fb_loc(adev);
+		if (ret) {
+			DRM_ERROR("Failed to get mem train fb location.\n");
+			return ret;
+		}
 	} else {
 		amdgpu_atombios_scratch_regs_init(adev);
 		amdgpu_atombios_allocate_fb_scratch(adev);
diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_atomfirmware.c b/drivers/gpu/drm/amd/amdgpu/amdgpu_atomfirmware.c
index 39fd8ae5a822..ff4eb96bdfb5 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_atomfirmware.c
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_atomfirmware.c
@@ -27,6 +27,7 @@
 #include "amdgpu_atomfirmware.h"
 #include "atom.h"
 #include "atombios.h"
+#include "soc15_hw_ip.h"
 
 bool amdgpu_atomfirmware_gpu_supports_virtualization(struct amdgpu_device *adev)
 {
@@ -462,3 +463,138 @@ int amdgpu_atomfirmware_get_gfx_info(struct amdgpu_device *adev)
 	}
 	return -EINVAL;
 }
+
+/*
+ * Check if VBIOS supports GDDR6 training data save/restore
+ */
+static bool gddr6_mem_train_vbios_support(struct amdgpu_device *adev)
+{
+	uint16_t data_offset;
+	int index;
+
+	index = get_index_into_master_table(atom_master_list_of_data_tables_v2_1,
+					    firmwareinfo);
+	if (amdgpu_atom_parse_data_header(adev->mode_info.atom_context, index, NULL,
+					  NULL, NULL, &data_offset)) {
+		struct atom_firmware_info_v3_1 *firmware_info =
+			(struct atom_firmware_info_v3_1 *)(adev->mode_info.atom_context->bios +
+							   data_offset);
+
+		DRM_DEBUG("atom firmware capability:0x%08x.\n",
+			  le32_to_cpu(firmware_info->firmware_capability));
+
+		if (le32_to_cpu(firmware_info->firmware_capability) &
+		    ATOM_FIRMWARE_CAP_ENABLE_2STAGE_BIST_TRAINING)
+			return true;
+	}
+
+	return false;
+}
+
+static int gddr6_mem_train_support(struct amdgpu_device *adev)
+{
+	int ret;
+	uint32_t major, minor, revision, hw_v;
+
+	if (gddr6_mem_train_vbios_support(adev)) {
+		amdgpu_discovery_get_ip_version(adev, MP0_HWID, &major, &minor, &revision);
+		hw_v = HW_REV(major, minor, revision);
+		/*
+		 * treat 0 revision as a special case since register for MP0 and MMHUB is missing
+		 * for some Navi10 A0, preventing driver from discovering the hwip information since
+		 * none of the functions will be initialized, it should not cause any problems
+		 */
+		switch (hw_v) {
+		case HW_REV(11, 0, 0):
+		case HW_REV(11, 0, 5):
+			ret = 1;
+			break;
+		default:
+			DRM_ERROR("memory training vbios supports but psp hw(%08x)"
+				  " doesn't support!\n", hw_v);
+			ret = -1;
+			break;
+		}
+	} else {
+		ret = 0;
+		hw_v = -1;
+	}
+
+
+	DRM_DEBUG("mp0 hw_v %08x, ret:%d.\n", hw_v, ret);
+	return ret;
+}
+
+int amdgpu_atomfirmware_get_mem_train_fb_loc(struct amdgpu_device *adev)
+{
+	struct atom_context *ctx = adev->mode_info.atom_context;
+	unsigned char *bios = ctx->bios;
+	struct vram_reserve_block *reserved_block;
+	int index, block_number;
+	uint8_t frev, crev;
+	uint16_t data_offset, size;
+	uint32_t start_address_in_kb;
+	uint64_t offset;
+	int ret;
+
+	adev->fw_vram_usage.mem_train_support = false;
+
+	if (adev->asic_type != CHIP_NAVI10 &&
+	    adev->asic_type != CHIP_NAVI14)
+		return 0;
+
+	if (amdgpu_sriov_vf(adev))
+		return 0;
+
+	ret = gddr6_mem_train_support(adev);
+	if (ret == -1)
+		return -EINVAL;
+	else if (ret == 0)
+		return 0;
+
+	index = get_index_into_master_table(atom_master_list_of_data_tables_v2_1,
+					    vram_usagebyfirmware);
+	ret = amdgpu_atom_parse_data_header(ctx, index, &size, &frev, &crev,
+					    &data_offset);
+	if (ret == 0) {
+		DRM_ERROR("parse data header failed.\n");
+		return -EINVAL;
+	}
+
+	DRM_DEBUG("atom firmware common table header size:0x%04x, frev:0x%02x,"
+		  " crev:0x%02x, data_offset:0x%04x.\n", size, frev, crev, data_offset);
+	/* only support 2.1+ */
+	if (((uint16_t)frev << 8 | crev) < 0x0201) {
+		DRM_ERROR("frev:0x%02x, crev:0x%02x < 2.1 !\n", frev, crev);
+		return -EINVAL;
+	}
+
+	reserved_block = (struct vram_reserve_block *)
+		(bios + data_offset + sizeof(struct atom_common_table_header));
+	block_number = ((unsigned int)size - sizeof(struct atom_common_table_header))
+		/ sizeof(struct vram_reserve_block);
+	reserved_block += (block_number > 0) ? block_number-1 : 0;
+	DRM_DEBUG("block_number:0x%04x, last block: 0x%08xkb sz, %dkb fw, %dkb drv.\n",
+		  block_number,
+		  le32_to_cpu(reserved_block->start_address_in_kb),
+		  le16_to_cpu(reserved_block->used_by_firmware_in_kb),
+		  le16_to_cpu(reserved_block->used_by_driver_in_kb));
+	if (reserved_block->used_by_firmware_in_kb > 0) {
+		start_address_in_kb = le32_to_cpu(reserved_block->start_address_in_kb);
+		offset = (uint64_t)start_address_in_kb * ONE_KiB;
+		if ((offset & (ONE_MiB - 1)) < (4 * ONE_KiB + 1) ) {
+			offset -= ONE_MiB;
+		}
+
+		offset &= ~(ONE_MiB - 1);
+		adev->fw_vram_usage.mem_train_fb_loc = offset;
+		adev->fw_vram_usage.mem_train_support = true;
+		DRM_DEBUG("mem_train_fb_loc:0x%09llx.\n", offset);
+		ret = 0;
+	} else {
+		DRM_ERROR("used_by_firmware_in_kb is 0!\n");
+		ret = -EINVAL;
+	}
+
+	return ret;
+}
diff --git a/drivers/gpu/drm/amd/amdgpu/amdgpu_atomfirmware.h b/drivers/gpu/drm/amd/amdgpu/amdgpu_atomfirmware.h
index 53449fc7baf4..f871af5ea6f3 100644
--- a/drivers/gpu/drm/amd/amdgpu/amdgpu_atomfirmware.h
+++ b/drivers/gpu/drm/amd/amdgpu/amdgpu_atomfirmware.h
@@ -31,6 +31,7 @@ void amdgpu_atomfirmware_scratch_regs_init(struct amdgpu_device *adev);
 int amdgpu_atomfirmware_allocate_fb_scratch(struct amdgpu_device *adev);
 int amdgpu_atomfirmware_get_vram_info(struct amdgpu_device *adev,
 	int *vram_width, int *vram_type, int *vram_vendor);
+int amdgpu_atomfirmware_get_mem_train_fb_loc(struct amdgpu_device *adev);
 int amdgpu_atomfirmware_get_clock_info(struct amdgpu_device *adev);
 int amdgpu_atomfirmware_get_gfx_info(struct amdgpu_device *adev);
 bool amdgpu_atomfirmware_mem_ecc_supported(struct amdgpu_device *adev);