clk: at91: re-factor clocks suspend/resume

SAMA5D2 and SAMA7G5 have a special power saving mode (backup mode) where
most of the SoC's components are powered off (including PMC). Resuming
from this mode is done with the help of bootloader. Peripherals are not
aware of the power saving mode thus most of them are disabling clocks in
proper suspend API and re-enable them in resume API without taking into
account the previously setup rate. Moreover some of the peripherals are
acting as wakeup sources and are not disabling the clocks in this
scenario, when suspending. Since backup mode cuts the power for
peripherals, in resume part these clocks needs to be re-configured.

The initial PMC suspend/resume code was designed only for SAMA5D2's PMC
(as it was the only one supporting backup mode). SAMA7G supports also
backup mode and its PMC is different (few new functionalities, different
registers offsets, different offsets in registers for each
functionalities). To address both SAMA5D2 and SAMA7G5 PMC add
.save_context()/.resume_context() support to each clocks driver and call
this from PMC driver.

Signed-off-by: Claudiu Beznea <claudiu.beznea@microchip.com>
Link: https://lore.kernel.org/r/20211011112719.3951784-2-claudiu.beznea@microchip.com
Acked-by: Nicolas Ferre <nicolas.ferre@microchip.com>
Signed-off-by: Stephen Boyd <sboyd@kernel.org>

12 files changed, 558 insertions, 181 deletions

diff --git a/drivers/clk/at91/clk-generated.c b/drivers/clk/at91/clk-generated.c
index b656d25a9767..23cc8297ec4c 100644
--- a/drivers/clk/at91/clk-generated.c
+++ b/drivers/clk/at91/clk-generated.c
@@ -27,6 +27,7 @@ struct clk_generated {
 	u32 id;
 	u32 gckdiv;
 	const struct clk_pcr_layout *layout;
+	struct at91_clk_pms pms;
 	u8 parent_id;
 	int chg_pid;
 };
@@ -34,25 +35,35 @@ struct clk_generated {
 #define to_clk_generated(hw) \
 	container_of(hw, struct clk_generated, hw)
 
-static int clk_generated_enable(struct clk_hw *hw)
+static int clk_generated_set(struct clk_generated *gck, int status)
 {
-	struct clk_generated *gck = to_clk_generated(hw);
 	unsigned long flags;
-
-	pr_debug("GCLK: %s, gckdiv = %d, parent id = %d\n",
-		 __func__, gck->gckdiv, gck->parent_id);
+	unsigned int enable = status ? AT91_PMC_PCR_GCKEN : 0;
 
 	spin_lock_irqsave(gck->lock, flags);
 	regmap_write(gck->regmap, gck->layout->offset,
 		     (gck->id & gck->layout->pid_mask));
 	regmap_update_bits(gck->regmap, gck->layout->offset,
 			   AT91_PMC_PCR_GCKDIV_MASK | gck->layout->gckcss_mask |
-			   gck->layout->cmd | AT91_PMC_PCR_GCKEN,
+			   gck->layout->cmd | enable,
 			   field_prep(gck->layout->gckcss_mask, gck->parent_id) |
 			   gck->layout->cmd |
 			   FIELD_PREP(AT91_PMC_PCR_GCKDIV_MASK, gck->gckdiv) |
-			   AT91_PMC_PCR_GCKEN);
+			   enable);
 	spin_unlock_irqrestore(gck->lock, flags);
+
+	return 0;
+}
+
+static int clk_generated_enable(struct clk_hw *hw)
+{
+	struct clk_generated *gck = to_clk_generated(hw);
+
+	pr_debug("GCLK: %s, gckdiv = %d, parent id = %d\n",
+		 __func__, gck->gckdiv, gck->parent_id);
+
+	clk_generated_set(gck, 1);
+
 	return 0;
 }
 
@@ -245,6 +256,23 @@ static int clk_generated_set_rate(struct clk_hw *hw,
 	return 0;
 }
 
