MLK-16750-5: arm: imx: support using psci to handle power stuff

Support using PSCI to handle Power stuff.
Use PSCI to differentiate secure/non-secure kernel.
i.MX7 LPSR mode not implemented now.

Signed-off-by: Peng Fan <peng.fan@nxp.com>
Reviewed-by: Anson Huang <Anson.Huang@nxp.com>

6 files changed, 338 insertions, 145 deletions

diff --git a/arch/arm/mach-imx/cpuidle-imx6sx.c b/arch/arm/mach-imx/cpuidle-imx6sx.c
index be73709ebce5..b16bee20fa14 100644
--- a/arch/arm/mach-imx/cpuidle-imx6sx.c
+++ b/arch/arm/mach-imx/cpuidle-imx6sx.c
@@ -13,6 +13,7 @@
 #include <linux/genalloc.h>
 #include <linux/interrupt.h>
 #include <linux/module.h>
+#include <linux/psci.h>
 #include <asm/cacheflush.h>
 #include <asm/cpuidle.h>
 #include <asm/fncpy.h>
@@ -21,6 +22,8 @@
 #include <asm/suspend.h>
 #include <asm/tlb.h>
 
+#include <uapi/linux/psci.h>
+
 #include "common.h"
 #include "cpuidle.h"
 #include "hardware.h"
