fix(imx8m): fix the rank to rank space issue

update umctl2's setting based on phy training CDD value
to workaround the rank-to-rank space issue.

Signed-off-by: Jacky Bai <ping.bai@nxp.com>
Reviewed-by: Anson Huang <anson.huang@nxp.com>
Change-Id: I0fab18cdc378fda760daa0f89c4dd84eb46f7e11
diff --git a/plat/imx/imx8m/ddr/dram.c b/plat/imx/imx8m/ddr/dram.c
index 53605cd..1fea69d 100644
--- a/plat/imx/imx8m/ddr/dram.c
+++ b/plat/imx/imx8m/ddr/dram.c
@@ -44,6 +44,25 @@
 	}
 }
 
+static void save_rank_setting(void)
+{
+	uint32_t i, offset;
+	uint32_t pstate_num = dram_info.num_fsp;
+
+	for (i = 0U; i < pstate_num; i++) {
+		offset = i ? (i + 1) * 0x1000 : 0U;
+		dram_info.rank_setting[i][0] = mmio_read_32(DDRC_DRAMTMG2(0) + offset);
+		if (dram_info.dram_type != DDRC_LPDDR4) {
+			dram_info.rank_setting[i][1] = mmio_read_32(DDRC_DRAMTMG9(0) + offset);
+		}
+#if !defined(PLAT_imx8mq)
+		dram_info.rank_setting[i][2] = mmio_read_32(DDRC_RANKCTL(0) + offset);
+#endif
+	}
+#if defined(PLAT_imx8mq)
+	dram_info.rank_setting[0][2] = mmio_read_32(DDRC_RANKCTL(0));
+#endif
+}
 /* Restore the ddrc configs */
 void dram_umctl2_init(struct dram_timing_info *timing)
 {
@@ -150,6 +169,9 @@
 	}
 	dram_info.num_fsp = i;
 
+	/* save the DRAMTMG2/9 for rank to rank workaround */
+	save_rank_setting();
+
 	/* check if has bypass mode support */
 	if (dram_info.timing_info->fsp_table[idx] < 666) {
 		dram_info.bypass_mode = true;
diff --git a/plat/imx/imx8m/ddr/dram_retention.c b/plat/imx/imx8m/ddr/dram_retention.c
index 7d4f823..96591d1 100644
--- a/plat/imx/imx8m/ddr/dram_retention.c
+++ b/plat/imx/imx8m/ddr/dram_retention.c
@@ -1,5 +1,5 @@
 /*
- * Copyright 2018-2022 NXP
+ * Copyright 2018-2023 NXP
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -25,6 +25,28 @@
 
 #define DBGCAM_EMPTY		0x36000000
 
+static void rank_setting_update(void)
+{
+	uint32_t i, offset;
+	uint32_t pstate_num = dram_info.num_fsp;
+
+	for (i = 0U; i < pstate_num; i++) {
+		offset = i ? (i + 1) * 0x1000 : 0U;
+		mmio_write_32(DDRC_DRAMTMG2(0) + offset, dram_info.rank_setting[i][0]);
+		if (dram_info.dram_type != DDRC_LPDDR4) {
+			mmio_write_32(DDRC_DRAMTMG9(0) + offset, dram_info.rank_setting[i][1]);
+		}
+
+#if !defined(PLAT_imx8mq)
+		mmio_write_32(DDRC_RANKCTL(0) + offset,
+			dram_info.rank_setting[i][2]);
+#endif
+	}
+#if defined(PLAT_imx8mq)
+		mmio_write_32(DDRC_RANKCTL(0), dram_info.rank_setting[0][2]);
+#endif
+}
+
 void dram_enter_retention(void)
 {
 	/* Wait DBGCAM to be empty */
@@ -157,6 +179,9 @@
 	/* dram phy re-init */
 	dram_phy_init(dram_info.timing_info);
 
+	/* workaround for rank-to-rank issue */
+	rank_setting_update();
+
 	/* DWC_DDRPHYA_APBONLY0_MicroContMuxSel */
 	dwc_ddrphy_apb_wr(0xd0000, 0x0);
 	while (dwc_ddrphy_apb_rd(0x20097)) {
diff --git a/plat/imx/imx8m/include/dram.h b/plat/imx/imx8m/include/dram.h
index c9f18e8..db43352 100644
--- a/plat/imx/imx8m/include/dram.h
+++ b/plat/imx/imx8m/include/dram.h
@@ -59,6 +59,8 @@
 	struct dram_timing_info *timing_info;
 	/* mr, emr, emr2, emr3, mr11, mr12, mr22, mr14 */
 	uint32_t mr_table[3][8];
+	/* used for workaround for rank to rank issue */
+	uint32_t rank_setting[3][3];
 };
 
 extern struct dram_info dram_info;