+static int clk_generated_save_context(struct clk_hw *hw)
+{
+	struct clk_generated *gck = to_clk_generated(hw);
+
+	gck->pms.status = clk_generated_is_enabled(&gck->hw);
+
+	return 0;
+}
+
+static void clk_generated_restore_context(struct clk_hw *hw)
+{
+	struct clk_generated *gck = to_clk_generated(hw);
+
+	if (gck->pms.status)
+		clk_generated_set(gck, gck->pms.status);
+}
+
 static const struct clk_ops generated_ops = {
 	.enable = clk_generated_enable,
 	.disable = clk_generated_disable,
@@ -254,6 +282,8 @@ static const struct clk_ops generated_ops = {
 	.get_parent = clk_generated_get_parent,
 	.set_parent = clk_generated_set_parent,
 	.set_rate = clk_generated_set_rate,
+	.save_context = clk_generated_save_context,
+	.restore_context = clk_generated_restore_context,
 };
 
 /**
@@ -320,8 +350,6 @@ at91_clk_register_generated(struct regmap *regmap, spinlock_t *lock,
 	if (ret) {
 		kfree(gck);
 		hw = ERR_PTR(ret);
-	} else {
-		pmc_register_id(id);
 	}
 
 	return hw;
diff --git a/drivers/clk/at91/clk-main.c b/drivers/clk/at91/clk-main.c
index cfae2f59df66..8601b27c1ae0 100644
--- a/drivers/clk/at91/clk-main.c
+++ b/drivers/clk/at91/clk-main.c
@@ -28,6 +28,7 @@
 struct clk_main_osc {
 	struct clk_hw hw;
 	struct regmap *regmap;
+	struct at91_clk_pms pms;
 };
 
 #define to_clk_main_osc(hw) container_of(hw, struct clk_main_osc, hw)
@@ -37,6 +38,7 @@ struct clk_main_rc_osc {
 	struct regmap *regmap;
 	unsigned long frequency;
 	unsigned long accuracy;
+	struct at91_clk_pms pms;
 };
 
 #define to_clk_main_rc_osc(hw) container_of(hw, struct clk_main_rc_osc, hw)
@@ -51,6 +53,7 @@ struct clk_rm9200_main {
 struct clk_sam9x5_main {
 	struct clk_hw hw;
 	struct regmap *regmap;
+	struct at91_clk_pms pms;
 	u8 parent;
 };
 
@@ -120,10 +123,29 @@ static int clk_main_osc_is_prepared(struct clk_hw *hw)
 	return (status & AT91_PMC_MOSCS) && clk_main_parent_select(tmp);
 }
 
+static int clk_main_osc_save_context(struct clk_hw *hw)
+{
+	struct clk_main_osc *osc = to_clk_main_osc(hw);
+
+	osc->pms.status = clk_main_osc_is_prepared(hw);
+
+	return 0;
+}
+
+static void clk_main_osc_restore_context(struct clk_hw *hw)
+{
+	struct clk_main_osc *osc = to_clk_main_osc(hw);
+
+	if (osc->pms.status)
+		clk_main_osc_prepare(hw);
+}
+
 static const struct clk_ops main_osc_ops = {
 	.prepare = clk_main_osc_prepare,
 	.unprepare = clk_main_osc_unprepare,
 	.is_prepared = clk_main_osc_is_prepared,
+	.save_context = clk_main_osc_save_context,
+	.restore_context = clk_main_osc_restore_context,
 };
 
 struct clk_hw * __init
@@ -240,12 +262,31 @@ static unsigned long clk_main_rc_osc_recalc_accuracy(struct clk_hw *hw,
 	return osc->accuracy;
 }
 
+static int clk_main_rc_osc_save_context(struct clk_hw *hw)
+{
+	struct clk_main_rc_osc *osc = to_clk_main_rc_osc(hw);
+
+	osc->pms.status = clk_main_rc_osc_is_prepared(hw);
+
+	return 0;
+}
+
+static void clk_main_rc_osc_restore_context(struct clk_hw *hw)
+{
+	struct clk_main_rc_osc *osc = to_clk_main_rc_osc(hw);
+
+	if (osc->pms.status)
+		clk_main_rc_osc_prepare(hw);
+}
+
 static const struct clk_ops main_rc_osc_ops = {
 	.prepare = clk_main_rc_osc_prepare,
 	.unprepare = clk_main_rc_osc_unprepare,
 	.is_prepared = clk_main_rc_osc_is_prepared,
 	.recalc_rate = clk_main_rc_osc_recalc_rate,
 	.recalc_accuracy = clk_main_rc_osc_recalc_accuracy,
+	.save_context = clk_main_rc_osc_save_context,
+	.restore_context = clk_main_rc_osc_restore_context,
 };
 
 struct clk_hw * __init
@@ -465,12 +506,37 @@ static u8 clk_sam9x5_main_get_parent(struct clk_hw *hw)
 	return clk_main_parent_select(status);
 }
 
+static int clk_sam9x5_main_save_context(struct clk_hw *hw)
+{
+	struct clk_sam9x5_main *clkmain = to_clk_sam9x5_main(hw);
+
+	clkmain->pms.status = clk_main_rc_osc_is_prepared(&clkmain->hw);
+	clkmain->pms.parent = clk_sam9x5_main_get_parent(&clkmain->hw);
+
+	return 0;
+}
+
+static void clk_sam9x5_main_restore_context(struct clk_hw *hw)
+{
+	struct clk_sam9x5_main *clkmain = to_clk_sam9x5_main(hw);
+	int ret;
+
+	ret = clk_sam9x5_main_set_parent(hw, clkmain->pms.parent);
+	if (ret)
+		return;
+
+	if (clkmain->pms.status)
+		clk_sam9x5_main_prepare(hw);
+}
+
 static const struct clk_ops sam9x5_main_ops = {
 	.prepare = clk_sam9x5_main_prepare,
 	.is_prepared = clk_sam9x5_main_is_prepared,
 	.recalc_rate = clk_sam9x5_main_recalc_rate,
 	.set_parent = clk_sam9x5_main_set_parent,
 	.get_parent = clk_sam9x5_main_get_parent,
+	.save_context = clk_sam9x5_main_save_context,
+	.restore_context = clk_sam9x5_main_restore_context,
 };
 
 struct clk_hw * __init
diff --git a/drivers/clk/at91/clk-master.c b/drivers/clk/at91/clk-master.c
index a80427980bf7..f75549fff023 100644
--- a/drivers/clk/at91/clk-master.c
+++ b/drivers/clk/at91/clk-master.c
@@ -37,6 +37,7 @@ struct clk_master {
 	spinlock_t *lock;
 	const struct clk_master_layout *layout;
 	const struct clk_master_characteristics *characteristics;
+	struct at91_clk_pms pms;
 	u32 *mux_table;
 	u32 mckr;
 	int chg_pid;
@@ -112,10 +113,52 @@ static unsigned long clk_master_div_recalc_rate(struct clk_hw *hw,
 	return rate;
 }
 
+static int clk_master_div_save_context(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+	struct clk_hw *parent_hw = clk_hw_get_parent(hw);
+	unsigned long flags;
+	unsigned int mckr, div;
+
+	spin_lock_irqsave(master->lock, flags);
+	regmap_read(master->regmap, master->layout->offset, &mckr);
+	spin_unlock_irqrestore(master->lock, flags);
+
+	mckr &= master->layout->mask;
+	div = (mckr >> MASTER_DIV_SHIFT) & MASTER_DIV_MASK;
+	div = master->characteristics->divisors[div];
+
+	master->pms.parent_rate = clk_hw_get_rate(parent_hw);
+	master->pms.rate = DIV_ROUND_CLOSEST(master->pms.parent_rate, div);
+
+	return 0;
+}
+
+static void clk_master_div_restore_context(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+	unsigned long flags;
+	unsigned int mckr;
+	u8 div;
+
+	spin_lock_irqsave(master->lock, flags);
+	regmap_read(master->regmap, master->layout->offset, &mckr);
+	spin_unlock_irqrestore(master->lock, flags);
+
+	mckr &= master->layout->mask;
+	div = (mckr >> MASTER_DIV_SHIFT) & MASTER_DIV_MASK;
+	div = master->characteristics->divisors[div];
+
+	if (div != DIV_ROUND_CLOSEST(master->pms.parent_rate, master->pms.rate))
+		pr_warn("MCKR DIV not configured properly by firmware!\n");
+}
+
 static const struct clk_ops master_div_ops = {
 	.prepare = clk_master_prepare,
 	.is_prepared = clk_master_is_prepared,
 	.recalc_rate = clk_master_div_recalc_rate,
+	.save_context = clk_master_div_save_context,
+	.restore_context = clk_master_div_restore_context,
 };
 
 static int clk_master_div_set_rate(struct clk_hw *hw, unsigned long rate,
@@ -125,7 +168,9 @@ static int clk_master_div_set_rate(struct clk_hw *hw, unsigned long rate,
 	const struct clk_master_characteristics *characteristics =
 						master->characteristics;
 	unsigned long flags;
+	unsigned int mckr, tmp;
 	int div, i;
+	int ret;
 
 	div = DIV_ROUND_CLOSEST(parent_rate, rate);
 	if (div > ARRAY_SIZE(characteristics->divisors))
@@ -145,11 +190,24 @@ static int clk_master_div_set_rate(struct clk_hw *hw, unsigned long rate,
 		return -EINVAL;
 
 	spin_lock_irqsave(master->lock, flags);
-	regmap_update_bits(master->regmap, master->layout->offset,
-			   (MASTER_DIV_MASK << MASTER_DIV_SHIFT),
-			   (div << MASTER_DIV_SHIFT));
+	ret = regmap_read(master->regmap, master->layout->offset, &mckr);
+	if (ret)
+		goto unlock;
+
+	tmp = mckr & master->layout->mask;
+	tmp = (tmp >> MASTER_DIV_SHIFT) & MASTER_DIV_MASK;
+	if (tmp == div)
+		goto unlock;
+
+	mckr &= ~(MASTER_DIV_MASK << MASTER_DIV_SHIFT);
+	mckr |= (div << MASTER_DIV_SHIFT);
+	ret = regmap_write(master->regmap, master->layout->offset, mckr);
+	if (ret)
+		goto unlock;
+
 	while (!clk_master_ready(master))
 		cpu_relax();
+unlock:
 	spin_unlock_irqrestore(master->lock, flags);
 
 	return 0;
@@ -197,12 +255,25 @@ static int clk_master_div_determine_rate(struct clk_hw *hw,
 	return 0;
 }
 
+static void clk_master_div_restore_context_chg(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+	int ret;
+
+	ret = clk_master_div_set_rate(hw, master->pms.rate,
+				      master->pms.parent_rate);
+	if (ret)
+		pr_warn("Failed to restore MCK DIV clock\n");
+}
+
 static const struct clk_ops master_div_ops_chg = {
 	.prepare = clk_master_prepare,
 	.is_prepared = clk_master_is_prepared,
 	.recalc_rate = clk_master_div_recalc_rate,
 	.determine_rate = clk_master_div_determine_rate,
 	.set_rate = clk_master_div_set_rate,
+	.save_context = clk_master_div_save_context,
+	.restore_context = clk_master_div_restore_context_chg,
 };
 
 static void clk_sama7g5_master_best_diff(struct clk_rate_request *req,
@@ -272,7 +343,8 @@ static int clk_master_pres_set_rate(struct clk_hw *hw, unsigned long rate,
 {
 	struct clk_master *master = to_clk_master(hw);
 	unsigned long flags;
-	unsigned int pres;
+	unsigned int pres, mckr, tmp;
+	int ret;
 
 	pres = DIV_ROUND_CLOSEST(parent_rate, rate);
 	if (pres > MASTER_PRES_MAX)
@@ -284,15 +356,27 @@ static int clk_master_pres_set_rate(struct clk_hw *hw, unsigned long rate,
 		pres = ffs(pres) - 1;
 
 	spin_lock_irqsave(master->lock, flags);
-	regmap_update_bits(master->regmap, master->layout->offset,
-			   (MASTER_PRES_MASK << master->layout->pres_shift),
-			   (pres << master->layout->pres_shift));
+	ret = regmap_read(master->regmap, master->layout->offset, &mckr);
+	if (ret)
+		goto unlock;
+
+	mckr &= master->layout->mask;
+	tmp = (mckr >> master->layout->pres_shift) & MASTER_PRES_MASK;
+	if (pres == tmp)
+		goto unlock;
+
+	mckr &= ~(MASTER_PRES_MASK << master->layout->pres_shift);
+	mckr |= (pres << master->layout->pres_shift);
+	ret = regmap_write(master->regmap, master->layout->offset, mckr);
+	if (ret)
+		goto unlock;
 
 	while (!clk_master_ready(master))
 		cpu_relax();
+unlock:
 	spin_unlock_irqrestore(master->lock, flags);
 
-	return 0;
+	return ret;
 }
 
 static unsigned long clk_master_pres_recalc_rate(struct clk_hw *hw,
@@ -330,11 +414,68 @@ static u8 clk_master_pres_get_parent(struct clk_hw *hw)
 	return mckr & AT91_PMC_CSS;
 }
 
+static int clk_master_pres_save_context(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+	struct clk_hw *parent_hw = clk_hw_get_parent(hw);
+	unsigned long flags;
+	unsigned int val, pres;
+
+	spin_lock_irqsave(master->lock, flags);
+	regmap_read(master->regmap, master->layout->offset, &val);
+	spin_unlock_irqrestore(master->lock, flags);
+
+	val &= master->layout->mask;
+	pres = (val >> master->layout->pres_shift) & MASTER_PRES_MASK;
+	if (pres == MASTER_PRES_MAX && master->characteristics->have_div3_pres)
+		pres = 3;
+	else
+		pres = (1 << pres);
+
+	master->pms.parent = val & AT91_PMC_CSS;
+	master->pms.parent_rate = clk_hw_get_rate(parent_hw);
+	master->pms.rate = DIV_ROUND_CLOSEST_ULL(master->pms.parent_rate, pres);
+
+	return 0;
+}
+
+static void clk_master_pres_restore_context(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+	unsigned long flags;
+	unsigned int val, pres;
+
+	spin_lock_irqsave(master->lock, flags);
+	regmap_read(master->regmap, master->layout->offset, &val);
+	spin_unlock_irqrestore(master->lock, flags);
+
+	val &= master->layout->mask;
+	pres = (val >> master->layout->pres_shift) & MASTER_PRES_MASK;
+	if (pres == MASTER_PRES_MAX && master->characteristics->have_div3_pres)
+		pres = 3;
+	else
+		pres = (1 << pres);
+
+	if (master->pms.rate !=
+	    DIV_ROUND_CLOSEST_ULL(master->pms.parent_rate, pres) ||
+	    (master->pms.parent != (val & AT91_PMC_CSS)))
+		pr_warn("MCKR PRES was not configured properly by firmware!\n");
+}
+
+static void clk_master_pres_restore_context_chg(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+
+	clk_master_pres_set_rate(hw, master->pms.rate, master->pms.parent_rate);
+}
+
 static const struct clk_ops master_pres_ops = {
 	.prepare = clk_master_prepare,
 	.is_prepared = clk_master_is_prepared,
 	.recalc_rate = clk_master_pres_recalc_rate,
 	.get_parent = clk_master_pres_get_parent,
+	.save_context = clk_master_pres_save_context,
+	.restore_context = clk_master_pres_restore_context,
 };
 
 static const struct clk_ops master_pres_ops_chg = {
@@ -344,6 +485,8 @@ static const struct clk_ops master_pres_ops_chg = {
 	.recalc_rate = clk_master_pres_recalc_rate,
 	.get_parent = clk_master_pres_get_parent,
 	.set_rate = clk_master_pres_set_rate,
+	.save_context = clk_master_pres_save_context,
+	.restore_context = clk_master_pres_restore_context_chg,
 };
 
 static struct clk_hw * __init
@@ -539,20 +682,21 @@ static int clk_sama7g5_master_set_parent(struct clk_hw *hw, u8 index)
 	return 0;
 }
 
-static int clk_sama7g5_master_enable(struct clk_hw *hw)
+static void clk_sama7g5_master_set(struct clk_master *master,
+				   unsigned int status)
 {
-	struct clk_master *master = to_clk_master(hw);
 	unsigned long flags;
 	unsigned int val, cparent;
+	unsigned int enable = status ? PMC_MCR_EN : 0;
 
 	spin_lock_irqsave(master->lock, flags);
 
 	regmap_write(master->regmap, PMC_MCR, PMC_MCR_ID(master->id));
 	regmap_read(master->regmap, PMC_MCR, &val);
 	regmap_update_bits(master->regmap, PMC_MCR,
-			   PMC_MCR_EN | PMC_MCR_CSS | PMC_MCR_DIV |
+			   enable | PMC_MCR_CSS | PMC_MCR_DIV |
 			   PMC_MCR_CMD | PMC_MCR_ID_MSK,
-			   PMC_MCR_EN | (master->parent << PMC_MCR_CSS_SHIFT) |
+			   enable | (master->parent << PMC_MCR_CSS_SHIFT) |
 			   (master->div << MASTER_DIV_SHIFT) |
 			   PMC_MCR_CMD | PMC_MCR_ID(master->id));
 
@@ -563,6 +707,13 @@ static int clk_sama7g5_master_enable(struct clk_hw *hw)
 		cpu_relax();
 
 	spin_unlock_irqrestore(master->lock, flags);
+}
+
+static int clk_sama7g5_master_enable(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+
+	clk_sama7g5_master_set(master, 1);
 
 	return 0;
 }
@@ -620,6 +771,23 @@ static int clk_sama7g5_master_set_rate(struct clk_hw *hw, unsigned long rate,
 	return 0;
 }
 
+static int clk_sama7g5_master_save_context(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+
+	master->pms.status = clk_sama7g5_master_is_enabled(hw);
+
+	return 0;
+}
+
+static void clk_sama7g5_master_restore_context(struct clk_hw *hw)
+{
+	struct clk_master *master = to_clk_master(hw);
+
+	if (master->pms.status)
+		clk_sama7g5_master_set(master, master->pms.status);
+}
+
 static const struct clk_ops sama7g5_master_ops = {
 	.enable = clk_sama7g5_master_enable,
 	.disable = clk_sama7g5_master_disable,
@@ -629,6 +797,8 @@ static const struct clk_ops sama7g5_master_ops = {
 	.set_rate = clk_sama7g5_master_set_rate,
 	.get_parent = clk_sama7g5_master_get_parent,
 	.set_parent = clk_sama7g5_master_set_parent,
+	.save_context = clk_sama7g5_master_save_context,
+	.restore_context = clk_sama7g5_master_restore_context,
 };
 
 struct clk_hw * __init
diff --git a/drivers/clk/at91/clk-peripheral.c b/drivers/clk/at91/clk-peripheral.c
index 7a27ba8e0577..e14fa5ac734c 100644
--- a/drivers/clk/at91/clk-peripheral.c
+++ b/drivers/clk/at91/clk-peripheral.c
@@ -37,6 +37,7 @@ struct clk_sam9x5_peripheral {
 	u32 id;
 	u32 div;
 	const struct clk_pcr_layout *layout;
+	struct at91_clk_pms pms;
 	bool auto_div;
 	int chg_pid;
 };
@@ -155,10 +156,11 @@ static void clk_sam9x5_peripheral_autodiv(struct clk_sam9x5_peripheral *periph)
 	periph->div = shift;
 }
 
-static int clk_sam9x5_peripheral_enable(struct clk_hw *hw)
+static int clk_sam9x5_peripheral_set(struct clk_sam9x5_peripheral *periph,
+				     unsigned int status)
 {
-	struct clk_sam9x5_peripheral *periph = to_clk_sam9x5_peripheral(hw);
 	unsigned long flags;
+	unsigned int enable = status ? AT91_PMC_PCR_EN : 0;
 
 	if (periph->id < PERIPHERAL_ID_MIN)
 		return 0;
@@ -168,15 +170,21 @@ static int clk_sam9x5_peripheral_enable(struct clk_hw *hw)
 		     (periph->id & periph->layout->pid_mask));
 	regmap_update_bits(periph->regmap, periph->layout->offset,
 			   periph->layout->div_mask | periph->layout->cmd |
-			   AT91_PMC_PCR_EN,
+			   enable,
 			   field_prep(periph->layout->div_mask, periph->div) |
-			   periph->layout->cmd |
-			   AT91_PMC_PCR_EN);
+			   periph->layout->cmd | enable);
 	spin_unlock_irqrestore(periph->lock, flags);
 
 	return 0;
 }
 
+static int clk_sam9x5_peripheral_enable(struct clk_hw *hw)
+{
+	struct clk_sam9x5_peripheral *periph = to_clk_sam9x5_peripheral(hw);
+
+	return clk_sam9x5_peripheral_set(periph, 1);
+}
+
 static void clk_sam9x5_peripheral_disable(struct clk_hw *hw)
 {
 	struct clk_sam9x5_peripheral *periph = to_clk_sam9x5_peripheral(hw);
@@ -393,6 +401,23 @@ static int clk_sam9x5_peripheral_set_rate(struct clk_hw *hw,
 	return -EINVAL;
 }
 
+static int clk_sam9x5_peripheral_save_context(struct clk_hw *hw)
+{
+	struct clk_sam9x5_peripheral *periph = to_clk_sam9x5_peripheral(hw);
+
+	periph->pms.status = clk_sam9x5_peripheral_is_enabled(hw);
+
+	return 0;
+}
+
+static void clk_sam9x5_peripheral_restore_context(struct clk_hw *hw)
+{
+	struct clk_sam9x5_peripheral *periph = to_clk_sam9x5_peripheral(hw);
+
+	if (periph->pms.status)
+		clk_sam9x5_peripheral_set(periph, periph->pms.status);
+}
+
 static const struct clk_ops sam9x5_peripheral_ops = {
 	.enable = clk_sam9x5_peripheral_enable,
 	.disable = clk_sam9x5_peripheral_disable,
@@ -400,6 +425,8 @@ static const struct clk_ops sam9x5_peripheral_ops = {
 	.recalc_rate = clk_sam9x5_peripheral_recalc_rate,
 	.round_rate = clk_sam9x5_peripheral_round_rate,
 	.set_rate = clk_sam9x5_peripheral_set_rate,
+	.save_context = clk_sam9x5_peripheral_save_context,
+	.restore_context = clk_sam9x5_peripheral_restore_context,
 };
 
 static const struct clk_ops sam9x5_peripheral_chg_ops = {
@@ -409,6 +436,8 @@ static const struct clk_ops sam9x5_peripheral_chg_ops = {
 	.recalc_rate = clk_sam9x5_peripheral_recalc_rate,
 	.determine_rate = clk_sam9x5_peripheral_determine_rate,
 	.set_rate = clk_sam9x5_peripheral_set_rate,
+	.save_context = clk_sam9x5_peripheral_save_context,
+	.restore_context = clk_sam9x5_peripheral_restore_context,
 };
 
 struct clk_hw * __init
@@ -460,7 +489,6 @@ at91_clk_register_sam9x5_peripheral(struct regmap *regmap, spinlock_t *lock,
 		hw = ERR_PTR(ret);
 	} else {
 		clk_sam9x5_peripheral_autodiv(periph);
-		pmc_register_id(id);
 	}
 
 	return hw;
diff --git a/drivers/clk/at91/clk-pll.c b/drivers/clk/at91/clk-pll.c
index 6ed986d3eee0..249d6a53cedf 100644
--- a/drivers/clk/at91/clk-pll.c
+++ b/drivers/clk/at91/clk-pll.c
@@ -40,6 +40,7 @@ struct clk_pll {
 	u16 mul;
 	const struct clk_pll_layout *layout;
 	const struct clk_pll_characteristics *characteristics;
+	struct at91_clk_pms pms;
 };
 
 static inline bool clk_pll_ready(struct regmap *regmap, int id)
@@ -260,6 +261,42 @@ static int clk_pll_set_rate(struct clk_hw *hw, unsigned long rate,
 	return 0;
 }
 
+static int clk_pll_save_context(struct clk_hw *hw)
+{
+	struct clk_pll *pll = to_clk_pll(hw);
+	struct clk_hw *parent_hw = clk_hw_get_parent(hw);
+
+	pll->pms.parent_rate = clk_hw_get_rate(parent_hw);
+	pll->pms.rate = clk_pll_recalc_rate(&pll->hw, pll->pms.parent_rate);
+	pll->pms.status = clk_pll_ready(pll->regmap, PLL_REG(pll->id));
+
+	return 0;
+}
+
+static void clk_pll_restore_context(struct clk_hw *hw)
+{
+	struct clk_pll *pll = to_clk_pll(hw);
+	unsigned long calc_rate;
+	unsigned int pllr, pllr_out, pllr_count;
+	u8 out = 0;
+
+	if (pll->characteristics->out)
+		out = pll->characteristics->out[pll->range];
+
+	regmap_read(pll->regmap, PLL_REG(pll->id), &pllr);
+
+	calc_rate = (pll->pms.parent_rate / PLL_DIV(pllr)) *
+		     (PLL_MUL(pllr, pll->layout) + 1);
+	pllr_count = (pllr >> PLL_COUNT_SHIFT) & PLL_MAX_COUNT;
+	pllr_out = (pllr >> PLL_OUT_SHIFT) & out;
+
+	if (pll->pms.rate != calc_rate ||
+	    pll->pms.status != clk_pll_ready(pll->regmap, PLL_REG(pll->id)) ||
+	    pllr_count != PLL_MAX_COUNT ||
+	    (out && pllr_out != out))
+		pr_warn("PLLAR was not configured properly by firmware\n");
+}
+
 static const struct clk_ops pll_ops = {
 	.prepare = clk_pll_prepare,
 	.unprepare = clk_pll_unprepare,
@@ -267,6 +304,8 @@ static const struct clk_ops pll_ops = {
 	.recalc_rate = clk_pll_recalc_rate,
 	.round_rate = clk_pll_round_rate,
 	.set_rate = clk_pll_set_rate,
+	.save_context = clk_pll_save_context,
+	.restore_context = clk_pll_restore_context,
 };
 
 struct clk_hw * __init
diff --git a/drivers/clk/at91/clk-programmable.c b/drivers/clk/at91/clk-programmable.c
index fcf8f6a1c2c6..6c4b259d31d3 100644
--- a/drivers/clk/at91/clk-programmable.c
+++ b/drivers/clk/at91/clk-programmable.c
@@ -24,6 +24,7 @@ struct clk_programmable {
 	u32 *mux_table;
 	u8 id;
 	const struct clk_programmable_layout *layout;
+	struct at91_clk_pms pms;
 };
 
 #define to_clk_programmable(hw) container_of(hw, struct clk_programmable, hw)
@@ -177,12 +178,38 @@ static int clk_programmable_set_rate(struct clk_hw *hw, unsigned long rate,
 	return 0;
 }
 
+static int clk_programmable_save_context(struct clk_hw *hw)
+{
+	struct clk_programmable *prog = to_clk_programmable(hw);
+	struct clk_hw *parent_hw = clk_hw_get_parent(hw);
+
+	prog->pms.parent = clk_programmable_get_parent(hw);
+	prog->pms.parent_rate = clk_hw_get_rate(parent_hw);
+	prog->pms.rate = clk_programmable_recalc_rate(hw, prog->pms.parent_rate);
+
+	return 0;
+}
+
+static void clk_programmable_restore_context(struct clk_hw *hw)
+{
+	struct clk_programmable *prog = to_clk_programmable(hw);
+	int ret;
+
+	ret = clk_programmable_set_parent(hw, prog->pms.parent);
+	if (ret)
+		return;
+
+	clk_programmable_set_rate(hw, prog->pms.rate, prog->pms.parent_rate);
+}
+
 static const struct clk_ops programmable_ops = {
 	.recalc_rate = clk_programmable_recalc_rate,
 	.determine_rate = clk_programmable_determine_rate,
 	.get_parent = clk_programmable_get_parent,
 	.set_parent = clk_programmable_set_parent,
 	.set_rate = clk_programmable_set_rate,
+	.save_context = clk_programmable_save_context,
+	.restore_context = clk_programmable_restore_context,
 };
 
 struct clk_hw * __init
@@ -221,8 +248,6 @@ at91_clk_register_programmable(struct regmap *regmap,
 	if (ret) {
 		kfree(prog);
 		hw = ERR_PTR(ret);
-	} else {
-		pmc_register_pck(id);
 	}
 
 	return hw;
diff --git a/drivers/clk/at91/clk-sam9x60-pll.c b/drivers/clk/at91/clk-sam9x60-pll.c
index 34e3ab13741a..7020d3bf6e13 100644
--- a/drivers/clk/at91/clk-sam9x60-pll.c
+++ b/drivers/clk/at91/clk-sam9x60-pll.c
@@ -38,12 +38,14 @@ struct sam9x60_pll_core {
 
 struct sam9x60_frac {
 	struct sam9x60_pll_core core;
+	struct at91_clk_pms pms;
 	u32 frac;
 	u16 mul;
 };
 
 struct sam9x60_div {
 	struct sam9x60_pll_core core;
+	struct at91_clk_pms pms;
 	u8 div;
 };
 
@@ -75,9 +77,8 @@ static unsigned long sam9x60_frac_pll_recalc_rate(struct clk_hw *hw,
 		((u64)parent_rate * frac->frac >> 22));
 }
 
-static int sam9x60_frac_pll_prepare(struct clk_hw *hw)
+static int sam9x60_frac_pll_set(struct sam9x60_pll_core *core)
 {
-	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
 	struct sam9x60_frac *frac = to_sam9x60_frac(core);
 	struct regmap *regmap = core->regmap;
 	unsigned int val, cfrac, cmul;
@@ -141,6 +142,13 @@ unlock:
 	return 0;
 }
 
+static int sam9x60_frac_pll_prepare(struct clk_hw *hw)
+{
+	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
+
+	return sam9x60_frac_pll_set(core);
+}
+
 static void sam9x60_frac_pll_unprepare(struct clk_hw *hw)
 {
 	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
@@ -280,6 +288,25 @@ unlock:
 	return ret;
 }
 
+static int sam9x60_frac_pll_save_context(struct clk_hw *hw)
+{
+	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
+	struct sam9x60_frac *frac = to_sam9x60_frac(core);
+
+	frac->pms.status = sam9x60_pll_ready(core->regmap, core->id);
+
+	return 0;
+}
+
+static void sam9x60_frac_pll_restore_context(struct clk_hw *hw)
+{
+	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
+	struct sam9x60_frac *frac = to_sam9x60_frac(core);
+
+	if (frac->pms.status)
+		sam9x60_frac_pll_set(core);
+}
+
 static const struct clk_ops sam9x60_frac_pll_ops = {
 	.prepare = sam9x60_frac_pll_prepare,
 	.unprepare = sam9x60_frac_pll_unprepare,
@@ -287,6 +314,8 @@ static const struct clk_ops sam9x60_frac_pll_ops = {
 	.recalc_rate = sam9x60_frac_pll_recalc_rate,
 	.round_rate = sam9x60_frac_pll_round_rate,
 	.set_rate = sam9x60_frac_pll_set_rate,
+	.save_context = sam9x60_frac_pll_save_context,
+	.restore_context = sam9x60_frac_pll_restore_context,
 };
 
 static const struct clk_ops sam9x60_frac_pll_ops_chg = {
@@ -296,11 +325,12 @@ static const struct clk_ops sam9x60_frac_pll_ops_chg = {
 	.recalc_rate = sam9x60_frac_pll_recalc_rate,
 	.round_rate = sam9x60_frac_pll_round_rate,
 	.set_rate = sam9x60_frac_pll_set_rate_chg,
+	.save_context = sam9x60_frac_pll_save_context,
+	.restore_context = sam9x60_frac_pll_restore_context,
 };
 
-static int sam9x60_div_pll_prepare(struct clk_hw *hw)
+static int sam9x60_div_pll_set(struct sam9x60_pll_core *core)
 {
-	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
 	struct sam9x60_div *div = to_sam9x60_div(core);
 	struct regmap *regmap = core->regmap;
 	unsigned long flags;
@@ -334,6 +364,13 @@ unlock:
 	return 0;
 }
 
+static int sam9x60_div_pll_prepare(struct clk_hw *hw)
+{
+	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
+
+	return sam9x60_div_pll_set(core);
+}
+
 static void sam9x60_div_pll_unprepare(struct clk_hw *hw)
 {
 	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
@@ -482,6 +519,25 @@ unlock:
 	return 0;
 }
 
+static int sam9x60_div_pll_save_context(struct clk_hw *hw)
+{
+	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
+	struct sam9x60_div *div = to_sam9x60_div(core);
+
+	div->pms.status = sam9x60_div_pll_is_prepared(hw);
+
+	return 0;
+}
+
+static void sam9x60_div_pll_restore_context(struct clk_hw *hw)
+{
+	struct sam9x60_pll_core *core = to_sam9x60_pll_core(hw);
+	struct sam9x60_div *div = to_sam9x60_div(core);
+
+	if (div->pms.status)
+		sam9x60_div_pll_set(core);
+}
+
 static const struct clk_ops sam9x60_div_pll_ops = {
 	.prepare = sam9x60_div_pll_prepare,
 	.unprepare = sam9x60_div_pll_unprepare,
@@ -489,6 +545,8 @@ static const struct clk_ops sam9x60_div_pll_ops = {
 	.recalc_rate = sam9x60_div_pll_recalc_rate,
 	.round_rate = sam9x60_div_pll_round_rate,
 	.set_rate = sam9x60_div_pll_set_rate,
+	.save_context = sam9x60_div_pll_save_context,
+	.restore_context = sam9x60_div_pll_restore_context,
 };
 
 static const struct clk_ops sam9x60_div_pll_ops_chg = {
@@ -498,6 +556,8 @@ static const struct clk_ops sam9x60_div_pll_ops_chg = {
 	.recalc_rate = sam9x60_div_pll_recalc_rate,
 	.round_rate = sam9x60_div_pll_round_rate,
 	.set_rate = sam9x60_div_pll_set_rate_chg,
+	.save_context = sam9x60_div_pll_save_context,
+	.restore_context = sam9x60_div_pll_restore_context,
 };
 
 struct clk_hw * __init
diff --git a/drivers/clk/at91/clk-system.c b/drivers/clk/at91/clk-system.c
index f83ec0de86c3..80720fd1a9cf 100644
--- a/drivers/clk/at91/clk-system.c
+++ b/