17 files changed, 206 insertions, 73 deletions

diff --git a/drivers/power/reset/Kconfig b/drivers/power/reset/Kconfig
index f65ff49bb275..028e76504519 100644
--- a/drivers/power/reset/Kconfig
+++ b/drivers/power/reset/Kconfig
@@ -71,6 +71,15 @@ config POWER_RESET_HISI
 	help
 	  Reboot support for Hisilicon boards.
 
+config POWER_RESET_IMX
+	bool "IMX6 power-off driver"
+	depends on POWER_RESET && SOC_IMX6
+	help
+	  This driver support power off external PMIC by PMIC_ON_REQ on i.mx6
+	  boards.If you want to use other pin to control external power,please
+	  say N here or disable in dts to make sure pm_power_off never be
+	  overwrote wrongly by this driver.
+
 config POWER_RESET_MSM
 	bool "Qualcomm MSM power-off driver"
 	depends on ARCH_QCOM
diff --git a/drivers/power/reset/Makefile b/drivers/power/reset/Makefile
index 76ce1c59469b..1d4804d6b323 100644
--- a/drivers/power/reset/Makefile
+++ b/drivers/power/reset/Makefile
@@ -6,6 +6,7 @@ obj-$(CONFIG_POWER_RESET_BRCMSTB) += brcmstb-reboot.o
 obj-$(CONFIG_POWER_RESET_GPIO) += gpio-poweroff.o
 obj-$(CONFIG_POWER_RESET_GPIO_RESTART) += gpio-restart.o
 obj-$(CONFIG_POWER_RESET_HISI) += hisi-reboot.o
+obj-$(CONFIG_POWER_RESET_IMX) += imx-snvs-poweroff.o
 obj-$(CONFIG_POWER_RESET_MSM) += msm-poweroff.o
 obj-$(CONFIG_POWER_RESET_LTC2952) += ltc2952-poweroff.o
 obj-$(CONFIG_POWER_RESET_QNAP) += qnap-poweroff.o
