diff --git a/dsp/q6afe.c b/dsp/q6afe.c
index b65cc98f912d..eb3aa578f17c 100644
--- a/dsp/q6afe.c
+++ b/dsp/q6afe.c
@@ -97,7 +97,7 @@ enum {
 };
 
 struct wlock {
-	struct wakeup_source ws;
+	struct wakeup_source *ws;
 };
 
 static struct wlock wl;
@@ -8525,7 +8525,7 @@ static int afe_set_cal_fb_spkr_prot(int32_t cal_type, size_t data_size,
 		goto done;
 
 	if (cal_data->cal_info.mode == MSM_SPKR_PROT_CALIBRATION_IN_PROGRESS)
-		__pm_wakeup_event(&wl.ws, jiffies_to_msecs(WAKELOCK_TIMEOUT));
+		__pm_wakeup_event(wl.ws, jiffies_to_msecs(WAKELOCK_TIMEOUT));
 	mutex_lock(&this_afe.cal_data[AFE_FB_SPKR_PROT_CAL]->lock);
 	memcpy(&this_afe.prot_cfg, &cal_data->cal_info,
 		sizeof(this_afe.prot_cfg));
@@ -8731,7 +8731,7 @@ static int afe_get_cal_fb_spkr_prot(int32_t cal_type, size_t data_size,
 	}
 	this_afe.initial_cal = 0;
 	mutex_unlock(&this_afe.cal_data[AFE_FB_SPKR_PROT_CAL]->lock);
-	__pm_relax(&wl.ws);
+	__pm_relax(wl.ws);
 done:
 	return ret;
 }
@@ -9025,7 +9025,10 @@ int __init afe_init(void)
 	}
 	init_waitqueue_head(&this_afe.wait_wakeup);
 	init_waitqueue_head(&this_afe.lpass_core_hw_wait);
-	wakeup_source_init(&wl.ws, "spkr-prot");
+	wl.ws = wakeup_source_register(NULL, "spkr-prot");
+	if (!wl.ws)
+		return -ENOMEM;
+
 	ret = afe_init_cal_data();
 	if (ret)
 		pr_err("%s: could not init cal data! %d\n", __func__, ret);
@@ -9033,8 +9036,10 @@ int __init afe_init(void)
 	config_debug_fs_init();
 
 	this_afe.uevent_data = kzalloc(sizeof(*(this_afe.uevent_data)), GFP_KERNEL);
-	if (!this_afe.uevent_data)
+	if (!this_afe.uevent_data) {
+		wakeup_source_unregister(wl.ws);
 		return -ENOMEM;
+	}
 
 	/*
 	 * Set release function to cleanup memory related to kobject
@@ -9069,7 +9074,7 @@ void afe_exit(void)
 	config_debug_fs_exit();
 	mutex_destroy(&this_afe.afe_cmd_lock);
 	mutex_destroy(&this_afe.afe_apr_lock);
-	wakeup_source_trash(&wl.ws);
+	wakeup_source_unregister(wl.ws);
 }
 
 /*