imx: set CPU boot entry for partition reboot

With SPL running on OCRAM, when linux suspend, OCRAM
will lose power and if partition reboot is started from
SPL, system will hang as the OCRAM data lost, so for
partition reboot, the CPU boot entry can be set to
be from ATF BL31 entry directly, SCFW exposes such API
for this scenario.

Signed-off-by: Anson Huang <Anson.Huang@nxp.com>
(cherry picked from commit d4a0970cfea5f0dd6f88a5c202b90ce31e6baecc)
diff --git a/plat/imx/imx8qm/imx8qm_psci.c b/plat/imx/imx8qm/imx8qm_psci.c
index 1536255..e41baf4 100644
--- a/plat/imx/imx8qm/imx8qm_psci.c
+++ b/plat/imx/imx8qm/imx8qm_psci.c
@@ -300,6 +300,9 @@
 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;
 
@@ -321,5 +324,12 @@
 	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_psci.c b/plat/imx/imx8qx/imx8qx_psci.c
index 69addd1..b6b3a95 100644
--- a/plat/imx/imx8qx/imx8qx_psci.c
+++ b/plat/imx/imx8qx/imx8qx_psci.c
@@ -238,5 +238,11 @@
 	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;
 }