ASoC: wcd9xxx: Add support to compute hph L & R impedance

Add support to compute impedance of the headphone left
and right on Tapan codec. This can be used to adjust gain
according to the value of impedance so that user can have
consistent volume level with different headphones.

Change-Id: Id187bbfaf20c428d367bffa77538182554917e09
Signed-off-by: Phani Kumar Uppalapati <phaniu@codeaurora.org>
diff --git a/sound/soc/codecs/wcd9306.c b/sound/soc/codecs/wcd9306.c
index c509b08..a6ef992 100644
--- a/sound/soc/codecs/wcd9306.c
+++ b/sound/soc/codecs/wcd9306.c
@@ -45,6 +45,9 @@
 #define TAPAN_VDD_CX_OPTIMAL_UA 10000
 #define TAPAN_VDD_CX_SLEEP_UA 2000
 
+/* RX_HPH_CNP_WG_TIME increases by 0.24ms */
+#define TAPAN_WG_TIME_FACTOR_US  240
+
 static struct regulator *tapan_codec_find_regulator(
 	struct snd_soc_codec *codec,
 	const char *name);
@@ -97,6 +100,12 @@
 
 static enum tapan_codec_type codec_ver;
 
+/*
+ * Multiplication factor to compute impedance on Tapan
+ * This is computed from (Vx / (m*Ical)) = (10mV/(180*30uA))
+ */
+#define TAPAN_ZDET_MUL_FACTOR 1852
+
 enum {
 	AIF1_PB = 0,
 	AIF1_CAP,
@@ -251,6 +260,12 @@
 
 	/* pointers to regulators required for chargepump */
 	struct regulator *cp_regulators[CP_REG_MAX];
+
+	/*
+	 * list used to save/restore registers at start and
+	 * end of impedance measurement
+	 */
+	struct list_head reg_save_restore;
 };
 
 static const u32 comp_shift[] = {
@@ -4917,6 +4932,255 @@
 	return WCD9XXX_CDC_TYPE_TAPAN;
 }
 