diff --git a/drivers/power/reset/as3722-poweroff.c b/drivers/power/reset/as3722-poweroff.c
index 684971199bd3..60d0295fffb1 100644
--- a/drivers/power/reset/as3722-poweroff.c
+++ b/drivers/power/reset/as3722-poweroff.c
@@ -82,7 +82,6 @@ static int as3722_poweroff_remove(struct platform_device *pdev)
 static struct platform_driver as3722_poweroff_driver = {
 	.driver = {
 		.name = "as3722-power-off",
-		.owner = THIS_MODULE,
 	},
 	.probe = as3722_poweroff_probe,
 	.remove = as3722_poweroff_remove,
diff --git a/drivers/power/reset/at91-reset.c b/drivers/power/reset/at91-reset.c
index 3611806c9cfd..69a75d99ae92 100644
--- a/drivers/power/reset/at91-reset.c
+++ b/drivers/power/reset/at91-reset.c
@@ -19,8 +19,8 @@
 
 #include <asm/system_misc.h>
 
-#include <mach/at91sam9_ddrsdr.h>
-#include <mach/at91sam9_sdramc.h>
+#include <soc/at91/at91sam9_ddrsdr.h>
+#include <soc/at91/at91sam9_sdramc.h>
 
 #define AT91_RSTC_CR	0x00		/* Reset Controller Control Register */
 #define AT91_RSTC_PROCRST	BIT(0)		/* Processor Reset */
@@ -100,11 +100,11 @@ static void at91sam9g45_restart(enum reboot_mode mode, const char *cmd)
 		/* Disable SDRAM0 accesses */
 		"1:	str	%3, [%0, #" __stringify(AT91_DDRSDRC_RTR) "]\n\t"
 		/* Power down SDRAM0 */
-		"	str	%4, [%0, #" __stringify(AT91_DDRSDRC_RTR) "]\n\t"
+		"	str	%4, [%0, #" __stringify(AT91_DDRSDRC_LPR) "]\n\t"
 		/* Disable SDRAM1 accesses */
 		"	strne	%3, [%1, #" __stringify(AT91_DDRSDRC_RTR) "]\n\t"
 		/* Power down SDRAM1 */
-		"	strne	%4, [%1, #" __stringify(AT91_DDRSDRC_RTR) "]\n\t"
+		"	strne	%4, [%1, #" __stringify(AT91_DDRSDRC_LPR) "]\n\t"
 		/* Reset CPU */
 		"	str	%5, [%2, #" __stringify(AT91_RSTC_CR) "]\n\t"
 
diff --git a/drivers/power/reset/axxia-reset.c b/drivers/power/reset/axxia-reset.c
index 3b1f8d601784..4e4cd1c8fe50 100644
--- a/drivers/power/reset/axxia-reset.c
+++ b/drivers/power/reset/axxia-reset.c
@@ -19,14 +19,12 @@
 #include <linux/kernel.h>
 #include <linux/mfd/syscon.h>
 #include <linux/module.h>
+#include <linux/notifier.h>
 #include <linux/of.h>
 #include <linux/platform_device.h>
 #include <linux/reboot.h>
 #include <linux/regmap.h>
 
-#include <asm/system_misc.h>
-
-
 #define SC_CRIT_WRITE_KEY	0x1000
 #define SC_LATCH_ON_RESET	0x1004
 #define SC_RESET_CONTROL	0x1008
@@ -39,7 +37,8 @@
 
 static struct regmap *syscon;
 
-static void do_axxia_restart(enum reboot_mode reboot_mode, const char *cmd)
+static int axxia_restart_handler(struct notifier_block *this,
+				 unsigned long mode, void *cmd)
 {
 	/* Access Key (0xab) */
 	regmap_write(syscon, SC_CRIT_WRITE_KEY, 0xab);
@@ -50,11 +49,19 @@ static void do_axxia_restart(enum reboot_mode reboot_mode, const char *cmd)
 	/* Assert chip reset */
 	regmap_update_bits(syscon, SC_RESET_CONTROL,
 			   RSTCTL_RST_CHIP, RSTCTL_RST_CHIP);
+
+	return NOTIFY_DONE;
 }
 
+static struct notifier_block axxia_restart_nb = {
+	.notifier_call = axxia_restart_handler,
+	.priority = 128,
+};
+
 static int axxia_reset_probe(struct platform_device *pdev)
 {
 	struct device *dev = &pdev->dev;
+	int err;
 
 	syscon = syscon_regmap_lookup_by_phandle(dev->of_node, "syscon");
 	if (IS_ERR(syscon)) {
@@ -62,9 +69,11 @@ static int axxia_reset_probe(struct platform_device *pdev)
 		return PTR_ERR(syscon);
 	}
 
-	arm_pm_restart = do_axxia_restart;
+	err = register_restart_handler(&axxia_restart_nb);
+	if (err)
+		dev_err(dev, "cannot register restart handler (err=%d)\n", err);
 
-	return 0;
+	return err;
 }
 
 static const struct of_device_id of_axxia_reset_match[] = {
diff --git a/drivers/power/reset/brcmstb-reboot.c b/drivers/power/reset/brcmstb-reboot.c
index 3f236924742a..100606f9b3dc 100644
--- a/drivers/power/reset/brcmstb-reboot.c
+++ b/drivers/power/reset/brcmstb-reboot.c
@@ -16,6 +16,7 @@
 #include <linux/init.h>
 #include <linux/io.h>
 #include <linux/jiffies.h>
+#include <linux/notifier.h>
 #include <linux/of_address.h>
 #include <linux/of_irq.h>
 #include <linux/of_platform.h>
@@ -26,8 +27,6 @@
 #include <linux/smp.h>
 #include <linux/mfd/syscon.h>
 
-#include <asm/system_misc.h>
-
 #define RESET_SOURCE_ENABLE_REG 1
 #define SW_MASTER_RESET_REG 2
 
@@ -35,7 +34,8 @@ static struct regmap *regmap;
 static u32 rst_src_en;
 static u32 sw_mstr_rst;
 
-static void brcmstb_reboot(enum reboot_mode mode, const char *cmd)
+static int brcmstb_restart_handler(struct notifier_block *this,
+				   unsigned long mode, void *cmd)
 {
 	int rc;
 	u32 tmp;
@@ -43,31 +43,38 @@ static void brcmstb_reboot(enum reboot_mode mode, const char *cmd)
 	rc = regmap_write(regmap, rst_src_en, 1);
 	if (rc) {
 		pr_err("failed to write rst_src_en (%d)\n", rc);
-		return;
+		return NOTIFY_DONE;
 	}
 
 	rc = regmap_read(regmap, rst_src_en, &tmp);
 	if (rc) {
 		pr_err("failed to read rst_src_en (%d)\n", rc);
-		return;
+		return NOTIFY_DONE;
 	}
 
 	rc = regmap_write(regmap, sw_mstr_rst, 1);
 	if (rc) {
 		pr_err("failed to write sw_mstr_rst (%d)\n", rc);
-		return;
+		return NOTIFY_DONE;
 	}
 
 	rc = regmap_read(regmap, sw_mstr_rst, &tmp);
 	if (rc) {
 		pr_err("failed to read sw_mstr_rst (%d)\n", rc);
-		return;
+		return NOTIFY_DONE;
 	}
 
 	while (1)
 		;
+
+	return NOTIFY_DONE;
 }
 
+static struct notifier_block brcmstb_restart_nb = {
+	.notifier_call = brcmstb_restart_handler,
+	.priority = 128,
+};
+
 static int brcmstb_reboot_probe(struct platform_device *pdev)
 {
 	int rc;
@@ -93,9 +100,12 @@ static int brcmstb_reboot_probe(struct platform_device *pdev)
 		return -EINVAL;
 	}
 
-	arm_pm_restart = brcmstb_reboot;
+	rc = register_restart_handler(&brcmstb_restart_nb);
+	if (rc)
+		dev_err(&pdev->dev,
+			"cannot register restart handler (err=%d)\n", rc);
 
-	return 0;
+	return rc;
 }
 
 static const struct of_device_id of_match[] = {
@@ -107,7 +117,6 @@ static struct platform_driver brcmstb_reboot_driver = {
 	.probe = brcmstb_reboot_probe,
 	.driver = {
 		.name = "brcmstb-reboot",
-		.owner = THIS_MODULE,
 		.of_match_table = of_match,
 	},
 };
diff --git a/drivers/power/reset/gpio-poweroff.c b/drivers/power/reset/gpio-poweroff.c
index ce849bc9b269..e5332f1db8a7 100644
--- a/drivers/power/reset/gpio-poweroff.c
+++ b/drivers/power/reset/gpio-poweroff.c
@@ -99,7 +99,6 @@ static struct platform_driver gpio_poweroff_driver = {
 	.remove = gpio_poweroff_remove,
 	.driver = {
 		.name = "poweroff-gpio",
-		.owner = THIS_MODULE,
 		.of_match_table = of_gpio_poweroff_match,
 	},
 };
diff --git a/drivers/power/reset/gpio-restart.c b/drivers/power/reset/gpio-restart.c
index a76829b3f1cd..edb327efee8b 100644
--- a/drivers/power/reset/gpio-restart.c
+++ b/drivers/power/reset/gpio-restart.c
@@ -137,7 +137,6 @@ static struct platform_driver gpio_restart_driver = {
 	.remove = gpio_restart_remove,
 	.driver = {
 		.name = "restart-gpio",
-		.owner = THIS_MODULE,
 		.of_match_table = of_gpio_restart_match,
 	},
 };
diff --git a/drivers/power/reset/hisi-reboot.c b/drivers/power/reset/hisi-reboot.c
index 0c91d0231d36..5385460e23bb 100644
--- a/drivers/power/reset/hisi-reboot.c
+++ b/drivers/power/reset/hisi-reboot.c
@@ -14,27 +14,36 @@
 #include <linux/delay.h>
 #include <linux/io.h>
 #include <linux/module.h>
+#include <linux/notifier.h>
 #include <linux/of_address.h>
 #include <linux/platform_device.h>
 #include <linux/reboot.h>
 
 #include <asm/proc-fns.h>
-#include <asm/system_misc.h>
 
 static void __iomem *base;
 static u32 reboot_offset;
 
-static void hisi_restart(enum reboot_mode mode, const char *cmd)
+static int hisi_restart_handler(struct notifier_block *this,
+				unsigned long mode, void *cmd)
 {
 	writel_relaxed(0xdeadbeef, base + reboot_offset);
 
 	while (1)
 		cpu_do_idle();
+
+	return NOTIFY_DONE;
 }
 
+static struct notifier_block hisi_restart_nb = {
+	.notifier_call = hisi_restart_handler,
+	.priority = 128,
+};
+
 static int hisi_reboot_probe(struct platform_device *pdev)
 {
 	struct device_node *np = pdev->dev.of_node;
+	int err;
 
 	base = of_iomap(np, 0);
 	if (!base) {
@@ -47,9 +56,12 @@ static int hisi_reboot_probe(struct platform_device *pdev)
 		return -EINVAL;
 	}
 
-	arm_pm_restart = hisi_restart;
+	err = register_restart_handler(&hisi_restart_nb);
+	if (err)
+		dev_err(&pdev->dev, "cannot register restart handler (err=%d)\n",
+			err);
 
-	return 0;
+	return err;
 }
 
 static struct of_device_id hisi_reboot_of_match[] = {
diff --git a/drivers/power/reset/imx-snvs-poweroff.c b/drivers/power/reset/imx-snvs-poweroff.c
new file mode 100644
index 000000000000..ad6ce5020ea7
--- /dev/null
+++ b/drivers/power/reset/imx-snvs-poweroff.c
@@ -0,0 +1,66 @@
+/* Power off driver for i.mx6
+ * Copyright (c) 2014, FREESCALE CORPORATION.  All rights reserved.
+ *
+ * based on msm-poweroff.c
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License version 2 and
+ * only version 2 as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ */
+
+#include <linux/err.h>
+#include <linux/init.h>
+#include <linux/io.h>
+#include <linux/kernel.h>
+#include <linux/module.h>
+#include <linux/of.h>
+#include <linux/of_address.h>
+#include <linux/platform_device.h>
+
+static void __iomem *snvs_base;
+
+static void do_imx_poweroff(void)
+{
+	u32 value = readl(snvs_base);
+
+	/* set TOP and DP_EN bit */
+	writel(value | 0x60, snvs_base);
+}
+
+static int imx_poweroff_probe(struct platform_device *pdev)
+{
+	snvs_base = of_iomap(pdev->dev.of_node, 0);
+	if (!snvs_base) {
+		dev_err(&pdev->dev, "failed to get memory\n");
+		return -ENODEV;
+	}
+
+	pm_power_off = do_imx_poweroff;
+	return 0;
+}
+
+static const struct of_device_id of_imx_poweroff_match[] = {
+	{ .compatible = "fsl,sec-v4.0-poweroff", },
+	{},
+};
+MODULE_DEVICE_TABLE(of, of_imx_poweroff_match);
+
+static struct platform_driver imx_poweroff_driver = {
+	.probe = imx_poweroff_probe,
+	.driver = {
+		.name = "imx-snvs-poweroff",
+		.of_match_table = of_match_ptr(of_imx_poweroff_match),
+	},
+};
+
+static int __init imx_poweroff_init(void)
+{
+	return platform_driver_register(&imx_poweroff_driver);
+}
+device_initcall(imx_poweroff_init);
diff --git a/drivers/power/reset/keystone-reset.c b/drivers/power/reset/keystone-reset.c
index 408a18fd91cb..faedf16c8111 100644
--- a/drivers/power/reset/keystone-reset.c
+++ b/drivers/power/reset/keystone-reset.c
@@ -12,9 +12,9 @@
 
 #include <linux/io.h>
 #include <linux/module.h>
+#include <linux/notifier.h>
 #include <linux/reboot.h>
 #include <linux/regmap.h>
-#include <asm/system_misc.h>
 #include <linux/mfd/syscon.h>
 #include <linux/of_platform.h>
 
@@ -52,7 +52,8 @@ static inline int rsctrl_enable_rspll_write(void)
 				  RSCTRL_KEY_MASK, RSCTRL_KEY);
 }
 
-static void rsctrl_restart(enum reboot_mode mode, const char *cmd)
+static int rsctrl_restart_handler(struct notifier_block *this,
+				  unsigned long mode, void *cmd)
 {
 	/* enable write access to RSTCTRL */
 	rsctrl_enable_rspll_write();
@@ -60,8 +61,15 @@ static void rsctrl_restart(enum reboot_mode mode, const char *cmd)
 	/* reset the SOC */
 	regmap_update_bits(pllctrl_regs, rspll_offset + RSCTRL_RG,
 			   RSCTRL_RESET_MASK, 0);
+
+	return NOTIFY_DONE;
 }
 
+static struct notifier_block rsctrl_restart_nb = {
+	.notifier_call = rsctrl_restart_handler,
+	.priority = 128,
+};
+
 static struct of_device_id rsctrl_of_match[] = {
 	{.compatible = "ti,keystone-reset", },
 	{},
@@ -114,8 +122,6 @@ static int rsctrl_probe(struct platform_device *pdev)
 	if (ret)
 		return ret;
 
-	arm_pm_restart = rsctrl_restart;
-
 	/* disable a reset isolation for all module clocks */
 	ret = regmap_write(pllctrl_regs, rspll_offset + RSISO_RG, 0);
 	if (ret)
@@ -147,13 +153,16 @@ static int rsctrl_probe(struct platform_device *pdev)
 			return ret;
 	}
 
-	return 0;
+	ret = register_restart_handler(&rsctrl_restart_nb);
+	if (ret)
+		dev_err(dev, "cannot register restart handler (err=%d)\n", ret);
+
+	return ret;
 }
 
 static struct platform_driver rsctrl_driver = {
 	.probe = rsctrl_probe,
 	.driver = {
-		.owner = THIS_MODULE,
 		.name = KBUILD_MODNAME,
 		.of_match_table = rsctrl_of_match,
 	},
diff --git a/drivers/power/reset/ltc2952-poweroff.c b/drivers/power/reset/ltc2952-poweroff.c
index 116a1cef8f7b..34f38a3dc3ff 100644
--- a/drivers/power/reset/ltc2952-poweroff.c
+++ b/drivers/power/reset/ltc2952-poweroff.c
@@ -342,7 +342,6 @@ static struct platform_driver ltc2952_poweroff_driver = {
 	.remove = ltc2952_poweroff_remove,
 	.driver = {
 		.name = "ltc2952-poweroff",
-		.owner = THIS_MODULE,
 		.of_match_table = of_ltc2952_poweroff_match,
 	},
 	.suspend = ltc2952_poweroff_suspend,
diff --git a/drivers/power/reset/qnap-poweroff.c b/drivers/power/reset/qnap-poweroff.c
index a75db7f8a92f..2789a61cec68 100644
--- a/drivers/power/reset/qnap-poweroff.c
+++ b/drivers/power/reset/qnap-poweroff.c
@@ -129,7 +129,6 @@ static struct platform_driver qnap_power_off_driver = {
 	.probe	= qnap_power_off_probe,
 	.remove	= qnap_power_off_remove,
 	.driver	= {
-		.owner	= THIS_MODULE,
 		.name	= "qnap_power_off",
 		.of_match_table = of_match_ptr(qnap_power_off_of_match_table),
 	},
diff --git a/drivers/power/reset/restart-poweroff.c b/drivers/power/reset/restart-poweroff.c
index edd707ee7281..f46f2c2e4648 100644
--- a/drivers/power/reset/restart-poweroff.c
+++ b/drivers/power/reset/restart-poweroff.c
@@ -55,7 +55,6 @@ static struct platform_driver restart_poweroff_driver = {
 	.remove = restart_poweroff_remove,
 	.driver = {
 		.name = "poweroff-restart",
-		.owner = THIS_MODULE,
 		.of_match_table = of_restart_poweroff_match,
 	},
 };
diff --git a/drivers/power/reset/syscon-reboot.c b/drivers/power/reset/syscon-reboot.c
index 815b901822cf..c4049f45663f 100644
--- a/drivers/power/reset/syscon-reboot.c
+++ b/drivers/power/reset/syscon-reboot.c
@@ -68,7 +68,7 @@ static int syscon_reboot_probe(struct platform_device *pdev)
 		return -EINVAL;
 
 	ctx->restart_handler.notifier_call = syscon_restart_handle;
-	ctx->restart_handler.priority = 128;
+	ctx->restart_handler.priority = 192;
 	err = register_restart_handler(&ctx->restart_handler);
 	if (err)
 		dev_err(dev, "can't register restart notifier (err=%d)\n", err);
diff --git a/drivers/power/reset/vexpress-poweroff.c b/drivers/power/reset/vexpress-poweroff.c
index 4dc102e2b230..9dfc9cee3232 100644
--- a/drivers/power/reset/vexpress-poweroff.c
+++ b/drivers/power/reset/vexpress-poweroff.c
@@ -12,14 +12,14 @@
  */
 
 #include <linux/delay.h>
+#include <linux/notifier.h>
 #include <linux/of.h>
 #include <linux/of_device.h>
 #include <linux/platform_device.h>
+#include <linux/reboot.h>
 #include <linux/stat.h>
 #include <linux/vexpress.h>
 
-#include <asm/system_misc.h>
-
 static void vexpress_reset_do(struct device *dev, const char *what)
 {
 	int err = -ENOENT;
@@ -43,11 +43,19 @@ static void vexpress_power_off(void)
 
 static struct device *vexpress_restart_device;
 
-static void vexpress_restart(enum reboot_mode reboot_mode, const char *cmd)
+static int vexpress_restart(struct notifier_block *this, unsigned long mode,
+			     void *cmd)
 {
 	vexpress_reset_do(vexpress_restart_device, "restart");
+
+	return NOTIFY_DONE;
 }
 
+static struct notifier_block vexpress_restart_nb = {
+	.notifier_call = vexpress_restart,
+	.priority = 128,
+};
+
 static ssize_t vexpress_reset_active_show(struct device *dev,
 		struct device_attribute *attr, char *buf)
 {
@@ -86,12 +94,28 @@ static struct of_device_id vexpress_reset_of_match[] = {
 	{}
 };
 
+static int _vexpress_register_restart_handler(struct device *dev)
+{
+	int err;
+
+	vexpress_restart_device = dev;
+	err = register_restart_handler(&vexpress_restart_nb);
+	if (err) {
+		dev_err(dev, "cannot register restart handler (err=%d)\n", err);
+		return err;
+	}
+	device_create_file(dev, &dev_attr_active);
+
+	return 0;
+}
+
 static int vexpress_reset_probe(struct platform_device *pdev)
 {
 	enum vexpress_reset_func func;
 	const struct of_device_id *match =
 			of_match_device(vexpress_reset_of_match, &pdev->dev);
 	struct regmap *regmap;
+	int ret = 0;
 
 	if (match)
 		func = (enum vexpress_reset_func)match->data;
@@ -110,18 +134,14 @@ static int vexpress_reset_probe(struct platform_device *pdev)
 		break;
 	case FUNC_RESET:
 		if (!vexpress_restart_device)
-			vexpress_restart_device = &pdev->dev;
-		arm_pm_restart = vexpress_restart;
-		device_create_file(&pdev->dev, &dev_attr_active);
+			ret = _vexpress_register_restart_handler(&pdev->dev);
 		break;
 	case FUNC_REBOOT:
-		vexpress_restart_device = &pdev->dev;
-		arm_pm_restart = vexpress_restart;
-		device_create_file(&pdev->dev, &dev_attr_active);
+		ret = _vexpress_register_restart_handler(&pdev->dev);
 		break;
 	};
 
-	return 0;
+	return ret;
 }
 
 static const struct platform_device_id vexpress_reset_id_table[] = {
diff --git a/drivers/power/reset/xgene-reboot.c b/drivers/power/reset/xgene-reboot.c
index 6b49be6867ab..b0e5002f8deb 100644
--- a/drivers/power/reset/xgene-reboot.c
+++ b/drivers/power/reset/xgene-reboot.c
@@ -24,63 +24,67 @@
  * For system shutdown, this is board specify. If a board designer
  * implements GPIO shutdown, use the gpio-poweroff.c driver.
  */
+#include <linux/delay.h>
 #include <linux/io.h>
+#include <linux/notifier.h>
 #include <linux/of_device.h>
 #include <linux/of_address.h>
 #include <linux/platform_device.h>
+#include <linux/reboot.h>
 #include <linux/stat.h>
 #include <linux/slab.h>
-#include <asm/system_misc.h>
 
 struct xgene_reboot_context {
-	struct platform_device *pdev;
+	struct device *dev;
 	void *csr;
 	u32 mask;
+	struct notifier_block restart_handler;
 };
 
-static struct xgene_reboot_context *xgene_restart_ctx;
-
-static void xgene_restart(enum reboot_mode mode, const char *cmd)
+static int xgene_restart_handler(struct notifier_block *this,
+				 unsigned long mode, void *cmd)
 {
-	struct xgene_reboot_context *ctx = xgene_restart_ctx;
-	unsigned long timeout;
+	struct xgene_reboot_context *ctx =
+		container_of(this, struct xgene_reboot_context,
+			     restart_handler);
 
 	/* Issue the reboot */
-	if (ctx)
-		writel(ctx->mask, ctx->csr);
+	writel(ctx->mask, ctx->csr);
+
+	mdelay(1000);
 
-	timeout = jiffies + HZ;
-	while (time_before(jiffies, timeout))
-		cpu_relax();
+	dev_emerg(ctx->dev, "Unable to restart system\n");
 
-	dev_emerg(&ctx->pdev->dev, "Unable to restart system\n");
+	return NOTIFY_DONE;
 }
 
 static int xgene_reboot_probe(struct platform_device *pdev)
 {
 	struct xgene_reboot_context *ctx;
+	struct device *dev = &pdev->dev;
+	int err;
 
-	ctx = devm_kzalloc(&pdev->dev, sizeof(*ctx), GFP_KERNEL);
-	if (!ctx) {
-		dev_err(&pdev->dev, "out of memory for context\n");
-		return -ENODEV;
-	}
+	ctx = devm_kzalloc(dev, sizeof(*ctx), GFP_KERNEL);
+	if (!ctx)
+		return -ENOMEM;
 
-	ctx->csr = of_iomap(pdev->dev.of_node, 0);
+	ctx->csr = of_iomap(dev->of_node, 0);
 	if (!ctx->csr) {
-		devm_kfree(&pdev->dev, ctx);
-		dev_err(&pdev->dev, "can not map resource\n");
+		dev_err(dev, "can not map resource\n");
 		return -ENODEV;
 	}
 
-	if (of_property_read_u32(pdev->dev.of_node, "mask", &ctx->mask))
+	if (of_property_read_u32(dev->of_node, "mask", &ctx->mask))
 		ctx->mask = 0xFFFFFFFF;
 
-	ctx->pdev = pdev;
-	arm_pm_restart = xgene_restart;
-	xgene_restart_ctx = ctx;
+	ctx->dev = dev;
+	ctx->restart_handler.notifier_call = xgene_restart_handler;
+	ctx->restart_handler.priority = 128;
+	err = register_restart_handler(&ctx->restart_handler);
+	if (err)
+		dev_err(dev, "cannot register restart handler (err=%d)\n", err);
 
-	return 0;
+	return err;
 }
 
 static struct of_device_id xgene_reboot_of_match[] = {