return 0;
 }
 
+static int am33xx_do_sram_idle(u32 wfi_flags)
+{
+       int ret = 0;
+
+       if (!m3_ipc || !pm_ops)
+               return 0;
+
+       if (wfi_flags & WFI_FLAG_WAKE_M3)
+               ret = m3_ipc->ops->prepare_low_power(m3_ipc, WKUP_M3_IDLE);
+
+       return pm_ops->cpu_suspend(am33xx_do_wfi_sram, wfi_flags);
+}
+
 static int __init am43xx_map_gic(void)
 {
        gic_dist_base = ioremap(AM43XX_GIC_DIST_BASE, SZ_4K);
                rtc_only_idle = 0;
        }
 
+       pm_ops->begin_suspend();
+
        switch (state) {
        case PM_SUSPEND_MEM:
                ret = m3_ipc->ops->prepare_low_power(m3_ipc, WKUP_M3_DEEPSLEEP);
        }
 
        rtc_only_idle = 0;
+
+       pm_ops->finish_suspend();
 }
 
 static int am33xx_pm_valid(suspend_state_t state)
        suspend_wfi_flags |= WFI_FLAG_WAKE_M3;
 #endif /* CONFIG_SUSPEND */
 
-       ret = pm_ops->init(NULL);
+       ret = pm_ops->init(am33xx_do_sram_idle);
        if (ret) {
                dev_err(dev, "Unable to call core pm init!\n");
                ret = -ENODEV;
 
 static int am33xx_pm_remove(struct platform_device *pdev)
 {
+       if (pm_ops->deinit)
+               pm_ops->deinit();
        suspend_set_ops(NULL);
        wkup_m3_ipc_put(m3_ipc);
        am33xx_pm_free_sram();