+static void wcd9xxx_prepare_hph_pa(struct wcd9xxx_mbhc *mbhc,
+				   struct list_head *lh)
+{
+	int i;
+	struct snd_soc_codec *codec = mbhc->codec;
+	u32 delay;
+
+	const struct wcd9xxx_reg_mask_val reg_set_paon[] = {
+		{WCD9XXX_A_CDC_CLSH_B1_CTL, 0x0F, 0x00},
+		{WCD9XXX_A_RX_HPH_CHOP_CTL, 0xFF, 0xA4},
+		{WCD9XXX_A_RX_HPH_OCP_CTL, 0xFF, 0x67},
+		{WCD9XXX_A_RX_HPH_L_TEST, 0x1, 0x0},
+		{WCD9XXX_A_RX_HPH_R_TEST, 0x1, 0x0},
+		{WCD9XXX_A_RX_HPH_BIAS_WG_OCP, 0xFF, 0x1A},
+		{WCD9XXX_A_RX_HPH_CNP_WG_CTL, 0xFF, 0xDB},
+		{WCD9XXX_A_RX_HPH_CNP_WG_TIME, 0xFF, 0x2A},
+		{TAPAN_A_CDC_CONN_RX2_B2_CTL, 0xFF, 0x10},
+		{WCD9XXX_A_CDC_CLK_OTHR_CTL, 0xFF, 0x05},
+		{WCD9XXX_A_CDC_RX1_B6_CTL, 0xFF, 0x81},
+		{WCD9XXX_A_CDC_CLK_RX_B1_CTL, 0x03, 0x03},
+		{WCD9XXX_A_RX_HPH_L_GAIN, 0xFF, 0x2C},
+		{WCD9XXX_A_CDC_RX2_B6_CTL, 0xFF, 0x81},
+		{WCD9XXX_A_RX_HPH_R_GAIN, 0xFF, 0x2C},
+		{WCD9XXX_A_BUCK_CTRL_CCL_4, 0xFF, 0x50},
+		{WCD9XXX_A_BUCK_CTRL_VCL_1, 0xFF, 0x08},
+		{WCD9XXX_A_BUCK_CTRL_CCL_1, 0xFF, 0x5B},
+		{WCD9XXX_A_NCP_CLK, 0xFF, 0x9C},
+		{WCD9XXX_A_NCP_CLK, 0xFF, 0xFC},
+		{WCD9XXX_A_BUCK_MODE_3, 0xFF, 0xCE},
+		{WCD9XXX_A_BUCK_CTRL_CCL_3, 0xFF, 0x6B},
+		{WCD9XXX_A_BUCK_CTRL_CCL_3, 0xFF, 0x6F},
+		{TAPAN_A_RX_BUCK_BIAS1, 0xFF, 0x62},
+		{TAPAN_A_RX_HPH_BIAS_PA, 0xFF, 0x7A},
+		{TAPAN_A_CDC_CLK_RDAC_CLK_EN_CTL, 0xFF, 0x02},
+		{TAPAN_A_CDC_CLK_RDAC_CLK_EN_CTL, 0xFF, 0x06},
+		{WCD9XXX_A_RX_COM_BIAS, 0xFF, 0x80},
+		{WCD9XXX_A_BUCK_MODE_3, 0xFF, 0xC6},
+		{WCD9XXX_A_BUCK_MODE_4, 0xFF, 0xE6},
+		{WCD9XXX_A_BUCK_MODE_5, 0xFF, 0x02},
+		{WCD9XXX_A_BUCK_MODE_1, 0xFF, 0xA1},
+		/* Delay 1ms */
+		{WCD9XXX_A_NCP_EN, 0xFF, 0xFF},
+		/* Delay 1ms */
+		{WCD9XXX_A_BUCK_MODE_5, 0xFF, 0x03},
+		{WCD9XXX_A_BUCK_MODE_5, 0xFF, 0x7B},
+		{WCD9XXX_A_CDC_CLSH_B1_CTL, 0xFF, 0xE6},
+		{WCD9XXX_A_RX_HPH_L_DAC_CTL, 0xFF, 0x40},
+		{WCD9XXX_A_RX_HPH_L_DAC_CTL, 0xFF, 0xC0},
+		{WCD9XXX_A_RX_HPH_R_DAC_CTL, 0xFF, 0x40},
+		{WCD9XXX_A_RX_HPH_R_DAC_CTL, 0xFF, 0xC0},
+		{WCD9XXX_A_NCP_STATIC, 0xFF, 0x08},
+		{WCD9XXX_A_RX_HPH_L_DAC_CTL, 0x03, 0x01},
+		{WCD9XXX_A_RX_HPH_R_DAC_CTL, 0x03, 0x01},
+	};
+
+	/*
+	 * Configure PA in class-AB, -18dB gain,
+	 * companding off, OCP off, Chopping ON
+	 */
+	for (i = 0; i < ARRAY_SIZE(reg_set_paon); i++) {
+		/*
+		 * Some of the codec registers like BUCK_MODE_1
+		 * and NCP_EN requires 1ms wait time for them
+		 * to take effect. Other register writes for
+		 * PA configuration do not require any wait time.
+		 */
+		if (reg_set_paon[i].reg == WCD9XXX_A_BUCK_MODE_1 ||
+		    reg_set_paon[i].reg == WCD9XXX_A_NCP_EN)
+			delay = 1000;
+		else
+			delay = 0;
+		wcd9xxx_soc_update_bits_push(codec, lh,
+					     reg_set_paon[i].reg,
+					     reg_set_paon[i].mask,
+					     reg_set_paon[i].val, delay);
+	}
+	pr_debug("%s: PAs are prepared\n", __func__);
+	return;
+}
+
+static int wcd9xxx_enable_static_pa(struct wcd9xxx_mbhc *mbhc, bool enable)
+{
+	struct snd_soc_codec *codec = mbhc->codec;
+	int wg_time = snd_soc_read(codec, WCD9XXX_A_RX_HPH_CNP_WG_TIME) *
+				   TAPAN_WG_TIME_FACTOR_US;
+	/*
+	 * Tapan requires additional time to enable PA.
+	 * It is observed during experiments that we need
+	 * an additional wait time about 0.35 times of
+	 * the WG_TIME
+	 */
+	wg_time += (int) (wg_time * 35) / 100;
+
+	snd_soc_update_bits(codec, WCD9XXX_A_RX_HPH_CNP_EN, 0x30,
+			    enable ? 0x30 : 0x0);
+	/* Wait for wave gen time to avoid pop noise */
+	usleep_range(wg_time, wg_time + WCD9XXX_USLEEP_RANGE_MARGIN_US);
+	pr_debug("%s: PAs are %s as static mode (wg_time %d)\n", __func__,
+		 enable ? "enabled" : "disabled", wg_time);
+	return 0;
+}
+
+static int tapan_setup_zdet(struct wcd9xxx_mbhc *mbhc,
+			    enum mbhc_impedance_detect_stages stage)
+{
+
+	int ret = 0;
+	struct snd_soc_codec *codec = mbhc->codec;
+	struct tapan_priv *tapan = snd_soc_codec_get_drvdata(codec);
+	const int mux_wait_us = 25;
+
+	switch (stage) {
+
+	case PRE_MEAS:
+		INIT_LIST_HEAD(&tapan->reg_save_restore);
+		/* Configure PA */
+		wcd9xxx_prepare_hph_pa(mbhc, &tapan->reg_save_restore);
+
+#define __wr(reg, mask, value)						  \
+	do {								  \
+		ret = wcd9xxx_soc_update_bits_push(codec,		  \
+						   &tapan->reg_save_restore, \
+						   reg, mask, value, 0);  \
+		if (ret < 0)						  \
+			return ret;					  \
+	} while (0)
+
+		/* Setup MBHC */
+		__wr(WCD9XXX_A_MBHC_SCALING_MUX_1, 0x7F, 0x40);
+		__wr(WCD9XXX_A_MBHC_SCALING_MUX_2, 0xFF, 0xF0);
+		__wr(WCD9XXX_A_TX_7_MBHC_TEST_CTL, 0xFF, 0x78);
+		__wr(WCD9XXX_A_TX_7_MBHC_EN, 0xFF, 0xEC);
+		__wr(WCD9XXX_A_CDC_MBHC_TIMER_B4_CTL, 0xFF, 0x45);
+		__wr(WCD9XXX_A_CDC_MBHC_TIMER_B5_CTL, 0xFF, 0x80);
+
+		__wr(WCD9XXX_A_CDC_MBHC_CLK_CTL, 0xFF, 0x0A);
+		snd_soc_write(codec, WCD9XXX_A_CDC_MBHC_EN_CTL, 0x2);
+		__wr(WCD9XXX_A_CDC_MBHC_CLK_CTL, 0xFF, 0x02);
+
+		/* Enable Impedance Detection */
+		__wr(WCD9XXX_A_MBHC_HPH, 0xFF, 0xC8);
+
+		/*
+		 * CnP setup for 0mV
+		 * Route static data as input to noise shaper
+		 */
+		__wr(TAPAN_A_CDC_RX1_B3_CTL, 0xFF, 0x02);
+		__wr(TAPAN_A_CDC_RX2_B3_CTL, 0xFF, 0x02);
+
+		snd_soc_update_bits(codec, WCD9XXX_A_RX_HPH_L_TEST,
+				    0x02, 0x00);
+		snd_soc_update_bits(codec, WCD9XXX_A_RX_HPH_R_TEST,
+				    0x02, 0x00);
+
+		/* Reset the HPHL static data pointer */
+		__wr(TAPAN_A_CDC_RX1_B2_CTL, 0xFF, 0x00);
+		/* Four consecutive writes to set 0V as static data input */
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B1_CTL, 0x00);
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B1_CTL, 0x00);
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B1_CTL, 0x00);
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B1_CTL, 0x00);
+
+		/* Reset the HPHR static data pointer */
+		__wr(TAPAN_A_CDC_RX2_B2_CTL, 0xFF, 0x00);
+		/* Four consecutive writes to set 0V as static data input */
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B1_CTL, 0x00);
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B1_CTL, 0x00);
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B1_CTL, 0x00);
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B1_CTL, 0x00);
+
+		/* Enable the HPHL and HPHR PA */
+		wcd9xxx_enable_static_pa(mbhc, true);
+		break;
+	case POST_MEAS:
+		/* Turn off ICAL */
+		snd_soc_write(codec, WCD9XXX_A_MBHC_SCALING_MUX_2, 0xF0);
+
+		wcd9xxx_enable_static_pa(mbhc, false);
+
+		/*
+		 * Setup CnP wavegen to ramp to the desired
+		 * output using a 40ms ramp
+		 */
+
+		/* CnP wavegen current to 0.5uA */
+		snd_soc_write(codec, WCD9XXX_A_RX_HPH_BIAS_WG_OCP, 0x1A);
+		/* Set the current division ratio to 2000 */
+		snd_soc_write(codec, WCD9XXX_A_RX_HPH_CNP_WG_CTL, 0xDF);
+		/* Set the wavegen timer to max (60msec) */
+		snd_soc_write(codec, WCD9XXX_A_RX_HPH_CNP_WG_TIME, 0xA0);
+		/* Set the CnP reference current to sc_bias */
+		snd_soc_write(codec, WCD9XXX_A_RX_HPH_OCP_CTL, 0x6D);
+
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B2_CTL, 0x00);
+		/* Four consecutive writes to set -10mV as static data input */
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B1_CTL, 0x00);
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B1_CTL, 0x1F);
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B1_CTL, 0x19);
+		snd_soc_write(codec, TAPAN_A_CDC_RX1_B1_CTL, 0xAA);
+
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B2_CTL, 0x00);
+		/* Four consecutive writes to set -10mV as static data input */
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B1_CTL, 0x00);
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B1_CTL, 0x1F);
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B1_CTL, 0x19);
+		snd_soc_write(codec, TAPAN_A_CDC_RX2_B1_CTL, 0xAA);
+
+		snd_soc_update_bits(codec, WCD9XXX_A_RX_HPH_L_TEST,
+				    0x02, 0x02);
+		snd_soc_update_bits(codec, WCD9XXX_A_RX_HPH_R_TEST,
+				    0x02, 0x02);
+		/* Enable the HPHL and HPHR PA and wait for 60mS */
+		wcd9xxx_enable_static_pa(mbhc, true);
+
+		snd_soc_update_bits(codec, WCD9XXX_A_MBHC_SCALING_MUX_1,
+				    0x7F, 0x40);
+		usleep_range(mux_wait_us,
+				mux_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
+		break;
+	case PA_DISABLE:
+		wcd9xxx_enable_static_pa(mbhc, false);
+		wcd9xxx_restore_registers(codec, &tapan->reg_save_restore);
+		break;
+	}
+#undef __wr
+
+	return ret;
+}
+
+static void tapan_compute_impedance(s16 *l, s16 *r, uint32_t *zl, uint32_t *zr)
+{
+	int zln, zld;
+	int zrn, zrd;
+	int rl = 0, rr = 0;
+
+	zln = (l[1] - l[0]) * TAPAN_ZDET_MUL_FACTOR;
+	zld = (l[2] - l[0]);
+	if (zld)
+		rl = zln / zld;
+
+	zrn = (r[1] - r[0]) * TAPAN_ZDET_MUL_FACTOR;
+	zrd = (r[2] - r[0]);
+	if (zrd)
+		rr = zrn / zrd;
+
+	*zl = rl;
+	*zr = rr;
+}
+
 static const struct wcd9xxx_mbhc_cb mbhc_cb = {
 	.enable_mux_bias_block = tapan_enable_mux_bias_block,
 	.cfilt_fast_mode = tapan_put_cfilt_fast_mode,
@@ -4926,6 +5190,8 @@
 	.select_cfilt = tapan_select_cfilt,
 	.free_irq = tapan_free_irq,
 	.get_cdc_type = tapan_get_cdc_type,
+	.setup_zdet = tapan_setup_zdet,
+	.compute_impedance = tapan_compute_impedance,
 };
 
 int tapan_hs_detect(struct snd_soc_codec *codec,
@@ -5000,7 +5266,8 @@
 		rco_clk_rate = TAPAN_MCLK_CLK_9P6MHZ;
 
 	ret = wcd9xxx_mbhc_init(&tapan->mbhc, &tapan->resmgr, codec, NULL,
-				&mbhc_cb, rco_clk_rate, false);
+				&mbhc_cb, rco_clk_rate,
+				TAPAN_CDC_ZDET_SUPPORTED);
 	if (ret)
 		pr_err("%s: mbhc init failed %d\n", __func__, ret);
 	else
