imx: keep boot device info for partition reboot

The boot device info should be kept during partition reboot,
the boot device ownership is changed after partition management
in ATF, so calling sc_pm_set_boot_parm() with boot device
parameter will fail in PSCI initialization phase, moving it
to bl31_early_platform_setup2() can make it work, correct them.

Signed-off-by: Anson Huang <Anson.Huang@nxp.com>
(cherry picked from commit d0c4a4e2c2d2618bef7c469c11ebeeb98174630a)
diff --git a/plat/imx/imx8qm/imx8qm_bl31_setup.c b/plat/imx/imx8qm/imx8qm_bl31_setup.c
index 11af036..ed5df85 100644
--- a/plat/imx/imx8qm/imx8qm_bl31_setup.c
+++ b/plat/imx/imx8qm/imx8qm_bl31_setup.c
@@ -405,6 +405,9 @@
 void bl31_early_platform_setup2(u_register_t arg0, u_register_t arg1,
 				u_register_t arg2, u_register_t arg3)
 {
+	u_register_t mpidr = read_mpidr_el1();
+	unsigned int cluster_id = MPIDR_AFFLVL1_VAL(mpidr);
+	sc_rsrc_t boot_dev;
 #if DEBUG_CONSOLE
 	static console_lpuart_t console;
 
@@ -413,6 +416,11 @@
 	if (sc_ipc_open(&ipc_handle, SC_IPC_BASE) != SC_ERR_NONE)
 		panic();
 
+	/* set primary CPU boot entry to BL31_BASE for partition reboot */
+	sc_misc_get_boot_dev(ipc_handle, &boot_dev);
+	sc_pm_set_boot_parm(ipc_handle, cluster_id == 0 ? SC_R_A53_0 : SC_R_A72_0,
+		BL31_BASE, SC_R_MU_0A, boot_dev);
+
 #if DEBUG_CONSOLE_A53
 	sc_pm_set_resource_power_mode(ipc_handle, SC_R_UART_0, SC_PM_PW_MODE_ON);
 	sc_pm_clock_rate_t rate = 80000000;
diff --git a/plat/imx/imx8qm/imx8qm_psci.c b/plat/imx/imx8qm/imx8qm_psci.c
index e41baf4..1536255 100644
--- a/plat/imx/imx8qm/imx8qm_psci.c
+++ b/plat/imx/imx8qm/imx8qm_psci.c
@@ -300,9 +300,6 @@
 int plat_setup_psci_ops(uintptr_t sec_entrypoint,
 			const plat_psci_ops_t **psci_ops)
 {
-	u_register_t mpidr = read_mpidr_el1();
-	unsigned int cluster_id = MPIDR_AFFLVL1_VAL(mpidr);
-
 	imx_mailbox_init(sec_entrypoint);
 	*psci_ops = &imx_plat_psci_ops;
 
@@ -324,12 +321,5 @@
 	sc_pm_req_sys_if_power_mode(ipc_handle, SC_R_A72, SC_PM_SYS_IF_INTERCONNECT,
 		SC_PM_PW_MODE_ON, SC_PM_PW_MODE_ON);
 
-	/*
-	 * set partition reboot address for primary CPU, boot device is NOT owned
-	 * by ATF, so pass 0 here
-	 */
-	sc_pm_set_boot_parm(ipc_handle, cluster_id == 0 ? SC_R_A53_0 : SC_R_A72_0,
-		BL31_BASE, SC_R_MU_0A, 0);
-
 	return 0;
 }
diff --git a/plat/imx/imx8qx/imx8qx_bl31_setup.c b/plat/imx/imx8qx/imx8qx_bl31_setup.c
index ef81738..259c683 100644
--- a/plat/imx/imx8qx/imx8qx_bl31_setup.c
+++ b/plat/imx/imx8qx/imx8qx_bl31_setup.c
@@ -383,6 +383,7 @@
 void bl31_early_platform_setup2(u_register_t arg0, u_register_t arg1,
 				u_register_t arg2, u_register_t arg3)
 {
+	sc_rsrc_t boot_dev;
 #if DEBUG_CONSOLE
 	static console_lpuart_t console;
 
@@ -391,6 +392,10 @@
 	if (sc_ipc_open(&ipc_handle, SC_IPC_BASE) != SC_ERR_NONE)
 		panic();
 
+	/* set primary CPU boot entry to BL31_BASE for partition reboot */
+	sc_misc_get_boot_dev(ipc_handle, &boot_dev);
+	sc_pm_set_boot_parm(ipc_handle, SC_R_A35_0, BL31_BASE, SC_R_MU_0A, boot_dev);
+
 #if DEBUG_CONSOLE_A35
 	sc_pm_set_resource_power_mode(ipc_handle, SC_R_UART_0, SC_PM_PW_MODE_ON);
 	sc_pm_clock_rate_t rate = 80000000;
diff --git a/plat/imx/imx8qx/imx8qx_psci.c b/plat/imx/imx8qx/imx8qx_psci.c
index b6b3a95..69addd1 100644
--- a/plat/imx/imx8qx/imx8qx_psci.c
+++ b/plat/imx/imx8qx/imx8qx_psci.c
@@ -238,11 +238,5 @@
 	sc_pm_req_sys_if_power_mode(ipc_handle, SC_R_A35, SC_PM_SYS_IF_INTERCONNECT,
 		SC_PM_PW_MODE_ON, SC_PM_PW_MODE_ON);
 
-	/*
-	 * set partition reboot address for primary CPU, boot device is NOT owned
-	 * by ATF, so pass 0 here
-	 */
-	sc_pm_set_boot_parm(ipc_handle, SC_R_A35_0, BL31_BASE, SC_R_MU_0A, 0);
-
 	return 0;
 }