1 files changed, 51 insertions, 19 deletions

diff --git a/lib/utils/suspend/fdt_suspend_andes_atcsmu.c b/lib/utils/suspend/fdt_suspend_andes_atcsmu.c
index 072accc0..86cdf955 100644
--- a/lib/utils/suspend/fdt_suspend_andes_atcsmu.c
+++ b/lib/utils/suspend/fdt_suspend_andes_atcsmu.c
@@ -12,22 +12,37 @@
 #include <sbi/sbi_ecall_interface.h>
 #include <sbi/sbi_hart.h>
 #include <sbi/sbi_system.h>
+#include <sbi/sbi_timer.h>
 #include <sbi_utils/cache/fdt_cmo_helper.h>
 #include <sbi_utils/fdt/fdt_driver.h>
 #include <sbi_utils/fdt/fdt_helper.h>
 #include <sbi_utils/hsm/fdt_hsm_andes_atcsmu.h>
 
-static int check_secondary_harts_sleep(u32 hartid, bool deep_sleep)
+#define HART_SLEEP_TIMEOUT_MS 1000
+
+static int wait_secondary_harts_sleep(u32 hartid, bool deep_sleep)
 {
 	const struct sbi_domain *dom = &root;
 	unsigned long i;
 	u32 target;
+	struct atcsmu_sleep_arg arg;
+
+	arg.deep_sleep = deep_sleep;
 
-	/* Ensure the secondary harts entering the corresponding sleep state */
+	/* Wait for the secondary harts entering the corresponding sleep state */
 	sbi_hartmask_for_each_hartindex(i, dom->possible_harts) {
 		target = sbi_hartindex_to_hartid(i);
-		if (target != hartid && !atcsmu_pcs_is_sleep(target, deep_sleep))
-			return SBI_EFAIL;
+		if (target == hartid)
+			continue;
+
+		arg.hartid = target;
+		if (!sbi_timer_waitms_until(atcsmu_hart_is_sleep, &arg,
+					    HART_SLEEP_TIMEOUT_MS)) {
+			sbi_printf("ATCSMU: hart%u (PCS%u): timed out waiting for %s sleep\n",
+				   target, target + 3,
+				   deep_sleep ? "deep" : "light");
+			return SBI_ETIMEOUT;
+		}
 	}
 
 	return SBI_OK;
@@ -42,39 +57,45 @@ static int ae350_system_suspend_check(u32 sleep_type)
 static int ae350_system_suspend(u32 sleep_type, unsigned long addr)
 {
 	u32 hartid = current_hartid();
+	unsigned long saved_mie;
 	int rc;
 
 	/* Prevent the core leaving the WFI mode unexpectedly */
+	saved_mie = csr_read(CSR_MIE);
 	csr_write(CSR_MIE, 0);
 
-	/*
-	 * Only allow the S-mode external interrupts (UART2 and RTC alarm) to
-	 * wake up the primary hart
-	 */
-	csr_set(CSR_SIE, MIP_SEIP);
+	/* SMU wakes the primary hart on RTC alarm / UART2 */
 	atcsmu_set_wakeup_events(PCS_WAKEUP_RTC_ALARM_MASK | PCS_WAKEUP_UART2_MASK, hartid);
 
 	if (sleep_type == SBI_SUSP_AE350_LIGHT_SLEEP) {
-		rc = check_secondary_harts_sleep(hartid, false);
+		rc = wait_secondary_harts_sleep(hartid, false);
 		if (rc)
-			return rc;
+			goto err_restore_mie;
 
+		/* Clock-gated only: enable SEI to resume past the WFI */
+		csr_set(CSR_MIE, MIP_SEIP);
 		atcsmu_set_command(LIGHT_SLEEP_CMD, hartid);
 	} else if (sleep_type == SBI_SUSP_SLEEP_TYPE_SUSPEND) {
-		rc = check_secondary_harts_sleep(hartid, true);
+		rc = wait_secondary_harts_sleep(hartid, true);
 		if (rc)
-			return rc;
+			goto err_restore_mie;
 
-		atcsmu_set_command(DEEP_SLEEP_CMD, hartid);
 		rc = atcsmu_set_reset_vector((ulong)ae350_enable_coherency_warmboot, hartid);
 		if (rc)
-			return rc;
+			goto err_restore_mie;
 
 		ae350_non_ret_save(sbi_scratch_thishart_ptr());
-		fdt_cmo_llc_enable(false);
+
+		/* No LLC is fine; only fail on real errors */
+		rc = fdt_cmo_llc_enable(false);
+		if (rc && rc != SBI_ENODEV)
+			goto err_discard_save;
+
 		rc = fdt_cmo_llc_flush_all();
-		if (rc)
-			return rc;
+		if (rc && rc != SBI_ENODEV)
+			goto err_enable_llc;
+
+		atcsmu_set_command(DEEP_SLEEP_CMD, hartid);
 	}
 
 	ae350_disable_coherency();
@@ -84,6 +105,15 @@ static int ae350_system_suspend(u32 sleep_type, unsigned long addr)
 	ae350_enable_coherency();
 
 	return SBI_OK;
+
+err_enable_llc:
+	fdt_cmo_llc_enable(true);
+err_discard_save:
+	ae350_non_ret_discard(sbi_scratch_thishart_ptr());
+err_restore_mie:
+	csr_write(CSR_MIE, saved_mie);
+
+	return rc;
 }
 
 static void ae350_system_resume(void)
@@ -91,8 +121,10 @@ static void ae350_system_resume(void)
 	u32 hartid = current_hartid();
 	u32 sleep_type = atcsmu_get_sleep_type(hartid);
 
-	if (sleep_type == SBI_SUSP_SLEEP_TYPE_SUSPEND)
+	if (sleep_type == SBI_SUSP_SLEEP_TYPE_SUSPEND) {
 		fdt_cmo_llc_enable(true);
+		ae350_non_ret_restore(sbi_scratch_thishart_ptr());
+	}
 }
 
 static struct sbi_system_suspend_device suspend_andes_atcsmu = {