@@ -5203,7 +5470,8 @@
 		rco_clk_rate = TAPAN_MCLK_CLK_9P6MHZ;
 
 	ret = wcd9xxx_mbhc_init(&tapan->mbhc, &tapan->resmgr, codec, NULL,
-				&mbhc_cb, rco_clk_rate, false);
+				&mbhc_cb, rco_clk_rate,
+				TAPAN_CDC_ZDET_SUPPORTED);
 
 	if (ret) {
 		pr_err("%s: mbhc init failed %d\n", __func__, ret);
diff --git a/sound/soc/codecs/wcd9306.h b/sound/soc/codecs/wcd9306.h
index fdd62d1..c4788cc 100644
--- a/sound/soc/codecs/wcd9306.h
+++ b/sound/soc/codecs/wcd9306.h
@@ -1,4 +1,4 @@
-/* Copyright (c) 2012, The Linux Foundation. All rights reserved.
+/* Copyright (c) 2012-2013, The Linux Foundation. All rights reserved.
  *
  * 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
@@ -24,6 +24,8 @@
 
 #define TAPAN_REG_VAL(reg, val)		{reg, 0, val}
 
+#define TAPAN_CDC_ZDET_SUPPORTED  true
+
 extern const u8 tapan_reg_readable[TAPAN_CACHE_SIZE];
 extern const u8 tapan_reset_reg_defaults[TAPAN_CACHE_SIZE];
 struct tapan_codec_dai_data {
diff --git a/sound/soc/codecs/wcd9320.c b/sound/soc/codecs/wcd9320.c
index 3621879..e52eb9d 100644
--- a/sound/soc/codecs/wcd9320.c
+++ b/sound/soc/codecs/wcd9320.c
@@ -49,6 +49,9 @@
 
 #define DAPM_MICBIAS2_EXTERNAL_STANDALONE "MIC BIAS2 External Standalone"
 
+/* RX_HPH_CNP_WG_TIME increases by 0.24ms */
+#define TAIKO_WG_TIME_FACTOR_US	240
+
 static atomic_t kp_taiko_priv;
 static int spkr_drv_wrnd_param_set(const char *val,
 				   const struct kernel_param *kp);
@@ -509,6 +512,12 @@
 
 	int (*machine_codec_event_cb)(struct snd_soc_codec *codec,
 			enum wcd9xxx_codec_event);
+
+	/*
+	 * list used to save/restore registers at start and
+	 * end of impedance measurement
+	 */
+	struct list_head reg_save_restore;
 };
 
 static const u32 comp_shift[] = {
@@ -6239,6 +6248,198 @@
 	return 0;
 }
 
