fix: power: ace: Move HST domain suspend before IMR context save

This patch addresses an issue in the ACE platform power management code
where the HST domain suspend operation was performed after the IMR
context save. This resulted in the power management context being
restored with outdated values upon wake-up from D3 state, leading to a
failure to resume the HST domain correctly.

By moving the `pm_device_runtime_put(INTEL_ADSP_HST_DOMAIN_DEV)` call
before the IMR context save, we ensure that the HST domain is suspended
with the current context, and upon resume, the power management context
has the correct information to restore the HST domain state.

Signed-off-by: Tomasz Leman <[email protected]>

soc/intel/intel_adsp/ace/power.c

@@ -296,6 +296,8 @@ void pm_state_set(enum pm_state state, uint8_t substate_id)
 		DSPCS.bootctl[cpu].bctl &= ~DSPBR_BCTL_WAITIPCG;
 		if (cpu == 0) {
 			soc_cpus_active[cpu] = false;
+			ret = pm_device_runtime_put(INTEL_ADSP_HST_DOMAIN_DEV);
+			__ASSERT_NO_MSG(ret == 0);
 #ifdef CONFIG_ADSP_IMR_CONTEXT_SAVE
 			/* save storage and restore information to imr */
 			__ASSERT_NO_MSG(global_imr_ram_storage != NULL);
@@ -349,8 +351,6 @@ void pm_state_set(enum pm_state state, uint8_t substate_id)
 #else
 #define HPSRAM_MASK false
 #endif /* CONFIG_ADSP_POWER_DOWN_HPSRAM */
-			ret = pm_device_runtime_put(INTEL_ADSP_HST_DOMAIN_DEV);
-			__ASSERT_NO_MSG(ret == 0);
 			/* do power down - this function won't return */
 			power_down(true, HPSRAM_MASK, true);
 		} else {