@@ -86,6 +89,11 @@ struct imx6_cpuidle_pm_info {
 
 static void (*imx6sx_wfi_in_iram_fn)(void __iomem *iram_vbase);
 
+#define MX6SX_POWERDWN_IDLE_PARAM	\
+	((1 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_POWER_DOWN << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
 static int imx6_idle_finish(unsigned long val)
 {
 	/*
@@ -97,7 +105,11 @@ static int imx6_idle_finish(unsigned long val)
 	 * just call flush_cache_all() is fine.
 	 */
 	flush_cache_all();
-	imx6sx_wfi_in_iram_fn(wfi_iram_base);
+	if (psci_ops.cpu_suspend)
+		psci_ops.cpu_suspend(MX6SX_POWERDWN_IDLE_PARAM,
+				     __pa(cpu_resume));
+	else
+		imx6sx_wfi_in_iram_fn(wfi_iram_base);
 
 	return 0;
 }
diff --git a/arch/arm/mach-imx/cpuidle-imx6ul.c b/arch/arm/mach-imx/cpuidle-imx6ul.c
index 7708d878615c..4f22b8f0d02b 100644
--- a/arch/arm/mach-imx/cpuidle-imx6ul.c
+++ b/arch/arm/mach-imx/cpuidle-imx6ul.c
@@ -12,6 +12,7 @@
 #include <linux/delay.h>
 #include <linux/module.h>
 #include <linux/platform_device.h>
+#include <linux/psci.h>
 #include <linux/regulator/consumer.h>
 #include <linux/regulator/driver.h>
 #include <linux/regulator/machine.h>
@@ -20,6 +21,8 @@
 #include <asm/proc-fns.h>
 #include <asm/suspend.h>
 
+#include <uapi/linux/psci.h>
+
 #include "common.h"
 #include "cpuidle.h"
 #include "hardware.h"
@@ -83,9 +86,18 @@ static const u32 imx6ul_mmdc_io_offset[] __initconst = {
 
 static void (*imx6ul_wfi_in_iram_fn)(void __iomem *iram_vbase);
 
+#define MX6UL_POWERDWN_IDLE_PARAM	\
+	((1 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_POWER_DOWN << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
 static int imx6ul_idle_finish(unsigned long val)
 {
-	imx6ul_wfi_in_iram_fn(wfi_iram_base);
+	if (psci_ops.cpu_suspend)
+		psci_ops.cpu_suspend(MX6UL_POWERDWN_IDLE_PARAM,
+				     __pa(cpu_resume));
+	else
+		imx6ul_wfi_in_iram_fn(wfi_iram_base);
 
 	return 0;
 }
diff --git a/arch/arm/mach-imx/cpuidle-imx7d.c b/arch/arm/mach-imx/cpuidle-imx7d.c
index 36ad21a33dfa..9b0fd7d301f3 100644
--- a/arch/arm/mach-imx/cpuidle-imx7d.c
+++ b/arch/arm/mach-imx/cpuidle-imx7d.c
@@ -16,6 +16,7 @@
 #include <linux/of.h>
 #include <linux/of_address.h>
 #include <linux/of_device.h>
+#include <linux/psci.h>
 #include <asm/cp15.h>
 #include <asm/cpuidle.h>
 #include <asm/fncpy.h>
@@ -24,6 +25,8 @@
 #include <asm/suspend.h>
 #include <asm/tlb.h>
 
+#include <uapi/linux/psci.h>
+
 #include "common.h"
 #include "cpuidle.h"
 #include "hardware.h"
@@ -77,11 +80,22 @@ struct imx7_cpuidle_pm_info {
 	struct imx7_pm_base gic_dist_base;
 } __aligned(8);
 
+static atomic_t master_lpi = ATOMIC_INIT(0);
 static atomic_t master_wait = ATOMIC_INIT(0);
 
 static void (*imx7d_wfi_in_iram_fn)(void __iomem *iram_vbase);
 static struct imx7_cpuidle_pm_info *cpuidle_pm_info;
 
+#define MX7D_POWERDWN_IDLE_PARAM	\
+	((1 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_POWER_DOWN << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
+#define MX7D_STANDBY_IDLE_PARAM	\
+	((1 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_STANDBY << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
 /* Mapped for the kernel, unlike cpuidle_pm_info->gic_dist_base.vbase */
 static void __iomem *imx7d_cpuidle_gic_base;
 
@@ -119,7 +133,11 @@ static void imx_pen_unlock(int cpu)
 
 static int imx7d_idle_finish(unsigned long val)
 {
-	imx7d_wfi_in_iram_fn(wfi_iram_base);
+	if (psci_ops.cpu_suspend)
+		psci_ops.cpu_suspend(MX7D_POWERDWN_IDLE_PARAM, __pa(cpu_resume));
+	else
+		imx7d_wfi_in_iram_fn(wfi_iram_base);
+
 	return 0;
 }
 
@@ -133,6 +151,7 @@ static bool imx7d_gic_sgis_pending(void)
 		readl_relaxed(sgip_base + 0xc));
 }
 
+static DEFINE_SPINLOCK(psci_lock);
 static int imx7d_enter_low_power_idle(struct cpuidle_device *dev,
 			    struct cpuidle_driver *drv, int index)
 {
@@ -148,41 +167,72 @@ static int imx7d_enter_low_power_idle(struct cpuidle_device *dev,
 		atomic_dec(&master_wait);
 		imx_gpcv2_set_lpm_mode(WAIT_CLOCKED);
 	} else {
-		imx_pen_lock(dev->cpu);
-		++cpuidle_pm_info->num_lpi_cpus;
-		cpu_pm_enter();
-		if (cpuidle_pm_info->num_lpi_cpus ==
-				cpuidle_pm_info->num_online_cpus) {
-			/*
-			 * GPC will not wake on SGIs so check for them
-			 * manually here. At this point we know the other cpu
-			 * is in wfi or waiting for the lock and can't send
-			 * any additional IPIs.
-			 */
-			if (imx7d_gic_sgis_pending()) {
-				index = -1;
-				goto skip_lpi_flow;
+		if (psci_ops.cpu_suspend) {
+			cpu_pm_enter();
+			spin_lock(&psci_lock);
+			if (atomic_inc_return(&master_lpi) == num_online_cpus()) {
+				if (imx7d_gic_sgis_pending()) {
+					index = -1;
+					goto psci_skip_lpi_flow;
+				}
+
+				imx_gpcv2_set_lpm_mode(WAIT_UNCLOCKED);
+				imx_gpcv2_set_cpu_power_gate_in_idle(true);
+
+				cpu_cluster_pm_enter();
 			}
-			imx_gpcv2_set_lpm_mode(WAIT_UNCLOCKED);
-			imx_gpcv2_set_cpu_power_gate_in_idle(true);
-			cpu_cluster_pm_enter();
+			spin_unlock(&psci_lock);
+
+			cpu_suspend(0, imx7d_idle_finish);
+
+			spin_lock(&psci_lock);
+			if (atomic_read(&master_lpi) == num_online_cpus()) {
+				cpu_cluster_pm_exit();
+				imx_gpcv2_set_cpu_power_gate_in_idle(false);
+				imx_gpcv2_set_lpm_mode(WAIT_CLOCKED);
+			}
+
+			atomic_dec(&master_lpi);
+psci_skip_lpi_flow:
+			spin_unlock(&psci_lock);
+			cpu_pm_exit();
 		} else {
-			imx_set_cpu_jump(dev->cpu, ca7_cpu_resume);
-		}
+			imx_pen_lock(dev->cpu);
+			++cpuidle_pm_info->num_lpi_cpus;
+			cpu_pm_enter();
+			if (cpuidle_pm_info->num_lpi_cpus ==
+					cpuidle_pm_info->num_online_cpus) {
+				/*
+				 * GPC will not wake on SGIs so check for them
+				 * manually here. At this point we know the other cpu
+				 * is in wfi or waiting for the lock and can't send
+				 * any additional IPIs.
+				 */
+				if (imx7d_gic_sgis_pending()) {
+					index = -1;
+					goto skip_lpi_flow;
+				}
+				imx_gpcv2_set_lpm_mode(WAIT_UNCLOCKED);
+				imx_gpcv2_set_cpu_power_gate_in_idle(true);
+				cpu_cluster_pm_enter();
+			} else {
+				imx_set_cpu_jump(dev->cpu, ca7_cpu_resume);
+			}
 
-		cpu_suspend(0, imx7d_idle_finish);
+			cpu_suspend(0, imx7d_idle_finish);
 
-		if (cpuidle_pm_info->num_lpi_cpus ==
-				cpuidle_pm_info->num_online_cpus) {
-			cpu_cluster_pm_exit();
-			imx_gpcv2_set_cpu_power_gate_in_idle(false);
-			imx_gpcv2_set_lpm_mode(WAIT_CLOCKED);
-		}
+			if (cpuidle_pm_info->num_lpi_cpus ==
+					cpuidle_pm_info->num_online_cpus) {
+				cpu_cluster_pm_exit();
+				imx_gpcv2_set_cpu_power_gate_in_idle(false);
+				imx_gpcv2_set_lpm_mode(WAIT_CLOCKED);
+			}
 
 skip_lpi_flow:
-		cpu_pm_exit();
-		--cpuidle_pm_info->num_lpi_cpus;
-		imx_pen_unlock(dev->cpu);
+			cpu_pm_exit();
+			--cpuidle_pm_info->num_lpi_cpus;
+			imx_pen_unlock(dev->cpu);
+		}
 	}
 
 	return index;
@@ -347,10 +397,12 @@ int __init imx7d_cpuidle_init(void)
 	register_hotcpu_notifier(&cpu_hotplug_notifier);
 #endif
 	/* code size should include cpuidle_pm_info size */
-	imx7d_wfi_in_iram_fn = (void *)fncpy(wfi_iram_base +
-		sizeof(*cpuidle_pm_info),
-		&imx7d_low_power_idle,
-		MX7_CPUIDLE_OCRAM_SIZE - sizeof(*cpuidle_pm_info));
+	if (!psci_ops.cpu_suspend) {
+		imx7d_wfi_in_iram_fn = (void *)fncpy(wfi_iram_base +
+			sizeof(*cpuidle_pm_info),
+			&imx7d_low_power_idle,
+			MX7_CPUIDLE_OCRAM_SIZE - sizeof(*cpuidle_pm_info));
+	}
 
 	return cpuidle_register(&imx7d_cpuidle_driver, NULL);
 }
diff --git a/arch/arm/mach-imx/pm-imx6.c b/arch/arm/mach-imx/pm-imx6.c
index 41e148ee0b3c..7b1f7acddfd6 100644
--- a/arch/arm/mach-imx/pm-imx6.c
+++ b/arch/arm/mach-imx/pm-imx6.c
@@ -21,6 +21,7 @@
 #include <linux/of_address.h>
 #include <linux/of_fdt.h>
 #include <linux/of_platform.h>
+#include <linux/psci.h>
 #include <linux/regmap.h>
 #include <linux/slab.h>
 #include <linux/suspend.h>
@@ -31,6 +32,8 @@
 #include <asm/suspend.h>
 #include <asm/tlb.h>
 
+#include <uapi/linux/psci.h>
+
 #include "common.h"
 #include "hardware.h"
 
@@ -747,8 +750,18 @@ int imx6_set_lpm(enum mxc_cpu_pwr_mode mode)
 	return 0;
 }
 
+#define MX6Q_SUSPEND_PARAM	\
+	((0 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_POWER_DOWN << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
 static int imx6q_suspend_finish(unsigned long val)
 {
+	if (psci_ops.cpu_suspend) {
+		return psci_ops.cpu_suspend(MX6Q_SUSPEND_PARAM,
+					    __pa(cpu_resume));
+	}
+
 	if (!imx6_suspend_in_ocram_fn) {
 		cpu_do_idle();
 	} else {
@@ -993,6 +1006,11 @@ void __init imx6_pm_map_io(void)
 	 */
 	WARN_ON(of_scan_flat_dt(imx6_dt_find_lpsram, NULL));
 
+	/*
+	 * We moved suspend/resume and lowpower idle to TEE,
+	 * But busfreq now still in Linux, this table is still needed
+	 * If we later decide to move busfreq to TEE, we could drop this.
+	 */
 	/* Return if no IRAM space is allocated for suspend/resume code. */
 	if (!iram_tlb_base_addr) {
 		pr_warn("No IRAM/OCRAM memory allocated for suspend/resume \
@@ -1068,6 +1086,12 @@ static int __init imx6q_suspend_init(const struct imx6_pm_socdata *socdata)
 		return -EINVAL;
 	}
 
+	if (psci_ops.cpu_suspend) {
+		/* TODO: seems not needed */
+		/* of_node_put(node); */
+		return ret;
+	}
+
 	/*
 	 * 16KB is allocated for IRAM TLB, but only up 8k is for kernel TLB,
 	 * The lower 8K is not used, so use the lower 8K for IRAM code and
diff --git a/arch/arm/mach-imx/pm-imx7.c b/arch/arm/mach-imx/pm-imx7.c
index 540bc8effecf..e9f4d86f21aa 100644
--- a/arch/arm/mach-imx/pm-imx7.c
+++ b/arch/arm/mach-imx/pm-imx7.c
@@ -19,6 +19,7 @@
 #include <linux/of_device.h>
 #include <linux/of_fdt.h>
 #include <linux/of_irq.h>
+#include <linux/psci.h>
 #include <linux/slab.h>
 #include <linux/suspend.h>
 #include <linux/genalloc.h>
@@ -37,6 +38,8 @@
 #include <asm/suspend.h>
 #include <asm/tlb.h>
 
+#include <uapi/linux/psci.h>
+
 #include "common.h"
 #include "hardware.h"
 #include "cpuidle.h"
@@ -666,8 +669,29 @@ static void imx7_console_io_restore(void)
 	writel_relaxed(uart1_io[3], iomuxc_base + UART_TX_PAD);
 }
 
+#define MX7D_SUSPEND_POWERDWN_PARAM	\
+	((0 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_POWER_DOWN << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
+#define MX7D_SUSPEND_STANDBY_PARAM	\
+	((0 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_STANDBY << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
 static int imx7_suspend_finish(unsigned long val)
 {
+	u32 state;
+
+	if (val == 0)
+		state = MX7D_SUSPEND_POWERDWN_PARAM;
+	else
+		state = MX7D_SUSPEND_STANDBY_PARAM;
+
+	if (psci_ops.cpu_suspend) {
+		return psci_ops.cpu_suspend(state, __pa(cpu_resume));
+	}
+
 	if (!imx7_suspend_in_ocram_fn) {
 		cpu_do_idle();
 	} else {
@@ -726,7 +750,10 @@ static int imx7_pm_enter(suspend_state_t state)
 		imx_gpcv2_pre_suspend(false);
 
 		/* Zzz ... */
-		imx7_suspend_in_ocram_fn(suspend_ocram_base);
+		if (psci_ops.cpu_suspend)
+			cpu_suspend(1, imx7_suspend_finish);
+		else
+			imx7_suspend_in_ocram_fn(suspend_ocram_base);
 
 		imx_anatop_post_resume();
 		imx_gpcv2_post_resume();
@@ -886,6 +913,7 @@ void __init imx7_pm_map_io(void)
 		return;
 	}
 
+	/* TODO: Handle M4 in TEE? */
 	/* Set all entries to 0 except first 3 words reserved for M4. */
 	memset((void *)(iram_tlb_base_addr + M4_OCRAMS_RESERVED_SIZE),
 		0, MX7_IRAM_TLB_SIZE - M4_OCRAMS_RESERVED_SIZE);
@@ -974,7 +1002,13 @@ static int __init imx7_suspend_init(const struct imx7_pm_socdata *socdata)
 	/* Get the virtual address of the suspend code. */
 	suspend_ocram_base = (void *)IMX_IO_P2V(iram_paddr);
 
-	pm_info = suspend_ocram_base;
+	if (psci_ops.cpu_suspend) {
+		pm_info = kmalloc(sizeof(*pm_info), GFP_KERNEL);
+		if (!pm_info)
+			return -ENOMEM;
+	} else {
+		pm_info = suspend_ocram_base;
+	}
 	/* pbase points to iram_paddr. */
 	pm_info->pbase = iram_paddr;
 	pm_info->resume_addr = virt_to_phys(ca7_cpu_resume);
@@ -1057,6 +1091,9 @@ static int __init imx7_suspend_init(const struct imx7_pm_socdata *socdata)
 				ddrc_phy_offset_array[i][1];
 	}
 
+	if (psci_ops.cpu_suspend)
+		goto put_node;
+
 	imx7_suspend_in_ocram_fn = fncpy(
 		suspend_ocram_base + sizeof(*pm_info),
 		&imx7_suspend,
@@ -1123,6 +1160,9 @@ void __init imx7d_pm_init(void)
 	if (of_get_property(np, "fsl,enable-lpsr", NULL))
 		lpsr_enabled = true;
 
+	if (psci_ops.cpu_suspend)
+		lpsr_enabled = false;
+
 	if (lpsr_enabled) {
 		pr_info("LPSR mode enabled, DSM will go into LPSR mode!\n");
 		lpm_ocram_base = of_iomap(np, 0);
diff --git a/arch/arm/mach-imx/pm-imx7ulp.c b/arch/arm/mach-imx/pm-imx7ulp.c
index 15aa46837fb0..e5ac4229d0d4 100644
--- a/arch/arm/mach-imx/pm-imx7ulp.c
+++ b/arch/arm/mach-imx/pm-imx7ulp.c
@@ -21,6 +21,7 @@
 #include <linux/of_address.h>
 #include <linux/of_fdt.h>
 #include <linux/of_irq.h>
+#include <linux/psci.h>
 #include <linux/of_platform.h>
 #include <linux/regmap.h>
 #include <linux/slab.h>
@@ -32,6 +33,8 @@
 #include <asm/suspend.h>
 #include <asm/tlb.h>
 
+#include <uapi/linux/psci.h>
+
 #include "common.h"
 #include "hardware.h"
 
@@ -429,8 +432,29 @@ int imx7ulp_set_lpm(enum imx7ulp_cpu_pwr_mode mode)
 	return 0;
 }
 
+#define MX7ULP_SUSPEND_POWERDWN_PARAM	\
+	((0 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_POWER_DOWN << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
+#define MX7ULP_SUSPEND_STANDBY_PARAM	\
+	((0 << PSCI_0_2_POWER_STATE_ID_SHIFT) | \
+	 (1 << PSCI_0_2_POWER_STATE_AFFL_SHIFT) | \
+	 (PSCI_POWER_STATE_TYPE_STANDBY << PSCI_0_2_POWER_STATE_TYPE_SHIFT))
+
 static int imx7ulp_suspend_finish(unsigned long val)
 {
+	u32 state;
+
+	if (val == 0)
+		state = MX7ULP_SUSPEND_POWERDWN_PARAM;
+	else
+		state = MX7ULP_SUSPEND_STANDBY_PARAM;
+
+	if (psci_ops.cpu_suspend) {
+		return psci_ops.cpu_suspend(state, __pa(cpu_resume));
+	}
+
 	imx7ulp_suspend_in_ocram_fn(suspend_ocram_base);
 
 	return 0;
@@ -440,38 +464,48 @@ static int imx7ulp_pm_enter(suspend_state_t state)
 {
 	switch (state) {
 	case PM_SUSPEND_STANDBY:
-		imx7ulp_set_lpm(VLPS);
-		writel_relaxed(
-			readl_relaxed(pmc1_base + PMC_VLPS) | BM_VLPS_RBBEN,
-			pmc1_base + PMC_VLPS);
-
-		/* Zzz ... */
-		cpu_suspend(0, imx7ulp_suspend_finish);
-
-		writel_relaxed(
-			readl_relaxed(pmc1_base + PMC_VLPS) & ~BM_VLPS_RBBEN,
-			pmc1_base + PMC_VLPS);
-		imx7ulp_set_lpm(RUN);
+		if (psci_ops.cpu_suspend) {
+			/* Zzz ... */
+			cpu_suspend(1, imx7ulp_suspend_finish);
+		} else {
+			imx7ulp_set_lpm(VLPS);
+			writel_relaxed(
+				readl_relaxed(pmc1_base + PMC_VLPS) | BM_VLPS_RBBEN,
+				pmc1_base + PMC_VLPS);
+
+			/* Zzz ... */
+			cpu_suspend(0, imx7ulp_suspend_finish);
+
+			writel_relaxed(
+				readl_relaxed(pmc1_base + PMC_VLPS) & ~BM_VLPS_RBBEN,
+				pmc1_base + PMC_VLPS);
+			imx7ulp_set_lpm(RUN);
+		}
 		break;
 	case PM_SUSPEND_MEM:
-		imx7ulp_gpio_save();
-		imx7ulp_scg1_save();
-		imx7ulp_pcc2_save();
-		imx7ulp_pcc3_save();
-		imx7ulp_tpm_save();
-		imx7ulp_lpuart_save();
-		imx7ulp_iomuxc_save();
-		imx7ulp_set_lpm(VLLS);
-
-		/* Zzz ... */
-		cpu_suspend(0, imx7ulp_suspend_finish);
-
-		imx7ulp_pcc2_restore();
-		imx7ulp_pcc3_restore();
-		imx7ulp_lpuart_restore();
-		imx7ulp_set_dgo(0);
-		imx7ulp_tpm_restore();
-		imx7ulp_set_lpm(RUN);
+		if (psci_ops.cpu_suspend) {
+			/* Zzz ... */
+			cpu_suspend(0, imx7ulp_suspend_finish);
+		} else {
+			imx7ulp_gpio_save();
+			imx7ulp_scg1_save();
+			imx7ulp_pcc2_save();
+			imx7ulp_pcc3_save();
+			imx7ulp_tpm_save();
+			imx7ulp_lpuart_save();
+			imx7ulp_iomuxc_save();
+			imx7ulp_set_lpm(VLLS);
+
+			/* Zzz ... */
+			cpu_suspend(0, imx7ulp_suspend_finish);
+
+			imx7ulp_pcc2_restore();
+			imx7ulp_pcc3_restore();
+			imx7ulp_lpuart_restore();
+			imx7ulp_set_dgo(0);
+			imx7ulp_tpm_restore();
+			imx7ulp_set_lpm(RUN);
+	}
 		break;
 	default:
 		return -EINVAL;
@@ -531,6 +565,9 @@ static int __init imx7ulp_dt_find_lpsram(unsigned long node, const char *uname,
 
 void __init imx7ulp_pm_map_io(void)
 {
+	if (psci_ops.cpu_suspend) {
+		return;
+	}
 	/*
 	 * Get the address of IRAM or OCRAM to be used by the low
 	 * power code from the device tree.
@@ -557,63 +594,71 @@ void __init imx7ulp_pm_common_init(const struct imx7ulp_pm_socdata
 	unsigned long i, j;
 	int ret;
 
-	/*
-	 * Make sure the IRAM virtual address has a mapping in the IRAM
-	 * page table.
-	 *
-	 * Only use the top 12 bits [31-20] when storing the physical
-	 * address in the page table as only these bits are required
-	 * for 1M mapping.
-	 */
-	j = ((iram_tlb_base_addr >> 20) << 2) / 4;
-	*((unsigned long *)iram_tlb_base_addr + j) =
-		(iram_tlb_phys_addr & ADDR_1M_MASK) |
-		TT_ATTRIB_NON_CACHEABLE_1M;
-	/*
-	 * Make sure the AIPS1 virtual address has a mapping in the
-	 * IRAM page table.
-	 */
-	aips1_base = ioremap(MX7ULP_AIPS1_BASE_ADDR, SZ_1M);
-	j = (((u32)aips1_base >> 20) << 2) / 4;
-	*((unsigned long *)iram_tlb_base_addr + j) =
-		((MX7ULP_AIPS1_BASE_ADDR) & ADDR_1M_MASK) |
-		TT_ATTRIB_NON_CACHEABLE_1M;
-	/*
-	 * Make sure the AIPS2 virtual address has a mapping in the
-	 * IRAM page table.
-	 */
-	aips2_base = ioremap(MX7ULP_AIPS2_BASE_ADDR, SZ_1M);
-	j = (((u32)aips2_base >> 20) << 2) / 4;
-	*((unsigned long *)iram_tlb_base_addr + j) =
-		((MX7ULP_AIPS2_BASE_ADDR) & ADDR_1M_MASK) |
-		TT_ATTRIB_NON_CACHEABLE_1M;
-	/*
-	 * Make sure the AIPS3 virtual address has a mapping in the
-	 * IRAM page table.
-	 */
-	aips3_base = ioremap(MX7ULP_AIPS3_BASE_ADDR, SZ_1M);
-	j = (((u32)aips3_base >> 20) << 2) / 4;
-	*((unsigned long *)iram_tlb_base_addr + j) =
-		((MX7ULP_AIPS3_BASE_ADDR) & ADDR_1M_MASK) |
-		TT_ATTRIB_NON_CACHEABLE_1M;
-	/*
-	 * Make sure the AIPS4 virtual address has a mapping in the
-	 * IRAM page table.
-	 */
-	aips4_base = ioremap(MX7ULP_AIPS4_BASE_ADDR, SZ_1M);
-	j = (((u32)aips4_base >> 20) << 2) / 4;
-	*((unsigned long *)iram_tlb_base_addr + j) =
-		((MX7ULP_AIPS4_BASE_ADDR) & ADDR_1M_MASK) |
-		TT_ATTRIB_NON_CACHEABLE_1M;
-	/*
-	 * Make sure the AIPS5 virtual address has a mapping in the
-	 * IRAM page table.
-	 */
-	aips5_base = ioremap(MX7ULP_AIPS5_BASE_ADDR, SZ_1M);
-	j = (((u32)aips5_base >> 20) << 2) / 4;
-	*((unsigned long *)iram_tlb_base_addr + j) =
-		((MX7ULP_AIPS5_BASE_ADDR) & ADDR_1M_MASK) |
-		TT_ATTRIB_NON_CACHEABLE_1M;
+	if (psci_ops.cpu_suspend) {
+		aips1_base = ioremap(MX7ULP_AIPS1_BASE_ADDR, SZ_1M);
+		aips2_base = ioremap(MX7ULP_AIPS2_BASE_ADDR, SZ_1M);
+		aips3_base = ioremap(MX7ULP_AIPS3_BASE_ADDR, SZ_1M);
+		aips4_base = ioremap(MX7ULP_AIPS4_BASE_ADDR, SZ_1M);
+		aips5_base = ioremap(MX7ULP_AIPS5_BASE_ADDR, SZ_1M);
+	} else {
+		/*
+		 * Make sure the IRAM virtual address has a mapping in the IRAM
+		 * page table.
+		 *
+		 * Only use the top 12 bits [31-20] when storing the physical
+		 * address in the page table as only these bits are required
+		 * for 1M mapping.
+		 */
+		j = ((iram_tlb_base_addr >> 20) << 2) / 4;
+		*((unsigned long *)iram_tlb_base_addr + j) =
+			(iram_tlb_phys_addr & ADDR_1M_MASK) |
+			TT_ATTRIB_NON_CACHEABLE_1M;
+		/*
+		 * Make sure the AIPS1 virtual address has a mapping in the
+		 * IRAM page table.
+		 */
+		aips1_base = ioremap(MX7ULP_AIPS1_BASE_ADDR, SZ_1M);
+		j = (((u32)aips1_base >> 20) << 2) / 4;
+		*((unsigned long *)iram_tlb_base_addr + j) =
+			((MX7ULP_AIPS1_BASE_ADDR) & ADDR_1M_MASK) |
+			TT_ATTRIB_NON_CACHEABLE_1M;
+		/*
+		 * Make sure the AIPS2 virtual address has a mapping in the
+		 * IRAM page table.
+		 */
+		aips2_base = ioremap(MX7ULP_AIPS2_BASE_ADDR, SZ_1M);
+		j = (((u32)aips2_base >> 20) << 2) / 4;
+		*((unsigned long *)iram_tlb_base_addr + j) =
+			((MX7ULP_AIPS2_BASE_ADDR) & ADDR_1M_MASK) |
+			TT_ATTRIB_NON_CACHEABLE_1M;
+		/*
+		 * Make sure the AIPS3 virtual address has a mapping in the
+		 * IRAM page table.
+		 */
+		aips3_base = ioremap(MX7ULP_AIPS3_BASE_ADDR, SZ_1M);
+		j = (((u32)aips3_base >> 20) << 2) / 4;
+		*((unsigned long *)iram_tlb_base_addr + j) =
+			((MX7ULP_AIPS3_BASE_ADDR) & ADDR_1M_MASK) |
+			TT_ATTRIB_NON_CACHEABLE_1M;
+		/*
+		 * Make sure the AIPS4 virtual address has a mapping in the
+		 * IRAM page table.
+		 */
+		aips4_base = ioremap(MX7ULP_AIPS4_BASE_ADDR, SZ_1M);
+		j = (((u32)aips4_base >> 20) << 2) / 4;
+		*((unsigned long *)iram_tlb_base_addr + j) =
+			((MX7ULP_AIPS4_BASE_ADDR) & ADDR_1M_MASK) |
+			TT_ATTRIB_NON_CACHEABLE_1M;
+		/*
+		 * Make sure the AIPS5 virtual address has a mapping in the
+		 * IRAM page table.
+		 */
+		aips5_base = ioremap(MX7ULP_AIPS5_BASE_ADDR, SZ_1M);
+		j = (((u32)aips5_base >> 20) << 2) / 4;
+		*((unsigned long *)iram_tlb_base_addr + j) =
+			((MX7ULP_AIPS5_BASE_ADDR) & ADDR_1M_MASK) |
+			TT_ATTRIB_NON_CACHEABLE_1M;
+	}
 
 	np = of_find_compatible_node(NULL, NULL, "fsl,imx7ulp-smc1");
 	smc1_base = of_iomap(np, 0);
@@ -658,22 +703,28 @@ void __init imx7ulp_pm_common_init(const struct imx7ulp_pm_socdata
 		WARN_ON(!gpio_base[i]);
 	}
 
-	/*
-	 * 16KB is allocated for IRAM TLB, but only up 8k is for kernel TLB,
-	 * The lower 8K is not used, so use the lower 8K for IRAM code and
-	 * pm_info.
-	 *
-	 */
-	sram_paddr = iram_tlb_phys_addr;
-
-	/* Make sure sram_paddr is 8 byte aligned. */
-	if ((uintptr_t)(sram_paddr) & (FNCPY_ALIGN - 1))
-		sram_paddr += FNCPY_ALIGN - sram_paddr % (FNCPY_ALIGN);
-
-	/* Get the virtual address of the suspend code. */
-	suspend_ocram_base = (void *)IMX_IO_P2V(sram_paddr);
-
-	pm_info = suspend_ocram_base;
+	if (psci_ops.cpu_suspend) {
+		pm_info = kzalloc(SZ_16K, GFP_KERNEL);
+		if (!pm_info)
+			panic("pm info allocation failed\n");
+	} else {
+		/*
+		 * 16KB is allocated for IRAM TLB, but only up 8k is for kernel TLB,
+		 * The lower 8K is not used, so use the lower 8K for IRAM code and
+		 * pm_info.
+		 *
+		 */
+		sram_paddr = iram_tlb_phys_addr;
+
+		/* Make sure sram_paddr is 8 byte aligned. */
+		if ((uintptr_t)(sram_paddr) & (FNCPY_ALIGN - 1))
+			sram_paddr += FNCPY_ALIGN - sram_paddr % (FNCPY_ALIGN);
+
+		/* Get the virtual address of the suspend code. */
+		suspend_ocram_base = (void *)IMX_IO_P2V(sram_paddr);
+
+		pm_info = suspend_ocram_base;
+	}
 	pm_info->pbase = sram_paddr;
 	pm_info->resume_addr = virt_to_phys(imx7ulp_cpu_resume);
 	pm_info->pm_info_size = sizeof(*pm_info);
@@ -710,10 +761,12 @@ void __init imx7ulp_pm_common_init(const struct imx7ulp_pm_socdata
 	for (i = 0; i < pm_info->mmdc_num; i++)
 		pm_info->mmdc_val[i][1] = imx7ulp_lpddr3_script[i];
 
-	imx7ulp_suspend_in_ocram_fn = fncpy(
-		suspend_ocram_base + sizeof(*pm_info),
-		&imx7ulp_suspend,
-		MX7ULP_SUSPEND_OCRAM_SIZE - sizeof(*pm_info));
+	if (!psci_ops.cpu_suspend) {
+		imx7ulp_suspend_in_ocram_fn = fncpy(
+			suspend_ocram_base + sizeof(*pm_info),
+			&imx7ulp_suspend,
+			MX7ULP_SUSPEND_OCRAM_SIZE - sizeof(*pm_info));
+	}
 
 	if (IS_ENABLED(CONFIG_SUSPEND)) {
 		ret = imx7ulp_suspend_init();