struct device *dev;
        void __iomem *csr;
        u32 mask;
-       struct notifier_block restart_handler;
 };
 
-static int xgene_restart_handler(struct notifier_block *this,
-                                unsigned long mode, void *cmd)
+static int xgene_restart_handler(struct sys_off_data *data)
 {
-       struct xgene_reboot_context *ctx =
-               container_of(this, struct xgene_reboot_context,
-                            restart_handler);
+       struct xgene_reboot_context *ctx = data->cb_data;
 
        /* Issue the reboot */
        writel(ctx->mask, ctx->csr);
                ctx->mask = 0xFFFFFFFF;
 
        ctx->dev = dev;
-       ctx->restart_handler.notifier_call = xgene_restart_handler;
-       ctx->restart_handler.priority = 128;
-       err = register_restart_handler(&ctx->restart_handler);
+       err = devm_register_sys_off_handler(dev, SYS_OFF_MODE_RESTART, 128,
+                                           xgene_restart_handler, ctx);
        if (err)
                dev_err(dev, "cannot register restart handler (err=%d)\n", err);