+static int wcd9xxx_prepare_static_pa(struct wcd9xxx_mbhc *mbhc,
+				     struct list_head *lh)
+{
+	int i;
+	struct snd_soc_codec *codec = mbhc->codec;
+
+	const struct wcd9xxx_reg_mask_val reg_set_paon[] = {
+		{WCD9XXX_A_RX_HPH_OCP_CTL, 0x18, 0x00},
+		{WCD9XXX_A_RX_HPH_L_TEST, 0x1, 0x0},
+		{WCD9XXX_A_RX_HPH_R_TEST, 0x1, 0x0},
+		{WCD9XXX_A_RX_HPH_BIAS_WG_OCP, 0xff, 0x1A},
+		{WCD9XXX_A_RX_HPH_CNP_WG_CTL, 0xff, 0xDB},
+		{WCD9XXX_A_RX_HPH_CNP_WG_TIME, 0xff, 0x15},
+		{WCD9XXX_A_CDC_RX1_B6_CTL, 0xff, 0x81},
+		{WCD9XXX_A_CDC_CLK_RX_B1_CTL, 0x01, 0x01},
+		{WCD9XXX_A_RX_HPH_CHOP_CTL, 0xff, 0xA4},
+		{WCD9XXX_A_RX_HPH_L_GAIN, 0xff, 0x2C},
+		{WCD9XXX_A_CDC_RX2_B6_CTL, 0xff, 0x81},
+		{WCD9XXX_A_CDC_CLK_RX_B1_CTL, 0x02, 0x02},
+		{WCD9XXX_A_RX_HPH_R_GAIN, 0xff, 0x2C},
+		{WCD9XXX_A_NCP_CLK, 0xff, 0xFC},
+		{WCD9XXX_A_BUCK_CTRL_CCL_3, 0xff, 0x60},
+		{WCD9XXX_A_RX_COM_BIAS, 0xff, 0x80},
+		{WCD9XXX_A_BUCK_MODE_3, 0xff, 0xC6},
+		{WCD9XXX_A_BUCK_MODE_4, 0xff, 0xE6},
+		{WCD9XXX_A_BUCK_MODE_5, 0xff, 0x02},
+		{WCD9XXX_A_BUCK_MODE_1, 0xff, 0xA1},
+		{WCD9XXX_A_NCP_EN, 0xff, 0xFF},
+		{WCD9XXX_A_BUCK_MODE_5, 0xff, 0x7B},
+		{WCD9XXX_A_CDC_CLSH_B1_CTL, 0xff, 0xE6},
+		{WCD9XXX_A_RX_HPH_L_DAC_CTL, 0xff, 0xC0},
+		{WCD9XXX_A_RX_HPH_R_DAC_CTL, 0xff, 0xC0},
+	};
+
+	for (i = 0; i < ARRAY_SIZE(reg_set_paon); i++)
+		wcd9xxx_soc_update_bits_push(codec, lh,
+					     reg_set_paon[i].reg,
+					     reg_set_paon[i].mask,
+					     reg_set_paon[i].val, 0);
+	pr_debug("%s: PAs are prepared\n", __func__);
+
+	return 0;
+}
+
+static int wcd9xxx_enable_static_pa(struct wcd9xxx_mbhc *mbhc, bool enable)
+{
+	struct snd_soc_codec *codec = mbhc->codec;
+	const int wg_time = snd_soc_read(codec, WCD9XXX_A_RX_HPH_CNP_WG_TIME) *
+			    TAIKO_WG_TIME_FACTOR_US;
+
+	snd_soc_update_bits(codec, WCD9XXX_A_RX_HPH_CNP_EN, 0x30,
+			    enable ? 0x30 : 0x0);
+	/* Wait for wave gen time to avoid pop noise */
+	usleep_range(wg_time, wg_time + WCD9XXX_USLEEP_RANGE_MARGIN_US);
+	pr_debug("%s: PAs are %s as static mode (wg_time %d)\n", __func__,
+		 enable ? "enabled" : "disabled", wg_time);
+	return 0;
+}
+
+static int taiko_setup_zdet(struct wcd9xxx_mbhc *mbhc,
+			    enum mbhc_impedance_detect_stages stage)
+{
+	int ret = 0;
+	struct snd_soc_codec *codec = mbhc->codec;
+	struct taiko_priv *taiko = snd_soc_codec_get_drvdata(codec);
+	const int ramp_wait_us = 18 * 1000;
+
+#define __wr(reg, mask, value)						  \
+	do {								  \
+		ret = wcd9xxx_soc_update_bits_push(codec,		  \
+						   &taiko->reg_save_restore, \
+						   reg, mask, value, 0);  \
+		if (ret < 0)						  \
+			return ret;					  \
+	} while (0)
+
+	switch (stage) {
+
+	case PRE_MEAS:
+		INIT_LIST_HEAD(&taiko->reg_save_restore);
+		wcd9xxx_prepare_static_pa(mbhc, &taiko->reg_save_restore);
+		wcd9xxx_enable_static_pa(mbhc, true);
+
+		/*
+		 * save old value of registers and write the new value to
+		 * restore old value back, WCD9XXX_A_CDC_PA_RAMP_B{1,2,3,4}_CTL
+		 * registers don't need to be restored as those are solely used
+		 * by impedance detection.
+		 */
+		/* Phase 1 */
+		/* Reset the PA Ramp */
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B1_CTL, 0x1C);
+		/*
+		 * Connect the PA Ramp to PA chain and release reset with
+		 * keep it connected.
+		 */
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B1_CTL, 0x1F);
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B1_CTL, 0x03);
+		/*
+		 * Program the PA Ramp to FS_48K, L shift 1 and sample
+		 * num to 24
+		 */
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B3_CTL,
+			      0x3 << 4 | 0x6);
+		/* 0x56 for 10mv.  0xC0 is for 50mv */
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B4_CTL, 0xC0);
+		/* Enable MBHC MUX, Set MUX current to 37.5uA and ADC7 */
+		__wr(WCD9XXX_A_MBHC_SCALING_MUX_1, 0xFF, 0xC0);
+		__wr(WCD9XXX_A_MBHC_SCALING_MUX_2, 0xFF, 0xF0);
+		__wr(WCD9XXX_A_TX_7_MBHC_TEST_CTL, 0xFF, 0x78);
+		__wr(WCD9XXX_A_TX_7_MBHC_EN, 0xFF, 0x8C);
+		/* Change NSA and NAVG */
+		__wr(WCD9XXX_A_CDC_MBHC_TIMER_B4_CTL, 0x4 << 4, 0x4 << 4);
+		__wr(WCD9XXX_A_CDC_MBHC_TIMER_B5_CTL, 0xFF, 0x10);
+		/* Reset MBHC and set it up for STA */
+		__wr(WCD9XXX_A_CDC_MBHC_CLK_CTL, 0xFF, 0x0A);
+		__wr(WCD9XXX_A_CDC_MBHC_EN_CTL, 0xFF, 0x02);
+		__wr(WCD9XXX_A_CDC_MBHC_CLK_CTL, 0xFF, 0x02);
+
+		/* Set HPH_MBHC for zdet */
+		__wr(WCD9XXX_A_MBHC_HPH, 0xB3, 0x80);
+		break;
+	case POST_MEAS:
+		/* Phase 2 */
+		/* Start the PA ramp on HPH L and R */
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B2_CTL, 0x05);
+		/* Ramp generator takes ~17ms */
+		usleep_range(ramp_wait_us,
+				ramp_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
+
+		/* Disable Ical */
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B2_CTL, 0x00);
+		/* Ramp generator takes ~17ms */
+		usleep_range(ramp_wait_us,
+				ramp_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
+		break;
+	case PA_DISABLE:
+		/* Ramp HPH L & R back to Zero */
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B2_CTL, 0x0A);
+		/* Ramp generator takes ~17ms */
+		usleep_range(ramp_wait_us,
+				ramp_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B2_CTL, 0x00);
+
+		/* Clean up starts */
+		/* Turn off PA ramp generator */
+		snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B1_CTL, 0x0);
+		wcd9xxx_enable_static_pa(mbhc, false);
+		wcd9xxx_restore_registers(codec, &taiko->reg_save_restore);
+		break;
+	}
+#undef __wr
+
+	return ret;
+}
+
+static void taiko_compute_impedance(s16 *l, s16 *r, uint32_t *zl, uint32_t *zr)
+{
+
+	int64_t rl, rr = 0; /* milliohm */
+	const int alphal = 364; /* 0.005555 * 65536 = 364.05 */
+	const int alphar = 364; /* 0.005555 * 65536 = 364.05 */
+	const int beta = 3855; /* 0.011765 * 5 * 65536 = 3855.15 */
+	const int rref = 11333; /* not scaled up */
+	const int shift = 16;
+
+	rl = (int)(l[0] - l[1]) * 1000 / (l[0] - l[2]);
+	rl = rl * rref * alphal;
+	rl = rl >> shift;
+	rl = rl * beta;
+	rl = rl >> shift;
+	*zl = rl;
+
+	rr = (int)(r[0] - r[1]) * 1000 / (r[0] - r[2]);
+	rr = rr * rref  * alphar;
+	rr = rr >> shift;
+	rr = rr * beta;
+	rr = rr >> shift;
+	*zr = rr;
+}
+
+static enum wcd9xxx_cdc_type taiko_get_cdc_type(void)
+{
+	return WCD9XXX_CDC_TYPE_TAIKO;
+}
+
+static const struct wcd9xxx_mbhc_cb mbhc_cb = {
+	.get_cdc_type = taiko_get_cdc_type,
+	.setup_zdet = taiko_setup_zdet,
+	.compute_impedance = taiko_compute_impedance,
+};
+
 static int taiko_post_reset_cb(struct wcd9xxx *wcd9xxx)
 {
 	int ret = 0;
@@ -6284,7 +6485,7 @@
 
 		ret = wcd9xxx_mbhc_init(&taiko->mbhc, &taiko->resmgr, codec,
 					taiko_enable_mbhc_micbias,
-					NULL, rco_clk_rate, true);
+					&mbhc_cb, rco_clk_rate, true);
 		if (ret)
 			pr_err("%s: mbhc init failed %d\n", __func__, ret);
 		else
