diff --git a/drivers/power/p938x_charger/p938x_charger.c b/drivers/power/p938x_charger/p938x_charger.c
index a67a85294b87..cc93350e6a3f 100644
--- a/drivers/power/p938x_charger/p938x_charger.c
+++ b/drivers/power/p938x_charger/p938x_charger.c
@@ -303,7 +303,7 @@ struct p938x_charger {
 
 	unsigned long flags;
 
-	struct wakeup_source wls_wake_source;
+	struct wakeup_source *wls_wake_source;
 	bool	epp_mode;
 
 	u32 bpp_current_limit;
@@ -391,10 +391,10 @@ static void p938x_reset(struct p938x_charger *chip)
 static void p938x_pm_set_awake(struct p938x_charger *chip, int awake)
 {
 	if(!test_bit(WLS_FLAG_KEEP_AWAKE, &chip->flags) && awake) {
-		__pm_stay_awake(&chip->wls_wake_source);
+		__pm_stay_awake(chip->wls_wake_source);
 		set_bit(WLS_FLAG_KEEP_AWAKE, &chip->flags);
 	} else if(test_bit(WLS_FLAG_KEEP_AWAKE, &chip->flags) && !awake) {
-		__pm_relax(&chip->wls_wake_source);
+		__pm_relax(chip->wls_wake_source);
 		clear_bit(WLS_FLAG_KEEP_AWAKE, &chip->flags);
 	}
 }
@@ -2942,7 +2942,13 @@ static int p938x_charger_probe(struct i2c_client *client,
 	mutex_init(&chip->disconnect_lock);
 	mutex_init(&chip->txmode_lock);
 
-	wakeup_source_init(&chip->wls_wake_source, "p938x wireless charger");
+	chip->wls_wake_source = wakeup_source_register(chip->dev, "p938x wireless charger");
+
+	if (!chip->wls_wake_source) {
+		p938x_err(chip, "failed to allocate wakeup source\n");
+		rc = -ENOMEM;
+		goto free_psy;
+	}
 
 	/* In case we are already powered on */
 	p938x_check_status(chip);
@@ -3012,7 +3018,7 @@ static int p938x_charger_remove(struct i2c_client *client)
 {
 	struct p938x_charger *chip = i2c_get_clientdata(client);
 
-	wakeup_source_trash(&chip->wls_wake_source);
+	wakeup_source_unregister(chip->wls_wake_source);
 	sysfs_remove_groups(&chip->dev->kobj, p938x_groups);
 	cancel_delayed_work_sync(&chip->heartbeat_work);
 	cancel_delayed_work_sync(&chip->tx_mode_work);
diff --git a/drivers/power/qpnp-smbcharger-mmi/qpnp-smbcharger-mmi.c b/drivers/power/qpnp-smbcharger-mmi/qpnp-smbcharger-mmi.c
index 3c4be375dccc..8ccb9ea4a67f 100644
--- a/drivers/power/qpnp-smbcharger-mmi/qpnp-smbcharger-mmi.c
+++ b/drivers/power/qpnp-smbcharger-mmi/qpnp-smbcharger-mmi.c
@@ -372,7 +372,7 @@ struct smb_mmi_charger {
 	struct mmi_sm_params	sm_param[3];
 	int			base_fv_mv;
 	int			vfloat_comp_mv;
-	struct wakeup_source	smb_mmi_hb_wake_source;
+	struct wakeup_source	*smb_mmi_hb_wake_source;
 	struct alarm		heartbeat_alarm;
 	int			suspended;
 	bool			awake;
@@ -2243,7 +2243,7 @@ static enum alarmtimer_restart mmi_heartbeat_alarm_cb(struct alarm *alarm,
 
 	mmi_dbg(chip, "HB alarm fired\n");
 
-	__pm_stay_awake(&chip->smb_mmi_hb_wake_source);
+	__pm_stay_awake(chip->smb_mmi_hb_wake_source);
 	cancel_delayed_work(&chip->heartbeat_work);
 	/* Delay by 500 ms to allow devices to resume. */
 	schedule_delayed_work(&chip->heartbeat_work,
@@ -3430,7 +3430,7 @@ sch_hb:
 
 	kfree(chrg_rate_string);
 
-	__pm_relax(&chip->smb_mmi_hb_wake_source);
+	__pm_relax(chip->smb_mmi_hb_wake_source);
 }
 
 static int mmi_psy_notifier_call(struct notifier_block *nb, unsigned long val,
@@ -4121,8 +4121,13 @@ static int smb_mmi_probe(struct platform_device *pdev)
 
 	INIT_DELAYED_WORK(&chip->heartbeat_work, mmi_heartbeat_work);
 	INIT_DELAYED_WORK(&chip->weakcharger_work, mmi_weakcharger_work);
-	wakeup_source_init(&chip->smb_mmi_hb_wake_source,
-			   "smb_mmi_hb_wake");
+
+	chip->smb_mmi_hb_wake_source = wakeup_source_register(chip->dev, "smb_mmi_hb_wake");
+	if (!chip->smb_mmi_hb_wake_source) {
+		mmi_err(chip, "failed to allocate wakeup source\n");
+		return -ENOMEM;
+	}
+
 	alarm_init(&chip->heartbeat_alarm, ALARM_BOOTTIME,
 		   mmi_heartbeat_alarm_cb);
 
@@ -4444,7 +4449,7 @@ static void smb_mmi_shutdown(struct platform_device *pdev)
 	if (chip->max_flip_psy)
 		power_supply_put(chip->max_flip_psy);
 
-	wakeup_source_trash(&chip->smb_mmi_hb_wake_source);
+	wakeup_source_unregister(chip->smb_mmi_hb_wake_source);
 
 	return;
 }