.init_ack_masked = true,
 };
 
-static struct i2c_client *rk808_i2c_client;
-
-static void rk808_pm_power_off(void)
+static int rk808_power_off(struct sys_off_data *data)
 {
+       struct rk808 *rk808 = data->cb_data;
        int ret;
        unsigned int reg, bit;
-       struct rk808 *rk808 = i2c_get_clientdata(rk808_i2c_client);
 
        switch (rk808->variant) {
        case RK805_ID:
                bit = DEV_OFF;
                break;
        default:
-               return;
+               return NOTIFY_DONE;
        }
        ret = regmap_update_bits(rk808->regmap, reg, bit, bit);
        if (ret)
-               dev_err(&rk808_i2c_client->dev, "Failed to shutdown device!\n");
+               dev_err(&rk808->i2c->dev, "Failed to shutdown device!\n");
+
+       return NOTIFY_DONE;
 }
 
-static int rk808_restart_notify(struct notifier_block *this, unsigned long mode, void *cmd)
+static int rk808_restart(struct sys_off_data *data)
 {
-       struct rk808 *rk808 = i2c_get_clientdata(rk808_i2c_client);
+       struct rk808 *rk808 = data->cb_data;
        unsigned int reg, bit;
        int ret;
 
        }
        ret = regmap_update_bits(rk808->regmap, reg, bit, bit);
        if (ret)
-               dev_err(&rk808_i2c_client->dev, "Failed to restart device!\n");
+               dev_err(&rk808->i2c->dev, "Failed to restart device!\n");
 
        return NOTIFY_DONE;
 }
 
-static struct notifier_block rk808_restart_handler = {
-       .notifier_call = rk808_restart_notify,
-       .priority = 192,
-};
-
 static void rk8xx_shutdown(struct i2c_client *client)
 {
        struct rk808 *rk808 = i2c_get_clientdata(client);
                return -EINVAL;
        }
 
-       ret = regmap_add_irq_chip(rk808->regmap, client->irq,
-                                 IRQF_ONESHOT, -1,
-                                 rk808->regmap_irq_chip, &rk808->irq_data);
+       ret = devm_regmap_add_irq_chip(&client->dev, rk808->regmap, client->irq,
+                                      IRQF_ONESHOT, -1,
+                                      rk808->regmap_irq_chip, &rk808->irq_data);
        if (ret) {
                dev_err(&client->dev, "Failed to add irq_chip %d\n", ret);
                return ret;
                              regmap_irq_get_domain(rk808->irq_data));
        if (ret) {
                dev_err(&client->dev, "failed to add MFD devices %d\n", ret);
-               goto err_irq;
+               return ret;
        }
 
        if (of_property_read_bool(np, "rockchip,system-power-controller")) {
-               rk808_i2c_client = client;
-               pm_power_off = rk808_pm_power_off;
+               ret = devm_register_sys_off_handler(&client->dev,
+                                   SYS_OFF_MODE_POWER_OFF_PREPARE, SYS_OFF_PRIO_HIGH,
+                                   &rk808_power_off, rk808);
+               if (ret)
+                       return dev_err_probe(&client->dev, ret,
+                                            "failed to register poweroff handler\n");
 
                switch (rk808->variant) {
                case RK809_ID:
                case RK817_ID:
-                       ret = register_restart_handler(&rk808_restart_handler);
+                       ret = devm_register_sys_off_handler(&client->dev,
+                                                           SYS_OFF_MODE_RESTART, SYS_OFF_PRIO_HIGH,
+                                                           &rk808_restart, rk808);
                        if (ret)
                                dev_warn(&client->dev, "failed to register rst handler, %d\n", ret);
                        break;
        }
 
        return 0;
-
-err_irq:
-       regmap_del_irq_chip(client->irq, rk808->irq_data);
-       return ret;
-}
-
-static void rk808_remove(struct i2c_client *client)
-{
-       struct rk808 *rk808 = i2c_get_clientdata(client);
-
-       regmap_del_irq_chip(client->irq, rk808->irq_data);
-
-       /**
-        * pm_power_off may points to a function from another module.
-        * Check if the pointer is set by us and only then overwrite it.
-        */
-       if (pm_power_off == rk808_pm_power_off)
-               pm_power_off = NULL;
-
-       unregister_restart_handler(&rk808_restart_handler);
 }
 
 static int __maybe_unused rk8xx_suspend(struct device *dev)
                .pm = &rk8xx_pm_ops,
        },
        .probe_new = rk808_probe,
-       .remove   = rk808_remove,
        .shutdown = rk8xx_shutdown,
 };