From e8a327329a55124a553e75491b776871a7a82f50 Mon Sep 17 00:00:00 2001
From: Adam Kondraciuk <adam.kondraciuk@nordicsemi.no>
Date: Tue, 9 Jul 2024 10:02:46 +0200
Subject: [PATCH] drivers: spi: nrfx_spim: Add support for device runtime PM

Enable the device runtime power management on the SPIM shim.

Signed-off-by: Adam Kondraciuk <adam.kondraciuk@nordicsemi.no>
Signed-off-by: Nikodem Kastelik <nikodem.kastelik@nordicsemi.no>
---
 drivers/spi/spi_nrfx_spim.c | 73 ++++++++++++++++++-------------------
 1 file changed, 36 insertions(+), 37 deletions(-)

diff --git a/drivers/spi/spi_nrfx_spim.c b/drivers/spi/spi_nrfx_spim.c
index 3fb5468b542a8b..f8d5b60330ee56 100644
--- a/drivers/spi/spi_nrfx_spim.c
+++ b/drivers/spi/spi_nrfx_spim.c
@@ -8,6 +8,7 @@
 #include <zephyr/drivers/spi/rtio.h>
 #include <zephyr/cache.h>
 #include <zephyr/pm/device.h>
+#include <zephyr/pm/device_runtime.h>
 #include <zephyr/drivers/pinctrl.h>
 #include <zephyr/mem_mgmt/mem_attr.h>
 #include <soc.h>
@@ -90,6 +91,8 @@ static inline void finalize_spi_transaction(const struct device *dev, bool deact
 	if (NRF_SPIM_IS_320MHZ_SPIM(reg) && !(dev_data->ctx.config->operation & SPI_HOLD_ON_CS)) {
 		nrfy_spim_disable(reg);
 	}
+
+	pm_device_runtime_put_async(dev, K_NO_WAIT);
 }
 
 static inline uint32_t get_nrf_spim_frequency(uint32_t frequency)
@@ -460,6 +463,7 @@ static int transceive(const struct device *dev,
 	void *reg = dev_config->spim.p_reg;
 	int error;
 
+	pm_device_runtime_get(dev);
 	spi_context_lock(&dev_data->ctx, asynchronous, cb, userdata, spi_cfg);
 
 	error = configure(dev, spi_cfg);
@@ -571,56 +575,49 @@ static DEVICE_API(spi, spi_nrfx_driver_api) = {
 	.release = spi_nrfx_release,
 };
 
-#ifdef CONFIG_PM_DEVICE
-static int spim_nrfx_pm_action(const struct device *dev,
-			       enum pm_device_action action)
+static void spim_resume(const struct device *dev)
 {
-	int ret = 0;
-	struct spi_nrfx_data *dev_data = dev->data;
 	const struct spi_nrfx_config *dev_config = dev->config;
 
-	switch (action) {
-	case PM_DEVICE_ACTION_RESUME:
-		ret = pinctrl_apply_state(dev_config->pcfg,
-					  PINCTRL_STATE_DEFAULT);
-		if (ret < 0) {
-			return ret;
-		}
-		/* nrfx_spim_init() will be called at configuration before
-		 * the next transfer.
-		 */
+	(void)pinctrl_apply_state(dev_config->pcfg, PINCTRL_STATE_DEFAULT);
+	/* nrfx_spim_init() will be called at configuration before
+	 * the next transfer.
+	 */
 
 #ifdef CONFIG_SOC_NRF54H20_GPD
-		nrf_gpd_retain_pins_set(dev_config->pcfg, false);
+	nrf_gpd_retain_pins_set(dev_config->pcfg, false);
 #endif
+}
 
-		break;
+static void spim_suspend(const struct device *dev)
+{
+	const struct spi_nrfx_config *dev_config = dev->config;
+	struct spi_nrfx_data *dev_data = dev->data;
 
-	case PM_DEVICE_ACTION_SUSPEND:
-		if (dev_data->initialized) {
-			nrfx_spim_uninit(&dev_config->spim);
-			dev_data->initialized = false;
-		}
+	if (dev_data->initialized) {
+		nrfx_spim_uninit(&dev_config->spim);
+		dev_data->initialized = false;
+	}
 
 #ifdef CONFIG_SOC_NRF54H20_GPD
-		nrf_gpd_retain_pins_set(dev_config->pcfg, true);
+	nrf_gpd_retain_pins_set(dev_config->pcfg, true);
 #endif
 
-		ret = pinctrl_apply_state(dev_config->pcfg,
-					  PINCTRL_STATE_SLEEP);
-		if (ret < 0) {
-			return ret;
-		}
-		break;
+	(void)pinctrl_apply_state(dev_config->pcfg, PINCTRL_STATE_SLEEP);
+}
 
-	default:
-		ret = -ENOTSUP;
+static int spim_nrfx_pm_action(const struct device *dev, enum pm_device_action action)
+{
+	if (action == PM_DEVICE_ACTION_RESUME) {
+		spim_resume(dev);
+	} else if (IS_ENABLED(CONFIG_PM_DEVICE) && (action == PM_DEVICE_ACTION_SUSPEND)) {
+		spim_suspend(dev);
+	} else {
+		return -ENOTSUP;
 	}
 
-	return ret;
+	return 0;
 }
-#endif /* CONFIG_PM_DEVICE */
-
 
 static int spi_nrfx_init(const struct device *dev)
 {
@@ -655,10 +652,12 @@ static int spi_nrfx_init(const struct device *dev)
 	spi_context_unlock_unconditionally(&dev_data->ctx);
 
 #ifdef CONFIG_SOC_NRF52832_ALLOW_SPIM_DESPITE_PAN_58
-	return anomaly_58_workaround_init(dev);
-#else
-	return 0;
+	err = anomaly_58_workaround_init(dev);
+	if (err < 0) {
+		return err;
+	}
 #endif
+	return pm_device_driver_init(dev, spim_nrfx_pm_action);
 }
 /*
  * We use NODELABEL here because the nrfx API requires us to call