@@ -6473,7 +6674,7 @@
 	/* init and start mbhc */
 	ret = wcd9xxx_mbhc_init(&taiko->mbhc, &taiko->resmgr, codec,
 				taiko_enable_mbhc_micbias,
-				NULL, rco_clk_rate, true);
+				&mbhc_cb, rco_clk_rate, true);
 	if (ret) {
 		pr_err("%s: mbhc init failed %d\n", __func__, ret);
 		goto err_init;
diff --git a/sound/soc/codecs/wcd9xxx-common.c b/sound/soc/codecs/wcd9xxx-common.c
index d00b843..05f2191 100644
--- a/sound/soc/codecs/wcd9xxx-common.c
+++ b/sound/soc/codecs/wcd9xxx-common.c
@@ -11,6 +11,7 @@
  */
 
 #include <linux/module.h>
+#include <linux/slab.h>
 #include <sound/soc.h>
 #include <linux/kernel.h>
 #include <linux/delay.h>
@@ -32,6 +33,8 @@
 
 #define MAX_IMPED_PARAMS 13
 
+#define USLEEP_RANGE_MARGIN_US 100
+
 struct wcd9xxx_imped_val {
 	u32 imped_val;
 	u8 index;
@@ -614,6 +617,46 @@
 	}
 }
 
+
+int wcd9xxx_soc_update_bits_push(struct snd_soc_codec *codec,
+					struct list_head *list,
+					uint16_t reg, uint8_t mask,
+					uint8_t value, int delay)
+{
+	int rc;
+	struct wcd9xxx_register_save_node *node;
+
+	node = kmalloc(sizeof(*node), GFP_KERNEL);
+	if (unlikely(!node)) {
+		pr_err("%s: Not enough memory\n", __func__);
+		return -ENOMEM;
+	}
+	node->reg = reg;
+	node->value = snd_soc_read(codec, reg);
+	list_add(&node->lh, list);
+	if (mask == 0xFF)
+		rc = snd_soc_write(codec, reg, value);
+	else
+		rc = snd_soc_update_bits(codec, reg, mask, value);
+	if (delay)
+		usleep_range(delay, delay + USLEEP_RANGE_MARGIN_US);
+	return rc;
+}
+EXPORT_SYMBOL(wcd9xxx_soc_update_bits_push);
+
+void wcd9xxx_restore_registers(struct snd_soc_codec *codec,
+			       struct list_head *lh)
+{
+	struct wcd9xxx_register_save_node *node, *nodetmp;
+
+	list_for_each_entry_safe(node, nodetmp, lh, lh) {
+		snd_soc_write(codec, node->reg, node->value);
+		list_del(&node->lh);
+		kfree(node);
+	}
+}
+EXPORT_SYMBOL(wcd9xxx_restore_registers);
+
 static void wcd9xxx_enable_buck_mode(struct snd_soc_codec *codec,
 		u8 buck_vref)
 {
diff --git a/sound/soc/codecs/wcd9xxx-common.h b/sound/soc/codecs/wcd9xxx-common.h
index 654964e..9fba8e3 100644
--- a/sound/soc/codecs/wcd9xxx-common.h
+++ b/sound/soc/codecs/wcd9xxx-common.h
@@ -81,4 +81,16 @@
 	WCD9XXX_CODEC_EVENT_CODEC_UP = 0,
 };
 
