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;
 }