+struct wcd9xxx_register_save_node {
+	struct list_head lh;
+	u16 reg;
+	u16 value;
+};
+
+extern int wcd9xxx_soc_update_bits_push(struct snd_soc_codec *codec,
+					struct list_head *lh,
+					uint16_t reg, uint8_t mask,
+					uint8_t value, int delay);
+extern void wcd9xxx_restore_registers(struct snd_soc_codec *codec,
+				      struct list_head *lh);
 #endif
diff --git a/sound/soc/codecs/wcd9xxx-mbhc.c b/sound/soc/codecs/wcd9xxx-mbhc.c
index cc3c19c..4e95fa9 100644
--- a/sound/soc/codecs/wcd9xxx-mbhc.c
+++ b/sound/soc/codecs/wcd9xxx-mbhc.c
@@ -142,12 +142,6 @@
 	enum wcd9xxx_mbhc_plug_type _type;
 };
 
-struct wcd9xxx_register_save_node {
-	struct list_head lh;
-	u16 reg;
-	u16 value;
-};
-
 enum meas_type {
 	STA = 0,
 	DCE,
@@ -4216,107 +4210,6 @@
 	return ret;
 }
 
-static int wcd9xxx_soc_update_bits_push(struct snd_soc_codec *codec,
-					struct list_head *list,
-					uint16_t reg, uint8_t mask,
-					uint8_t value)
-{
-	int rc;
-	struct wcd9xxx_register_save_node *node;
-
-	node = kmalloc(sizeof(*node), GFP_KERNEL);
-	if (unlikely(!node)) {
-		pr_err("%s: Not enough memory\n", __func__);
-		return -ENOMEM;
-	}
-	node->reg = reg;
-	node->value = snd_soc_read(codec, reg);
-	list_add(&node->lh, list);
-	if (mask == 0xFF)
-		rc = snd_soc_write(codec, reg, value);
-	else
-		rc = snd_soc_update_bits(codec, reg, mask, value);
-	return rc;
-}
-
-static int wcd9xxx_prepare_static_pa(struct wcd9xxx_mbhc *mbhc,
-				     struct list_head *lh)
-{
-	int i;
-	struct snd_soc_codec *codec = mbhc->codec;
-
-	const struct wcd9xxx_reg_mask_val reg_set_paon[] = {
-		{WCD9XXX_A_RX_HPH_OCP_CTL, 0x18, 0x00},
-		{WCD9XXX_A_RX_HPH_L_TEST, 0x1, 0x0},
-		{WCD9XXX_A_RX_HPH_R_TEST, 0x1, 0x0},
-		{WCD9XXX_A_RX_HPH_BIAS_WG_OCP, 0xff, 0x1A},
-		{WCD9XXX_A_RX_HPH_CNP_WG_CTL, 0xff, 0xDB},
-		{WCD9XXX_A_RX_HPH_CNP_WG_TIME, 0xff, 0x15},
-		{WCD9XXX_A_CDC_RX1_B6_CTL, 0xff, 0x81},
-		{WCD9XXX_A_CDC_CLK_RX_B1_CTL, 0x01, 0x01},
-		{WCD9XXX_A_RX_HPH_CHOP_CTL, 0xff, 0xA4},
-		{WCD9XXX_A_RX_HPH_L_GAIN, 0xff, 0x2C},
-		{WCD9XXX_A_CDC_RX2_B6_CTL, 0xff, 0x81},
-		{WCD9XXX_A_CDC_CLK_RX_B1_CTL, 0x02, 0x02},
-		{WCD9XXX_A_RX_HPH_R_GAIN, 0xff, 0x2C},
-		{WCD9XXX_A_NCP_CLK, 0xff, 0xFC},
-		{WCD9XXX_A_BUCK_CTRL_CCL_3, 0xff, 0x60},
-		{WCD9XXX_A_RX_COM_BIAS, 0xff, 0x80},
-		{WCD9XXX_A_BUCK_MODE_3, 0xff, 0xC6},
-		{WCD9XXX_A_BUCK_MODE_4, 0xff, 0xE6},
-		{WCD9XXX_A_BUCK_MODE_5, 0xff, 0x02},
-		{WCD9XXX_A_BUCK_MODE_1, 0xff, 0xA1},
-		{WCD9XXX_A_NCP_EN, 0xff, 0xFF},
-		{WCD9XXX_A_BUCK_MODE_5, 0xff, 0x7B},
-		{WCD9XXX_A_CDC_CLSH_B1_CTL, 0xff, 0xE6},
-		{WCD9XXX_A_RX_HPH_L_DAC_CTL, 0xff, 0xC0},
-		{WCD9XXX_A_RX_HPH_R_DAC_CTL, 0xff, 0xC0},
-	};
-
-	for (i = 0; i < ARRAY_SIZE(reg_set_paon); i++)
-		wcd9xxx_soc_update_bits_push(codec, lh,
-					     reg_set_paon[i].reg,
-					     reg_set_paon[i].mask,
-					     reg_set_paon[i].val);
-	pr_debug("%s: PAs are prepared\n", __func__);
-
-	return 0;
-}
-
-static void wcd9xxx_restore_registers(struct wcd9xxx_mbhc *mbhc,
-				      struct list_head *lh)
-{
-	struct wcd9xxx_register_save_node *node, *nodetmp;
-	struct snd_soc_codec *codec = mbhc->codec;
-
-	list_for_each_entry_safe(node, nodetmp, lh, lh) {
-		snd_soc_write(codec, node->reg, node->value);
-		list_del(&node->lh);
-		kfree(node);
-	}
-}
-
-static void wcd9xxx_unprepare_static_pa(struct wcd9xxx_mbhc *mbhc,
-					struct list_head *lh)
-{
-	wcd9xxx_restore_registers(mbhc, lh);
-}
-
-static int wcd9xxx_enable_static_pa(struct wcd9xxx_mbhc *mbhc, bool enable)
-{
-	struct snd_soc_codec *codec = mbhc->codec;
-	const int wg_time = snd_soc_read(codec, WCD9XXX_A_RX_HPH_CNP_WG_TIME) *
-			    WCD9XXX_WG_TIME_FACTOR_US;
-
-	snd_soc_update_bits(codec, WCD9XXX_A_RX_HPH_CNP_EN, 0x30,
-			    enable ? 0x30 : 0x0);
-	/* Wait for wave gen time to avoid pop noise */
-	usleep_range(wg_time, wg_time + WCD9XXX_USLEEP_RANGE_MARGIN_US);
-	pr_debug("%s: PAs are %s as static mode (wg_time %d)\n", __func__,
-		 enable ? "enabled" : "disabled", wg_time);
-	return 0;
-}
-
 static int wcd9xxx_detect_impedance(struct wcd9xxx_mbhc *mbhc, uint32_t *zl,
 				    uint32_t *zr)
 {
@@ -4326,17 +4219,8 @@
 	s16 *z[] = {
 		&l[0], &r[0], &r[1], &l[1], &l[2], &r[2],
 	};
-	LIST_HEAD(lh);
-	LIST_HEAD(lhpa);
 	struct snd_soc_codec *codec = mbhc->codec;
-	const int ramp_wait_us = 18 * 1000;
 	const int mux_wait_us = 25;
-	const int alphal = 364; /* 0.005555 * 65536 = 364.05 */
-	const int alphar = 364; /* 0.005555 * 65536 = 364.05 */
-	const int beta = 3855; /* 0.011765 * 5 * 65536 = 3855.15 */
-	const int rref = 11333; /* not scaled up */
-	const int shift = 16;
-	int64_t rl, rr = 0; /* milliohm */
 	const struct wcd9xxx_reg_mask_val reg_set_mux[] = {
 		/* Phase 1 */
 		/* Set MBHC_MUX for HPHL without ical */
@@ -4357,6 +4241,11 @@
 	pr_debug("%s: enter\n", __func__);
 	WCD9XXX_BCL_ASSERT_LOCKED(mbhc->resmgr);
 
+	if (!mbhc->mbhc_cb || !mbhc->mbhc_cb->setup_zdet ||
+	    !mbhc->mbhc_cb->compute_impedance || !zl ||
+	    !zr)
+		return -EINVAL;
+
 	/*
 	 * Impedance detection is an intrusive function as it mutes RX paths,
 	 * enable PAs and etc.  Therefore codec drvier including ALSA
@@ -4375,127 +4264,66 @@
 	wcd9xxx_resmgr_get_clk_block(mbhc->resmgr, WCD9XXX_CLK_RCO);
 	WCD9XXX_BG_CLK_UNLOCK(mbhc->resmgr);
 
+	wcd9xxx_turn_onoff_override(codec, true);
 	pr_debug("%s: Setting impedance detection\n", __func__);
-	wcd9xxx_prepare_static_pa(mbhc, &lhpa);
-	wcd9xxx_enable_static_pa(mbhc, true);
 
-	/*
-	 * save old value of registers and write the new value to restore old
-	 * value back, WCD9XXX_A_CDC_PA_RAMP_B{1,2,3,4}_CTL registers don't
-	 * need to be restored as those are solely used by impedance detection.
-	 */
-#define __w(reg, mask, value)						  \
-	do {								  \
-		ret = wcd9xxx_soc_update_bits_push(codec, &lh, reg, mask, \
-						   value);		  \
-		if (ret < 0)						  \
-			return ret;					  \
-	} while (0)
-
-	/* Reset the PA Ramp */
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B1_CTL, 0x1C);
-	/*
-	 * Connect the PA Ramp to PA chain and release reset with keep it
-	 * connected.
-	 */
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B1_CTL, 0x1F);
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B1_CTL, 0x03);
-	/* Program the PA Ramp to FS_48K, L shift 1 and sample num to 24 */
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B3_CTL, 0x3 << 4 | 0x6);
-	/* 0x56 for 10mv.  0xC0 is for 50mv */
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B4_CTL, 0xC0);
-	/* Enable MBHC MUX, Set MUX current to 37.5uA and ADC7 */
-	__w(WCD9XXX_A_MBHC_SCALING_MUX_1, 0xFF, 0xC0);
-	__w(WCD9XXX_A_MBHC_SCALING_MUX_2, 0xFF, 0xF0);
-	__w(WCD9XXX_A_TX_7_MBHC_TEST_CTL, 0xFF, 0x78);
-	__w(WCD9XXX_A_TX_7_MBHC_EN, 0xFF, 0x8C);
-	/* Change NSA and NAVG */
-	__w(WCD9XXX_A_CDC_MBHC_TIMER_B4_CTL, 0x4 << 4, 0x4 << 4);
-	__w(WCD9XXX_A_CDC_MBHC_TIMER_B5_CTL, 0xFF, 0x10);
-	/* Reset MBHC and set it up for STA */
-	__w(WCD9XXX_A_CDC_MBHC_CLK_CTL, 0xFF, 0x0A);
-	__w(WCD9XXX_A_CDC_MBHC_EN_CTL, 0xFF, 0x02);
-	__w(WCD9XXX_A_CDC_MBHC_CLK_CTL, 0xFF, 0x02);
-
-	/* Set HPH_MBHC for zdet */
-	__w(WCD9XXX_A_MBHC_HPH, 0xB3, 0x80);
+	/* Codec specific setup for L0, R0, L1 and R1 measurements */
+	mbhc->mbhc_cb->setup_zdet(mbhc, PRE_MEAS);
 
 	pr_debug("%s: Performing impedance detection\n", __func__);
-	/* Phase 1 */
 	for (i = 0; i < ARRAY_SIZE(reg_set_mux) - 2; i++) {
-		__w(reg_set_mux[i].reg, reg_set_mux[i].mask,
-		    reg_set_mux[i].val);
+		snd_soc_update_bits(codec, reg_set_mux[i].reg,
+				    reg_set_mux[i].mask,
+				    reg_set_mux[i].val);
+		if (mbhc->mbhc_cb && mbhc->mbhc_cb->enable_mux_bias_block)
+			mbhc->mbhc_cb->enable_mux_bias_block(codec);
+		else
+			snd_soc_update_bits(codec,
+					    WCD9XXX_A_MBHC_SCALING_MUX_1,
+					    0x80, 0x80);
 		/* 25us is required after mux change to settle down */
 		usleep_range(mux_wait_us,
 			     mux_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
-		*(z[i]) = __wcd9xxx_codec_sta_dce(mbhc, 0, false, false);
+		*(z[i]) = __wcd9xxx_codec_sta_dce(mbhc, 0, true, false);
 	}
 
-	/* Phase 2 */
-	/* Start the PA ramp on HPH L and R */
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B2_CTL, 0x05);
-	/* Ramp generator takes ~17ms */
-	usleep_range(ramp_wait_us,
-		     ramp_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
+	/* Codec specific setup for L2 and R2 measurements */
+	mbhc->mbhc_cb->setup_zdet(mbhc, POST_MEAS);
 
-	/* Disable Ical */
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B2_CTL, 0x00);
-	/* Ramp generator takes ~17ms */
-	usleep_range(ramp_wait_us,
-		     ramp_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
 	for (; i < ARRAY_SIZE(reg_set_mux); i++) {
-		__w(reg_set_mux[i].reg, reg_set_mux[i].mask,
-		    reg_set_mux[i].val);
+		snd_soc_update_bits(codec, reg_set_mux[i].reg,
+				    reg_set_mux[i].mask,
+				    reg_set_mux[i].val);
+		if (mbhc->mbhc_cb && mbhc->mbhc_cb->enable_mux_bias_block)
+			mbhc->mbhc_cb->enable_mux_bias_block(codec);
+		else
+			snd_soc_update_bits(codec,
+					    WCD9XXX_A_MBHC_SCALING_MUX_1,
+					    0x80, 0x80);
 		/* 25us is required after mux change to settle down */
 		usleep_range(mux_wait_us,
 			     mux_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
-		*(z[i]) = __wcd9xxx_codec_sta_dce(mbhc, 0, false, false);
+		*(z[i]) = __wcd9xxx_codec_sta_dce(mbhc, 0, true, false);
 	}
 
-	/* Ramp HPH L & R back to Zero */
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B2_CTL, 0x0A);
-	/* Ramp generator takes ~17ms */
-	usleep_range(ramp_wait_us,
-		     ramp_wait_us + WCD9XXX_USLEEP_RANGE_MARGIN_US);
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B2_CTL, 0x00);
-#undef __w
-
-	/* Clean up starts */
-	/* Turn off PA ramp generator */
-	snd_soc_write(codec, WCD9XXX_A_CDC_PA_RAMP_B1_CTL, 0x0);
-	wcd9xxx_enable_static_pa(mbhc, false);
-	wcd9xxx_restore_registers(mbhc, &lh);
-	wcd9xxx_unprepare_static_pa(mbhc, &lhpa);
+	mbhc->mbhc_cb->setup_zdet(mbhc, PA_DISABLE);
 
 	mutex_unlock(&codec->mutex);
 
 	WCD9XXX_BG_CLK_LOCK(mbhc->resmgr);
-	wcd9xxx_resmgr_put_bandgap(mbhc->resmgr, WCD9XXX_BANDGAP_MBHC_MODE);
+	wcd9xxx_resmgr_put_bandgap(mbhc->resmgr, WCD9XXX_BANDGAP_AUDIO_MODE);
 	wcd9xxx_resmgr_put_clk_block(mbhc->resmgr, WCD9XXX_CLK_RCO);
 	WCD9XXX_BG_CLK_UNLOCK(mbhc->resmgr);
 
-	rl = (int)(l[0] - l[1]) * 1000 / (l[0] - l[2]);
-	rl = rl * rref * alphal;
-	rl = rl >> shift;
-	rl = rl * beta;
-	rl = rl >> shift;
-	*zl = rl;
+	wcd9xxx_turn_onoff_override(codec, false);
+	mbhc->mbhc_cb->compute_impedance(l, r, zl, zr);
 
-	rr = (int)(r[0] - r[1]) * 1000 / (r[0] - r[2]);
-	rr = rr * rref  * alphar;
-	rr = rr >> shift;
-	rr = rr * beta;
-	rr = rr >> shift;
-	*zr = rr;
-
-	pr_debug("%s: L0: 0x%x(%d), L1: 0x%x(%d), L2: 0x%x(%d), rl: %lld\n",
+	pr_debug("%s: L0: 0x%x(%d), L1: 0x%x(%d), L2: 0x%x(%d)\n",
 		 __func__,
-		 l[0] & 0xffff, l[0], l[1] & 0xffff, l[1], l[2] & 0xffff, l[2],
-		 rl);
-	pr_debug("%s: R0: 0x%x(%d), R1: 0x%x(%d), R2: 0x%x(%d), rr: %lld\n",
+		 l[0] & 0xffff, l[0], l[1] & 0xffff, l[1], l[2] & 0xffff, l[2]);
+	pr_debug("%s: R0: 0x%x(%d), R1: 0x%x(%d), R2: 0x%x(%d)\n",
 		 __func__,
-		 r[0] & 0xffff, r[0], r[1] & 0xffff, r[1], r[2] & 0xffff, r[2],
-		 rr);
+		 r[0] & 0xffff, r[0], r[1] & 0xffff, r[1], r[2] & 0xffff, r[2]);
 	pr_debug("%s: RL %d milliohm, RR %d milliohm\n", __func__, *zl, *zr);
 	pr_debug("%s: Impedance detection completed\n", __func__);
 
diff --git a/sound/soc/codecs/wcd9xxx-mbhc.h b/sound/soc/codecs/wcd9xxx-mbhc.h
index 3040bc4..e5259c3 100644
--- a/sound/soc/codecs/wcd9xxx-mbhc.h
+++ b/sound/soc/codecs/wcd9xxx-mbhc.h
@@ -19,6 +19,8 @@
 #define WCD9XXX_CFILT_EXT_PRCHG_EN 0x30
 #define WCD9XXX_CFILT_EXT_PRCHG_DSBL 0x00
 
+#define WCD9XXX_USLEEP_RANGE_MARGIN_US 100
+
 struct mbhc_micbias_regs {
 	u16 cfilt_val;
 	u16 cfilt_ctl;
@@ -40,6 +42,12 @@
 	MBHC_CAL_NUM,
 };
 
+enum mbhc_impedance_detect_stages {
+	PRE_MEAS,
+	POST_MEAS,
+	PA_DISABLE,
+};
+
 /* Data used by MBHC */
 struct mbhc_internal_cal_data {
 	u16 dce_z;
@@ -239,6 +247,9 @@
 	void (*free_irq) (struct wcd9xxx_mbhc *);
 	enum wcd9xxx_cdc_type (*get_cdc_type) (void);
 	void (*enable_clock_gate) (struct snd_soc_codec *, bool);
+	int (*setup_zdet) (struct wcd9xxx_mbhc *,
+			   enum mbhc_impedance_detect_stages stage);
+	void (*compute_impedance) (s16 *, s16 *, uint32_t *, uint32_t *);
 };
 
 struct wcd9xxx_mbhc {